RMSNorm and SwiGLU (Llama's MLP Block)

easy
Unit 2 · Transformers & Training0/3 to milestone

Llama, Mistral and most modern LLMs swapped two classic pieces for cheaper or better ones:

  • RMSNorm replaces LayerNorm: no mean subtraction, no bias. RMSNorm(x)=xmean(x2)+ϵ⊙g\text{RMSNorm}(x) = \frac{x}{\sqrt{\text{mean}(x^2) + \epsilon}} \odot g, computed over the last axis.
  • SwiGLU replaces the ReLU/GELU MLP: SwiGLU(x)=(SiLU(xW1)⊙xW3)W2\text{SwiGLU}(x) = \big(\text{SiLU}(xW_1) \odot xW_3\big) W_2, where SiLU(z)=z⋅σ(z)\text{SiLU}(z) = z \cdot \sigma(z).

Implement silu(z), rms_norm(x, g, eps=1e-6) and swiglu_mlp(x, W1, W3, W2), then llama_block_mlp(x, g, W1, W3, W2), the residual sub-block x+SwiGLU(RMSNorm(x))x + \text{SwiGLU}(\text{RMSNorm}(x)).

#transformers#normalization#activations2 visible + 2 hidden tests
solution.py · Python 3.13
Loading editor…
Run your code to see results. ⌘↵ to run · ⌘⇧↵ to submit
RMSNorm and SwiGLU (Llama's MLP Block) · TechGradify