gemma2

gemma2

Gemma 2 on vLLM (vllm.model_executor.models.gemma2).

Llama's names and a fused block, with Gemma-2's sandwich norms: what the block adds is each sublayer's post norm's output, not the module's, as on transformers, so the contributions point at the sibling norms.

Two things are not transformers': there is no lm_head module (the logits processor unembeds with embed_tokens' weight, which the checkpoint ties), so project_on_vocab is defined here; and the sqrt(hidden_size) scaling of the embeddings is applied by the model after embed_tokens, not inside it, so token_embeddings is the unscaled lookup and layers[0].layer_input the scaled stream. The final logit softcapping is applied by vLLM's logits processor, which both logits and project_on_vocab go through.

RENAME module-attribute

RENAME = {'model.embed_tokens': 'embed_tokens', 'model.layers': 'layers', 'model.norm': 'norm'}

ENVOYS module-attribute

ENVOYS = {Gemma2DecoderLayer: Layer, Gemma2Attention: Attention, Gemma2MLP: Mlp}

Layer

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

Bases: FusedLayer

Gemma-2's decoder block: forward(positions, hidden_states, residual) -> (hidden_states, residual).

layer_mid class-attribute instance-attribute

layer_mid = summed('pre_feedforward_layernorm.inputs', description='The residual stream after the attention add: hidden_states + residual, [1, tokens, hidden_size]')

mlp_norm_output

mlp_norm_output(value) -> Residual

project_on_vocab

project_on_vocab(model: StandardizedVLLM, hidden: Tensor) -> Tensor

The final norm, then vLLM's logits processor over embed_tokens, which is the unembedding here.