AI Engineer Exhaustive
HardPlusGemma 3

Full Gemma 3 Block

Gemma 3

Hard

Complete Gemma 3 transformer block

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

A Gemma 3 layer is a pre-norm / post-norm sandwich around grouped-query attention and a gated MLP. The technical report says the models “follow the same general decoder-only transformer architecture” as Gemma 1/2 and that they “use Grouped-Query Attention (GQA) with post-norm and pre-norm with RMSNorm.” It also records the Gemma 3-specific attention changes: QK-norm instead of soft-capping, and 5:1 local/global interleaving with a short local window. The residual algebra is not written out. The official Block in gemma/gm/nn/_modules.py is the definition.

Every released Gemma 3 size sets use_post_attn_norm=True, use_post_ffw_norm=True, use_qk_norm=True, and attn_logits_soft_cap=None. The block therefore has four RMSNorms on the residual stream (pre-attention, post-attention, pre-MLP, post-MLP) plus the two Q/K norms inside attention. Final-logit soft-cap is also None; that lives on the transformer, not the block.

How it works

Let (x \in \mathbb{R}^{B \times T \times D}) be the residual stream. Block.__call__ is

\begin{aligned} a &= \mathrm{Attn}\!\left(\mathrm{RMSNorm}_{\mathrm{pre\text{-}attn}}(x)\right),\\ a &\leftarrow \mathrm{RMSNorm}_{\mathrm{post\text{-}attn}}(a),\\ h &= x + a,\\ m &= \mathrm{FFN}\!\left(\mathrm{RMSNorm}_{\mathrm{pre\text{-}ffn}}(h)\right),\\ m &\leftarrow \mathrm{RMSNorm}_{\mathrm{post\text{-}ffn}}(m),\\ y &= h + m. \end{aligned}

Each RMSNorm is the official last-axis form with (\varepsilon=10^{-6}) and a zeros-init scale applied as ((1+\gamma)) (gemma/gm/nn/_layers.py). Post-norms sit on the sublayer output, before the residual add — the Gemma 2-style sandwich, not a post-norm-only Transformer.

Attention. Attn is the module in the attention-block note: GQA projections, QK-norm, RoPE, query scale, local or global mask, softmax, output projection. The block does not choose the mask itself. Transformer.setup constructs each Block with an attn_type from the tiled GEMMA3_ATTENTION_PATTERN and the matching RoPE base/scale. Local layers get (\theta=10^4) and window (W\in{512,1024}); global layers get (\theta=10^6) and, on 4B/12B/27B, RoPE scale 8. query_pre_attn_scalar is (d_h^{-1/2}) except on 27B, where it is ((D/H)^{-1/2}).

MLP. FeedForward is a GeGLU. A gating einsum produces a tensor of shape [B, T, 2, D_{\mathrm{ff}}]. Gemma 3 sets transpose_gating_einsum=True, so the weight is stored as (2, D_ff, D) and contracted with '...F,NHF->...NH'. Then

\mathrm{FFN}(u)=\bigl(\mathrm{GELU}(uW_g)\odot (uW_u)\bigr)W_d,

where (W_g,W_u) are the two slices of that gating tensor and (W_d) is linear of shape (D_ff, D). nn.gelu here is the Flax default (the approximate tanh form). Hidden widths in _gemma.py are 2048 (270M), (6D) (1B), and (4D) on 4B/12B/27B (hidden_dim = embed_dim * 8 // 2).

What the paper specifies vs. what the code fills in. Specified: decoder-only GQA, pre- and post-RMSNorm, QK-norm, 5:1 local/global, local window 1024, RoPE bases (10^4) / (10^6). Inferred: residual-add after each post-norm, GeGLU with transposed gating weights, the four residual-stream norms plus two Q/K norms, the 27B query-scale exception, and the 512-token window on 270M/1B. The report never names GeGLU; the feed-forward is visible only in the library.

A full forward pass over the model is then: embed (times (\sqrt{D})), run (N) such blocks, final RMSNorm, untie-less decode through the embedding table. Image tokens, when present, are merged before the first block via a frozen SigLIP encoder and are outside this note.

Official code

Watch-outs

Sources