All the Transformer Math You Need to Know | How To Scale Your Model

[jacobaustin123]

Discussing the Transformer architecture!

[main-horse]

a scant few minutes ago, the mathjax for this chapter's introduction was not rendering, but it appears that it was fixed live 😅

[main-horse]

[addressed]

[jacobaustin123]

Good catch, I've updated the text!

[spacewander]

an MoE is a dense model with E MLP blocks per layer,

Should be 'sparse model'?

[jacobaustin123]

I'm trying to make the point that it's just like a normal dense model, but with each MLP block replaced with E copies.

[spacewander]

I see, thanks for your explanation.

[burichh]

Two remarks:

1. In the section What Should You Take Away from this Section?, you write "The parameter count of the MLP block dominates the total parameter count and the MLP block also dominates the FLOPs budget until the sequence length $T > 8D$."
   Question: Don't you mean "as long as the sequence length $T<8D$."? Just because when $T>8D$ then computing the attention matrix becomes the leading term in FLOPs.
2. In the Appendix A: How does Flash Attention work? at the beginning you say While it’s true that the attention QK product has shape $[B, S, T, H]$ where B is the batch size, S and T are the Q and K sequence dims, and H is the number of heads..." but above you defined $H$ as the attention head dimension, not the num of heads.
   Question: Don't't you mean $N$? That's the number of attention heads.

[jacobaustin123]

(1) Good point, I've reworded this as you suggested. (2) Yes. Fixed. Thank you!

[rahulbatra85]

In the transformer decoder architecture figure, the term K is overloaded(K=XWk and also K=number of KV heads).

[jacobaustin123]

@levskaya I have been running into this too, maybe we pick another symbol? N for total, M for KV heads maybe, as in "MQA". Maybe it's fine as is too, it's not that confusing

[kerrywang]

In the transformer decoder architecture plot. G is query heads per kv heads. I wonder should it instead be reversed? i.e. KV heads per Query Heads? Given group query attention we are grouping the query so query heads should always be smaller than KV heads?

[jacobaustin123]

No, the key is to have fewer KV heads than query heads. We group query heads that attend to the same KV head.

[FL33TW00D]

From a hardware standpoint, this lets us fit our chunk of Q into VMEM (what the algorithm above calls on-chip SRAM) so we only have to load the KV chunks on each iteration, reducing the arithmetic intensity. We can also keep the running statistics in VMEM.

I might be incorrect here, but I was under the impression flash attention increased arithmetic intensity by reducing the number of memory accesses?

[jacobaustin123]

I think in this case it's ultimately the same thing: if you load the same Qs or Ks multiple times, you're increasing the bytes loaded (reducing arithmetic intensity) and adding more, smaller memory accesses.

[maximilianmbeck]

typo in appendix?
In the first equation (where you define $S_{ij}$). Shouldn't be the index of $k$ in the denominator also $k$ (i.e. the sum index)?
I think in your last equation (where you describe the local contraction) it is correct. Thanks for this great book!

[ohodson]

Great resource, thank you.

The second reference on the Transformer Math page looks like it should have the author "Shazeer, N." rather than "Noam, S.".

[hammersam]

Thanks for sharing! Great resource for learning LLM.

[wenchenvincent]

In the section "What Should You Take Away from this Section?", the Training FLOPs per layer for Vocab is shown as 12BTDV. I think it should be 6BTDV, and it should be total.

[jacobaustin123]

Agree it's total. It's 12BTDV because the embedding is used for both embedding and un-embedding (the output projection). We call this weight tying. This means each parameter is used twice.

[wenchenvincent]

That makes sense. In this case, you might want to include both embedding and un-embedding in the Section "Global FLOPs and Params Calculation: Other Operations" to make it consistent.