OLMo 2 Reordered Post-Norm

§ 1 · Premise

Why anyone would revisit Post-Norm in 2024

The 2020-era settlement (see norm-placement entry) was Pre-Norm everywhere: keep the residual stream clean of normalization, and the gradient through depth collapses to nearly $\mathbf{I}$ instead of the $\Theta(\sqrt{L})$ growth that Post-Norm exhibited (Xiong et al. arXiv 2002.04745 Theorem 1). Every open dense LLM from Llama 1 through Llama 3, DeepSeek V3, and Qwen 3 was Pre-Norm.

Five years of frontier training have surfaced a second failure mode that Pre-Norm trades into. Because no normalization touches the residual stream itself, the L2 norm of $\mathbf{x}_\ell$ grows monotonically across depth: each sublayer adds a perturbation of bounded variance, so $\mathbb{E}[\|\mathbf{x}_\ell\|_2^2] \approx \|\mathbf{x}_0\|_2^2 + \ell\,\sigma^2_\text{block}$, and $\|\mathbf{x}_\ell\|_2 \sim \sqrt{\ell}$. At depths of 60–120 layers this growth is observable and interacts badly with three concrete production failure modes:

1. Attention logit blow-up. Q and K projections of an unboundedly-growing residual stream produce vectors with unboundedly-growing norms, which dominate the attention softmax (the QK-Norm entry is the proximate fix for this).
2. FFN gating saturation. SwiGLU and GeGLU FFNs gate by $\sigma$ or GELU; large inputs saturate the gate and stall the gradient through the FFN’s down-projection.
3. Periodic loss spikes during long training runs. OLMo 1 (the Pre-Norm predecessor) reported these spikes in its training writeup (arXiv 2402.00838 §3.3) — sudden order-of-magnitude jumps in loss that required restarting from earlier checkpoints. The OLMo 1 team attributed the spikes to interactions between residual growth and the AdamW optimizer’s second-moment estimate.

OLMo 2 (arXiv 2501.00656) asks: at frontier scale, with modern stabilizers (QK-Norm, Z-loss, careful init), can a Post-Norm-flavored layout actually train more reliably than Pre-Norm? The answer, per the paper’s 5T-token training-stability study (§3.1), is yes — provided the placement change is paired with the rest of the stability package.

§ 2 · Derivation

RMSNorm after each sublayer, inside the residual

OLMo 2’s block, in the convention of the paper (§3.1, Eq. 1):

Compare to standard Pre-Norm:

The RMSNorm has moved from the sublayer’s input to the sublayer’s output, but stays inside the residual branch. The identity path $\mathbf{x}_\ell \to \mathbf{x}_{\ell+1}$ — the residual connection itself — remains unnormalized.

Why this preserves Pre-Norm’s gradient flow. The forward identity term is the same: $\mathbf{x}_{\ell+1} = \mathbf{x}_\ell + (\text{stuff})$ with no normalization on $\mathbf{x}_\ell$ itself. The backward Jacobian inherits the dominant identity:

with $\partial \mathbf{h}_\ell / \partial \mathbf{x}_\ell = \mathbf{I} + \mathbf{J}_{\mathrm{RMSNorm}}(\mathrm{Attn}(\mathbf{x}_\ell))\, \mathbf{J}_{\mathrm{Attn}}(\mathbf{x}_\ell)$ having the same identity-plus-correction shape. The depth product collapses to nearly $\mathbf{I}$ at initialization for the same reason Pre-Norm’s does — random sublayer initialization makes the correction term small, and the $\mathrm{RMSNorm}$ Jacobian is $O(1)$ rather than the $O(1/\sqrt{L})$-suppressing factor that classical Post-Norm pays.

What the placement changes about the forward pass. The sublayer output $\mathrm{Attn}(\mathbf{x}_\ell)$ is now explicitly rescaled to unit RMS before being added to the residual. Each block contributes a perturbation of fixed RMS magnitude, regardless of what the sublayer learned to output. The residual stream norm grows as $\sqrt{\ell}$ exactly, with a constant set by the RMSNorm’s learnable gain, instead of with the unbounded constant that Pre-Norm allows. Empirically, this caps the activation magnitudes that downstream attention and FFN sublayers see.

Relation to Sandwich-LN. Sandwich-LN (see Sandwich-LN entry) places norms both before and after each sublayer (two norms per sublayer, four per block). OLMo 2’s variant drops the input norm (one norm per sublayer, two per block). Algebraically:

Both put a norm on the sublayer output — that is the load-bearing piece for bounding the residual contribution. OLMo 2’s choice trades the redundant input normalization (the sublayer projections already produce an output whose statistics are controlled by their weight init, and Pre-Norm-style input normalization was about gradient flow, which OLMo 2’s identity-path residual already provides) for half the normalization compute. Both choices live in the same bounded-perturbation design space; the empirical decision between them has not been ablated in public.

Why the bare placement change alone is insufficient. Reordered Post-Norm at 13B without other interventions would still hit the failure modes Pre-Norm was originally chosen to avoid. The OLMo 2 tech report (§3.1) is explicit that the placement is part of a package, not a standalone improvement. The supporting cast:

• QK-Norm (per-head RMSNorm on Q and K before the attention dot product) caps the attention-logit magnitude regardless of residual-stream growth. Detailed in the QK-Norm entry. The Henry et al. 2020 paper (arXiv 2010.04245) is the original; OLMo 2 uses the RMSNorm variant.
• Output Z-loss — an auxiliary loss penalizing $(\log Z)^2$ where $Z = \sum_v e^{z_v}$ is the softmax partition function over the vocabulary. Keeps the unembedding logits from drifting toward extreme scales, which would compound any residual-stream growth at the network’s top.
• Conservative init — truncated normal at $\sigma_\text{init} = 0.02$ tuned for the specific block layout, so the initial per-block gradient magnitude matches what a Pre-Norm baseline at the same depth would produce.

Parameter and FLOP cost. Identical to Pre-Norm: two RMSNorms per block, $2 L d$ extra parameters total, $O(L T d)$ extra FLOPs per forward pass. The placement change is free in compute; the cost is paid in coordinating the three-stabilizer package.

§ 3 · Reference implementation

Reference implementation

# OLMo 2 reordered Post-Norm block.
def block_olmo2(x, attn, ffn, n1, n2):
    # x: [B, T, d]
    # n1, n2: RMSNorm modules with learnable gain [d]
    x = x + n1(attn(x))         # RMSNorm AFTER attention, inside residual
    x = x + n2(ffn(x))          # RMSNorm AFTER ffn, inside residual
    return x

# QK-Norm sits inside the attention sublayer itself.
def attn_with_qk_norm(x, w_qkv, w_o, qn, kn, n_heads, head_dim):
    B, T, d = x.shape
    qkv = x @ w_qkv
    q, k, v = qkv.split(d, dim=-1)
    q = q.view(B, T, n_heads, head_dim).transpose(1, 2)  # [B, H, T, d_h]
    k = k.view(B, T, n_heads, head_dim).transpose(1, 2)
    v = v.view(B, T, n_heads, head_dim).transpose(1, 2)
    q = qn(q); k = kn(k)        # per-head RMSNorm — see qk-norm entry
    logits = (q @ k.transpose(-2, -1)) / (head_dim ** 0.5)
    return (logits.softmax(-1) @ v).transpose(1, 2).reshape(B, T, d) @ w_o

The two changes vs Pre-Norm are mechanical: the RMSNorm calls migrate from the input of each sublayer to its output, inside the residual addition. The Z-loss term is added to the training objective, not to the block’s forward pass; it does not appear in this sketch.

§ 4 · Empirical evidence

Empirical evidence

The OLMo 2 stability study (Walsh et al. 2024, §3.1). The most complete public ablation of the placement choice at LLM scale. OLMo 2 trains 7B and 13B variants comparing four configurations: (a) OLMo 1’s Pre-Norm baseline, (b) Pre-Norm + QK-Norm, (c) reordered Post-Norm without QK-Norm, (d) full OLMo 2 stack (reordered Post-Norm + QK-Norm + Z-loss). Figure 4 reports loss spike counts per 1000 training steps over a 5T-token run; the full stack reduces spikes to roughly 1/10 of OLMo 1 baseline frequency, while (b) and (c) each give a partial reduction. The paper is careful to note that no individual component of the package fully recovers stability — the placement change alone (configuration c) still spikes; QK-Norm alone (configuration b) still spikes; only the full package matches the spike-free regime.

OLMo 1 baseline failures. The OLMo 1 paper (arXiv 2402.00838, §3.3) documents the Pre-Norm failure mode that motivated the OLMo 2 redesign: periodic loss spikes during the OLMo 1B and 7B training runs, attributed to interactions between residual-stream growth and AdamW’s second-moment estimate. The cumulative engineering cost of these spikes — restarts from checkpoints, lost compute — was substantial enough to motivate a top-to-bottom architecture review for OLMo 2.

Carryover to OLMo 3 (AI2, 2025). The OLMo 3 release (arXiv 2512.13961) inherits the reordered Post-Norm, QK-Norm, and Z-loss package without modification for the 32B dense flagship. The OLMo 3 team’s choice not to revisit the placement is a soft signal that the OLMo 2 setup generalizes to the 64-layer / 32B regime; the OLMo 3 paper does not present a fresh ablation against Pre-Norm at the 32B size.

Cross-team adoption. No other open-weight LLM family in this knowledge base has adopted OLMo 2’s specific layout. The Gemma 2/3 family uses Sandwich-LN (the two-norm variant), which lives in the same bounded-perturbation design space but is a different point in it. DeepNet (arXiv 2203.00555) demonstrates a related modified-Post-Norm approach for deep encoder-decoders but with different scaling choices and has not been picked up by any production open dense LLM either. Llama, DeepSeek, Qwen, Mistral, Hunyuan, Kimi, and MiniMax all continue with plain Pre-Norm. The honest reading is that the OLMo 2 choice is empirically defensible at the 13B/32B scale where AI2 trains it, but the broader open frontier has not converged on it — possibly because the three-component stability package is harder to port than a single architectural swap.

No public cross-team head-to-head. No paper has published a fixed-budget, fixed-data comparison of OLMo 2 reordered Post-Norm vs Pre-Norm vs Sandwich-LN at constant parameter count. The strongest case for the OLMo 2 layout is the AI2 internal ablation in Figure 4; the strongest case against it is its non-adoption outside AI2.

@article{arxiv2501_00656,
  title  = {2 OLMo 2 Furious},
  author = {Allen Institute for AI (AI2)},
  year   = {2024},
  eprint = {2501.00656},
  archivePrefix = {arXiv},
  url    = {https://arxiv.org/abs/2501.00656}
}

@article{arxiv_2501_00656,
  title         = {2 OLMo 2 Furious},
  author        = {Allen Institute for AI (AI2)},
  year          = {2024},
  eprint        = {2501.00656},
  archivePrefix = {arXiv},
  url           = {https://arxiv.org/abs/2501.00656}
}

{
  "id": "arxiv_2501_00656",
  "type": "article-journal",
  "title": "2 OLMo 2 Furious",
  "author": [
    {
      "literal": "Allen Institute for AI (AI2)"
    }
  ],
  "issued": {
    "date-parts": [
      [
        2024
      ]
    ]
  },
  "URL": "https://arxiv.org/abs/2501.00656",
  "number": "2501.00656",
  "source": "arXiv"
}

TY  - JOUR
TI  - 2 OLMo 2 Furious
AU  - Allen Institute for AI (AI2)
PY  - 2024
JO  - arXiv
AN  - arXiv:2501.00656
UR  - https://arxiv.org/abs/2501.00656
ER  -