f of x

Homemade Tinyllama - Easiest way to get from GPT2 to Llama

So you know how a GPT works and you also know how to make one from the GPT 2/3 era? And now you wanna step up the game?

In this post, I'm gonna take you from GPT2 to your homemade Tinyllama (decoder-only).

It seems like a big change and it definitely is one but going from GPT2 to a Llama is surprisingly easy!

We're gonna cover what changes in the architecture and the forward pass and we'll also talk about the intuition and math behind everything. It's just 4 major changes and if you understand all of them, you'll be ready to implement the model all by yourself.

What we will not talk about here is the code part as I've written all of that in my repo called lantern.

What GPT2 did:

If you understand almost all of this clearly, you're good to go and you can skip to the next heading now.

If not, i suggest you to go and study gpt2 first. In my opinion, the best way is to make it yourself. For me, the best resources were Andrej Karpathy's Zero to Hero series' videos about building a GPT step by step and the other one in which he tries to reproduce GPT2. You can also refer to the lantern code and dig through how everything's implemented.

What changes does this Tinyllama bring?

Every change here has an intuitive reason behind it so it'll be a lot of fun because we'll know exactly why we're doing a certain thing. The most tricky part might be RoPE (so we'll start with it first). MHA to GQA is all about efficiency and a few tensor gymnastics. In MHA, all the q heads have separate k and v heads while in MQA, some q heads will share common k and v heads. The change from LayerNorm to RMSNorm has one simple reason. And Swiglu is just a fancy way of deciding how much information should go ahead.

Rotary Position Encoding (RoPE)

Instead of learned position encodings as in GPT2, Llama models use RoPE on query and key just before computing attention scores.

q, k = rope(q), rope(k) # apply the same RoPE transformation to both q and k

Modern LLMs use RoPE which captures position of tokens relative to each other and works better than absolute encodings, especially for longer contexts. One token influences other tokens based on its position relative to the other tokens rather than solely its absolute position in the sequence. So, the information about a token's relative position matters a lot more than its absolute position. To achieve this kind of encoding, we just have to rotate the features by an angle that depends on the position of the token. This way, each token gets a different angle for rotation and this information is carried forward in the q and k dot product during attention score computation.

Now, the thing to keep in mind is that we're not gonna have the same angle for all the features within the same token. The features will have a different angle of rotation. But, overall, this angle will depend on the position of the token. We make pairs of features and then compute the angle using the given formula for inverse frequency.

If we have x = [[[x1, x2, x3, x4,...n_dim features],...]], we make pairs of features as (x1, x2), (x3, x4)... Here, the first pair has the index i = 0, the second one has an index of i = 1 and so on. RoPE makes use of the formula proposed in the original transformer paper -

inverse frequency = 1/10000^(2i/n_dim)
[i is the index of the pair and n_dim is the head_size of the attn head]

Now, we multiply the token's position with the inverse frequency to get the rotation angle or theta.

theta = token_position * inverse_frequency

Here, theta actually contains different rotation angles for the features i.e. the -1 dim.

But why involve features too? We need just the token positions, right? What does giving each feature pair a different frequency do?

Think of each feature pair as a container for any arbitrary kind of information about the token. Since, each pair has a different frequency and a token has a variety of these containers, the captured information can be of a wide range. The pairs with a lower inverse frequency (resulting in a smaller rotation angle) will be less sensitive to small range changes, meaning that they won't react much with the nearby tokens and can carry information over long distances in the context window. On the other hand, the ones with a higher inverse frequency will be more sensitive to the tokens immediately surrounding them.

We can get the cosine and the sine vectors using this theta and rotate the pairs to create rotary encodings.

To rotate the pairs, we could've used the 2D rotation matrix but modern LLMs approach this in a clever way.

In the Llama style implementation, we split the features (last dim of x) into two parts and then concatenate them as pairs of (-x2, x1). [x1, x2, x3, x4] would become [-x3, -x4, x1, x2]. Let's call this y. Then we can element-wise multiply x by cos and y by sin and add them. This results in the same tensor as in the case of multiplying each pair by the rotation matrix.

You can verify this by doing the math on the example features [x1, x2, x3, x4] by hand.

In models like GPTNeoX, y looks like [-x2, x1, -x4, x3]. The math is identical. The discrepancy in Llama is just because of lineage.

The final step written out is:

x = x * cos + y * sin

This avoids construction rotation matrices and is much faster on GPUs.

Grouped Query Attention (GQA)

The core of GQA is the same as MHA except that key and value heads are actually lesser in number compared to query heads to save memory and compute.

In GQA, we divide the query heads into n_kv_head groups and when calculating attention scores, one kv head is shared across n_head // n_kv_head query heads. Setting n_kv_head to one creates just one kv head that is used for all query heads, making this an MQA or Multi Query Attention which saves a lot of resources at the cost of model performance.

In super simple words, if, in a library, there is one shelf (kv head) for all the students (query heads), it's MQA (Multi Query Attention) i.e. one key head and one value head for all query heads but if we had x shelves where 1 < x < number of students, it's the kind of attention that llama models use i.e. GQA.

GQA with a setting like n_kv_head = 8 for 32 query heads is a nice sweet spot that gives almost MHA level performance and also saves compute like MQA.

RMSNorm (Root Mean Square Norm)

Let's start with what our good old layernorm does. It's essentially -

y_i = ((x_i  - x_mean)/std) * gamma_i + beta_i

We are centering around zero here and also adding a beta_i.

What RMSNorm does is just -

y_i = (x_i/RMS) * gamma_i

Here, RMS = sqrt([Σ(x_i^2)/n_dim] + epsilon) where epsilon is a really small constant.

See how simple that is? And it makes the whole thing faster too while giving almost identical results. So, many models use RMSNorm for this reason.

Swiglu

The MLP part in Tinyllama is a Swiglu.

Swish Gated Linear Unit or SwiGLU is made up of two things - a Swish Activation Function and a Gated Linear Unit (GLU).

Swish: x * sigmoid(beta * x) where beta is a scalar (usually 1) and x is the input

GLU: sigmoid(W1 * x + b1) * (W2 * x + b2)

SwiGLU is derived from a combination of these two and helps decide how much of information we want to keep i.e. if we want to amplify the signal or make it less powerful.

SwiGLU is implemented as follows:

That's it!

What does the forward pass look like in Tinyllama?:

x = x + GQA(input_RMSNorm(x))
x = x + SwigluMLP(post_attn_RMSNorm(x))

*Note that this model doesn't share weights as in GPT2.

If you've understood these changes, you can call yourself ready to implement the model from scratch in Pytorch. Make sure to have a look at the original model code For reference, it recommended that you checkout my lantern code for the model implementation (self-promotion tehee!).

May you run into bugs and logical errors so you get to learn more!

See ya in the next one!