vaultgemma¶
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.