Diagnose why a 48-layer model diverges right after a pre-norm to post-norm refactor

**The architectures differ at one line.** Pre-norm computes `x = x + sublayer(norm(x))`. Post-norm computes `x = norm(x + sublayer(x))`. The norm moved from before the sublayer to after the residual add. That one-line move changes the gradient story greatly. In pre-norm, the residual `+` provides an identity path from the loss back to every layer. Gradients flow back cleanly without depth-dependent amplification. In post-norm, the norm sits on top of the residual add. So the gradient through it depends on the activation scale at that layer. The cumulative effect across 48 layers makes gradients at the bottom either too big or too small at initialization. The first few optimizer steps with a cold model and no warmup make the loss diverge. Your pre-norm run was tolerant because the identity path kept gradients sane regardless of LR (learning rate). **Fix 1: add warmup.** The 2017 Transformer was post-norm and shipped with `lr = d_model^-0.5 * min(step^-0.5, step * warmup^-1.5)` because cold starts at full LR diverge. A linear warmup over 4000-10000 steps is the standard modern equivalent. Optionally lower the peak LR too. **Fix 2: residual scaling.** Scale initial sublayer output weights by `1/sqrt(2L)` (L layers).

**The mechanism.** Pre-norm keeps an identity path from loss to every layer because the residual `x + sublayer(norm(x))` passes gradients through the `+` unchanged. Post-norm's `norm(x + sublayer(x))` puts the norm on top of the add, so the gradient through it is scaled by the LayerNorm Jacobian at every layer. Stack 48 of these and the gradient at layer 1 picks up a depth-dependent scale factor that is sensitive to the activation magnitudes early in training. **Why the first 300 steps blow up.** At init, sublayer outputs have variance proportional to their projections. With standard Xavier or Kaiming init, the sum `x + sublayer(x)` has larger variance than `x`, so every post-norm LayerNorm rescales the identity path. After backpropagating through 48 such layers, the gradient scale depends strongly on depth, and parameter updates are large enough to overshoot. The first optimizer step lands in a region where the LayerNorm denominators are different from what the rest of the network expects, and the loss explodes. **Pre-norm tolerated the same schedule** because its identity gradient path makes the effective depth multiplicative factor close to 1 regardless of layer count. GPT-2's 48-layer model used pre-norm. **Fix 1: warmup.** Add 4000-10000 steps of linear LR warmup. This is what the 2017 paper's inverse-sqrt schedule does implicitly. Optionally lower the peak LR by 2-4x if you want extra margin. **Fix 2: depth-aware residual scaling.** Scale each sublayer's output projection by `1/sqrt(2L)` at init, or use DeepNet. Each block then starts close to identity.