gemma4_text¶

gemma4_text ¶

Gemma 4, text (Gemma4ForCausalLM, model_type gemma4_text): E2B, E4B, 26B-A4B and 31B.

model.{embed_tokens, layers[i].{input_layernorm, self_attn, post_attention_layernorm, pre_feedforward_layernorm, mlp, post_feedforward_layernorm}, norm} and lm_head: Gemma-3's sandwich tree. A gemma4 checkpoint (Gemma4ForConditionalGeneration, the published repos) 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.

The block, in order::

x1 = x  + post_attention_layernorm(self_attn(input_layernorm(x)))
x2 = x1 + post_feedforward_layernorm(mlp(pre_feedforward_layernorm(x1)) [+ experts])
p  = gelu(per_layer_input_gate(x2)) * per_layer_input
x3 = x2 + post_per_layer_input_norm(per_layer_projection(p))
out = x3 * layer_scalar          # in place, on x3

The contributions are the post-norms' outputs, as on Gemma-2/3. On a mixture-of-experts checkpoint (26B-A4B, enable_moe_block) the block also runs router and experts, direct children of the block, not of mlp: post_feedforward_layernorm_1(mlp(...)) + post_feedforward_layernorm_2(experts(pre_feedforward_layernorm_2(x1))), and post_feedforward_layernorm norms that sum, so mlp_output is the dense MLP and the experts together, what the block adds. The mixture has no module of its own, so mlp (a Moe) hosts its values, reading the block's router and experts through its parent. The router runs on x1, the stream after the attention's add (the block's residual at that point, not its input), with its own norm, and returns probabilities; router_logits are its projection's output (router.proj). shared_expert_output is the dense MLP's output, before its post-norm: mlp_output == post_feedforward_layernorm( post_feedforward_layernorm_1(shared_expert_output) + post_feedforward_layernorm_2(routed_output)). On a dense checkpoint every mixture value is unavailable.
Per-layer embeddings (E2B/E4B, hidden_size_per_layer_input): a third add, served as layers[i].per_layer_output (the post-per-layer-input norm's output). per_layer_input is the block's second argument, one slice of the model's per-layer embedding table mixed with a projection of the token embeddings. On a checkpoint without them the value is unavailable.
layer_scalar, a per-block buffer, multiplies the sum in place. The contributions are served unscaled, so the identity is (input + attention_output + mlp_output [+ per_layer_output]) * layer_scalar == layer_output.
Keys and values that are not the block's own. The last num_kv_shared_layers blocks have no k_proj/v_proj: each borrows the keys and values of the last earlier block of the same kind (sliding or full) before the sharing starts. On a checkpoint with attention_k_eq_v (26B-A4B, 31B) the full-attention blocks have no v_proj: their values are v_norm(k_proj(x)), the keys' projection before k_norm and the rotary embedding. attention_keys and attention_values are what the attention receives on every block, borrowed or not, served as copies private to the block: an edit, in place or by assignment, changes what that block attends with and nothing else, and the blocks that borrow a source block's tensors keep the originals.
Per-layer sizes. Sliding and full blocks differ in head_dim (and on 26B-A4B/31B in num_key_value_heads), and on E2B (use_double_wide_mlp) the KV-sharing blocks' MLP is twice intermediate_size wide; the config marks the attention sizes per-layer and refuses a plain config.head_dim. The root's head_dim and num_kv_heads report the config's top-level values (the sliding blocks'); each block's own are on layers[i].self_attn and layers[i].mlp, read off the module.

final_logit_softcapping is in the text config; the root's project_on_vocab reads it there.

On the wrapper the vision tower model.vision_tower is vision (a Vision) and model.embed_vision, the last module before the scatter, is projector. The tower is a ViT over each image's patches: the patch embedder's linear over the processor's flattened patches (patch_embedder.input_proj) is vision.patch_embed, so vision.patch_embeddings is before the 2D position embedding patch_embedder adds, and the blocks (encoder.layers) are vision.layers. Each block is the text block's sandwich, so its contributions are the post-norms' outputs too (VisionAttention, VisionMlp). The processor pads every image's patches to max_soft_tokens * pooling_kernel_size**2 rows with zero pixels at position (-1, -1); the padded rows run through every block (masked as keys only), so they are present in vision.layers[i].layer_output. There is no final norm: the tower's pooler zeroes the padded rows, average-pools the patches pooling_kernel_size by pooling_kernel_size and multiplies by sqrt(vision_hidden_size) (in float32); the tower then strips the padding and, where the vision config sets standardize (26B-A4B, 31B), subtracts std_bias and multiplies by std_scale. So vision.tower_output is the encoder's output (the last block's stream, padded) and projector.input is the pooled, scaled (and standardized) soft tokens, flat over the images. The audio tower (model.audio_tower, a Conformer) and its embedder model.embed_audio keep their native names.

RENAME `module-attribute` ¶

RENAME = {'model.embed_tokens': 'embed_tokens', 'model.layers': 'layers', 'model.norm': 'norm', 'model.language_model.embed_tokens': 'embed_tokens', 'model.language_model.layers': 'layers', 'model.language_model.norm': 'norm', 'model.vision_tower': 'vision', 'model.embed_vision': 'projector', 'patch_embedder.input_proj': 'patch_embed', 'encoder.layers': 'layers'}

ENVOYS `module-attribute` ¶

ENVOYS = {Gemma4Model: ImageScatter, Gemma4TextDecoderLayer: Layer, Gemma4TextAttention: Attention, Gemma4TextMLP: Mlp, Gemma4VisionModel: Vision, Gemma4VisionEncoderLayer: VisionLayer, Gemma4VisionAttention: VisionAttention, Gemma4VisionMLP: VisionMlp}

Mlp ¶

Mlp(*args: Any, **kwargs: Any)

Bases: Moe

Gemma-4's MLP: what reaches the residual stream is the post-feedforward norm's output (with the experts', on a mixture block).

On a mixture block it hosts the block's mixture of experts, whose router and experts it reads off its parent block; the dense MLP itself is the shared expert. Read order: shared_expert_output (the dense MLP runs first), then the routing values, then routed_output.

has_experts `property` ¶

has_experts: bool

Whether the block runs a mixture of experts beside this MLP (enable_moe_block).

router `property` ¶

router: Envoy

The block's router.

experts `property` ¶

experts: Envoy

The block's experts.

top_k `property` ¶

top_k: int

Experts each token is routed to: the router's config's top_k_experts.

no_mixture ¶

no_mixture() -> str | None

mlp_output ¶

mlp_output(value) -> Residual

router_logits ¶

router_logits(value) -> RouterLogits

The router's projection, [batch, seq, num_experts]: the router returns their softmax.

shared_expert_output ¶

shared_expert_output(value) -> Residual

head_dim ¶

head_dim(model: StandardizedTransformer) -> int

The config's top-level head_dim (the sliding blocks'), as stored; full blocks' are on layers[i].self_attn.

num_kv_heads ¶

num_kv_heads(model: StandardizedTransformer) -> int

The config's top-level num_key_value_heads (the sliding blocks'), as stored; full blocks' are on layers[i].self_attn.