0
$\begingroup$

I have had a look at LLamas model card, specifically the 7B parameter version: https://github.com/facebookresearch/llama/blob/main/MODEL_CARD.md

which I assume is an encoder only transformer similar to this:

But then I did some math.If the dimension of every Dense layer, including the one connecting to the Attention layer is 4096, the context length is 2048, the number of attention heads is 32 and the embedding size is 786, then the output size of the attention layer is 32 * 786 * 2048 and as such the number of weights to connect it to the dense layer is 32 * 766 * 2048 * 4096, which is 205B parameters, which is obviously far more than 7B. So how is this accomplished? How big is the ouput of the attention layer and how is it connected to the following Dense layers?

Improve this question
$\endgroup$

2 Answers 2

Reset to default
2
$\begingroup$

Let's start over and count the number of parameters, looking at the code in the same repository:

  • The Attention layer is defined here in llama/model.py, it defines four matrices wq, wk, wv and wo. Each matrix is dim $\times$ n_heads * head_dim. And head_dim is computed as head_dim = dim // n_heads.

  • The FeedForward bit is defined here in llama/model.py, it defines three matrices w1, w2, w3. Each matrix is dim $\times$ hidden_dim. And hidden_dim is computed as:

      hidden_dim = 4 * dim    # Call on line 186
  hidden_dim = int(2 * hidden_dim / 3)
  hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)

Now, dim is the embedding dimension - for 7B model it is 4096 (not 786).
Number of heads n_heads is 32. And multiple_of is some kind of align-to-power-of-2 factor, equals to 256.

Gathering the whole thing:

dim = 4096
n_heads = 32
multiple_of = 256
n_blocks = 32

head_dim = dim // n_heads
n_attention_weights = 4 * dim * (n_heads * head_dim)

hidden_dim = 4 * dim 
hidden_dim = int(2 * hidden_dim / 3)
hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)
n_feed_forward_weights = 3 * dim * hidden_dim

total_weights = n_blocks * (n_attention_weights + n_feed_forward_weights)

I'm getting total_weights = 6 476 005 376.

There are also input and output embeddings that add 2 * vocab_size * dim getting me to 6.9 billon for vocab_size = 50000.

Note that max context size does not enter the computation at all.

I'll leave computing the number of weights for larger models as an exercise.

Improve this answer
$\endgroup$
2
$\begingroup$

The dense layers do not include interaction between different timesteps. As such, each one is just 4096x4096, and every one of the 2048 4096-dimensional vectors is multiplied by this. (Some of the matrices are rectangular instead of square, but the important part is that there aren't parameters for the full 2048 x whatever input).

Improve this answer
$\endgroup$
5
  • $\begingroup$ Im not sure I understand. At every timestep, meaning new token being generated, the whole context is considered. And at some point, the whole context must be passed from the attention layer to the dense layers, which includes each variation each of the 32 attention heads generates. Or does it not work that way for LLMs? $\endgroup$
    – user2741831
    May 26 at 7:58
  • $\begingroup$ For a single attention head in LLaMA it works like this: 1. Each 4096-dim vector is projected to 3 128-dimensional vectors, the key, query, and value. (All of the reference to parameters is here) 2. The keys, queries, and values for all positions are used as inputs to the attention function. All the interaction between timesteps happens here, and it doesn't depend on parameters other than the projections into key/query/value. $\endgroup$
    – Davis Yoshida
    May 26 at 8:02
  • $\begingroup$ By 4096-dim vector, to you mean the embedding size or the context length? Because I'm pretty sure Key Query and Values are supposed to be matrices (representing the input text). But you are saying the get dimensionally reduced to a smaller dimension? At 128x128x32x4096 that would still be 2B Weights just to feed the attention output to the Dense Layers $\endgroup$
    – user2741831
    May 26 at 8:08
  • $\begingroup$ For a single attention head, K/Q/V are all sequence_length x head_size matrices yeah. Each column is a single vector corresponding to a particular timestep. The important thing to understand is that every single linear layer in a transformer is applied to each timestep independently. I'm not sure where you got 128x128x32x4096, but the first dense layer is something like 4096 x 11000. The output of attention is sequence_length x 4096 and applying that 4096 x 11000ish matrix outputs a sequence_length x 11000ish matrix (but really you should think of it as a sequence of vectors). $\endgroup$
    – Davis Yoshida
    May 26 at 9:05
  • 1
    $\begingroup$ I assumed the attention heads were calculated in parallel and then fed into a single set of Dense layers. But they are actually stacked on top of each other, so now it makes more sense to me $\endgroup$
    – user2741831
    May 26 at 9:38

You must log in to answer this question.

Not the answer you're looking for? Browse other questions tagged .