kempnerforge.model.transformer

Transformer model for KempnerForge.

Architecture: Llama-style pre-norm transformer.

Token Embedding → [TransformerBlock × N] → Final Norm → Output Head

Design choices:

• ModuleDict (not ModuleList) for layers — preserves FQNs for DCP checkpointing.

• Embedding and output head are optional (can be None for PP middle stages).

• Forward is a simple loop over blocks — pipeline-parallelism friendly.

Classes

Transformer	Full transformer model built from ModelConfig.
TransformerBlock	Single transformer block with pre-norm architecture.

class kempnerforge.model.transformer.TransformerBlock

Bases: Module

Single transformer block with pre-norm architecture.

Structure: norm → attention → residual, norm → mlp → residual

__init__(config, layer_idx)

Parameters:

• config (ModelConfig)

• layer_idx (int)

Return type:

None

forward(x, rope_cos, rope_sin, *, kv_cache=None, attn_mask=None)

Run attention and the MLP over the residual stream.

Parameters:

• x (torch.Tensor) – Residual stream, shape (batch, seq_len, dim).

• rope_cos (torch.Tensor) – RoPE cosine frequencies for this block’s positions.

• rope_sin (torch.Tensor) – RoPE sine frequencies for this block’s positions.

• kv_cache (KVCache | None) – Optional KV cache for incremental generation.

• attn_mask (torch.Tensor | BlockMask | None) – The one mask from kempnerforge.model.masking.build_attention_mask, built once per forward and shared by every block.

Returns:

The updated residual stream, shape (batch, seq_len, dim).

Return type:

torch.Tensor

class kempnerforge.model.transformer.Transformer

Bases: Module

Full transformer model built from ModelConfig.

Embedding → TransformerBlocks → Norm → Output Head

include_token_embedding=False / include_output_head=False build the model without that module: the attribute is None and registers no state-dict keys. Without the embedding, pass pre-embedded inputs via modality.inputs_embeds; without the head, forward returns the final hidden states. tie_embeddings=True requires both modules. generate() rejects a model built without either module. Tensor parallelism is not supported for such a model.

__init__(config, *, vlm_config=None, num_image_tokens=0, include_token_embedding=True, include_output_head=True)

Parameters:

• config (ModelConfig)

• vlm_config (VLMConfig | None)

• num_image_tokens (int)

• include_token_embedding (bool)

• include_output_head (bool)

Return type:

None

init_weights_and_freqs()

Initialize weights and RoPE frequencies after meta-device materialization.

Called after model.to_empty(device=...) to fill in parameter values and compute RoPE frequency table. Safe to call on already-initialized models (skips if freqs are already computed).

Return type:

None

set_moe_step(step, max_steps)

Set training step on all MoE routers for adaptive bias scheduling.

Reaches every sigmoid router, including each MoT modality’s.

Parameters:

• step (int)

• max_steps (int)

Return type:

None

get_moe_aux_loss()

Sum the auxiliary losses of every MoE router. Returns 0 if dense.

A MoT layer contributes one term per modality, so an image + text MoT model sums twice as many terms as a text model with the same MoE layers.

Return type:

torch.Tensor

get_moe_router_z_loss()

Sum the router z-losses (ST-MoE) of every MoE router. Returns 0 if dense.

Like get_moe_aux_loss, a MoT layer contributes one term per modality.

Return type:

torch.Tensor

get_expert_counts()

Collect per-layer expert utilization for flat MoE layers.

Returns {layer_idx: (num_experts,) tensor} for layers whose MLP is a MoEMLP (the standard, single-pool MoE). Returns {} for dense models, for MoT and for MoMa, whose per-modality routers surface through get_mot_expert_counts and get_moma_expert_counts instead.

Return type:

dict[int, torch.Tensor]

get_mot_expert_counts()

Collect per-layer, per-modality expert utilization for MoT MoE layers.

Returns {layer_idx: {modality: (num_experts,) tensor}} for every MoT layer whose per-modality FFNs are MoEMLP``s; returns ``{} otherwise. Each tensor holds the token count per expert from that modality’s router in the most recent forward, so the image and text counts of a layer sum to different totals.

Return type:

dict[int, dict[str, torch.Tensor]]

get_moma_expert_counts()

Collect per-layer, per-modality expert utilization for MoMa layers.

Returns {layer_idx: {modality: (num_experts_for_modality,) tensor}} for every layer whose MLP is a MoMaFFN; returns {} otherwise. Each modality’s expert count tensor reflects the most recent forward through that layer’s expert-choice router (paper Figure 5-style utilization). Counts on a fresh model (no forward yet) are the router’s init zeros.

Return type:

dict[int, dict[str, torch.Tensor]]

forward(tokens=None, *, modality=None, kv_caches=None, doc_ids=None, position_ids=None, attention_group_ids=None)

Forward pass.

Exactly one of tokens or modality.inputs_embeds must be provided. Modality-injection routes (prefix_embeds, output_slice, image_features, image_mask, modality_ids, inline_embeds) are grouped on the optional ModalityContext arg; see kempnerforge/model/modality.py for the full intra-context invariant table. inline_embeds are written over the embedded tokens at (inline_rows, inline_positions) before the first layer, which is how a packed row carries each document’s image tokens before its own text, and output_rows / output_positions pick the positions the head runs on, giving (1, M, vocab) logits.

Parameters:

• tokens (torch.Tensor | None) – Integer token ids, shape (batch, seq_len).

• modality (ModalityContext | None) – Optional ModalityContext bundling pre-embedded inputs, prefix embeds, output slicing, image features, and modality routing tags for VLM arches. None is the pure text-only forward. It may also carry position_ids and attention_group_ids (see those arguments) or the rotation itself as rope_cos and rope_sin, (batch, seq_len, head_dim // 2) over the residual, or (1, seq_len, head_dim // 2) shared by every row, for position schemes that are not one integer per token; the rotation excludes position_ids, covers the appended tokens under kv_caches as position_ids does, is applied without the model’s own tables, and is unsupported for MoT.

• kv_caches (list[KVCache] | None) – Optional list of KVCache (one per layer) for generation. The tokens of this forward are appended after the cached ones; without caller positions or a rotation their RoPE positions continue from the cache fill level. Cross-arg invariant: kv_caches forbids modality.prefix_embeds, modality.output_slice, modality.image_features, modality.modality_ids, modality.key_padding_mask, doc_ids and attention_group_ids, by argument or as modality.attention_group_ids (all training-only).

• doc_ids (torch.Tensor | None) – Optional per-token document IDs for packed sequences, shape (batch, seq_len) over the residual stream. Enables block-diagonal causal attention that isolates documents within packed sequences.

• position_ids (torch.Tensor | None) – Optional explicit per-token RoPE positions, shape (batch, seq_len) matching the residual reaching attention (including any modality.prefix_embeds), dtype int64, values in [0, model.max_seq_len). Rows may differ, so positions can e.g. restart at zero at each document of a packed sequence. None (default) assigns contiguous positions starting at the KV-cache fill level, exactly as before. With kv_caches the tensor covers the tokens of this forward alone, the ones being appended, so a caller that positions its own tokens does so at prefill and at every decode step alike. Unsupported for MoT, which derives per-modality positions internally. A ModalityContext may carry the same tensor as modality.position_ids instead, which is how a strategy chooses the positions of the residual it lays out; passing both raises.

• attention_group_ids (torch.Tensor | None) – Optional per-token visibility group ids over the residual stream, shape (batch, seq_len), int64; negative means no group. Tokens sharing a nonnegative id see each other bidirectionally (e.g. prefix-LM spans): visible(q, k) = same_document AND (k <= q OR same_nonnegative_group). Without doc_ids the row is one document. None leaves attention unchanged; all-negative ids give the same visibility as None but still build a mask, so None is the unmasked fast path. On MoT, doc_ids and group ids need vlm.mot_modalities in residual order, ("image", "text"). May likewise come from modality.attention_group_ids, not both.

Returns:

Logits tensor of shape (batch, out_seq_len, vocab_size) where out_seq_len == seq_len normally or the sliced length when modality.output_slice is set, or (1, M, vocab_size) over the M positions modality.output_rows and output_positions pick. Built with include_output_head=False, returns the final hidden states (batch, out_seq_len, dim) instead.

Return type:

torch.Tensor