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 are worth noticing before any code. 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)

gemma_270M/model.py

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. Masks are built per forward pass at sequence length T rather than registered as full-context buffers; at a 32,768-token context, a cached full mask is real memory, and the comment in the code says exactly why it is not kept.

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; accumulating in fp32 is the standard fix, made explicit here rather than hidden inside a fused kernel. 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

gemma_270M/model.py

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

The transformer block, gemma_270M/model.py

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

GeGLU feed-forward, condensed, gemma_270M/model.py

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.