# Gradient of llama2 computed by Zygote seems to be incorrect

**URL:** <https://discourse.julialang.org/t/gradient-of-llama2-computed-by-zygote-seems-to-be-incorrect/103810>\
**Category:** Machine Learning\
**Created:** [September 13, 2023, 12:20pm UTC](https://discourse.julialang.org/t/gradient-of-llama2-computed-by-zygote-seems-to-be-incorrect/103810 "2023-09-13T12:20:05Z")\
**Posts on this page:** 7\
**Page:** 1

<div class="post-metadata">

**Author:** ![Tomas\_Pevny](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/tomas_pevny/32/25466_2.png) [@Tomas\_Pevny](https://discourse.julialang.org/u/Tomas_Pevny)\
**Post date:** [September 13, 2023, 12:20pm UTC](https://discourse.julialang.org/t/gradient-of-llama2-computed-by-zygote-seems-to-be-incorrect/103810/1 "2023-09-13T12:20:05Z")

</div>

Hi,

I have tested gradient of llama2-7b model with respect its input and it gives me wrong results (tested only on GPU). I have created a MWE as follows

```julia
using ProfileSummarizer
using Transformers
using Flux
using TextEncodeBase
using NeuralAttentionlib
using HuggingFaceApi
using FiniteDifferences
using Zygote
using CUDA
using StatsBase

CUDA.device!(0)
textenc = HuggingFace.load_tokenizer("meta-llama/Llama-2-7b-chat-hf"; auth_token = access_token);
model = f32(HuggingFace.load_model("llama", "meta-llama/Llama-2-7b-chat-hf", "forCausallm", auth_token = access_token));
model = gpu(model);
embeddings = model.model.embedding.token.embeddings;
decoder = model.model.decoder;

# We compute randomly gradients with respect to few selected elements and compare them
# to true values. This is a sanity check, because otherwise the test would take ages.

let begin
	tokens = TextEncodeBase.encode(textenc, "the most important text for this test").token
	θ = cpu(embeddings * tokens)
	ii = sample(eachindex(θ), 1000, replace = false)
	ii = [23, 5641, 18928, 4354, 16801, 294, 18929, 17793, 1277, 17098]
	sub_θ = θ[ii]
	function sub_f(sub_θ)
		θ[ii] = sub_θ
		f(θ)
	end

	function f(θ)
		hidden_state = gpu(θ)
		sum(decoder((;hidden_state)).hidden_state)
	end

	fin_gs = grad(central_fdm(5, 1), sub_f, sub_θ)[1]
	zyg_gs = Zygote.gradient(f, θ)[1][ii]
	hcat(fin_gs, zyg_gs)
	sort(abs.(fin_gs .- zyg_gs), rev = true)
end

```

which returns

```julia
10-element Vector{Float32}:
 4.4233932
 2.523777
 1.9660205
 1.4594455
 0.67348194
 0.589447
 0.5284171
 0.5228882
 0.36881256
 0.18504143

```

which seems to me off by a large margin.

I use Julia 1.9.2 and my `Pkg.status()` says

```julia
  [7d9f7c33] Accessors v0.1.32
  [6e4b80f9] BenchmarkTools v1.3.2
⌃ [052768ef] CUDA v4.4.0
  [d360d2e6] ChainRulesCore v1.16.0
  [a93c6f00] DataFrames v1.6.1
  [26cc04aa] FiniteDifferences v0.12.30
⌅ [587475ba] Flux v0.13.17
  [d9f16b24] Functors v0.4.5
  [3cc741c3] HuggingFaceApi v0.1.0
  [f1d291b0] MLUtils v0.4.3
  [5da4648a] NVTX v0.3.2
⌃ [12afc1b8] NeuralAttentionlib v0.2.11
  [0b1bfda6] OneHotArrays v0.2.4
  [6099a3de] PythonCall v0.9.14
  [2913bbd2] StatsBase v0.34.0
  [354b36f9] StringViews v1.3.3
⌃ [899adc3e] TensorBoardLogger v0.1.21
  [f92c20c0] TextEncodeBase v0.6.0
  [21ca0261] Transformers v0.2.7
⌃ [e88e6eb3] Zygote v0.6.63

```

I will try to narrow down the problem to test individual layers, but this is going to take bit of time.

---

<div class="post-metadata">

**Author:** ![Tomas\_Pevny](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/tomas_pevny/32/25466_2.png) [@Tomas\_Pevny](https://discourse.julialang.org/u/Tomas_Pevny)\
**Post date:** [September 13, 2023, 12:47pm UTC](https://discourse.julialang.org/t/gradient-of-llama2-computed-by-zygote-seems-to-be-incorrect/103810/2 "2023-09-13T12:47:49Z")

</div>

It seems to me to be the effect of normalization layers and `Float32`, which was my suspect. I redid my mwe as follows

```julia
using ProfileSummarizer
using Transformers
using Flux
using TextEncodeBase
using NeuralAttentionlib
using HuggingFaceApi
using FiniteDifferences
using Zygote
using CUDA
using StatsBase

CUDA.device!(0)
textenc = HuggingFace.load_tokenizer("meta-llama/Llama-2-7b-chat-hf"; auth_token = access_token);
model = f32(HuggingFace.load_model("llama", "meta-llama/Llama-2-7b-chat-hf", "forCausallm", auth_token = access_token));
model = gpu(model);
embeddings = model.model.embedding.token.embeddings;
decoder = model.model.decoder;

function test_layer(layer;togpu = cpu, float_precision = f64, use_hidden = true)
	layer = togpu(float_precision(layer))
	tokens = TextEncodeBase.encode(textenc, "the most important text for this test").token
	θ = cpu(float_precision(embeddings * tokens))
	ii = [23, 5641, 18928, 4354, 16801, 294, 18929, 17793, 1277, 17098]
	sub_θ = θ[ii]
	function sub_f(sub_θ)
		θ[ii] = sub_θ
		f(θ)
	end

	f(θ)= use_hidden ? sum(layer((;hidden_state = togpu(θ))).hidden_state) : sum(layer(togpu(θ)))
	fin_gs = grad(central_fdm(5, 1), sub_f, sub_θ)[1];
	zyg_gs = Zygote.gradient(f, θ)[1][ii];
	abs.(fin_gs .- zyg_gs)
end

test_layer(decoder.layers[1][1].feedforward.layer;togpu = gpu, float_precision = f32) |> sort
test_layer(decoder.layers[1][1].feedforward.layer;togpu = cpu) |> sort
test_layer(decoder.layers[1][1].feedforward.norm;togpu = gpu, float_precision = f32, use_hidden = false) |> sort
test_layer(decoder.layers[1][1].feedforward.norm;togpu = cpu, float_precision = f32, use_hidden = false) |> sort
test_layer(decoder.layers[1][1].feedforward.norm;togpu = cpu, float_precision = f64, use_hidden = false) |> sort

```

For brevity, the output of last two tests showing execution on `cpu` with different precision gives

```julia
julia> test_layer(decoder.layers[1][1].feedforward.norm;togpu = cpu, float_precision = f32, use_hidden = false) |> sort
10-element Vector{Float32}:
 2.2888184f-5
 3.2424927f-5
 3.4332275f-5
 3.6239624f-5
 0.00011920929
 0.0010175705
 0.0033392906
 0.0042181015
 0.004609585
 0.051587105

julia> test_layer(decoder.layers[1][1].feedforward.norm;togpu = cpu, float_precision = f64, use_hidden = false) |> sort
10-element Vector{Float64}:
 1.602273869139026e-12
 6.565414878423326e-12
 7.716494110354688e-12
 8.521183758603001e-12
 9.283240842705709e-12
 6.4617158352930915e-6
 6.750580402670181e-6
 9.646580691580198e-6
 0.00020662468795151767
 0.0010416488879068098

```

Increasing precision therefore greatly increases accuracy of the gradient.

---

<div class="post-metadata">

**Author:** ![Tomas\_Pevny](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/tomas_pevny/32/25466_2.png) [@Tomas\_Pevny](https://discourse.julialang.org/u/Tomas_Pevny)\
**Post date:** [September 13, 2023, 3:56pm UTC](https://discourse.julialang.org/t/gradient-of-llama2-computed-by-zygote-seems-to-be-incorrect/103810/3 "2023-09-13T15:56:12Z")

</div>

So the problem is in `RMSLayerNorm`. Not that it is wrong, but the finite difference might have large error, when the vectors have small norm.  
This was my suspect from the beginning. I mark this as a solution.

---

<div class="post-metadata">

**Author:** ![Palli](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/palli/32/3380_2.png) [@Palli](https://discourse.julialang.org/u/Palli)\
**Post date:** [September 13, 2023, 4:36pm UTC](https://discourse.julialang.org/t/gradient-of-llama2-computed-by-zygote-seems-to-be-incorrect/103810/4 "2023-09-13T16:36:13Z")

</div>

It seems you figured out and solved the problem?

I want to dive into using these models (with Julia, it seemed simple enough for you, had it worked the first time; or by using Python).

I want to ask why did you choose that model (and that size), and use Julia? Would Python have been easier? It gets 45.3 on MMLU (the largest 70B gets 68.9, so I guess you chose the smallest for resource limitations, what your GPU can handle).

Better and almost as good as the largest Llama 2:

> **[Open-Orca/OpenOrca-Platypus2-13B · Hugging Face](https://huggingface.co/Open-Orca/OpenOrca-Platypus2-13B)**
>
> We’re on a journey to advance and democratize artificial intelligence through open source and open science.

Can read about it and try the demo here: [https://platypus-llm.github.io/](https://platypus-llm.github.io/)

> **[Platypus: Quick, Cheap, and Powerful Refinement of LLMs](https://arxiv.org/abs/2308.07317)**
>
> We present $\\textbf{Platypus}$, a family of fine-tuned and merged Large Language Models (LLMs) that achieves the strongest performance and currently stands at first place in HuggingFace's Open LLM Leaderboard as of the release date of this work. In...

You can get over 51 (for non-huge) with e.g.

> **[lvkaokao/llama2-7b-hf-chat-lora-v2 · Hugging Face](https://huggingface.co/lvkaokao/llama2-7b-hf-chat-lora-v2)**
>
> We’re on a journey to advance and democratize artificial intelligence through open source and open science.

> **[posicube/Llama2-chat-AYT-13B at main](https://huggingface.co/posicube/Llama2-chat-AYT-13B/tree/main)**
>
> We’re on a journey to advance and democratize artificial intelligence through open source and open science.

Depending on the metric, I think these might be best:

> **[Riiid/sheep-duck-llama-2 · Hugging Face](https://huggingface.co/Riiid/sheep-duck-llama-2)**
>
> We’re on a journey to advance and democratize artificial intelligence through open source and open science.

> MMLU (5-shot) 70.74

> **[uni-tianyan/Uni-TianYan · Hugging Face](https://huggingface.co/uni-tianyan/Uni-TianYan)**
>
> We’re on a journey to advance and democratize artificial intelligence through open source and open science.

> Uni-TianYan is a finetuned model from LLaMA2.  
> […]
> 
> | MMLU (5-shot) | 69.91 |
> | --- | --- |
> | TruthfulQA (0-shot) | 65.81 |
> | Avg. | 73.81 |

You can get 64.75 MMLU on this one from 3 weeks ago with at least the Llama 2 licence, yes, it’s a 35.3 GB file (also available down to 26.78 GB), maybe why you do not use that large:

> **[TheBloke/Lemur-70B-Chat-v1-GPTQ · Hugging Face](https://huggingface.co/TheBloke/Lemur-70B-Chat-v1-GPTQ)**
>
> We’re on a journey to advance and democratize artificial intelligence through open source and open science.

> Open large language models (LLMs) have traditionally been tailored for either textual or code-related tasks, with limited ability to effectively balance both. However, […]  
> In this work, we introduce Lemur-70B-v1 and Lemur-70B-chat-v1, the state-of-the-art open pretrained and supervised fine-tuned large language models balancing text and code intelligence.  
> […]
> 
> - **[23 August, 2023]:** 🎉 We release the weights of [`OpenLemur/lemur-70b-v1`](https://huggingface.co/OpenLemur/lemur-70b-v1), and [`OpenLemur/lemur-70b-chat-v1`](https://huggingface.co/OpenLemur/lemur-70b-chat-v1)! Check it out in [HuggingFace Hub](https://huggingface.co/OpenLemur).

> **[stabilityai/StableBeluga-7B · Hugging Face](https://huggingface.co/stabilityai/StableBeluga-7B)**
>
> We’re on a journey to advance and democratize artificial intelligence through open source and open science.

> Model type: Stable Beluga 7B is an auto-regressive language model fine-tuned on Llama2 7B.  
> […]  
> License: Fine-tuned checkpoints (Stable Beluga 7B) is licensed under the STABLE BELUGA NON-COMMERCIAL COMMUNITY LICENSE AGREEMENT

> **[JosephusCheung/Qwen-LLaMAfied-7B-Chat · Hugging Face](https://huggingface.co/JosephusCheung/Qwen-LLaMAfied-7B-Chat)**
>
> We’re on a journey to advance and democratize artificial intelligence through open source and open science.

> This is the LLaMAfied version of [Qwen/Qwen-7B-Chat](https://huggingface.co/Qwen/Qwen-7B-Chat), recalibrated to fit the original LLaMA/LLaMA-2-like model structure.
> 
> You can use LlamaForCausalLM for model inference, which is the same as LLaMA/LLaMA-2 models (using GPT2Tokenizer converted from the original tiktoken, by [vonjack](https://huggingface.co/vonjack)).

> **[chargoddard/llama-2-26b-trenchcoat-stack · Hugging Face](https://huggingface.co/chargoddard/llama-2-26b-trenchcoat-stack)**
>
> We’re on a journey to advance and democratize artificial intelligence through open source and open science.

> Llama 2 13b is a pretty decent language model. You know what’s probably better? _Two_ Llama 2 13b models. In a trenchcoat.  
> Produced by [`bakllama.py`](https://github.com/cg123/mergekit/blob/main/bakllama.py) with this config file:  
> […]  
> This tells us two very important things:
> 
> 1. TruthfulQA is a perfect benchmark in every way.
> 2. Llama models are amazingly robust to being fed their own output.

You can get slightly lower 43.99 MMLU  
#Parms  
1.13

> **[TheBloke/Llama-2-7B-GPTQ · Hugging Face](https://huggingface.co/TheBloke/Llama-2-7B-GPTQ)**
>
> We’re on a journey to advance and democratize artificial intelligence through open source and open science.

> **[meta-llama/Llama-2-7b-chat-hf · Hugging Face](https://huggingface.co/meta-llama/Llama-2-7b-chat-hf)**
>
> We’re on a journey to advance and democratize artificial intelligence through open source and open science.

This one is very intriguing:

> **[ajibawa-2023/scarlett-33b · Hugging Face](https://huggingface.co/ajibawa-2023/scarlett-33b)**
>
> We’re on a journey to advance and democratize artificial intelligence through open source and open science.

> **[fangloveskari/ORCA\_LLaMA\_70B\_QLoRA · Hugging Face](https://huggingface.co/fangloveskari/ORCA_LLaMA_70B_QLoRA)**
>
> We’re on a journey to advance and democratize artificial intelligence through open source and open science.

> **mixed strategy: 100%Open-Platypus + ~1%Dolphin(GPT-4) + ~1%OpenOrca(GPT-4)**

This one gets 70.26 MMLU:

> **[TheBloke/Genz-70b-GPTQ · Hugging Face](https://huggingface.co/TheBloke/Genz-70b-GPTQ)**
>
> We’re on a journey to advance and democratize artificial intelligence through open source and open science.

> **[Open LLM Leaderboard - a Hugging Face Space by HuggingFaceH4](https://huggingface.co/spaces/HuggingFaceH4/open_llm_leaderboard)**
>
> Discover amazing ML apps made by the community

> **[Open LLM Leaderboard - a Hugging Face Space by gsaivinay](https://huggingface.co/spaces/gsaivinay/open_llm_leaderboard)**
>
> Discover amazing ML apps made by the community

> gpt4  
> uni-tianyan/Uni-TianYan

---

<div class="post-metadata">

**Author:** ![Tomas\_Pevny](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/tomas_pevny/32/25466_2.png) [@Tomas\_Pevny](https://discourse.julialang.org/u/Tomas_Pevny)\
**Post date:** [September 13, 2023, 7:14pm UTC](https://discourse.julialang.org/t/gradient-of-llama2-computed-by-zygote-seems-to-be-incorrect/103810/5 "2023-09-13T19:14:24Z")

</div>

Hi Palli,

thanks for the links.

Yes, I think I solved the problem.  
To answer your question, why I have chosen Julia instead of Python has a very simple answer.  
I like Julia. By solving issues, I learn a lot despite sometimes smaller progress. Transformers.jl is very good lib and serves me well. Moreover, I think we should not be afraid to use LLM with Julia.

---

<div class="post-metadata">

**Author:** ![Palli](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/palli/32/3380_2.png) [@Palli](https://discourse.julialang.org/u/Palli)\
**Post date:** [September 13, 2023, 7:25pm UTC](https://discourse.julialang.org/t/gradient-of-llama2-computed-by-zygote-seems-to-be-incorrect/103810/6 "2023-09-13T19:25:03Z")

</div>

> [@Tomas\_Pevny](#):
>
> Moreover, I think we should not be afraid to use LLM with Julia.

Well we shouldn’t, but I was a bit, and you inspire me to use Julia (it’s seems easy with the MWE). I also like Julia, I believe and understand Python ecosystem may be better in some ways, at least for learning, if docs are missing or you encounter bugs in Julia not in Python.

Would you have avoided this specific issue with Python? Do you have good experience with LLMs in both languages? Even if Julia might be less scalable to train, it’s good to know if you could at least use LLMs (for inference).

---

<div class="post-metadata">

**Author:** ![Tomas\_Pevny](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/tomas_pevny/32/25466_2.png) [@Tomas\_Pevny](https://discourse.julialang.org/u/Tomas_Pevny)\
**Post date:** [September 13, 2023, 7:35pm UTC](https://discourse.julialang.org/t/gradient-of-llama2-computed-by-zygote-seems-to-be-incorrect/103810/7 "2023-09-13T19:35:49Z")

</div>

To be honest, my experience with Python ecosystem is small. I delve into it only when I read and use reference implementation of some paper. In the case of this thread, there was no bug. The very same thing would happen in Python.

Transformers from huggingface have more features than Transformers.jl. Especially with stuff like mixed and low precision and distribution of models across multiple GPUs. But Julia’s Transformers.jl are super-fun to hack, because in nutshell, the lib is very transparent and that is why I enjoy the work with it. The biggest trouble I have faces with it that I sometimes do not know, how to download some model from huggingfacce.