SwiGLU FFN

LLaMA

Easy

Gated feedforward with Swish activation

Independent study note. Written from the public paper and official code. This is not TensorTonic Plus and does not reproduce their exercises, starter code, or tests. For the official version, subscribe on TensorTonic.

Overview

Every Llama 3 block’s position-wise network is SwiGLU (Shazeer, 2020): a SiLU-gated linear unit, not a single GELU/d_ff = 4d MLP. Table 3 of the herd paper lists the activation as SwiGLU and gives the inner widths 14,336 / 28,672 / 53,248. The paper does not write (W_1,W_2,W_3) or the SiLU formula. Those are the official FeedForward module, which is the same gated layout Llama has used since the 2023 papers the authors say they kept.

SwiGLU sits after ffn_norm and writes back through a residual add. It is the bulk of the parameter count per layer (three wide projections versus attention’s four).

How it works

Let (h \in \mathbb{R}^{B \times S \times d}) be the post-RMSNorm residual. Shazeer’s SwiGLU with a down-projection is

\mathrm{FFN}(h) = \bigl(\mathrm{SiLU}(h W_1) \odot (h W_3)\bigr) W_2,

where (\mathrm{SiLU}(z)=z,\sigma(z)) (the same function PyTorch calls F.silu, also known as Swish-1). (W_1,W_3 \in \mathbb{R}^{d \times d_{\mathrm{ff}}}) expand; (W_2 \in \mathbb{R}^{d_{\mathrm{ff}} \times d}) contracts. All three are bias-free in the official code (ColumnParallelLinear / RowParallelLinear with bias=False).

The one-line forward is w2(silu(w1(x)) * w3(x)). The gate (w1 + SiLU) and the value (w3) are independent projections of the same (h); their product is elementwise. There is no extra activation on (W_3 h) and none after (W_2).

Inner width. TransformerBlock constructs the FFN with hidden_dim=4*args.dim, then FeedForward.__init__ rewrites that integer:

1. (d_{\mathrm{ff}} \leftarrow \lfloor 2\cdot(4d)/3 \rfloor = \lfloor 8d/3 \rfloor).
2. If ffn_dim_multiplier is set, multiply and again take int.
3. Round up to a multiple of multiple_of (default 256):
   (d_{\mathrm{ff}} \leftarrow m \cdot \lceil d_{\mathrm{ff}} / m \rceil).

Table 3’s 14336 / 28672 / 53248 are the checkpoint widths after this recipe (and the per-size multiplier stored in params.json). They are not (4d) (16384 / 32768 / 65536). The (2/3) factor is the usual parameter-budget correction so that three projections cost about as much as a two-matrix (4d) MLP. The paper states the resulting sizes, not the rounding code.

Shapes stay (B\times S\times d) in and out, so the residual h + FFN(...) is legal. Sequence tokens do not mix; this is a pointwise MLP.

Specified vs inferred: “Activation Function: SwiGLU” and the three FFN dimensions are paper. SiLU vs Swish naming, the (W_1/W_3) gate-value split, the (2/3) + multiple_of constructor, and the lack of bias are official-code / Shazeer details.

Official code

llama/model.py — class FeedForward (w1, w2, w3, forward) and the allocation in TransformerBlock.__init__. Inner size for a released model is the ffn_dim_multiplier / multiple_of pair in that checkpoint’s params.json, loaded by Llama.build in llama/generation.py.

Watch-outs

• Gate is SiLU on (W_1), not on (W_3), and not GELU. Swapping the two weights or using GELU((W_1 h)) (\odot W_3 h) is a different FFN.
• The product is elementwise, then (W_2). A matmul between the gated tensor and (W_3) (as if it were a second sequence mix) is wrong rank and wrong math.
• Do not hard-code (d_{\mathrm{ff}}=4d). Llama 3 8B is 14336, not 16384; the extra columns will not match w1.weight.
• Parallel linear layers split (d_{\mathrm{ff}}) across ranks. A single-GPU reimplementation still needs the full Table 3 width, just without the column/row split.

Sources

• SwiGLU: Shazeer, “GLU Variants Improve Transformer,” 2020
• Paper: The Llama 3 Herd of Models (Llama Team, 2024), §3.2, Table 3
• Code: meta-llama/llama3 llama/model.py (FeedForward)