SwiGLU Feed-Forward Block (LLaMA MLP) - LLM Internals

SwiGLU Feed-Forward Block (LLaMA MLP)

What you'll build

The MLP half of a LLaMA decoder block: RMSNorm, a SwiGLU gated feed-forward network, and the residual connection around them. Three functions: rms_norm, swiglu_ffn and llama_mlp_block.

Background

The original Transformer MLP is ReLU(x W_1) W_2. LLaMA, Mistral, Qwen and PaLM use a gated MLP instead, SwiGLU (Shazeer, 2020): one projection goes through SiLU and gates a second, linear projection elementwise. At equal parameter count it gives lower loss. LLaMA also replaces LayerNorm with RMSNorm, which skips the mean subtraction and the bias: cheaper, and empirically just as good. About two thirds of a dense LLM's parameters live in these MLPs.

Formula

rms_norm(x)  = x / sqrt(mean(x^2, axis=-1) + eps) * weight

silu(z)      = z * sigmoid(z) = z / (1 + exp(-z))
gate         = silu(x @ W_gate)                  # (..., d_ff)
up           = x @ W_up                          # (..., d_ff)
swiglu_ffn(x) = (gate * up) @ W_down             # (..., d_model)

llama_mlp_block(x) = x + swiglu_ffn(rms_norm(x))

No biases anywhere (LLaMA uses none). With three weight matrices instead of two, LLaMA sets d_ff ≈ (2/3) · 4 · d_model to match the parameter count of a 4 · d_model ReLU MLP.

Inputs

• x: float array (..., d_model) (any number of leading axes, e.g. (batch, seq_len, d_model))
• weight / norm_weight: (d_model,)
• W_gate, W_up: (d_model, d_ff); W_down: (d_ff, d_model)
• eps: float, default 1e-6

Expected Return Format

• rms_norm(x, weight, eps=1e-6): same shape as x
• swiglu_ffn(x, W_gate, W_up, W_down): shape (..., d_model)
• llama_mlp_block(x, norm_weight, W_gate, W_up, W_down, eps=1e-6): same shape as x

Edge cases and conventions

• RMSNorm normalizes over the last axis only; eps goes inside the square root.
• The residual adds the un-normalized x (pre-norm): x + ffn(norm(x)), not norm(x) + ffn(norm(x)).
• The gate is SiLU, not sigmoid (that would be GLU) and not ReLU (ReGLU).

Constraints

NumPy only, vectorised.

Follow-up

• Count the parameters of the SwiGLU MLP and of a ReLU MLP with d_ff = 4 d_model. What d_ff makes them equal?
• Write the attention half of the LLaMA block (RMSNorm, RoPE, GQA, causal mask). Which earlier problems give you each piece?

Input: rms_norm(np.array([[3.0, 4.0]]), np.ones(2))
Output: [[0.8485, 1.1314]]

Input: swiglu_ffn(np.array([[1.0, -1.0]]), np.array([[1.0], [0.0]]), np.array([[2.0], [1.0]]), np.array([[1.0, 0.5]]))
Output: [[0.7311, 0.3655]]

import numpy as np

def rms_norm(x, weight, eps=1e-6):
    """
    RMSNorm over the last axis: x / sqrt(mean(x**2) + eps) * weight.

    Args:
        x: (..., d_model)
        weight: (d_model,) learned scale

    Returns:
        Array with the same shape as x
    """
    # Your code here
    pass


def swiglu_ffn(x, W_gate, W_up, W_down):
    """
    SwiGLU feed-forward network: (silu(x @ W_gate) * (x @ W_up)) @ W_down

    Args:
        x: (..., d_model)
        W_gate, W_up: (d_model, d_ff)
        W_down: (d_ff, d_model)

    Returns:
        Array of shape (..., d_model)
    """
    # Your code here
    pass


def llama_mlp_block(x, norm_weight, W_gate, W_up, W_down, eps=1e-6):
    """
    Pre-norm MLP sub-block with a residual: x + swiglu_ffn(rms_norm(x)).

    Returns:
        Array with the same shape as x
    """
    # Your code here
    pass