nnterp · encyclopedia gemma3_text · nnterp.families.gemma3_text

family gemma3_text

Gemma 3

Llama's module tree with a sandwich block, so what reaches the residual stream is a post-norm's output; queries and keys are normed per head before the rotary, and five sliding-window blocks precede each full one.

Sandwich normsQuery/key normsSliding windowScaled embeddings1 + weight norm gainMulti-query attention

The block

01 · the residual stream

One block of the stack. Hover any part for its nnterp name; click to pin the card. The slider picks a block; this family's blocks differ by layer_types.

embed_tokenstoken_embeddings
→
layers× 26
→
normfinal_norm_output
→
lm_headlm_head_output
→
logitsthe output
i = 0

identity layers[i].input + self_attn.attention_output + mlp.mlp_output == layer_output

The API

02 · support()

Values, by host

⚠︎ marks a value this checkpoint has only under some load: hover it for the condition.

model Gemma3ForCausalLMlayoutwhat it isread at
logits Logits [batch seq vocab_size] The logits the model returns, after anything it does past lm_head (a softcap, a scale)

the module's own output
token_embeddings Residual [batch seq hidden_size] The embedding module's output; layers[0].input is what enters the first block

the child module embed_tokens's output
final_norm_output Residual [batch seq hidden_size] The residual stream after the final norm, what lm_head reads

the child module norm's output
lm_head_output Logits [batch seq vocab_size] The head's logits before anything the model does past lm_head (a softcap, a scale)

the child module lm_head's output
next_token_probs NextTokenProbs [batch vocab_size] The next-token distribution at the last position; derived, read-only

the module's own output
input_ids Tokens [batch seq] The token ids the model was called with

the module's own input
attention_mask Tokens [batch seq] The attention mask the model was called with; zeros are padding

the module's own input
input_size scalar [batch, seq] of the current call; read-only

the module's own input
model.layers[i] Gemma3DecoderLayerlayoutwhat it isread at
layer_input Residual [batch seq hidden_size] The residual stream entering the block

the module's own input
attention_norm_output Residual [batch seq hidden_size] The normed stream entering the attention

the child module input_layernorm's output
layer_mid Residual [batch seq hidden_size] The residual stream after the attention add

the child module pre_feedforward_layernorm's input
mlp_norm_output Residual [batch seq hidden_size] The normed stream entering the MLP

the child module pre_feedforward_layernorm's output
mlp_input Residual [batch seq hidden_size] What the MLP is called with

the child module mlp's input
layer_output Residual [batch seq hidden_size] The residual stream leaving the block, a tensor even when the block returns a tuple

the module's own output
model.layers[i].self_attn Gemma3Attentionlayoutwhat it isread at
⚠︎attention_queries Queries [batch num_heads seq qk_head_dim] The queries entering attention

inside the forward: the arguments of attention_interface_1, element 1
⚠︎attention_keys Keys [batch num_kv_heads seq qk_head_dim] The keys entering attention

inside the forward: the arguments of attention_interface_1, element 2
⚠︎attention_values Values [batch num_kv_heads seq head_dim] The values entering attention

inside the forward: the arguments of attention_interface_1, element 3
⚠︎attention_scores Pattern [batch num_heads query key] The attention scores entering the softmax, masked

inside the forward: the input of nn_functional_softmax_0 inside attention_interface_1
attention_premix HeadOutputs [batch seq num_heads head_dim] Every head's output side by side, as the output projection takes it

the child module o_proj's input
attention_output Residual [batch seq hidden_size] What the attention adds to the residual stream: the post-attention norm's output

the sibling module post_attention_layernorm's output
⚠︎attention_probabilities Pattern [batch num_heads query key] The attention pattern the values are mixed with

inside the forward: the output of nn_functional_dropout_0 inside attention_interface_1
⚠︎attention_head_outputs HeadOutputs [batch seq num_heads head_dim] The per-head outputs before the output projection

inside the forward: the output of attention_interface_1, element 0
model.layers[i].mlp Gemma3MLPlayoutwhat it isread at
mlp_activation Neurons [batch seq intermediate_size] The activation function's output

the child module act_fn's output
mlp_neurons Neurons [batch seq intermediate_size] What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise

the child module down_proj's input
mlp_output Residual [batch seq hidden_size] What the MLP adds to the residual stream: the post-feedforward norm's output

the sibling module post_feedforward_layernorm's output

Printout

What print(model) shows on google/gemma-3-1b-pt: the native modules, standard names and every standard value with its layout.

Gemma3ForCausalLM(
  (lm_head): Linear(in_features=1152, out_features=262144, bias=False)
  (generator): Generator(
    (streamer): Streamer()
  )
  (embed_tokens): Gemma3TextScaledWordEmbedding(262144, 1152, padding_idx=0)
  (layers): ModuleList(
    (0-25): 26 x Gemma3DecoderLayer(
      (self_attn): Gemma3Attention(
        (q_proj): Linear(in_features=1152, out_features=1024, bias=False)
        (k_proj): Linear(in_features=1152, out_features=256, bias=False)
        (v_proj): Linear(in_features=1152, out_features=256, bias=False)
        (o_proj): Linear(in_features=1024, out_features=1152, bias=False)
        (q_norm): Gemma3RMSNorm((256,), eps=1e-06)
        (k_norm): Gemma3RMSNorm((256,), eps=1e-06)
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream: the post-attention norm's output
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_premix) -> HeadOutputs [batch seq num_heads head_dim]: Every head's output side by side, as the output projection takes it
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
      )
      (mlp): Gemma3MLP(
        (gate_proj): Linear(in_features=1152, out_features=6912, bias=False)
        (up_proj): Linear(in_features=1152, out_features=6912, bias=False)
        (down_proj): Linear(in_features=6912, out_features=1152, bias=False)
        (act_fn): GELUTanh()
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream: the post-feedforward norm's output
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
      )
      (input_layernorm): Gemma3RMSNorm((1152,), eps=1e-06)
      (post_attention_layernorm): Gemma3RMSNorm((1152,), eps=1e-06)
      (pre_feedforward_layernorm): Gemma3RMSNorm((1152,), eps=1e-06)
      (post_feedforward_layernorm): Gemma3RMSNorm((1152,), eps=1e-06)
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
  )
  (norm): Gemma3RMSNorm((1152,), eps=1e-06)
  (logits) -> Logits [batch seq vocab_size]: The logits the model returns, after anything it does past lm_head (a softcap, a scale)
  (token_embeddings) -> Residual [batch seq hidden_size]: The embedding module's output; layers[0].input is what enters the first block
  (final_norm_output) -> Residual [batch seq hidden_size]: The residual stream after the final norm, what lm_head reads
  (lm_head_output) -> Logits [batch seq vocab_size]: The head's logits before anything the model does past lm_head (a softcap, a scale)
  (next_token_probs) -> NextTokenProbs [batch vocab_size]: The next-token distribution at the last position; derived, read-only
  (input_ids) -> Tokens [batch seq]: The token ids the model was called with
  (attention_mask) -> Tokens [batch seq]: The attention mask the model was called with; zeros are padding
  (input_size): [batch, seq] of the current call; read-only
)

Config

model.num_layers26
model.hidden_size1152
model.num_heads4
model.num_kv_heads1
model.head_dim256
model.qk_head_dim256
model.vocab_size262144
model.intermediate_size6912
architecturesGemma3ForCausalLM
num_hidden_layers26
hidden_size1152
intermediate_size6912
num_attention_heads4
num_key_value_heads1
head_dim256
vocab_size262144
max_position_embeddings32768
sliding_window512
hidden_activationgelu_pytorch_tanh
rms_norm_eps1e-06
query_pre_attn_scalar256
tie_word_embeddingsTrue
dtypebfloat16
attention_biasFalse
use_bidirectional_attentionFalse

The Notes

03 · what to know before you run

The block, in order

a    = input_layernorm(x)
q, k = rope(q_norm(q_proj(a))), rope(k_norm(k_proj(a)))    # per head
h    = x + post_attention_layernorm(o_proj(attend(q, k, v_proj(a))))
out  = h + post_feedforward_layernorm(mlp(pre_feedforward_layernorm(h)))

Four RMSNorms on the block and two inside the attention. post_attention_layernorm is the norm after the attention, not Llama's pre-MLP norm; the MLP's input norm is pre_feedforward_layernorm. q_norm and k_norm act on each head's head_dim vector; the values are not normed.

The contributions are the post-norms' outputs

attention_output and mlp_output are what the block adds: the outputs of post_attention_layernorm and post_feedforward_layernorm. self_attn.output[0] and mlp.output are the tensors entering those norms. The identity holds exactly, in float32 and in bfloat16:

with model.trace(prompt):
    x = model.layers[5].input.save()
    attn = model.layers[5].self_attn.attention_output.save()
    mlp = model.layers[5].mlp.mlp_output.save()
    out = model.layers[5].layer_output.save()

torch.testing.assert_close(x + attn + mlp, out)

The post-norm sets the scale of what the stream receives, so scaling the module's output does almost nothing. On 1B, halving mlp.output at block 5 moves the logits by at most 0.3 in bfloat16 (0.03 in float32); halving mlp_output moves them by 2.9:

with model.trace(prompt):
    model.layers[5].mlp.mlp_output[:] *= 0.5   # halves what the block adds
    model.layers[5].mlp.output[:] *= 0.5       # nearly a no-op

The effect of halving mlp.output is not exactly zero: bfloat16 rounds, and the norm's eps matters where the module's output is small (the first position on 1B, RMS 0.004 at block 5). Ablations, steering and direct logit attribution of a sublayer go through attention_output and mlp_output.

Eager and sdpa compute the same attention

attn_logit_softcapping and final_logit_softcapping are null in every released config, and Gemma3Attention passes no softcap to the attention function in any case. A default (sdpa) load runs the model's own computation: on 1B in float32 its logits are within 2e-5 of an eager load's, on a 10-token and on a 701-token prompt. attn_implementation="eager" is needed only for the values read inside the attention function (the ones marked ⚠).

Queries and keys are read after their norms and the rotary

attention_queries and attention_keys are the tensors the attention function receives: rope(q_norm(q_proj(x))) and rope(k_norm(k_proj(x))), per head. The normed keys before the rotary are self_attn.k_norm.output ([batch, num_kv_heads, seq, head_dim]), the queries q_norm.output (num_heads wide); the norms' gain is 1 + weight, as on every Gemma 3 norm. attention_values is v_proj's output reshaped to heads.

The query scale is query_pre_attn_scalar ** -0.5. That is 256 ** -0.5 on 270M, 1B, 4B and 12B, where it equals head_dim ** -0.5; on 27B it is 168 ** -0.5 with head_dim 128. attention_scores are the scaled scores plus the mask, at the softmax's input.

Grouped-query attention

1B and 270M have 4 query heads and one key/value head: attention_keys and attention_values are [batch, 1, seq, 256], served before repeat_kv, and an edit to them reaches all four query heads. 4B, 12B and 27B have two query heads per key/value head (8/4, 16/8, 32/16).

Five sliding blocks to one full block

config.layer_types is full_attention on every sixth block (5, 11, 17, 23 on 1B's 26) and sliding_attention elsewhere. The window is 512 tokens on 270M and 1B, 1024 on 4B, 12B and 27B; a query in a sliding block attends to itself and the 511 tokens before it on 1B. On a prompt no longer than the window both masks are the same causal mask. A block's kind is model.layers[i].self_attn._module.sliding_window (the window, or None on a full block).

The two kinds also use different rotary embeddings, from config.rope_parameters: base 10,000 on sliding blocks and 1,000,000 on full blocks, with linear scaling by 8 on the full blocks of 4B, 12B and 27B. The model computes both cos/sin tables and hands each block the one for its kind, so attention_queries and attention_keys carry a different rotation in the two kinds of block.

The readout has no softcap

model.logits is model.lm_head.output, and project_on_vocab is the final norm and lm_head only, so a logit lens at the last block equals logits exactly:

with model.trace(prompt):
    resid = model.layers[12].layer_output[:, -1].save()

lens = model.project_on_vocab(resid)    # norm, then lm_head; no cap

The vocabulary is 262,144 tokens on 270M and 1B and 262,208 on 4B, 12B and 27B, so logits from the two groups are not the same width.

The norm gain is 1 + weight, and the weights are not small

Gemma3RMSNorm multiplies by 1 + weight. The stored weights are far from zero: on 1B the final norm's weights average 8.9 (maximum 50) and block 5's post_feedforward_layernorm reaches 246. Folding a norm into a neighbouring matrix uses 1 + model.norm._module.weight.

Embeddings are scaled and tied

embed_tokens multiplies its lookup by √hidden_size, cast to the weight's dtype first: 33.94 on 1B in float32 and 34.0 in bfloat16 (50.5 on 4B in bfloat16). token_embeddings is the scaled tensor and equals layers[0].input on a text prompt. The embedding weight is the unembedding (tie_word_embeddings), so an edit to embed_tokens.weight edits lm_head too.

4B, 12B and 27B are multimodal wrappers

270M and 1B are Gemma3ForCausalLM. 4B, 12B and 27B are gemma3 checkpoints: the text-generation task builds Gemma3ForConditionalGeneration, and nnterp picks this family from config.text_config and binds embed_tokens, layers and norm under model.language_model, with lm_head at the root. A text prompt runs the same block as above.

With an image, pass the processor's output to the trace (model.trace(dict(inputs))). The vision features replace the 256 image tokens' embeddings after embed_tokens, so at those positions token_embeddings is the placeholder token's embedding and layers[0].input is the image; at text positions the two are equal. With the processor's token_type_ids, the tokens of one image attend to each other in both directions, in sliding and full blocks alike. An image token still attends causally to earlier images and text, and text stays causal.

Sparse autoencoders: Gemma Scope 2

Gemma Scope 2 (google/gemma-scope-2-<size>-<pt|it>, every size from 270M to 27B) publishes SAEs, transcoders, crosscoders and, for 270M and 1B, cross-layer transcoders. Each config.json names its hook point as a transformers module path:

with model.trace(prompt):
    heads = model.layers[13].self_attn.o_proj.input.save()  # attn_out
    mlp_in = model.layers[13].mlp.input.save()       # transcoder input
    mlp_out = model.layers[13].mlp.mlp_output.save() # mlp_out, target
    resid = model.layers[13].layer_output.save()     # resid_post

The attn_out SAEs read o_proj's input, the heads' outputs concatenated (num_heads × head_dim, 1024 on 1B), not attention_output; that tensor is attention_head_outputs flattened over its last two axes, and o_proj.input is readable under any attention implementation. The mlp_out SAEs read post_feedforward_layernorm's output, which is mlp_output; the transcoders map pre_feedforward_layernorm's output (mlp.input) to it. The 4B, 12B and 27B configs name model.layers.<i> paths, which on the wrapper are model.language_model.layers.<i>: model.layers[i] either way.

From the family module

04 · nnterp/families/gemma3_text.py

Gemma 3, text (Gemma3ForCausalLM, model_type gemma3_text).

Llama's names plus a sandwich block: every sublayer is normed before and after, x + post_attention_layernorm(attn(input_layernorm(x))) and then + post_feedforward_layernorm(mlp(pre_feedforward_layernorm(x))). What the block adds is the post-norm's output, not the module's, so the contributions point at the sibling norms. A gemma3 checkpoint (Gemma3ForConditionalGeneration, gemma-3-4b/12b/27b) nests its config's text_config, of this type; the text-generation task builds the wrapper, whose text stack sits at model.language_model.{embed_tokens, layers, norm} with lm_head at the root, so RENAME carries both spellings and whichever the tree has binds. On the wrapper the SigLIP tower model.vision_tower is vision (a Vision; its post_layernorm over the patches is vision.norm) and the pooling projector model.multi_modal_projector, the last module before the scatter, is projector; image_features is read at the scatter in Gemma3Model's forward (ImageScatter). Loaded with task="image-text-to-text" (the processor), the tower serves vision.image_token_mask and vision.image_features. Only Gemma3ForCausalLM softcaps its logits where the config sets final_logit_softcapping (no released checkpoint does); the wrapper never does, and project_on_vocab follows the class.