$ cat ./blog/gen-ai/llm-fundamentals/ben-trong-transformer-block.md
seriesHow LLMs Workpart 8

Inside a Transformer Block: FFN, the Residual Stream, and Normalization

Part 8 of the How LLMs Work series: the other half of a Transformer block besides attention. The residual stream as a shared bus every layer reads from and writes to, the FFN as a key–value memory, why GELU and then SwiGLU replaced ReLU, where LayerNorm and RMSNorm go, and how to count the parameters and FLOPs of Llama 3.1 8B so they match the published numbers.

~/diagrams/ben-trong-transformer-block/residual-stream-bus-en.htmlinteractive
open ↗
FigThe middle row is the residual stream: a d_model-dimensional vector that runs straight from the embedding through every block. The two sub-layers of each block sit in the rows beside it: attention above, the FFN below. Each sub-layer reads the stream (normalizes it and projects it into its own space), computes, then writes its result back by addition. No sub-layer replaces the stream, so a later block always sees what earlier blocks wrote.
sub-layer that reads then writes (attention, FFN)residual stream state and later blocksgrey arrows: read from the stream, write back by addition

Part 7 (Vietnamese for now; the English version is on the way) stopped at section 09, with a diagram of a Transformer block made of four parts: attention, the MLP/FFN, the residual connection, and normalization. That post only went deep on the first one.

There is a detail that is easy to miss: attention is not the heaviest part of the block. In Llama 3.1 8B, each layer has about 218 million parameters, and less than 20% of them belong to attention. Almost all of the rest sits in the FFN, a component that many Transformer introductions mention in a single line.

This post fills in that other half. It answers four questions:

  1. How does information get from one layer to another? Through the residual stream.
  2. What does the FFN do to each token, and why do people call it a key–value memory?
  3. Why did the activation function move from ReLU to GELU and then to SwiGLU, and why does that make the d_ff of Llama 3 equal to 14,336 instead of 4 × 4096?
  4. Where does normalization go, and why does “where” matter more than “which kind”?

At the end, everything comes together in one concrete count: the parameters and FLOPs of a Llama 3.1 8B block, worked out by hand and matching the number Meta published.

Part 1 already has an overview section on the block. This post does not repeat it; it goes deeper and corrects a few familiar numbers that are out of date.


01The other half of the block: attention mixes, the FFN processes

A decoder block in a modern LLM has two sub-layers, and they do very different jobs:

Attention FFN
Looks at every earlier token in the context the current token only
Exchanges information between tokens yes no
Keys and values come from other tokens in the sentence, recomputed on every run fixed parameters learned during training
Characteristic cost grows with context length (the Q · Kᵀ step) parameter count

The most useful way to see it: attention is the communication step, the FFN is the computation step. Attention lets each token pull information from other tokens. The FFN then processes each position on its own, as if every token were passed through the same large transformation, independently of the others.

Because the FFN never looks at other tokens, everything it needs to know about the context has to be brought in by attention first. Conversely, attention only mixes and moves information; it applies no strong per-token non-linear transformation. A block without either one is incomplete.

The two sub-layers do not connect directly to each other. They both connect to a third thing: the residual stream.


02The residual stream: a shared bus every layer reads and writes

Part 1 covered the residual connection at a basic level: output = x + f(x), and why gradients vanish in a deep model without it. This section looks at the same addition from a different angle, the one Elhage et al. 2021 take in A Mathematical Framework for Transformer Circuits.

Write out the whole forward pass of an L-layer model in the Pre-LN style (the style Llama uses):

x_0 = Embedding(token)
x_{l+½} = x_l     + Attn_l( Norm(x_l) )
x_{l+1} = x_{l+½} + FFN_l ( Norm(x_{l+½}) )
logits  = lm_head( Norm(x_L) )

No line overwrites x. Each sub-layer only adds a vector. Expand all the additions:

x_L = x_0 + Attn_1(…) + FFN_1(…) + Attn_2(…) + … + FFN_L(…)

In other words, the hidden state at the end of the model is the sum of the initial embedding and the contributions of every sub-layer. Elhage et al. call this running sum the residual stream and describe it as a linear communication channel:

  • each sub-layer reads from the stream through a linear projection (the Q/K/V matrices of attention, the input matrix of the FFN);
  • and writes to the stream through another projection (the output matrix) followed by an addition.

Figure 1 at the top of the post draws exactly this: the stream is an axis running straight through every block, while attention and the FFN are stations that branch off to read and then merge back in.

Why the width of the stream never changes

Since every sub-layer adds into the same vector, every output must have the same width, d_model. For Llama 3.1 8B, d_model = 4096. Inside a sub-layer, the width can change:

Where Width in Llama 3.1 8B
Residual stream 4,096
Query, after projection 32 heads × 128 = 4,096
Key and value, after projection (GQA) 8 heads × 128 = 1,024
FFN hidden layer 14,336
Output of each sub-layer, before the addition 4,096

Every component has to project back to 4,096 before it writes. That is the bus’s “interface standard”.

What a shared bus means in practice

The “shared bus” idea has three concrete consequences:

  1. A later layer can read what every earlier layer wrote, not just the layer right before it. A feature that the FFN in layer 3 writes is still in the stream at layer 20, unless some layer deliberately writes an opposite vector to cancel it.
  2. Layers can write into different subspaces. The stream has 4,096 dimensions. Elhage et al. argue that sub-layers can use different directions in that space, so many kinds of information can coexist in one vector without overwriting each other.
  3. Gradients have a straight path. Because x_L is a sum that contains x_l, the derivative of x_L with respect to x_l always has an identity term. Seen from the stream, this is the technical reason residuals make deep models trainable.

From this angle, the question “what does the FFN do” becomes: what does the FFN read from the stream, and what does it write back?


03The FFN as a key–value memory

Write the classic FFN (GPT-2 style, two matrices, ignoring biases) for one token x:

FFN(x) = f(x · W_in) · W_out

x      : 1 × d_model
W_in   : d_model × d_ff
W_out  : d_ff × d_model
f      : activation function (ReLU, GELU, …)

Now put it next to the attention formula from part 7:

Attention(q, K, V) = softmax( q · Kᵀ ) · V
FFN(x)             =        f( x · K  ) · V     with K = W_in, V = W_out

The two formulas have the same shape. Geva et al. 2021 built the title of their paper on exactly this resemblance: Transformer Feed-Forward Layers Are Key-Value Memories.

Reading the FFN that way:

  • Each column of W_in is a key. There are d_ff keys: 3,072 per layer in GPT-2 small, 14,336 in Llama 3.1 8B.
  • x · W_in produces d_ff numbers, each measuring how well the current token matches the corresponding key.
  • The activation f keeps keys that match strongly and suppresses keys that do not. With ReLU, negative values are cut to exactly 0.
  • Each row of W_out is a value. The output is a sum of values, each weighted by the match coefficient of its key.
~/diagrams/ben-trong-transformer-block/ffn-key-value-memory-en.htmlinteractive
open ↗
FigBoth rows have the same shape: compare a query with a set of keys, pass the match scores through a function (softmax in attention, the activation function in the FFN), then take a weighted sum of values. The difference is where keys and values come from. Attention (top row) takes them from other tokens in the context, recomputed on every run. The FFN (bottom row) stores d_ff fixed key–value pairs in its two matrices W_in and W_out, learned during training. Both write their result into the residual stream.
computation steps (match, activate, mix)residual streampurple boxes: keys and values as parameters; blue box: query from the token

Two differences from attention are worth noting. First, the FFN’s keys and values do not depend on the input: they are parameters, the same for every sentence. Second, the FFN uses a non-linear activation instead of softmax, so the coefficients are not forced to sum to 1. Many keys can fire strongly at once, or none at all.

What keys catch, and what values push

Geva et al. 2021 did not stop at the formal comparison. They took individual keys in a 16-layer language model, found the training sentences that made each key fire most strongly, and asked human annotators what those sentences had in common. The results:

  • Keys catch patterns a human can recognize. In lower layers, the patterns tend to be shallow: sentences ending with the same word or the same n-gram. In upper layers, the patterns are more semantic: sentences about the same topic, such as TV shows.
  • Values push probability toward a group of output tokens. Projecting a value through the output embedding matrix gives a distribution over the vocabulary. In upper layers, that distribution often matches the token that comes next after the pattern the corresponding key catches.
  • The FFN output is a composition of many memories. A token typically activates hundreds of keys at once, and the output is the sum of those values, not a single value.

A follow-up paper from the same group (Geva et al. 2022) takes this further: each value can be read as an “update” to the token distribution in vocabulary space, and many values correspond to concepts humans can interpret.

Limits of this interpretation

From here it is easy to slide into “the FFN is where the model stores its knowledge”. That statement goes beyond the evidence:

  • Geva et al. 2021 is an observation over a selected set of keys annotated by hand. The paper does not prove that most knowledge lives in the FFN.
  • Methods like Knowledge Neurons (Dai et al. 2022) and ROME (Meng et al. 2022) find places in the FFN where an edit makes the model “remember” a different fact. That is evidence suggesting the FFN plays a large role in recalling facts.
  • But Hase et al. 2023 show that the place a localization method points to is not necessarily the most effective place to edit. Editing a different layer is sometimes as good or better. So “it can be localized there” and “the knowledge lives there” are two different claims.

The safer statement is: evidence suggests the FFN behaves like a key–value memory and plays a large role in recalling knowledge, while attention handles bringing in the right context first. The idea “the FFN is a knowledge store, so split the store into many experts” is exactly what motivates Mixture of Experts. Part 1 has a short section on MoE; the deep dive is saved for a later post in the series.


04Activation functions: from ReLU to GELU to SwiGLU

The function f in f(x · W_in) is the only thing that makes the FFN non-linear. Without it, x · W_in · W_out is just a matrix multiplication, and stacking any number of layers would still give a single linear transformation. The three generations of activation functions differ in how they decide how much to keep from each key.

ReLU: a hard cut

The original 2017 Transformer used ReLU(z) = max(0, z). A key with a negative match is dropped entirely; a key with a positive match is kept as is. Simple and cheap, but with two drawbacks: the derivative is 0 over the entire negative half, and there is a kink at 0.

GELU: multiply by the probability of keeping it

Hendrycks and Gimpel proposed GELU with a probabilistic intuition: instead of a hard “keep or drop” decision, multiply the input by the probability of keeping it, a probability that grows with the size of the input:

GELU(z) = z · Φ(z)        Φ is the cumulative distribution function of N(0, 1)

A few values show the shape:

z ReLU(z) GELU(z)
−3 0 −0.004
−1 0 −0.159
0 0 0
1 1 0.841
3 3 2.996

For large z, GELU is almost identical to ReLU. Around 0 it is smooth, and it gives a small negative value instead of cutting to zero. The paper also gives a tanh approximation that is faster to compute:

GELU(z) ≈ 0.5 · z · (1 + tanh( √(2/π) · (z + 0.044715 · z³) ))

BERT and GPT-2 use GELU. In GPT-2’s config.json, the line "activation_function": "gelu_new" is exactly this tanh approximation.

SwiGLU: add a gate

In 2020, Shazeer tested a different family of variants in GLU Variants Improve Transformer. The idea comes from the Gated Linear Unit (Dauphin et al. 2016): instead of one branch, send the input through two parallel projections and use one branch as a “gate” for the other:

FFN_SwiGLU(x) = ( Swish(x · W_gate) ⊙ (x · W_up) ) · W_down

Swish(z) = z · σ(z)        σ is the sigmoid
⊙        = element-wise multiplication

Swish with β = 1 is the same function as SiLU (Elfwing et al. 2017; Ramachandran et al. 2017). That is why the config.json of Llama and Qwen says "hidden_act": "silu", even though the architecture is called SwiGLU. The variant that uses GELU as the gate is called GEGLU. Gemma 2 uses that variant, so its config says gelu_pytorch_tanh.

~/diagrams/ben-trong-transformer-block/swiglu-ffn-en.htmlinteractive
open ↗
FigThe input after RMSNorm (4096 dimensions) goes through two parallel projections up to 14336 dimensions. The gate branch goes through SiLU and is multiplied element-wise with the up branch, so every dimension of the up branch is scaled by its own gate. W_down projects the result back to 4096 dimensions to be added to the residual stream. The three matrices give 3 × 4096 × 14336 = 176,160,768 parameters per layer.
parameter-free operations (SiLU, element-wise product)residual stream (in and out)purple boxes: the three weight matrices W_gate, W_up, W_down

In the key–value language of the previous section: the x · W_up branch is still “how well it matches a key”, while the gate branch decides which dimensions get through, and how much. Because the gate also depends on x, each token gets its own soft mask over the d_ff dimensions.

Shazeer measured perplexity at the same parameter and compute budget: GEGLU and SwiGLU beat ReLU and GELU. The paper ends with an often-quoted sentence:

“We offer no explanation as to why these architectures seem to work; we attribute their success, as all else, to divine benevolence.”

In other words, the author gives no theoretical reason why SwiGLU is better. It is an empirical choice, later adopted by LLaMA, Llama 2, Llama 3, Qwen, Mistral and many other models.

The 2/3 rule: why Llama 3 has d_ff = 14,336

SwiGLU has three matrices instead of two. Keeping d_ff = 4 · d_model as in the classic FFN would make the FFN 1.5 times larger. For a fair comparison, Shazeer shrinks d_ff to 2/3, so that 3 · d · (2/3 · 4d) = 8d², exactly the same as a two-matrix FFN with 4d. The LLaMA paper states it directly: SwiGLU with a dimension of 2/3 · 4d instead of 4d.

Meta’s reference code (llama/model.py) computes d_ff in three steps: take 2/3 · 4d, multiply by an optional factor ffn_dim_multiplier, then round up to a multiple of multiple_of:

Model d_model 2/3 · 4d Factor Actual d_ff d_ff / d_model
GPT-2 small (GELU, no gate) 768 — — 3,072 4×
Llama 2 7B 4,096 10,922 none 11,008 ≈ 2.7×
Llama 3.1 8B 4,096 10,922 1.3 14,336 3.5×
Qwen2.5-7B 3,584 — — 18,944 ≈ 5.3×

For Llama 3: 10,922 × 1.3 ≈ 14,198, rounded up to a multiple of 1,024, gives 14,336. Qwen2.5-7B does not follow Meta’s formula; its 18,944 comes straight from intermediate_size in the config.

The lesson: “d_ff = 4d” only holds for a GPT-2-style two-matrix FFN. In LLMs that use SwiGLU, the ratio can be 2.7×, 3.5× or 5.3×, depending on how many parameters the team wants to put into the FFN.


05Normalization: LayerNorm, RMSNorm, and where the norm goes

Each sub-layer adds a vector to the stream. After a few dozen additions, the magnitude of the stream can drift far from where it started, and different dimensions can end up on different scales. Normalization brings the vector back to a stable scale before it enters a magnitude-sensitive computation, such as the dot products in attention.

Part 1 briefly introduced the two kinds of norm. Here we write the formulas out in full, including ε, then focus on the more important question: where the norm goes.

LayerNorm and RMSNorm

For a d-dimensional vector x:

LayerNorm(x) = (x − μ) / √(σ² + ε) · γ + β
  μ  = mean over the dimensions of x
  σ² = variance over the dimensions of x

RMSNorm(x)   = x / √( mean(x²) + ε ) · γ

γ and β are learned parameters, each with d elements. ε is a small constant that prevents division by zero: Llama 3.1 uses 1e-5, Qwen2.5 uses 1e-6.

LayerNorm (Ba et al. 2016) does two things: re-centering (subtracting the mean) and re-scaling (dividing by the standard deviation). Zhang and Sennrich 2019 asked whether both are necessary. Their RMSNorm drops the mean subtraction and drops β as well, dividing only by the root mean square. The paper reports running time reductions of 7% to 64% depending on the model, with quality on par with LayerNorm.

A small example with x = [2, 4, 6], ignoring γ, β and ε:

Formula Result
LayerNorm μ = 4, σ ≈ 1.633 [−1.22, 0, 1.22]
RMSNorm RMS = √((4 + 16 + 36) / 3) ≈ 4.32 [0.46, 0.93, 1.39]

LayerNorm moves the vector to a zero center, so the information “every dimension is positive” disappears. RMSNorm only rescales, so the direction of the vector is preserved. Experiments show that re-centering is not needed for quality, so dropping it is an almost free saving. Llama, Qwen, Mistral and Gemma all use RMSNorm.

Post-LN and Pre-LN

The original 2017 Transformer put the norm after the residual addition. This is called Post-LN:

Post-LN:  x_{l+1} = Norm( x_l + Sub(x_l) )
Pre-LN:   x_{l+1} = x_l + Sub( Norm(x_l) )

The difference is whether the norm sits on the residual path:

  • With Post-LN, the norm sits right on the stream. Every signal (and every gradient) travelling from the last layer to the first has to pass through L norms. The stream is no longer a simple sum.
  • With Pre-LN, the norm sits only on the branch that enters the sub-layer. The stream remains a pure sum as in section 02, and the straight gradient path stays intact.
~/diagrams/ben-trong-transformer-block/norm-placement-en.htmlinteractive
open ↗
FigThe three rows are three ways to place the norm around the same sub-layer. Post-LN (top row) normalizes after the addition, so the norm sits on the residual stream itself. Pre-LN (middle row) normalizes only the input of the branch, and the dashed residual path goes straight to the addition without being normalized. Gemma 2 (bottom row) normalizes both the input and the output of the sub-layer, while keeping the straight residual path of Pre-LN.
sub-layer: attention or FFNresidual stream and the residual additionpink boxes: norm; dashed purple line: the residual path

Xiong et al. 2020 analyze both arrangements at initialization. With Post-LN, the gradients of the parameters near the output layer are very large. Using a large learning rate right away easily makes training diverge. That is why Post-LN needs learning rate warmup: start with a very small learning rate and increase it gradually. With Pre-LN, the gradients are “well-behaved” from the start, and the paper shows that warmup can be removed entirely, making training faster.

GPT-2 was one of the first large models to switch to Pre-LN. The GPT-2 paper says that LayerNorm was moved to the input of each sub-block, and that an extra LayerNorm was added after the final block. That extra norm at the end is not an accident. With Pre-LN, the stream is never normalized on the main path, so there has to be a final norm before lm_head. In the Hugging Face code for Llama, that is the model.norm module.

Norm placement is still unsettled

It is tempting to conclude “every modern LLM uses Pre-LN, end of story”. Reality is more complicated:

  • Post-LN can give better results, if it trains. Liu et al. 2020 observe that when Post-LN trains stably, it often reaches better final results than Pre-LN. They propose an initialization (Admin) that keeps this advantage without divergence.
  • Post-LN can go 1,000 layers deep. DeepNet (Wang et al. 2022) uses DeepNorm, a Post-LN variant that up-scales the shortcut x by a factor α > 1 before the addition and uses its own initialization, to train a 1,000-layer Transformer.
  • Gemma 2 normalizes both sides. The paper says: “use RMSNorm to normalize the input and output of each transformer sub-layer”. The formula becomes x + Norm(Sub(Norm(x))). The bottom row of Figure 4 draws this arrangement. Gemma 3 keeps it, using RMSNorm both before and after each sub-layer (Gemma Team 2025).
  • OLMo 2 normalizes the output of attention and the FFN instead of the input. Peri-LN (Kim et al. 2025) systematically analyzes placing the norm on both sides of a sub-layer and names it Peri-LN.

So the more accurate statement is: most open LLMs today use Pre-LN with RMSNorm, but where to put the norm is still an active research question. What the newer approaches have in common is that they all keep the residual path straight and do not put a norm on top of it.


06Counting the parameters and FLOPs of a block

We now have every piece we need to count. The goal is to compute the parameter count of a real model from its config.json, match the published number, and understand why the familiar formulas give different answers.

Kaplan’s formula and its assumptions

Kaplan et al. 2020 (Scaling Laws for Neural Language Models, Table 1) give a well-known estimate of the parameter count excluding embeddings:

N ≈ 2 · d_model · n_layer · (2 · d_attn + d_ff)
  ≈ 12 · n_layer · d_model²        when d_attn = d_model and d_ff = 4 · d_model

The 12 comes from: attention has four d × d matrices (Q, K, V, O), giving 4d²; the FFN has two d × 4d matrices, giving 8d². Together that is 12d² per layer.

The formula rests on three assumptions: MHA (K and V have as many heads as Q), a two-matrix FFN with d_ff = 4d, and ignoring biases, norms and embeddings. Modern LLMs break the first two.

GPT-2 small: where “2/3” comes from

For GPT-2 small (d = 768), exactly the setup Kaplan assumes:

Part Formula Parameters per layer
Attention 4 · 768² 2,359,296
FFN 2 · 768 · 3,072 = 8 · 768² 4,718,592
FFN / (Attention + FFN) 8 / 12 66.7%

This is where “the FFN holds about 2/3 of the parameters” comes from. It is true for a GPT-2-style architecture, and only for that architecture.

Llama 3.1 8B: counting matrix by matrix

The configuration from config.json: hidden_size = 4096, intermediate_size = 14336, num_hidden_layers = 32, num_attention_heads = 32, num_key_value_heads = 8, vocab_size = 128256, tie_word_embeddings = false, no biases. Each head has 4096 / 32 = 128 dimensions.

Part Formula Parameters
W_Q, W_O 2 · 4096 · 4096 33,554,432
W_K, W_V (GQA, 8 KV heads) 2 · 4096 · 1024 8,388,608
Attention 41,943,040
FFN (W_gate, W_up, W_down) 3 · 4096 · 14336 176,160,768
Two RMSNorms 2 · 4096 8,192
One block 218,112,000

The FFN is 80.8% of the block, not 66.7%. Two causes push in the same direction:

  • GQA makes attention lighter. K and V have only 8 heads instead of 32, so attention is 2.5d² instead of 4d².
  • SwiGLU at 3.5× makes the FFN heavier. Three d × 3.5d matrices give 10.5d² instead of 8d².

Each block totals 13d², against Kaplan’s 12d², with a very different split inside. Adding up the whole model:

Part Parameters
32 blocks × 218,112,000 6,979,584,000
Final RMSNorm 4,096
Embedding (128,256 × 4,096) 525,336,576
lm_head (weights not shared with the embedding) 525,336,576
Total 8,030,261,248 ≈ 8.03B

That matches the “8B” Meta published. The FFN accounts for about 70% of the whole model; most of the remaining 30% is the embedding and lm_head. The 12 · n_layer · d² formula would give 6.44B for the non-embedding part, about 8% below the real 6.98B.

Qwen2.5-7B for comparison

Qwen2.5-7B (d = 3584, d_ff = 18944, 28 layers, 28 Q heads / 4 KV heads, biases on Q, K and V) pushes the ratio further: the FFN is 87.4% of each block. The total comes to 7,615,616,512, and the non-embedding part to 6.53B. Qwen’s model card lists 7.61B and 6.53B, matching the count. The d_ff / d_model ≈ 5.3× ratio together with only 4 KV heads is why the FFN dominates even more than in Llama.

Counting it yourself in code

The code below uses only the Python standard library (run with Python 3.12, no external packages). The configurations are copied from each model’s config.json:

def count_params(cfg, swiglu=True, qkv_bias=False):
    d, ff, L = cfg["hidden_size"], cfg["intermediate_size"], cfg["num_hidden_layers"]
    head_dim = d // cfg["num_attention_heads"]
    kv_dim = cfg["num_key_value_heads"] * head_dim

    attn = 2 * d * d + 2 * d * kv_dim                 # W_Q, W_O, W_K, W_V
    if qkv_bias:
        attn += d + 2 * kv_dim                        # biases on Q, K, V
    ffn = (3 if swiglu else 2) * d * ff               # gate, up, down
    block = attn + ffn + 2 * d                        # + two RMSNorms

    embed = cfg["vocab_size"] * d
    head = 0 if cfg["tie_word_embeddings"] else embed
    total = L * block + d + embed + head              # + final norm
    return block, ffn / block, total, total - embed - head

llama31_8b = dict(hidden_size=4096, intermediate_size=14336, num_hidden_layers=32,
                  num_attention_heads=32, num_key_value_heads=8,
                  vocab_size=128256, tie_word_embeddings=False)
qwen25_7b = dict(hidden_size=3584, intermediate_size=18944, num_hidden_layers=28,
                 num_attention_heads=28, num_key_value_heads=4,
                 vocab_size=152064, tie_word_embeddings=False)

for name, cfg, bias in [("Llama 3.1 8B", llama31_8b, False), ("Qwen2.5-7B", qwen25_7b, True)]:
    block, ffn_share, total, non_emb = count_params(cfg, qkv_bias=bias)
    print(f"{name}: block={block:,}  FFN={ffn_share:.1%}  "
          f"total={total / 1e9:.2f}B  non-embedding={non_emb / 1e9:.2f}B")

Output:

Llama 3.1 8B: block=218,112,000  FFN=80.8%  total=8.03B  non-embedding=6.98B
Qwen2.5-7B: block=233,057,792  FFN=87.4%  total=7.62B  non-embedding=6.53B

For another model, replace the dict with the matching fields from its config.json. The code only covers dense Llama-style architectures: it does not apply as is to MoE models, models with biases in every layer, or models that share weights across layers.

FLOPs: 2N for the forward pass, 6N for training

A matrix multiplication (1 × a) · (a × b) takes a · b multiplications and a · b additions, so roughly 2 FLOPs per parameter. From that, Kaplan et al. estimate the forward cost per token:

C_forward ≈ 2N + 2 · n_layer · n_ctx · d_attn

The first term covers the multiplications by weight matrices. The second is the context-dependent part of attention: Q · Kᵀ and the multiplication by V, exactly the cost part 7, section 10 describes. In training, the backward pass costs about twice the forward pass, so the total is roughly 6N FLOPs per token.

Applied to Llama 3.1 8B, with N as the non-embedding parameter count (6.98B), as Kaplan defines it:

Part Formula GFLOPs per token
Matrices in the 32 blocks 2 · 6.98B ≈ 14.0
Context-dependent attention, n_ctx = 8,192 2 · 32 · 8,192 · 4,096 ≈ 2.1
lm_head 2 · 128,256 · 4,096 ≈ 1.05
Context-dependent attention, n_ctx = 131,072 2 · 32 · 131,072 · 4,096 ≈ 34.4
Training, 6N (N without embeddings) 6 · 6.98B ≈ 41.9
Training, 6N with N = 8.03B 6 · 8.03B ≈ 48.2

The table shows three things:

  1. At moderate context lengths, almost all FLOPs come from the matrices, and most of those from the FFN. At 8K tokens, the context-dependent part is only about 15% of 2N.
  2. At very long context lengths, the order flips. At 128K tokens, context-dependent attention (34.4 GFLOPs) is larger than all the matrix multiplications combined. This is the quadratic cost part 7 talks about.
  3. GQA does not reduce this n_ctx term. The query still has 32 heads × 128 = 4,096 dimensions, so Q · Kᵀ costs just as much. GQA reduces the parameters of W_K and W_V and shrinks the KV cache; it does not reduce the FLOPs of the attention operation itself.

The 6N figure is also what scaling papers such as Chinchilla use to convert between parameter count, token count and compute budget. When you see a “6N” number, check whether N includes embeddings: for Llama 3.1 8B, the difference is 42 versus 48 GFLOPs per token.


07Common misconceptions

Misconception Reality
“The FFN always holds about 2/3 of the parameters” 2/3 only holds for MHA with a two-matrix 4d FFN (GPT-2). With GQA and SwiGLU, the FFN is about 81% of a Llama 3.1 8B block and 87% of a Qwen2.5-7B block
“d_ff is always 4 · d_model” that is the convention for a two-matrix FFN. SwiGLU shrinks it to about 2/3 · 4d, then each team adjusts: Llama 2 7B uses 11,008, Llama 3 8B uses 14,336, Qwen2.5-7B uses 18,944
“Every modern LLM uses Pre-LN” most do, but Gemma 2 and Gemma 3 (Gemma Team 2024, 2025) normalize both before and after each sub-layer, OLMo 2 normalizes outputs, and Post-LN can still give better results if it trains
“All of a model’s knowledge lives in the FFN” evidence suggests the FFN plays a large role in recalling facts, but where a fact can be localized is not necessarily where it is stored or where it is best edited (Hase et al. 2023)
“RMSNorm exists only to run faster” RMSNorm is faster, but the key finding is that re-centering is not needed for quality. RMSNorm preserves the direction of the vector and only rescales its magnitude
“The residual connection is just a trick against vanishing gradients” that is one role. The other: residuals turn the hidden state into a shared bus that every layer reads and writes, so later layers can use the results of all earlier layers
“SwiGLU is better because theory proves it” Shazeer only measured it empirically and said outright that there is no explanation
“GQA makes attention cheaper in FLOPs” GQA reduces the K and V parameters and the KV cache size. The context-dependent Q · Kᵀ part costs the same as with MHA

Two halves of a block, and a bus that connects them

A Transformer block has two halves. Attention, the subject of part 7, is the communication step: each token pulls information from other tokens. This post covered the other half, and on closer inspection that is the half holding most of the parameters and compute.

What to take away:

  • the residual stream is a d_model-dimensional vector that runs straight through the whole model; each sub-layer reads by projecting and writes by adding, never overwriting;
  • the FFN has the same shape of formula as attention, f(x · K) · V, but its keys and values are fixed parameters, and there is evidence that keys catch text patterns while values push token probabilities;
  • “the FFN is the knowledge store” is an interpretation backed by evidence, not a theorem;
  • GELU smooths out ReLU; SwiGLU adds a gate branch, needs three matrices, and therefore shrinks d_ff to about 2/3 · 4d before each team adjusts it;
  • RMSNorm drops the mean subtraction without losing quality; Pre-LN keeps the residual path straight and trains stably, but norm placement is still being explored;
  • “FFN ≈ 2/3” and “d_ff = 4d” are GPT-2 numbers; in Llama 3.1 8B, the FFN is 80.8% of each block, and the 8.03B total can be counted by hand from config.json;
  • FLOPs are roughly 2N per token for the forward pass and 6N for training, plus a context-dependent attention term that overtakes 2N once the context is long enough.

After the last block, the stream goes through the final norm and into lm_head to become logits, one number per token in the vocabulary. Getting from logits to a concrete answer takes one more step: choosing the token. Part 9 covers that step, decoding and sampling: greedy, temperature, top-k, top-p, and why the same model can answer very differently depending only on how tokens are chosen.


References

Trương Đức Dũng

AI Engineer - Hanoi

Focused on LLMs, RAG, and AI systems that need to run reliably in production.