Gemma 3 270M, line by line

Gemma 3's 270M variant is the most instructive model in its family precisely because it is small: 18 layers, embedding dimension 640, and a design where almost every recent architecture idea appears once. Sliding-window attention five layers out of six, two RoPE frequency bases, QK-norm, normalization on both sides of every sublayer, and a GELU-gated feed-forward. My reimplementation is 252 lines of plain PyTorch, and this article walks the parts that differ from a vanilla decoder.

Two config choices shape the model. The vocabulary is enormous for the model size, 262,144 entries at dimension 640, so most of the parameters live in the embedding table; the file ends by printing the total parameter count with and without tying the output head to the embedding, which is the difference the release cares about at this scale. And the heads are wide: 4 query heads of dimension 256 each with a single KV head, so attention operates in a space wider than the residual stream itself.

One attention module, two personalities

Every sixth layer is a full-attention layer with RoPE base 1,000,000; the five layers between are sliding-window layers, window 512, with RoPE base 10,000. Which personality a layer gets is decided by the config's layer_type tuple, and inside the module it comes down to picking a theta and a mask.

if self.layer_type == "sliding_attention":
    theta = self.config.rope_local_base          # 10_000
    offset = -self.config.sliding_window
    mask = torch.triu(torch.tril(torch.ones((T, T), device=x.device)),
                      diagonal=offset)
elif self.layer_type == "full_attention":
    theta = self.config.rope_base                # 1_000_000
    mask = torch.tril(torch.ones((T, T), device=x.device))

cos, sin = self.get_rope_params(self.config.head_dim, theta, T, dtype=q.dtype)
q = self.apply_rope(q, cos, sin)
k = self.apply_rope(k, cos, sin)

The dual base is the interesting design decision: local layers get high-frequency rotations tuned for near context, and the sparse full layers get the long-wavelength base that makes 32k positions distinguishable. Each forward pass builds its mask at sequence length T. This sizes the mask to the current sequence and avoids retaining a full 32,768-token mask.

Attention accumulates in fp32

The model runs in bf16, but scores and softmax happen in float32 before casting back. At head dimension 256 the dot products are long sums, and bf16's 8 bits of mantissa lose real precision there; the implementation explicitly accumulates attention scores in fp32. QK-norm, an RMSNorm over the head dimension, runs on queries and keys before RoPE.

attn_scores = q.float() @ k.transpose(-1, -2).float()
attn_scores = attn_scores / math.sqrt(self.config.head_dim)
attn_scores = attn_scores.masked_fill(mask == 0, float("-inf"))
attn_scores = F.softmax(attn_scores, dim=-1)
attn_scores = attn_scores.to(dtype=self.config.dtype)   # back to bf16

Sandwich normalization and GeGLU

Gemma normalizes four times per block: before and after attention, before and after the feed-forward. The whole block is two lines. The feed-forward itself is GeGLU, the GELU-gated variant, where one projection gates the other before the down projection; the rest of my reimplementations use SiLU gating, so the contrast lives in one activation function.

def forward(self, x):
    x = x + self.n2(self.attn(self.n1(x)))
    x = x + self.n4(self.ffn(self.n3(x)))
    return x

def forward(self, x):
    x1 = F.gelu(self.w1(x))     # B, T, hidden_dim
    x2 = self.w2(x)             # B, T, hidden_dim
    return self.w3(x1 * x2)     # B, T, emb_dim

The small print

Two Gemma signatures round it out: token embeddings are scaled by the square root of the embedding dimension on the way in, and a final RMSNorm sits before the output head. Loss is computed in float32. Everything else, pre-norm residual blocks, broadcast grouped queries, a crop-and-sample generate(), follows the same skeleton as the other models in the repository, which is the point: hold the scaffold constant, and each architecture becomes a small diff against the previous one.