nnterp · encyclopedia vaultgemma · nnterp.families.vaultgemma

family vaultgemma

VaultGemma

Gemma 2's names on a pre-norm block: each sublayer is normed on the way in only, so what reaches the residual stream is the module's own output, and the MLP's input norm is pre_feedforward_layernorm.

Scaled embeddings1 + weight norm gain

The block

01 · the residual stream

One block of the stack. Hover any part for its nnterp name; click to pin the card.

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 VaultGemmaForCausalLMlayoutwhat 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] VaultGemmaDecoderLayerlayoutwhat 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 post_attention_layernorm's input
mlp_norm_output Residual [batch seq hidden_size] The normed stream entering the MLP

the child module post_attention_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 VaultGemmaAttentionlayoutwhat 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 module's own 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 VaultGemmaMLPlayoutwhat 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 module's own output

Printout

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

VaultGemmaForCausalLM(
  (lm_head): Linear(in_features=1152, out_features=256000, bias=False)
  (generator): Generator(
    (streamer): Streamer()
  )
  (embed_tokens): VaultGemmaTextScaledWordEmbedding(256000, 1152, padding_idx=0)
  (layers): ModuleList(
    (0-25): 26 x VaultGemmaDecoderLayer(
      (self_attn): VaultGemmaAttention(
        (q_proj): Linear(in_features=1152, out_features=1024, bias=False)
        (k_proj): Linear(in_features=1152, out_features=1024, bias=False)
        (v_proj): Linear(in_features=1152, out_features=1024, bias=False)
        (o_proj): Linear(in_features=1024, out_features=1152, bias=False)
        (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_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (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): VaultGemmaMLP(
        (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_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
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): VaultGemmaRMSNorm((1152,), eps=1e-06)
      (post_attention_layernorm/pre_feedforward_layernorm): VaultGemmaRMSNorm((1152,), eps=1e-06)
      (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
      (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
      (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): VaultGemmaRMSNorm((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_heads4
model.head_dim256
model.qk_head_dim256
model.vocab_size256000
model.intermediate_size6912
architecturesVaultGemmaForCausalLM
num_hidden_layers26
hidden_size1152
intermediate_size6912
num_attention_heads4
num_key_value_heads4
head_dim256
vocab_size256000
max_position_embeddings1024
sliding_window512
hidden_activationgelu_pytorch_tanh
rms_norm_eps1e-06
query_pre_attn_scalar256
tie_word_embeddingsTrue
dtypebfloat16
attention_biasFalse

The Notes

03 · what to know before you run

The block, in order

h   = x + self_attn(input_layernorm(x))
out = h + mlp(pre_feedforward_layernorm(h))

Two RMSNorms per block, each before its sublayer, none after. The pre-MLP norm keeps Gemma 2's name, pre_feedforward_layernorm, and nnterp aliases it post_attention_layernorm, the name Llama gives the same place: both reach one module. There is no norm after the attention, so a hook ported from Gemma 2 that reads post_attention_layernorm as the attention's post-norm lands on the MLP's input norm here.

The contributions are the modules' outputs

attention_output is self_attn.output[0] and mlp_output is mlp.output, added to the stream with nothing in between. The stream between the two adds is the pre-MLP norm's input. On the pinned checkpoint, in float32:

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

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

Scaling a module's output scales its contribution: mlp.output[:] *= 0.5 and mlp.mlp_output[:] *= 0.5 are the same edit.

The released checkpoint sets no softcap and no window

The modeling code has Gemma 2's score and logit caps and its sliding blocks: a score softcap inside the eager attention (attn_logit_softcapping), a logit softcap after lm_head (final_logit_softcapping) and config.layer_types with sliding blocks. On google/vaultgemma-1b both caps are null and all 26 blocks are full_attention (the config's sliding_window of 512 is read by no block), so model.logits equals model.lm_head.output, project_on_vocab is lm_head(norm(hidden)), and attention_scores are the plain scaled scores. The pinned tiny checkpoint sets both caps (50 and 30) and alternates sliding and full blocks: there logits is 30 · tanh(lm_head.output / 30), and project_on_vocab applies the same cap.

Attention

1B has 4 query heads and 4 key/value heads, each 256 wide, so the heads are 1024 wide against a hidden_size of 1152. The query scale is query_pre_attn_scalar ** -0.5, which is 256 ** -0.5, equal to head_dim ** -0.5 on 1B. With no score cap the default sdpa load and attn_implementation="eager" run the same attention on 1B; eager serves the interior values. On a config that sets attn_logit_softcapping, only eager applies the cap: transformers' sdpa path takes no softcap.

Embeddings and the norm gain

embed_tokens multiplies its lookup by √hidden_size, cast to the weight's dtype (33.94 on 1B, 34.0 in bfloat16), so token_embeddings is the scaled tensor and equals layers[0].input. Every VaultGemmaRMSNorm multiplies by 1 + weight: folding the final norm into the unembedding uses 1 + model.norm._module.weight. On 1B lm_head shares the weight of embed_tokens (tie_word_embeddings), so an edit to the embedding matrix edits the readout.

From the family module

04 · nnterp/families/vaultgemma.py

VaultGemma (VaultGemmaForCausalLM).

Gemma-2's tree without the post-norms: a pre-norm block, x + attn(input_layernorm(x)) and then + mlp(pre_feedforward_layernorm(x)), so the contributions are the modules' outputs and the base holds. pre_feedforward_layernorm is the pre-MLP norm, aliased post_attention_layernorm (the Llama meaning; VaultGemma has no norm after the attention). The embedding module scales its own output by sqrt(hidden_size), so token_embeddings is layers[0].input. The attention runs the shared eager forward with Gemma-2's score softcapping inside it where attn_logit_softcapping is set (attention_scores are then the capped scores entering the softmax), and final_logit_softcapping, where set, applies to the logits after lm_head, which the root's project_on_vocab follows. The released google/vaultgemma-1b sets both caps to null and every block to full attention, so it caps nothing; the pinned tiny sets both.