Transformer Block

LLaMA

Hard

Full Llama 3 transformer block

Independent study note. Written from the public paper and official code. This is not TensorTonic Plus and does not reproduce their exercises, starter code, or tests. For the official version, subscribe on TensorTonic.

Overview

A Llama 3 layer is a pre-norm residual Transformer block: RMSNorm, grouped-query attention with RoPE, a residual add, another RMSNorm, a SwiGLU feed-forward, and a second residual add. Section 3.2 of The Llama 3 Herd of Models (arXiv:2407.21783) says the architecture is a standard dense Transformer that “does not deviate significantly” from Llama and Llama 2. The paper lists the ingredients (GQA, SwiGLU, RoPE with (\theta=500{,}000), Table 3 widths) but does not write the residual equations. Those come from the official TransformerBlock in llama/model.py, which matches the Llama 1/2 pre-norm layout.

The block is the unit that is stacked (L) times (32 / 80 / 126 for 8B / 70B / 405B). Everything that is “the model” besides embeddings, a final RMSNorm, and the output projection lives inside this unit. Neighboring notes cover the attention, RoPE, and FFN pieces in isolation; this one is the wiring.

How it works

Let (x \in \mathbb{R}^{B \times S \times d}) be the residual stream for a batch of (B) sequences of length (S), with model width (d) (args.dim: 4096 / 8192 / 16384). The official forward is exactly two residual branches:

h = x + \mathrm{Attention}\bigl(\mathrm{RMSNorm}(x)\bigr),\qquad y = h + \mathrm{FFN}\bigl(\mathrm{RMSNorm}(h)\bigr).

Both norms are independent RMSNorm modules (attention_norm, ffn_norm), each with its own gain vector of length (d) and (\varepsilon=) norm_eps (default (10^{-5}) on ModelArgs). There is no bias on any projection and no dropout in the public inference graph.

Attention receives four extra arguments the block merely threads through:

• start_pos — cache write offset for incremental decoding.
• freqs_cis — the RoPE complex table slice for positions ([\texttt{start_pos},,\texttt{start_pos}+S)).
• mask — optional additive attention mask, already shaped for the current query length versus the cached key length.

Inside attention, (x) is projected to queries, keys, and values; RoPE rotates (Q) and (K); keys/values are written into a cache and repeated to match the query-head count (GQA); then scaled-dot-product attention and an output projection return a tensor of shape (B \times S \times d). The FFN is SwiGLU: (\mathrm{SiLU}(W_1 h)\odot W_3 h), then (W_2). Hidden width is not (4d); Table 3 gives 14336 / 28672 / 53248. The constructor starts from hidden_dim=4*args.dim and then applies the (2/3) shrink, optional ffn_dim_multiplier, and multiple_of rounding described in the SwiGLU note.

Paper-only details that do not appear in TransformerBlock: the document-boundary attention mask used in packed pre-training, and the 128K long-context recipe. The released block is the inference form — causal (plus cache offset) masking supplied by Transformer.forward.

Shapes that must line up: head_dim = d / n_heads (128 for every size in Table 3), KV heads (n_{kv}=8), and FFN inner size as above. A mismatch in any of those breaks the residual add, which requires the attention and FFN outputs to stay in (\mathbb{R}^{B \times S \times d}).

Official code

llama/model.py — class TransformerBlock. Construction allocates Attention, FeedForward, attention_norm, and ffn_norm. The two-line residual body is forward. Transformer stacks these in a ModuleList. Checkpoint hyperparameters arrive via params.json loaded in llama/generation.py (Llama.build).

Watch-outs

• Pre-norm, not post-norm. Normalizing after the residual (GPT-2 style) or sharing one norm for both branches changes every trained gain vector.
• The two RMSNorms are not tied. Reusing one module for attention and FFN is a different parameterization than the checkpoint.
• Attention and FFN must preserve (d). A GQA projection that emits (n_{kv}\cdot\text{head_dim}) on the residual path cannot add back to (x).
• start_pos / freqs_cis / mask are not optional in the official signature. Dropping the cache offset applies RoPE as if every call started at position 0.

Sources

• Paper: The Llama 3 Herd of Models (Llama Team, 2024), §3.2 and Table 3
• Code: meta-llama/llama3 llama/model.py (TransformerBlock)
• Lineage: Touvron et al., Llama / Llama 2 (2023), for the pre-norm residual layout the Llama 3 paper inherits but does not restate