From f59615c41c464dd48d27f91d70e027d5f8728a12 Mon Sep 17 00:00:00 2001 From: XiangAn Date: Sun, 17 May 2026 23:34:28 +0800 Subject: [PATCH 1/6] feat: register Gemma4 VL model family Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- .../core/transformer/transformer_config.py | 47 ++++++++++++++++++- aiak_training_llm/models/__init__.py | 1 + aiak_training_llm/train/__init__.py | 4 +- aiak_training_llm/train/arguments.py | 27 +++++++++++ aiak_training_llm/utils/constants.py | 1 + 5 files changed, 77 insertions(+), 3 deletions(-) diff --git a/aiak_megatron/megatron/core/transformer/transformer_config.py b/aiak_megatron/megatron/core/transformer/transformer_config.py index e41c3eeb..87f52bc7 100644 --- a/aiak_megatron/megatron/core/transformer/transformer_config.py +++ b/aiak_megatron/megatron/core/transformer/transformer_config.py @@ -3,7 +3,7 @@ """tranformer config""" import warnings -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Callable, List, Optional, Tuple, Union import torch @@ -530,6 +530,51 @@ class TransformerConfig(ModelParallelConfig): Using fp32 or fp64 can improve stability especially when the number of experts is large (e.g. finegrained-moe). None means no changes for dtype.""" + final_logit_softcapping: Optional[float] = None + """If set, apply ``cap * tanh(logits / cap)`` to the language-model output logits. + Used by Gemma4 (cap=30.0). None disables the softcap.""" + + use_layer_scalar: bool = False + """If True, each transformer layer learns a scalar gate (init 1.0) applied to the + layer output before the residual add. Used by Gemma4.""" + + partial_rotary_factor: float = 1.0 + """Fraction of the head dim that receives RoPE rotation. Gemma4 uses 0.25. + A value of 1.0 (default) keeps full RoPE and is a no-op for existing models.""" + + sliding_window: Optional[int] = None + """Window size used by sliding-window attention layers. None means no sliding mask.""" + + rotary_base_sliding: Optional[int] = None + """Optional separate RoPE base for sliding-window attention layers. + If None, sliding layers share ``rotary_base`` with global layers.""" + + layer_pattern: list = field(default_factory=list) + """Per-layer attention type as a list of strings ('sliding' / 'global'), + length == num_layers. Empty list (default) means all layers use the same + config-level head_dim / num_query_groups (no hybrid behavior).""" + + per_layer_kv_channels: dict = field(default_factory=dict) + """Optional per-layer-type override of ``kv_channels`` (a.k.a. head_dim). + Keyed by layer-type string ('sliding' / 'global'). Empty dict (default) + means use ``kv_channels`` for every layer.""" + + per_layer_num_query_groups: dict = field(default_factory=dict) + """Optional per-layer-type override of ``num_query_groups``. + Empty dict (default) means use ``num_query_groups`` for every layer.""" + + attention_k_eq_v: bool = False + """If True, Gemma4 global layers use K=V tying according to ``kv_tied_layers``. + Existing models keep the default False behavior.""" + + kv_tied_layers: list = field(default_factory=list) + """Layer indices for which K and V projections share the same weight tensor + (used by Gemma4 global layers). Empty list (default) means no tying.""" + + scale_emb_by_sqrt_hidden: bool = False + """If True, multiply token embeddings by ``sqrt(hidden_size)`` before the + transformer (Gemma family convention).""" + def __post_init__(self): """Python dataclass method that is used to modify attributes after initialization. See https://docs.python.org/3/library/dataclasses.html#post-init-processing for more diff --git a/aiak_training_llm/models/__init__.py b/aiak_training_llm/models/__init__.py index dbf99330..86d393e8 100644 --- a/aiak_training_llm/models/__init__.py +++ b/aiak_training_llm/models/__init__.py @@ -5,6 +5,7 @@ from .qwen_vl import qwen2_vl_config, qwen2_vl_provider from .llava_onevision1_5 import llava_onevision1_5_config, llava_onevision1_5_provider from .llava_onevision2 import llava_onevision2_config, llava_onevision2_provider +from .gemma4_vl import gemma4_vl_config, gemma4_vl_provider from .factory import ( get_support_model_archs, diff --git a/aiak_training_llm/train/__init__.py b/aiak_training_llm/train/__init__.py index f7fe5b2a..7d940744 100644 --- a/aiak_training_llm/train/__init__.py +++ b/aiak_training_llm/train/__init__.py @@ -1,8 +1,8 @@ """aiak train module""" from .arguments import parse_train_args -from .pretrain import pretrain_llava_onevision2, pretrain_llm, pretrain_qwen2_vl -from .sft import sft_llava_onevision2, sft_llava_onevision1_5, sft_llm, sft_qwen2_vl +from .pretrain import pretrain_gemma4_vl, pretrain_llava_onevision2, pretrain_llm, pretrain_qwen2_vl +from .sft import sft_gemma4_vl, sft_llava_onevision1_5, sft_llava_onevision2, sft_llm, sft_qwen2_vl from .trainer_builder import build_model_trainer diff --git a/aiak_training_llm/train/arguments.py b/aiak_training_llm/train/arguments.py index 6d0aff9a..4930b8f7 100644 --- a/aiak_training_llm/train/arguments.py +++ b/aiak_training_llm/train/arguments.py @@ -126,6 +126,33 @@ def _add_extra_model_args(parser: argparse.ArgumentParser): "this option, the head dimensions will be aligned by padding, so that fa can be used." "Deprecated: use --attention-backend=flash") + # Gemma4 architecture fields. These are normally populated from the + # registered model config in _validate_extra_model_args; defining them here + # keeps fixed-architecture validation and TransformerConfig construction in + # sync with the Gemma4VLConfig dataclass. + group.add_argument('--final-logit-softcapping', type=float, default=None, + help='Gemma4 final logits softcap. Usually set by --model-name.') + group.add_argument('--use-layer-scalar', action='store_true', default=False, + help='Enable Gemma4 per-layer scalar buffers. Usually set by --model-name.') + group.add_argument('--sliding-window', type=int, default=None, + help='Gemma4 sliding attention window. Usually set by --model-name.') + group.add_argument('--rotary-base-sliding', type=int, default=None, + help='Gemma4 sliding-layer RoPE base. Usually set by --model-name.') + group.add_argument('--partial-rotary-factor', type=float, default=1.0, + help='Gemma4 proportional RoPE factor. Usually set by --model-name.') + group.add_argument('--layer-pattern', default=[], + help='Gemma4 per-layer attention pattern. Usually set by --model-name.') + group.add_argument('--per-layer-kv-channels', default={}, + help='Gemma4 per-layer-type head dim overrides. Usually set by --model-name.') + group.add_argument('--per-layer-num-query-groups', default={}, + help='Gemma4 per-layer-type KV head overrides. Usually set by --model-name.') + group.add_argument('--attention-k-eq-v', action='store_true', default=False, + help='Enable Gemma4 K=V tying. Usually set by --model-name.') + group.add_argument('--kv-tied-layers', default=[], + help='Gemma4 K=V tied layer indices. Usually set by --model-name.') + group.add_argument('--scale-emb-by-sqrt-hidden', action='store_true', default=False, + help='Enable Gemma-style embedding scaling. Usually set by --model-name.') + return parser diff --git a/aiak_training_llm/utils/constants.py b/aiak_training_llm/utils/constants.py index eefdfe64..a822e581 100644 --- a/aiak_training_llm/utils/constants.py +++ b/aiak_training_llm/utils/constants.py @@ -94,3 +94,4 @@ class VisionLanguageModelFamilies(_BaseFamilies): QWEN2_5_VL = "qwen2_5_vl" LLAVA_ONEVISION1_5 = "llava_onevision1_5" LLAVA_ONEVISION2 = "llava_onevision2" + GEMMA4_VL = "gemma4_vl" From 24a3d1d8bfc9d96d50a86bf881bce8f2218e3b14 Mon Sep 17 00:00:00 2001 From: XiangAn Date: Sun, 17 May 2026 23:34:57 +0800 Subject: [PATCH 2/6] feat: add Gemma4 VL language model core Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- .../models/gemma4_vl/__init__.py | 7 + .../models/gemma4_vl/gemma4_attention.py | 258 +++++++++++++++ .../models/gemma4_vl/gemma4_mlp.py | 169 ++++++++++ .../models/gemma4_vl/gemma4_model.py | 296 ++++++++++++++++++ .../models/gemma4_vl/gemma4_moe_layer.py | 121 +++++++ .../gemma4_vl/gemma4_proportional_rotary.py | 109 +++++++ .../models/gemma4_vl/gemma4_router.py | 189 +++++++++++ .../gemma4_vl/gemma4_transformer_layer.py | 127 ++++++++ .../models/gemma4_vl/gemma4_vl_config.py | 242 ++++++++++++++ .../models/gemma4_vl/gemma4_vl_layer_spec.py | 170 ++++++++++ .../models/gemma4_vl/gemma4_vl_provider.py | 141 +++++++++ 11 files changed, 1829 insertions(+) create mode 100644 aiak_training_llm/models/gemma4_vl/__init__.py create mode 100644 aiak_training_llm/models/gemma4_vl/gemma4_attention.py create mode 100644 aiak_training_llm/models/gemma4_vl/gemma4_mlp.py create mode 100644 aiak_training_llm/models/gemma4_vl/gemma4_model.py create mode 100644 aiak_training_llm/models/gemma4_vl/gemma4_moe_layer.py create mode 100644 aiak_training_llm/models/gemma4_vl/gemma4_proportional_rotary.py create mode 100644 aiak_training_llm/models/gemma4_vl/gemma4_router.py create mode 100644 aiak_training_llm/models/gemma4_vl/gemma4_transformer_layer.py create mode 100644 aiak_training_llm/models/gemma4_vl/gemma4_vl_config.py create mode 100644 aiak_training_llm/models/gemma4_vl/gemma4_vl_layer_spec.py create mode 100644 aiak_training_llm/models/gemma4_vl/gemma4_vl_provider.py diff --git a/aiak_training_llm/models/gemma4_vl/__init__.py b/aiak_training_llm/models/gemma4_vl/__init__.py new file mode 100644 index 00000000..f3faeeff --- /dev/null +++ b/aiak_training_llm/models/gemma4_vl/__init__.py @@ -0,0 +1,7 @@ +"""Gemma4-VL model implementation.""" + +from . import gemma4_vl_config, gemma4_vl_provider +from .gemma4_vl_model import Gemma4VL + + +__all__ = ["Gemma4VL", "gemma4_vl_config", "gemma4_vl_provider"] diff --git a/aiak_training_llm/models/gemma4_vl/gemma4_attention.py b/aiak_training_llm/models/gemma4_vl/gemma4_attention.py new file mode 100644 index 00000000..54005081 --- /dev/null +++ b/aiak_training_llm/models/gemma4_vl/gemma4_attention.py @@ -0,0 +1,258 @@ +"""Gemma4 self-attention with hybrid head_dim, K=V tying, and parameter-free V LayerNorm. + +This module subclasses Megatron's SelfAttention to implement Gemma4-specific +behaviors WITHOUT modifying megatron/core/transformer/attention.py: + + 1. Hybrid per-layer head_dim and num_kv_heads (sliding vs global layers). + Implemented by cloning TransformerConfig and overriding kv_channels + + num_query_groups before super().__init__. + + 2. K=V tying on global layers (the K projection is also reused as V). + Implemented by rebuilding self.linear_qkv with reduced fan_out + (q + 1*kv instead of q + 2*kv) and overriding get_query_key_value_tensors + to perform single-KV split with value = key. + + 3. Parameter-free V LayerNorm (F.layer_norm with no weight/bias). + Applied at the end of get_query_key_value_tensors. + + 4. Dual-RoPE dispatch: when ``rotary_pos_emb`` is a ``dict`` mapping + layer-type strings ('sliding'|'global') to per-type RoPE tensors, + this layer selects the tensor matching its own ``layer_pattern[layer_number-1]`` + entry before calling the parent forward. This lets ``Gemma4Model`` ship two + independent ``RotaryEmbedding`` instances (sliding theta=1e4 vs global + theta=1e6 + proportional zero-pad) through the unmodified ``TransformerBlock`` + pass-through path. Pre-existing single-tensor callers are unaffected. + + 5. Sandwich post-attention RMSNorm: HF Gemma4DecoderLayer applies + ``post_attention_layernorm`` between ``self_attn(...)`` output and the + residual add (modeling_gemma4.py:1387). Megatron's stock pre-norm + ``TransformerLayer`` has no slot for this norm, so it is owned here and + applied after ``super().forward`` returns ``(output, bias)``. The state_dict + key becomes ``self_attention.post_attention_layernorm.weight`` and is + written by the HF→MG converter from HF ``post_attention_layernorm.weight``. +""" + +import copy +from typing import Optional + +import torch +import torch.nn.functional as F # noqa: N812 +from megatron.core.transformer.attention import SelfAttention, SelfAttentionSubmodules +from megatron.core.transformer.enums import AttnMaskType +from megatron.core.transformer.spec_utils import build_module +from megatron.core.transformer.transformer_config import TransformerConfig + +from aiak_training_llm.models.dispatch import multiacc_modules +from aiak_training_llm.utils import is_te_min_version + + +try: + from megatron.core.extensions.transformer_engine import SplitAlongDim +except ImportError: + SplitAlongDim = None + + +def _resolve_per_layer_overrides( + config: TransformerConfig, layer_number: int +) -> TransformerConfig: + """Per-layer override of kv_channels / num_query_groups / window_size. + + Reads ``config.layer_pattern[layer_number - 1]`` ('sliding' or 'global') and + selects per-layer kv_channels and num_query_groups from + ``config.per_layer_kv_channels`` / ``config.per_layer_num_query_groups``. + + Critical contract with TE backend: also writes ``cloned.window_size``, which + ``TEDotProductAttention`` reads (megatron/core/extensions/transformer_engine.py + line 738) and forwards to TE FlashAttention as ``window_size=(left, right)``. + HF's ``sliding_window=1024`` ("1024 tokens visible") maps to TE's ``(1023, 0)`` + because TE's window is inclusive: ``[i-1023, i+0]`` = exactly 1024 tokens. + Global layers get ``None`` (pure causal). If ``layer_pattern`` is empty, the + config is returned unchanged (no-op for non-Gemma4 models). + + Caller MUST always receive the cloned config (never early-return after pattern + is set) — otherwise sliding layers silently fall back to global causal. + """ + if not config.layer_pattern: + return config + idx = layer_number - 1 + if idx < 0 or idx >= len(config.layer_pattern): + return config + layer_type = config.layer_pattern[idx] + new_kv_channels = (config.per_layer_kv_channels or {}).get(layer_type) + new_num_query_groups = (config.per_layer_num_query_groups or {}).get(layer_type) + sliding_w = getattr(config, "sliding_window", None) + cloned = copy.copy(config) + if new_kv_channels is not None: + cloned.kv_channels = new_kv_channels + if new_num_query_groups is not None: + cloned.num_query_groups = new_num_query_groups + cloned.window_size = (sliding_w - 1, 0) if (layer_type == "sliding" and sliding_w is not None) else None + return cloned + + +class Gemma4SelfAttention(SelfAttention): + """Self-attention layer for Gemma4 with hybrid head_dim, K=V tying, and V LN. + + Args mirror SelfAttention exactly. The class inspects `config.layer_pattern` + and `config.kv_tied_layers` to decide per-layer behavior. + """ + + def __init__( + self, + config: TransformerConfig, + submodules: SelfAttentionSubmodules, + layer_number: int, + attn_mask_type: AttnMaskType = AttnMaskType.padding, + cp_comm_type: Optional[str] = None, + ): + effective_config = _resolve_per_layer_overrides(config, layer_number) + super().__init__( + config=effective_config, + submodules=submodules, + layer_number=layer_number, + attn_mask_type=attn_mask_type, + cp_comm_type=cp_comm_type, + ) + + self.kv_tied = (layer_number - 1) in (config.kv_tied_layers or []) + + if self.kv_tied: + self.linear_qkv = build_module( + submodules.linear_qkv, + self.config.hidden_size, + self.query_projection_size + self.kv_projection_size, + config=self.config, + init_method=self.config.init_method, + gather_output=False, + bias=self.config.add_bias_linear or self.config.add_qkv_bias, + skip_bias_add=False, + is_expert=False, + tp_comm_buffer_name='qkv', + ) + + norm_cls = multiacc_modules.TENorm if is_te_min_version("1.9.0") else multiacc_modules.LocalNorm + self.post_attention_layernorm = norm_cls( + config=self.config, + hidden_size=self.config.hidden_size, + eps=self.config.layernorm_epsilon, + ) + + def forward( + self, + hidden_states, + attention_mask, + attn_mask_type=None, + key_value_states=None, + inference_params=None, + rotary_pos_emb=None, + rotary_pos_cos=None, + rotary_pos_sin=None, + attention_bias=None, + packed_seq_params=None, + sequence_len_offset=None, + ): + if isinstance(rotary_pos_emb, dict): + layer_type = self.config.layer_pattern[self.layer_number - 1] + try: + rotary_pos_emb = rotary_pos_emb[layer_type] + except KeyError as exc: + raise KeyError( + f"Gemma4SelfAttention(layer_number={self.layer_number}): " + f"layer_pattern entry '{layer_type}' missing from rotary_pos_emb " + f"dict keys {list(rotary_pos_emb.keys())}." + ) from exc + output, bias = super().forward( + hidden_states=hidden_states, + attention_mask=attention_mask, + attn_mask_type=attn_mask_type, + key_value_states=key_value_states, + inference_params=inference_params, + rotary_pos_emb=rotary_pos_emb, + rotary_pos_cos=rotary_pos_cos, + rotary_pos_sin=rotary_pos_sin, + attention_bias=attention_bias, + packed_seq_params=packed_seq_params, + sequence_len_offset=sequence_len_offset, + ) + if bias is not None: + output = output + bias + bias = None + output = self.post_attention_layernorm(output) + return output, bias + + def get_query_key_value_tensors(self, hidden_states, key_value_states=None): + mixed_qkv, _ = self.linear_qkv(hidden_states) + + if self.kv_tied: + new_tensor_shape = mixed_qkv.size()[:-1] + ( + self.num_query_groups_per_partition, + ( + (self.num_attention_heads_per_partition // self.num_query_groups_per_partition + 1) + * self.hidden_size_per_attention_head + ), + ) + mixed_qkv = mixed_qkv.view(*new_tensor_shape) + + split_arg_list = [ + ( + self.num_attention_heads_per_partition + // self.num_query_groups_per_partition + * self.hidden_size_per_attention_head + ), + self.hidden_size_per_attention_head, + ] + + if SplitAlongDim is not None: + (query, key) = SplitAlongDim(mixed_qkv, 3, split_arg_list) + else: + (query, key) = torch.split(mixed_qkv, split_arg_list, dim=3) + + # K=V: value shares the K tensor. No clone — downstream kernels (TE/flash-attn) + # treat K and V as read-only. Saves one tensor allocation per layer. + value = key + else: + new_tensor_shape = mixed_qkv.size()[:-1] + ( + self.num_query_groups_per_partition, + ( + (self.num_attention_heads_per_partition // self.num_query_groups_per_partition + 2) + * self.hidden_size_per_attention_head + ), + ) + mixed_qkv = mixed_qkv.view(*new_tensor_shape) + + split_arg_list = [ + ( + self.num_attention_heads_per_partition + // self.num_query_groups_per_partition + * self.hidden_size_per_attention_head + ), + self.hidden_size_per_attention_head, + self.hidden_size_per_attention_head, + ] + + if SplitAlongDim is not None: + (query, key, value) = SplitAlongDim(mixed_qkv, 3, split_arg_list) + else: + (query, key, value) = torch.split(mixed_qkv, split_arg_list, dim=3) + + query = query.reshape(query.size(0), query.size(1), -1, self.hidden_size_per_attention_head) + + if self.q_layernorm is not None: + query = self.q_layernorm(query) + + if self.k_layernorm is not None: + key = self.k_layernorm(key) + + # Parameter-free V LayerNorm: normalize over the head_dim with weight=None bias=None. + # Equivalent to LayerNorm with weight=1 and bias=0 that is NOT learned. + # Applied AFTER any K=V tying so V normalization does not mutate K. + if self.kv_tied: + # Re-derive value as a normalized copy of key so K is untouched. + value = F.layer_norm(key, [self.hidden_size_per_attention_head]) + else: + value = F.layer_norm(value, [self.hidden_size_per_attention_head]) + + if self.config.test_mode: + self.run_realtime_tests() + + return query, key, value diff --git a/aiak_training_llm/models/gemma4_vl/gemma4_mlp.py b/aiak_training_llm/models/gemma4_vl/gemma4_mlp.py new file mode 100644 index 00000000..a3176ae1 --- /dev/null +++ b/aiak_training_llm/models/gemma4_vl/gemma4_mlp.py @@ -0,0 +1,169 @@ +"""Gemma4 parallel dense + MoE FFN block. + +Each Gemma4 transformer layer runs TWO FFN branches in parallel from the same +*raw post-attention residual* and sums their outputs. The HF reference +(``transformers/models/gemma4/modeling_gemma4.py``, +``Gemma4TextDecoderLayer.forward`` lines 1374-1409) reads:: + + residual = h_after_attn + hidden_states = pre_feedforward_layernorm(residual) + hidden_states = self.mlp(hidden_states) # dense + if enable_moe_block: + hidden_states_1 = post_feedforward_layernorm_1(hidden_states) + hidden_states_flat = residual.reshape(-1, H) + _, top_w, top_i = self.router(hidden_states_flat) # router on RAW residual + hidden_states_2 = pre_feedforward_layernorm_2(hidden_states_flat) + hidden_states_2 = self.experts(hidden_states_2, top_i, top_w) + hidden_states_2 = post_feedforward_layernorm_2(hidden_states_2) + hidden_states = hidden_states_1 + hidden_states_2 + hidden_states = post_feedforward_layernorm(hidden_states) # done in Gemma4TransformerLayer + hidden_states = residual + hidden_states # done by Megatron mlp_bda + +Megatron's standard ``TransformerLayer`` always applies +``self.pre_mlp_layernorm`` and feeds *that* to ``self.mlp``. To preserve the +HF requirement that **dense, experts and router each see** the raw residual +(through their own independent norms — or, for the router, no norm at all +beyond its own internal RMSNorm), the Gemma4 layer spec sets +``pre_mlp_layernorm = IdentityOp`` and this block holds three pre-merge norms: + + - ``pre_feedforward_layernorm`` (in front of dense, with weight) + - ``post_feedforward_layernorm_1`` (after dense MLP, with weight) + - ``post_feedforward_layernorm_2`` (after experts, with weight) + +The fourth pre-merge norm (``pre_feedforward_layernorm_2``, the experts' +input norm) lives on :class:`~aiak_training_llm.models.gemma4_vl.gemma4_moe_layer.Gemma4MoELayer` +because :class:`MoELayer.forward` feeds the same tensor to both ``self.router`` +(which must see RAW residual) and ``self.token_dispatcher`` (which must see +the normed residual). Splitting the two requires an override that lives next +to that forward, not in this composer. + +The fifth norm ``post_feedforward_layernorm`` (post-merge, pre-residual-add) +lives on :class:`~aiak_training_llm.models.gemma4_vl.gemma4_transformer_layer.Gemma4TransformerLayer` +because Megatron's ``mlp_bda`` is what performs the residual add and there is +no slot between ``self.mlp(...)`` and ``mlp_bda(...)`` that we can populate +without subclassing the layer. + +Forward contract matches Megatron's ``MLP`` / ``MoELayer``:: + + forward(hidden_states) -> (output, bias_or_None) + +Bias handling: Gemma4 sets ``add_bias_linear=False`` everywhere, so both +branches return ``bias=None``. We assert this and return ``None`` for the +composed bias to refuse silently-wrong fallbacks. +""" + +from dataclasses import dataclass +from typing import Optional, Union + +import torch +from megatron.core.transformer.module import MegatronModule +from megatron.core.transformer.spec_utils import ModuleSpec, build_module +from megatron.core.transformer.transformer_config import TransformerConfig + + +@dataclass +class Gemma4ParallelDenseMoESubmodules: + """Submodule spec for :class:`Gemma4ParallelDenseMoE`.""" + + dense: Union[ModuleSpec, type] = None + moe: Union[ModuleSpec, type] = None + pre_feedforward_layernorm: Union[ModuleSpec, type] = None + post_feedforward_layernorm_1: Union[ModuleSpec, type] = None + post_feedforward_layernorm_2: Union[ModuleSpec, type] = None + + +class Gemma4ParallelDenseMoE(MegatronModule): + """Parallel dense + MoE FFN composition for Gemma4. + + Receives the raw post-attention residual (Megatron's ``pre_mlp_layernorm`` + is forced to ``IdentityOp`` in the Gemma4 layer spec). The dense branch + norms locally; the MoE branch passes the raw residual straight through to + :class:`Gemma4MoELayer`, which owns the router-vs.-experts norm split. + """ + + def __init__( + self, + config: TransformerConfig, + submodules: Gemma4ParallelDenseMoESubmodules, + layer_number: Optional[int] = None, + ): + super().__init__(config=config) + self.config = config + self.layer_number = layer_number + + self.pre_feedforward_layernorm = build_module( + submodules.pre_feedforward_layernorm, + config=config, + hidden_size=config.hidden_size, + eps=config.layernorm_epsilon, + ) + self.post_feedforward_layernorm_1 = build_module( + submodules.post_feedforward_layernorm_1, + config=config, + hidden_size=config.hidden_size, + eps=config.layernorm_epsilon, + ) + self.post_feedforward_layernorm_2 = build_module( + submodules.post_feedforward_layernorm_2, + config=config, + hidden_size=config.hidden_size, + eps=config.layernorm_epsilon, + ) + + self.dense = build_module(submodules.dense, config=config) + self.moe = build_module(submodules.moe, config=config) + if layer_number is not None and hasattr(self.moe, "set_layer_number"): + self.moe.set_layer_number(layer_number) + + def set_layer_number(self, layer_number: int): + self.layer_number = layer_number + if hasattr(self.moe, "set_layer_number"): + self.moe.set_layer_number(layer_number) + + def forward(self, hidden_states: torch.Tensor): + residual = hidden_states + + dense_in = self.pre_feedforward_layernorm(residual) + dense_out, dense_bias = self.dense(dense_in) + dense_out = self.post_feedforward_layernorm_1(dense_out) + + moe_out, moe_bias = self.moe(residual) + moe_out = self.post_feedforward_layernorm_2(moe_out) + + assert dense_bias is None, ( + "Gemma4ParallelDenseMoE: dense branch returned non-None bias; " + "Gemma4 expects add_bias_linear=False everywhere." + ) + assert moe_bias is None, ( + "Gemma4ParallelDenseMoE: MoE branch returned non-None mlp_bias; " + "Gemma4 expects add_bias_linear=False everywhere." + ) + return dense_out + moe_out, None + + @property + def use_shared_expert(self) -> bool: + return getattr(self.moe, "use_shared_expert", False) + + @property + def shared_expert_overlap(self) -> bool: + return getattr(self.moe, "shared_expert_overlap", False) + + @property + def shared_experts(self): + return getattr(self.moe, "shared_experts", None) + + @property + def router(self): + return self.moe.router + + @property + def experts(self): + return self.moe.experts + + @property + def local_experts(self): + return getattr(self.moe.experts, "local_experts", self.moe.experts) + + @property + def token_dispatcher(self): + return self.moe.token_dispatcher diff --git a/aiak_training_llm/models/gemma4_vl/gemma4_model.py b/aiak_training_llm/models/gemma4_vl/gemma4_model.py new file mode 100644 index 00000000..6caa1c08 --- /dev/null +++ b/aiak_training_llm/models/gemma4_vl/gemma4_model.py @@ -0,0 +1,296 @@ +"""Gemma4 LLM backbone wrapper (Megatron LanguageModule subclass). + +Mirrors :class:`aiak_training_llm.models.qwen.qwen_model.QwenModel` and adds three +Gemma4-specific behaviors to ``forward``: + + 1. Embedding scaling: decoder_input *= sqrt(hidden_size). HF reference: + ``Gemma4TextScaledWordEmbedding.forward`` multiplies token embeddings by + ``embed_scale = config.hidden_size ** 0.5``. + + 2. Final logit soft-capping: ``logits = softcap * tanh(logits / softcap)``, + gated by ``config.final_logit_softcapping`` (=30.0 for Gemma4-26B-A4B). + Disabled when the value is None or <= 0. + + 3. Dual RoPE: when ``config.layer_pattern`` is non-empty (Gemma4 hybrid + attention), construct two independent rotary embeddings — sliding layers + get the stock :class:`RotaryEmbedding` (``head_dim=256``, ``theta=1e4``, + full rotation), global layers get :class:`Gemma4ProportionalRotaryEmbedding` + (``head_dim=512``, ``theta=1e6``, ``partial_rotary_factor=0.25`` with HF + proportional zero-padding). Both are evaluated each step and shipped as + a ``dict[str, Tensor]`` keyed by layer-type. The unmodified + :class:`TransformerBlock` passes the dict through verbatim; + :class:`Gemma4SelfAttention.forward` selects the matching tensor based on + its own ``self.layer_number`` and ``config.layer_pattern``. +""" + +import math +from collections import OrderedDict +from typing import Literal, Optional + +import torch +from megatron.core import InferenceParams, tensor_parallel +from megatron.core.config_logger import has_config_logger_enabled, log_config_to_disk +from megatron.core.dist_checkpointing.mapping import ShardedStateDict +from megatron.core.models.common.embeddings.language_model_embedding import LanguageModelEmbedding +from megatron.core.models.common.embeddings.rotary_pos_embedding import RotaryEmbedding +from megatron.core.models.common.language_module.language_module import LanguageModule +from megatron.core.packed_seq_params import PackedSeqParams +from megatron.core.transformer.enums import AttnMaskType, ModelType +from megatron.core.transformer.spec_utils import ModuleSpec +from megatron.core.transformer.transformer_block import TransformerBlock +from megatron.core.transformer.transformer_config import TransformerConfig +from torch import Tensor + +from aiak_training_llm.models.gemma4_vl.gemma4_proportional_rotary import ( + Gemma4ProportionalRotaryEmbedding, +) +from aiak_training_llm.models.qwen.qwen_model import _load_state_dict_hook_ignore_extra_state + + +class Gemma4Model(LanguageModule): + """Gemma4 transformer language model backbone.""" + + def __init__( + self, + config: TransformerConfig, + transformer_layer_spec: ModuleSpec, + vocab_size: int, + max_sequence_length: int, + pre_process: bool = True, + post_process: bool = True, + fp16_lm_cross_entropy: bool = False, + parallel_output: bool = True, + share_embeddings_and_output_weights: bool = True, + position_embedding_type: Literal["learned_absolute", "rope"] = "rope", + rotary_percent: float = 1.0, + rotary_base: int = 10000, + rotary_base_sliding: Optional[int] = None, + rope_scaling: bool = False, + rope_scaling_factor: float = 8.0, + scatter_embedding_sequence_parallel: bool = True, + seq_len_interpolation_factor: Optional[float] = None, + ) -> None: + super().__init__(config=config) + + if has_config_logger_enabled(config): + log_config_to_disk(config, locals(), prefix=type(self).__name__) + + self.transformer_layer_spec: ModuleSpec = transformer_layer_spec + self.vocab_size = vocab_size + self.max_sequence_length = max_sequence_length + self.pre_process = pre_process + self.post_process = post_process + self.fp16_lm_cross_entropy = fp16_lm_cross_entropy + self.parallel_output = parallel_output + self.share_embeddings_and_output_weights = share_embeddings_and_output_weights + self.position_embedding_type = position_embedding_type + + self.model_type = ModelType.encoder_or_decoder + + self.max_position_embeddings = max_sequence_length + self.rotary_percent = rotary_percent + self.rotary_base = rotary_base + self.rotary_scaling = rope_scaling + + # Gemma4 embedding scale factor: applied after embedding lookup. + # See HF Gemma4TextScaledWordEmbedding.embed_scale. + self._embed_scale = math.sqrt(self.config.hidden_size) + + # Final logit soft-capping; off when <= 0 / None. + self._final_logit_softcapping = getattr( + self.config, "final_logit_softcapping", None + ) + + if self.pre_process: + self.embedding = LanguageModelEmbedding( + config=self.config, + vocab_size=self.vocab_size, + max_sequence_length=self.max_sequence_length, + position_embedding_type=position_embedding_type, + scatter_to_sequence_parallel=scatter_embedding_sequence_parallel, + ) + + if self.position_embedding_type == "rope" and not self.config.multi_latent_attention: + sliding_kv_channels = ( + self.config.per_layer_kv_channels.get("sliding") + if self.config.per_layer_kv_channels + else None + ) or self.config.kv_channels + global_kv_channels = ( + self.config.per_layer_kv_channels.get("global") + if self.config.per_layer_kv_channels + else None + ) or self.config.kv_channels + + sliding_base = rotary_base_sliding if rotary_base_sliding is not None else rotary_base + global_base = rotary_base + + partial_global = getattr(self.config, "partial_rotary_factor", 1.0) or 1.0 + + self.rotary_pos_emb_sliding = RotaryEmbedding( + kv_channels=sliding_kv_channels, + rotary_percent=rotary_percent, + rotary_interleaved=self.config.rotary_interleaved, + seq_len_interpolation_factor=seq_len_interpolation_factor, + rotary_base=sliding_base, + rope_scaling=rope_scaling, + rope_scaling_factor=rope_scaling_factor, + use_cpu_initialization=self.config.use_cpu_initialization, + ) + + if self.config.layer_pattern and "global" in self.config.layer_pattern: + self.rotary_pos_emb_global = Gemma4ProportionalRotaryEmbedding( + head_dim=global_kv_channels, + partial_rotary_factor=partial_global, + rotary_base=global_base, + rotary_interleaved=self.config.rotary_interleaved, + seq_len_interpolation_factor=seq_len_interpolation_factor, + use_cpu_initialization=self.config.use_cpu_initialization, + ) + else: + self.rotary_pos_emb_global = None + + self.rotary_pos_emb_cache = {} + + self.decoder = TransformerBlock( + config=self.config, + spec=transformer_layer_spec, + pre_process=self.pre_process, + post_process=self.post_process, + ) + + if post_process: + if self.config.defer_embedding_wgrad_compute: + self.embedding_activation_buffer = [] + self.grad_output_buffer = [] + else: + self.embedding_activation_buffer = None + self.grad_output_buffer = None + + self.output_layer = tensor_parallel.ColumnParallelLinear( + config.hidden_size, + self.vocab_size, + config=config, + init_method=config.init_method, + bias=False, + skip_bias_add=False, + gather_output=not self.parallel_output, + skip_weight_param_allocation=self.pre_process + and self.share_embeddings_and_output_weights, + embedding_activation_buffer=self.embedding_activation_buffer, + grad_output_buffer=self.grad_output_buffer, + ) + + if self.pre_process or self.post_process: + self.setup_embeddings_and_output_layer() + + if has_config_logger_enabled(self.config): + log_config_to_disk( + self.config, self.state_dict(), prefix=f"{type(self).__name__}_init_ckpt" + ) + + self.register_load_state_dict_post_hook(_load_state_dict_hook_ignore_extra_state) + + def set_input_tensor(self, input_tensor: Tensor) -> None: + """See :meth:`LanguageModule.set_input_tensor`.""" + if not isinstance(input_tensor, list): + input_tensor = [input_tensor] + assert len(input_tensor) == 1, "input_tensor should only be length 1" + self.decoder.set_input_tensor(input_tensor[0]) + + def forward( + self, + input_ids: Tensor, + position_ids: Tensor, + attention_mask: Tensor, + attn_mask_type: Optional[AttnMaskType] = None, + decoder_input: Tensor = None, + labels: Tensor = None, + rotary_pos_emb: Tensor = None, + inference_params: InferenceParams = None, + packed_seq_params: PackedSeqParams = None, + extra_block_kwargs: Optional[dict] = None, + runtime_gather_output: Optional[bool] = None, + ) -> Tensor: + if decoder_input is None: + if self.pre_process: + decoder_input = self.embedding(input_ids=input_ids, position_ids=position_ids) + # Gemma4 embedding scaling: in-dtype multiply preserves the + # downstream casting policy (do NOT promote to fp32 here). + decoder_input = decoder_input * self._embed_scale + else: + decoder_input = None + + if ( + rotary_pos_emb is None + and self.position_embedding_type == "rope" + and not self.config.multi_latent_attention + ): + rotary_seq_len = self.rotary_pos_emb_sliding.get_rotary_seq_len( + inference_params, self.decoder, decoder_input, self.config, packed_seq_params + ) + packed = ( + packed_seq_params is not None and packed_seq_params.qkv_format == "thd" + ) + sliding_emb = self.rotary_pos_emb_sliding(rotary_seq_len, packed_seq=packed) + if self.rotary_pos_emb_global is not None: + global_emb = self.rotary_pos_emb_global(rotary_seq_len, packed_seq=packed) + rotary_pos_emb = {"sliding": sliding_emb, "global": global_emb} + else: + rotary_pos_emb = sliding_emb + + hidden_states = self.decoder( + hidden_states=decoder_input, + attention_mask=attention_mask, + attn_mask_type=attn_mask_type, + inference_params=inference_params, + rotary_pos_emb=rotary_pos_emb, + packed_seq_params=packed_seq_params, + **(extra_block_kwargs or {}), + ) + + if not self.post_process: + return hidden_states + + output_weight = None + if self.share_embeddings_and_output_weights: + output_weight = self.shared_embedding_or_output_weight() + + logits, _ = self.output_layer( + hidden_states, weight=output_weight, runtime_gather_output=runtime_gather_output + ) + + # Gemma4 final-logit soft-capping: logits = c * tanh(logits / c). + # Bounds extreme logits before cross-entropy. Off when c is None / <= 0. + if self._final_logit_softcapping and self._final_logit_softcapping > 0: + cap = float(self._final_logit_softcapping) + logits = cap * torch.tanh(logits / cap) + + if has_config_logger_enabled(self.config): + payload = OrderedDict( + { + "input_ids": input_ids, + "position_ids": position_ids, + "attention_mask": attention_mask, + "decoder_input": decoder_input, + "logits": logits, + } + ) + log_config_to_disk(self.config, payload, prefix="input_and_logits") + + if labels is None: + return logits.transpose(0, 1).contiguous() + + loss = self.compute_language_model_loss(labels, logits) + return loss + + def sharded_state_dict( + self, prefix: str = "", sharded_offsets: tuple = (), metadata: Optional[dict] = None + ) -> ShardedStateDict: + sharded_state_dict = super().sharded_state_dict(prefix, sharded_offsets, metadata) + output_layer_extra_state_key = f"{prefix}output_layer._extra_state" + output_extra_state = sharded_state_dict.pop(output_layer_extra_state_key, None) + assert not ( + output_extra_state and output_extra_state.data + ), f"Expected output layer extra state to be empty, got: {output_extra_state}" + return sharded_state_dict diff --git a/aiak_training_llm/models/gemma4_vl/gemma4_moe_layer.py b/aiak_training_llm/models/gemma4_vl/gemma4_moe_layer.py new file mode 100644 index 00000000..2c1afa74 --- /dev/null +++ b/aiak_training_llm/models/gemma4_vl/gemma4_moe_layer.py @@ -0,0 +1,121 @@ +"""Gemma4-VL MoE layer (subclass of Megatron's :class:`MoELayer`). + +Two reasons we subclass :class:`MoELayer`: + +1. **Router rebind.** Megatron's :class:`MoELayer.__init__` hard-codes + ``self.router = TopKRouter(config)`` and ignores anything you might pass + through :class:`MoESubmodules`. Gemma4 needs its own router with a + parameter-free RMSNorm + ``scale[H]`` + ``per_expert_scale[E]`` + (see ``gemma4_router.py``), so we rebind ``self.router`` after the parent + has finished wiring the experts / dispatcher / shared experts. + +2. **Router input contract.** HF Gemma4 routes on the *raw post-attention + residual* but feeds the experts a separately-normed view of that same + residual (see ``modeling_gemma4.py`` lines 1394-1405):: + + hidden_states_flat = residual.reshape(-1, H) # raw + _, top_w, top_i = self.router(hidden_states_flat) # router on raw + hidden_states_2 = pre_feedforward_layernorm_2(hidden_states_flat) + hidden_states_2 = self.experts(hidden_states_2, top_i, top_w) + + Megatron's :class:`MoELayer.forward` instead feeds the same tensor to both + ``self.router`` and ``self.token_dispatcher.token_permutation`` — making it + impossible to reproduce HF's split-input behavior from a wrapper that lives + *outside* :class:`MoELayer`. We therefore own ``pre_feedforward_layernorm_2`` + here and override :meth:`forward` to thread raw vs. normed correctly. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Union + +import torch +from megatron.core import tensor_parallel +from megatron.core.transformer.moe.moe_layer import MoELayer, MoESubmodules +from megatron.core.transformer.spec_utils import ModuleSpec, build_module +from megatron.core.transformer.transformer_config import TransformerConfig +from megatron.training import utils + +from aiak_training_llm.models.gemma4_vl.gemma4_router import Gemma4Router + + +@dataclass +class Gemma4MoESubmodules(MoESubmodules): + """Extends :class:`MoESubmodules` with Gemma4's expert pre-norm slot.""" + + pre_feedforward_layernorm_2: Union[ModuleSpec, type] = None + + +class Gemma4MoELayer(MoELayer): + """:class:`MoELayer` with Gemma4 router + split router/expert input norm.""" + + def __init__( + self, + config: TransformerConfig, + submodules: Gemma4MoESubmodules = None, + layer_number: int | None = None, + ) -> None: + super().__init__(config=config, submodules=submodules, layer_number=layer_number) + self.router = Gemma4Router(config=self.config) + if layer_number is not None: + self.router.set_layer_number(layer_number) + + if submodules is None or submodules.pre_feedforward_layernorm_2 is None: + raise ValueError( + "Gemma4MoELayer requires submodules.pre_feedforward_layernorm_2 to be set; " + "the layer spec must wire this norm in." + ) + self.pre_feedforward_layernorm_2 = build_module( + submodules.pre_feedforward_layernorm_2, + config=config, + hidden_size=config.hidden_size, + eps=config.layernorm_epsilon, + ) + + def forward(self, hidden_states: torch.Tensor): + if ( + self.training + and self.config.tensor_model_parallel_size > 1 + and not self.config.sequence_parallel + ): + raise ValueError( + "During training, performance may degrade if MoE and tensor parallelism " + "are enabled without also enabling sequence parallelism." + ) + + def custom_forward(hidden_states): + probs, routing_map = self.router(hidden_states) + + expert_input = self.pre_feedforward_layernorm_2(hidden_states) + (dispatched_input, tokens_per_expert) = self.token_dispatcher.token_permutation( + expert_input, probs, routing_map + ) + + if self.config.enable_mem_monitor: + utils.monitor_max_memory_usage() + utils.monitor_max_dispatcher_tokens(tokens_per_expert) + if self.config.mem_monitor_log is not None: + utils.write_monitor_data_to_file( + self.config.mem_monitor_log, self.config.print_mem_monitor_interval + ) + + expert_output, mlp_bias = self.experts(dispatched_input, tokens_per_expert) + output, mlp_bias = self.token_dispatcher.token_unpermutation(expert_output, mlp_bias) + + if self.use_shared_expert and not self.shared_expert_overlap: + raise RuntimeError( + "Gemma4MoELayer was built with shared experts enabled; HF Gemma4 has none. " + "Check that moe_shared_expert_intermediate_size is unset (None)." + ) + return output, mlp_bias + + if self.moe_layer_recompute: + output, mlp_bias = tensor_parallel.checkpoint(custom_forward, False, hidden_states) + else: + output, mlp_bias = custom_forward(hidden_states) + + return output, mlp_bias + + +__all__ = ["Gemma4MoELayer", "Gemma4MoESubmodules"] diff --git a/aiak_training_llm/models/gemma4_vl/gemma4_proportional_rotary.py b/aiak_training_llm/models/gemma4_vl/gemma4_proportional_rotary.py new file mode 100644 index 00000000..1a1e2143 --- /dev/null +++ b/aiak_training_llm/models/gemma4_vl/gemma4_proportional_rotary.py @@ -0,0 +1,109 @@ +"""Gemma4 proportional RoPE — global-layer rotary with HF zero-padded inv_freq. + +HF Gemma4 global (full_attention) layers use ``rope_type='proportional'`` with +``partial_rotary_factor=0.25`` and ``rope_theta=1e6``. This is **not** the same +as Megatron's stock ``rotary_percent`` — the two differ in two ways that +together push logits off by O(1e-3) on a real micro-batch: + + 1. Denominator. HF proportional uses the **full head_dim** as the denominator + when computing inv_freq exponents: + + inv_freq_rotated[k] = 1 / base^( 2*k / head_dim ) for k in [0, rope_angles) + + where ``rope_angles = int(rope_proportion * head_dim // 2)``. Megatron's + ``RotaryEmbedding`` instead uses ``dim = int(kv_channels * rotary_percent)`` + as the denominator: ``inv_freq[k] = 1 / base^(2*k / dim)``. With + ``head_dim=512`` and ``partial_rotary_factor=0.25``, HF divides by 512 while + stock Megatron would divide by 128 — a 4× larger exponent — producing a + completely different frequency spectrum. + + 2. Padding vs. truncation. HF returns an inv_freq of length ``head_dim // 2`` + where the first ``rope_angles`` entries are non-zero and the remainder is + zero-padded; the cat-(freqs, freqs) in the forward then produces a + ``head_dim``-long cos/sin pair whose tail is identity (cos=1, sin=0). + Megatron's stock partial-RoPE truncates instead: it returns inv_freq of + length ``dim // 2`` and the cos/sin pair is ``dim``-long, with the + attention head splitting Q/K into rotated + un-rotated sub-vectors. The + attention math is **incompatible** between the two without code changes + elsewhere; we match HF by zero-padding inside ``__init__``. + +This subclass therefore reimplements ``__init__``'s inv_freq computation: +construct a length-``head_dim // 2`` buffer whose first ``rope_angles`` entries +are HF-equivalent and the rest are zero. The parent ``forward`` is left intact +because (a) ``cat((freqs, freqs), dim=-1)`` correctly emits a length-``head_dim`` +cos/sin pair, and (b) zeros in inv_freq produce cos=1, sin=0 — the desired +identity rotation on the un-rotated tail. + +For sliding (``rope_type='default'`` with no ``partial_rotary_factor``) we +**don't need this class** — the stock ``RotaryEmbedding(kv_channels=256, +rotary_percent=1.0, rotary_base=1e4)`` is bit-for-bit equivalent to HF. + +References: + - HF ``_compute_proportional_rope_parameters`` in + ``transformers/src/transformers/modeling_rope_utils.py:187-254`` (commit c9de109). + - HF ``Gemma4TextRotaryEmbedding`` in + ``transformers/src/transformers/models/gemma4/modeling_gemma4.py:1046-1130``. +""" + +from __future__ import annotations + +import torch +from megatron.core.models.common.embeddings.rotary_pos_embedding import RotaryEmbedding + + +class Gemma4ProportionalRotaryEmbedding(RotaryEmbedding): + """Gemma4 global-layer RoPE with HF-equivalent zero-padded inv_freq. + + Args: + head_dim: Full attention head dimension (e.g. 512 for Gemma4 global). + This is the **denominator** in the HF inv_freq exponent and the + length of the emitted cos/sin (= 2 * len(inv_freq)). + partial_rotary_factor: Fraction of the head_dim to actually rotate + (e.g. 0.25 for Gemma4 global → rotate first 64 angle pairs of 256). + rotary_base: RoPE base / theta (e.g. 1_000_000 for Gemma4 global). + rotary_interleaved: Forwarded to parent. Gemma4 uses False (cat layout). + seq_len_interpolation_factor: Forwarded to parent. + use_cpu_initialization: Forwarded to parent. + """ + + def __init__( + self, + head_dim: int, + partial_rotary_factor: float, + rotary_base: float, + rotary_interleaved: bool = False, + seq_len_interpolation_factor: float | None = None, + use_cpu_initialization: bool = False, + ) -> None: + # Parent __init__ computes self.inv_freq using kv_channels + rotary_percent. + # We pass kv_channels=head_dim, rotary_percent=1.0 so the parent emits an + # inv_freq of length head_dim//2 with the **correct denominator** (head_dim). + # Then we zero out the tail to match HF proportional padding. + super().__init__( + kv_channels=head_dim, + rotary_percent=1.0, + rotary_interleaved=rotary_interleaved, + seq_len_interpolation_factor=seq_len_interpolation_factor, + rotary_base=rotary_base, + rope_scaling=False, + rope_scaling_factor=1.0, + use_cpu_initialization=use_cpu_initialization, + ) + + rope_angles = int(partial_rotary_factor * head_dim // 2) + if rope_angles < 0 or rope_angles > head_dim // 2: + raise ValueError( + f"Gemma4ProportionalRotaryEmbedding: rope_angles={rope_angles} out of " + f"range [0, {head_dim // 2}] for head_dim={head_dim}, " + f"partial_rotary_factor={partial_rotary_factor}." + ) + + # Zero-pad the inv_freq tail so cos[..., rope_angles:]=1 and sin[..., rope_angles:]=0 + # in the parent forward. Use in-place fill on the registered buffer / parameter view. + # self.inv_freq is a Tensor (set by parent as a plain attribute, not a buffer). + if rope_angles < head_dim // 2: + with torch.no_grad(): + self.inv_freq[rope_angles:].zero_() + + +__all__ = ["Gemma4ProportionalRotaryEmbedding"] diff --git a/aiak_training_llm/models/gemma4_vl/gemma4_router.py b/aiak_training_llm/models/gemma4_vl/gemma4_router.py new file mode 100644 index 00000000..0572f766 --- /dev/null +++ b/aiak_training_llm/models/gemma4_vl/gemma4_router.py @@ -0,0 +1,189 @@ +"""Gemma4-VL router (real subclass of Megatron's :class:`Router`). + +Faithfully reproduces HF ``Gemma4TextRouter`` semantics inside Megatron's MoE +plumbing. The HF router is a self-contained module — no algebraic fusion with +any neighbouring norm is possible (see plan §0 point 2 / §1.4 / Change Log v4 +#1) — so we reimplement it verbatim here and override the ``Router`` interface +to return the ``(scores, routing_map)`` tuple Megatron's token dispatcher +expects. + +HF reference (``transformers/models/gemma4/modeling_gemma4.py``, +transformers==5.7.0):: + + class Gemma4TextRouter(nn.Module): + def __init__(self, config): + self.norm = Gemma4RMSNorm(hidden_size, eps, with_scale=False) + self.proj = nn.Linear(hidden_size, num_experts, bias=False) + self.scale = nn.Parameter(torch.ones(hidden_size)) + self.per_expert_scale = nn.Parameter(torch.ones(num_experts)) + + def forward(self, hidden_states): + h = self.norm(hidden_states) + h = h * self.scale * (hidden_size ** -0.5) + scores = self.proj(h) + probs = F.softmax(scores, dim=-1) + top_w, top_i = torch.topk(probs, k=top_k_experts, dim=-1) + top_w = top_w / top_w.sum(dim=-1, keepdim=True) + top_w = top_w * self.per_expert_scale[top_i] + return probs, top_w, top_i + +Notes on the Megatron mapping: + +* Megatron's :class:`Router` base class allocates ``self.weight`` of shape + ``[num_experts, hidden_size]``. We reuse that buffer as HF's + ``proj.weight`` (same shape, same role) so the conversion script can copy + it verbatim and we avoid duplicating a Parameter. +* ``self.scale`` and ``self.per_expert_scale`` are extra Parameters added on + top of the base class. +* ``self.norm`` is the HF ``Gemma4RMSNorm(with_scale=False)`` — a parameter- + free RMSNorm with the eps **inside** the sqrt (``1/sqrt(mean(x^2)+eps)``), + which differs from Megatron's default ``MegatronRMSNorm`` + (``1/sqrt(mean(x^2))+eps``). We implement it inline as a private nn.Module + so the numerics match HF byte-for-byte and there is no learnable weight to + convert. +* The :meth:`forward` override returns ``(scores, routing_map)`` where: + + - ``scores`` is a dense ``[num_tokens, num_experts]`` tensor whose + non-zero entries equal HF's ``top_w`` (so the dispatcher's + ``out = scores * dispatched_input`` reproduces HF's gating exactly); + - ``routing_map`` is a dense ``[num_tokens, num_experts]`` bool mask + with ``top_k_experts`` True per row. + + The base class's ``gating()`` / ``routing()`` plumbing is bypassed entirely + because (a) we need a custom pre-projection norm + scalar scale, and + (b) HF's per-expert-scale post-multiplication has no equivalent in the + ``moe_router_topk_scaling_factor`` scalar slot. + +Conversion (P2): direct per-Parameter copy ``HF -> MG``:: + + proj.weight -> router.weight [num_experts, hidden] + scale -> router.scale [hidden] + per_expert_scale -> router.per_expert_scale [num_experts] +""" + +from __future__ import annotations + +import torch +from megatron.core.transformer.moe.router import Router +from megatron.core.transformer.transformer_config import TransformerConfig + + +class _Gemma4ParameterFreeRMSNorm(torch.nn.Module): + """HF ``Gemma4RMSNorm(with_scale=False)`` — parameter-free, eps-in-sqrt. + + Computes ``x * pow(mean(x^2) + eps, -0.5)`` in float32 and casts back to + the input dtype. Matches HF byte-for-byte; intentionally different from + Megatron's ``MegatronRMSNorm`` (which puts eps outside the sqrt). + """ + + def __init__(self, eps: float): + super().__init__() + self.eps = eps + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + in_dtype = hidden_states.dtype + x = hidden_states.float() + mean_sq = x.pow(2).mean(-1, keepdim=True) + self.eps + out = x * torch.pow(mean_sq, -0.5) + return out.to(in_dtype) + + +class Gemma4Router(Router): + """Self-contained router matching HF ``Gemma4TextRouter`` semantics. + + Inherits Megatron's :class:`Router` so token dispatchers can consume the + standard ``(scores, routing_map)`` tuple, but completely overrides the + base ``forward`` (and does not use the base ``gating`` / ``routing`` + helpers). + """ + + def __init__(self, config: TransformerConfig) -> None: + super().__init__(config=config) + # Base class already created self.weight: [num_experts, hidden_size] + # — we reuse it as HF's proj.weight. Sanity-check the shape. + assert self.weight.shape == (config.num_moe_experts, config.hidden_size), ( + f"unexpected router weight shape {tuple(self.weight.shape)}; " + f"expected ({config.num_moe_experts}, {config.hidden_size})" + ) + + self.hidden_size = config.hidden_size + self.scalar_root_size = self.hidden_size ** -0.5 + self.topk = config.moe_router_topk + + # HF: Gemma4RMSNorm(hidden_size, eps=rms_norm_eps, with_scale=False) + # `layernorm_epsilon` is Megatron's RMSNorm eps field name. + self.norm = _Gemma4ParameterFreeRMSNorm(eps=config.layernorm_epsilon) + + # HF: scale = nn.Parameter(torch.ones(hidden_size)) + self.scale = torch.nn.Parameter( + torch.ones(self.hidden_size, dtype=config.params_dtype) + ) + setattr(self.scale, "sequence_parallel", config.sequence_parallel) + + # HF: per_expert_scale = nn.Parameter(torch.ones(num_experts)) + self.per_expert_scale = torch.nn.Parameter( + torch.ones(self.num_experts, dtype=config.params_dtype) + ) + setattr(self.per_expert_scale, "sequence_parallel", config.sequence_parallel) + + # The abstract ``routing`` is required by the base class but unused + # (forward is fully overridden). Implement as a no-op assert so accidental + # callers fail loudly instead of silently misbehaving. + def routing(self, logits: torch.Tensor, ori_dtype=None): # type: ignore[override] + raise NotImplementedError( + "Gemma4Router overrides forward() entirely; routing() is unused." + ) + + def forward(self, input: torch.Tensor): # noqa: A002 - mirror base name + """Compute HF-equivalent gating and emit Megatron-shaped outputs. + + Args: + input: ``[seq_len, batch, hidden_size]`` activations from the + post-attention residual (raw — *not* pre-FFN-normed; see plan + §1.4 "self-contained router" bullet). + + Returns: + (scores, routing_map): + - ``scores: [num_tokens, num_experts]`` dense, non-zero only + at top-k positions, equal to HF's ``top_w`` after the + ``per_expert_scale`` multiplication. + - ``routing_map: [num_tokens, num_experts]`` bool, True at + top-k positions. + """ + # Move weights to GPU on first call (mirror base class behaviour). + if self.weight.device.type == "cpu": + self.weight.data = self.weight.data.to(device=torch.cuda.current_device()) + + # HF: norm -> *scale -> *scalar_root_size -> linear(proj.weight) + h = self.norm(input) + h = h * self.scale * self.scalar_root_size + # self.weight: [E, H], h: [..., H] -> logits: [..., E] + logits = torch.nn.functional.linear(h, self.weight) + + # Flatten leading dims for the dispatcher contract. + logits = logits.view(-1, self.num_experts) + + # HF: softmax then topk (post-softmax topk). + probs = torch.nn.functional.softmax(logits, dim=-1, dtype=torch.float32) + top_w, top_i = torch.topk(probs, k=self.topk, dim=-1) + # Re-normalize so top-k weights sum to 1 per token. + top_w = top_w / top_w.sum(dim=-1, keepdim=True) + # Apply per-expert scale (cast index gather to float32 to match probs). + top_w = top_w * self.per_expert_scale.float()[top_i] + top_w = top_w.to(input.dtype) + + # Build dense (scores, routing_map) of shape [num_tokens, num_experts]. + num_tokens = logits.shape[0] + scores = torch.zeros( + num_tokens, self.num_experts, dtype=top_w.dtype, device=logits.device + ) + scores.scatter_(1, top_i, top_w) + routing_map = torch.zeros( + num_tokens, self.num_experts, dtype=torch.bool, device=logits.device + ) + routing_map.scatter_(1, top_i, True) + + return scores, routing_map + + +__all__ = ["Gemma4Router"] diff --git a/aiak_training_llm/models/gemma4_vl/gemma4_transformer_layer.py b/aiak_training_llm/models/gemma4_vl/gemma4_transformer_layer.py new file mode 100644 index 00000000..6953e886 --- /dev/null +++ b/aiak_training_llm/models/gemma4_vl/gemma4_transformer_layer.py @@ -0,0 +1,127 @@ +"""Gemma4-specific transformer layer subclass. + +Wraps Megatron's :class:`TransformerLayer` to attach two Gemma4-only mechanisms +that have no slot in the upstream layer: + +1. ``post_feedforward_layernorm`` (with weight) — applied to the merged + ``dense + experts`` output **before** the residual add. In HF + (``Gemma4TextDecoderLayer.forward``) this lives between the merged sum + and the residual add:: + + hidden_states = self.post_feedforward_layernorm(hidden_states) + hidden_states = residual + hidden_states + + Megatron's stock layer goes ``mlp -> mlp_bda(mlp_out, residual)`` with no + hookable position in between, so we override ``forward`` and inject the + norm by wrapping ``self.mlp(...)`` instead of patching ``mlp_bda``. + +2. ``layer_scalar`` — non-trainable buffer of shape ``(1,)`` initialised to + ``1.0`` and multiplied into the layer output at the very end. Loaded + from the HF checkpoint (per-layer distinct scalars after pretraining). + Gated by ``config.use_layer_scalar`` so the buffer is not registered when + the flag is off. + +Both mechanisms are no-ops when their gates are off, which is required so +this subclass can be used by any Gemma4-VL spec without hard-coding Gemma4 +into upstream Megatron. +""" + +from dataclasses import dataclass +from typing import Union + +import torch +from megatron.core.transformer.spec_utils import ModuleSpec, build_module +from megatron.core.transformer.transformer_config import TransformerConfig +from megatron.core.transformer.transformer_layer import ( + TransformerLayer, + TransformerLayerSubmodules, +) + + +@dataclass +class Gemma4TransformerLayerSubmodules(TransformerLayerSubmodules): + """Adds the ``post_feedforward_layernorm`` slot used by Gemma4. + + Default ``None`` means "no post-FFN norm"; ``Gemma4TransformerLayer`` only + wraps ``self.mlp`` when this slot is populated, so other consumers of + :class:`TransformerLayerSubmodules` are unaffected. + """ + + post_feedforward_layernorm: Union[ModuleSpec, type, None] = None + + +class Gemma4TransformerLayer(TransformerLayer): + """``TransformerLayer`` plus Gemma4's post-merge norm and ``layer_scalar``. + + The base ``TransformerLayer.forward`` calls ``self.mlp(...)`` then + ``mlp_bda(mlp_out, residual)`` (see ``transformer_layer.py:550-562``). + To inject ``post_feedforward_layernorm`` between them without forking the + full forward, we monkey-patch ``self.mlp`` at construction time so that + its ``__call__`` returns ``(post_ffn_norm(mlp_out), bias)``. The base + forward then sees the post-normed tensor and feeds it to ``mlp_bda``, + which performs the residual add — exactly the HF order. + + ``layer_scalar`` is post-multiplied into the layer's final output and + handles both the single-tensor and ``(output, context)`` return shapes + of the base forward. + """ + + def __init__( + self, + config: TransformerConfig, + submodules, + layer_number: int = 1, + hidden_dropout: float | None = None, + ): + super().__init__( + config=config, + submodules=submodules, + layer_number=layer_number, + hidden_dropout=hidden_dropout, + ) + + post_ffn_spec = getattr(submodules, "post_feedforward_layernorm", None) + if post_ffn_spec is not None: + self.post_feedforward_layernorm = build_module( + post_ffn_spec, + config=config, + hidden_size=config.hidden_size, + eps=config.layernorm_epsilon, + ) + self._wrap_mlp_with_post_ffn_norm() + + if getattr(self.config, "use_layer_scalar", False): + self.register_buffer("layer_scalar", torch.ones(1)) + + def _wrap_mlp_with_post_ffn_norm(self): + original_mlp = self.mlp + post_ffn = self.post_feedforward_layernorm + + class _MlpWithPostFfnNorm(torch.nn.Module): + def __init__(inner_self): + super().__init__() + inner_self._inner_mlp = original_mlp + inner_self._post_ffn = post_ffn + + def __getattr__(inner_self, name): + if name in ("_inner_mlp", "_post_ffn"): + return super().__getattr__(name) + try: + return super().__getattr__(name) + except AttributeError: + return getattr(inner_self._inner_mlp, name) + + def forward(inner_self, hidden_states, *args, **kwargs): + output, bias = inner_self._inner_mlp(hidden_states, *args, **kwargs) + output = inner_self._post_ffn(output) + return output, bias + + self.mlp = _MlpWithPostFfnNorm() + + def forward(self, *args, **kwargs): + out = super().forward(*args, **kwargs) + if not hasattr(self, "layer_scalar"): + return out + if isinstance(out, tuple): + return (out[0] * self.layer_scalar, *out[1:]) + return out * self.layer_scalar diff --git a/aiak_training_llm/models/gemma4_vl/gemma4_vl_config.py b/aiak_training_llm/models/gemma4_vl/gemma4_vl_config.py new file mode 100644 index 00000000..bb084cbc --- /dev/null +++ b/aiak_training_llm/models/gemma4_vl/gemma4_vl_config.py @@ -0,0 +1,242 @@ +"""Configuration dataclasses and registrations for Gemma4-VL family. + +Mirrors the structure of ``aiak_training_llm.models.llava_onevision2.llava_onevision2_config`` +and adds the Gemma4-specific knobs: + +- ``final_logit_softcapping`` (final logits tanh-cap, gated) +- ``use_layer_scalar`` (per-layer scalar buffer multiplied at end of layer forward) +- ``per_layer_kv_channels`` / ``layer_pattern`` (sliding vs global hybrid attention) +- ``attention_k_eq_v`` (Gemma4 K=V tying flag — combined with ``layer_pattern[i]=='global'`` + this drives the no-V-projection runtime aliasing on global layers) +- ``kv_tied_layers`` (precomputed list of global layer indices for fast runtime lookup) +- ``scale_emb_by_sqrt_hidden`` (Gemma family: x = embed(ids) * sqrt(hidden)) +- ``sliding_window`` (sliding-attention window size) + +Note: this checkpoint has dual RoPE — sliding layers use ``rope_type='default'`` +with ``rope_theta=10000`` and full rotation; global layers use HF +``rope_type='proportional'`` with ``rope_theta=1_000_000``, +``partial_rotary_factor=0.25``, and the proportional zero-pad scheme over +``global_head_dim=512`` (rotate the first 64 angle pairs of 256, leave the +remaining 192 as identity). The sliding/global split lives in +``layer_pattern``; ``Gemma4Model`` constructs two ``RotaryEmbedding`` instances +and ships them as a dict keyed by layer-type to ``Gemma4SelfAttention``. + +These fields land on ``TransformerConfig`` via the matching aiak_megatron patch +(P1.13A) so ``core_transformer_config_from_args`` can pump them straight from +the parsed args. +""" + +from dataclasses import dataclass, field + +import torch +from torch.nn.functional import gelu + +from aiak_training_llm.models.factory import register_model_config +from aiak_training_llm.utils.constants import VisionLanguageModelFamilies + + +@dataclass +class AdapterConfig: + """Configuration for the Gemma4-VL projector. + + The fields need to be consistent with the definitions in args. + """ + + normalization: str + activation_func: torch.nn.Module = gelu + add_bias_linear: bool = False + layernorm_epsilon: float = 1e-06 + + +@dataclass +class Gemma4VLConfig: + num_layers: int + hidden_size: int + ffn_hidden_size: int + num_attention_heads: int + group_query_attention: bool = True + num_query_groups: int = 8 + position_embedding_type: str = "rope" + add_position_embedding: bool = False + rotary_interleaved: bool = False + normalization: str = "RMSNorm" + swiglu: bool = True + attention_dropout: float = 0 + hidden_dropout: float = 0 + add_bias_linear: bool = False + add_qkv_bias: bool = False + qk_layernorm: bool = True + untie_embeddings_and_output_weights: bool = False + vocab_size_in_config_file: int = 262144 + make_vocab_size_divisible_by: int = 128 + norm_epsilon: float = 1e-06 + rotary_base: int = 1000000 + kv_channels: int = 256 + num_experts: int = None + moe_ffn_hidden_size: int = None + # ---- Gemma4-specific extensions (mirrored on TransformerConfig via P1.13A) ---- + final_logit_softcapping: float = None + use_layer_scalar: bool = False + sliding_window: int = None + rotary_base_sliding: int = None + partial_rotary_factor: float = 1.0 + # Per-layer hybrid attention overrides. ``layer_pattern`` is a list of + # ``"sliding"`` / ``"global"`` strings of length ``num_layers``; + # ``per_layer_kv_channels`` and ``per_layer_num_query_groups`` are dicts + # keyed by layer-type ("sliding"/"global") -> int. + layer_pattern: list = field(default_factory=list) + per_layer_kv_channels: dict = field(default_factory=dict) + per_layer_num_query_groups: dict = field(default_factory=dict) + # K=V tying configuration. ``attention_k_eq_v=True`` enables HF's "no V projection, + # alias V from K" mechanism on layers listed in ``kv_tied_layers`` (= global layers + # under HF's ``use_alternative_attention = attention_k_eq_v AND not is_sliding``). + attention_k_eq_v: bool = False + kv_tied_layers: list = field(default_factory=list) + scale_emb_by_sqrt_hidden: bool = False + + +@register_model_config( + model_family=VisionLanguageModelFamilies.GEMMA4_VL, + model_arch="gemma4-26b-a4b-vl", +) +def gemma4_26b_a4b_vl(): + """Gemma4-26B-A4B-it (instruction-tuned MoE VLM). + + Architecture (verified from /ov2/pretrain_models/google/gemma-4-26B-A4B-it/config.json): + - 30 layers with repeating ``[s,s,s,s,s,g]`` pattern (5 sliding + 1 global) + - hidden=2816, num_attention_heads=16 + - dense intermediate=2112, moe intermediate=704 (per routed expert) + - 128 routed experts top-k=8; dense MLP runs as a parallel branch from the + same post-attention residual (NOT a Megatron "shared expert" — see + Gemma4ParallelDenseMoE / plan v5 §531 R5) + - sliding layers: head_dim=256, num_kv=8, RoPE theta=10000, window=1024 + - global layers: head_dim=512, num_kv=2, RoPE theta=1000000, K=V tied (no V proj) + - full RoPE (no partial), qk_layernorm, parameter-free V LayerNorm + - final logit softcap 30.0, per-layer scalar buffer multiplied at layer end + - vocab=262144, ctx_len=262144 + - tie_word_embeddings=True + """ + layer_pattern = [] + for i in range(30): + layer_pattern.append("global" if (i + 1) % 6 == 0 else "sliding") + kv_tied_layers = [i for i, t in enumerate(layer_pattern) if t == "global"] + return Gemma4VLConfig( + num_layers=30, + hidden_size=2816, + ffn_hidden_size=2112, + num_attention_heads=16, + group_query_attention=True, + num_query_groups=8, + vocab_size_in_config_file=262144, + make_vocab_size_divisible_by=128, + qk_layernorm=True, + kv_channels=256, + add_qkv_bias=False, + rotary_base=1000000, + rotary_base_sliding=10000, + partial_rotary_factor=0.25, + sliding_window=1024, + num_experts=128, + moe_ffn_hidden_size=704, + final_logit_softcapping=30.0, + use_layer_scalar=True, + scale_emb_by_sqrt_hidden=True, + untie_embeddings_and_output_weights=False, + layer_pattern=layer_pattern, + per_layer_kv_channels={"sliding": 256, "global": 512}, + per_layer_num_query_groups={"sliding": 8, "global": 2}, + attention_k_eq_v=True, + kv_tied_layers=kv_tied_layers, + ) + + +@dataclass +class VisionConfig: + """Configuration for the Gemma4-VL vision tower. + + Verbatim mirror of HF ``Gemma4VisionConfig`` for ``gemma-4-26B-A4B-it`` + (see ``/ov2/pretrain_models/google/gemma-4-26B-A4B-it/config.json`` → + ``vision_config``). This is NOT a SigLIP tower — Gemma4 uses its own ViT + with: learned 2-D position embeddings via one-hot + matmul against a + ``[2, position_embedding_size, hidden_size]`` table; sandwich RMSNorm per + layer; v_proj followed by a parameterless RMSNorm (``with_scale=False``); + multidimensional RoPE applied per-axis on the head_dim/2 split; pooling via + a dynamic-kernel 2-D average pool whose kernel is derived from + ``pooling_kernel_size`` and the input sequence length; and an optional + output standardization stage gated by ``standardize``. + + Field semantics that DIFFER from OneVision-Encoder / SigLIP and MUST NOT + be silently aliased to the OV2 defaults: + + - ``patch_size=16`` (NOT 14) and ``in_channels=3`` give patch_embed input + ``in_features = 3 * 16 * 16 = 768``. + - ``num_attention_heads = num_key_value_heads = 16`` → MHA (not GQA); + ``head_dim=72`` → ``Q/K/V`` projection out_features ``= 16 * 72 = 1152``. + - ``hidden_activation='gelu_pytorch_tanh'`` is the *only* activation that + reproduces HF — exact ``F.gelu`` (erf form) silently diverges. + - ``use_clipped_linears=False`` for this checkpoint; the buffer branch in + :class:`Gemma4ClippableLinear` stays inactive. + - ``standardize=True`` → :class:`Gemma4VisionTower` registers persistent + ``std_bias``/``std_scale`` buffers and applies them after pooling. + - ``rope_theta=100.0`` (not 10000) for the vision multidimensional RoPE. + - ``position_embedding_size=10240`` is the per-axis position table length; + ``pooling_kernel_size=3`` controls the *output* sequence length via + ``pixel_values.shape[-2] // (pooling_kernel_size ** 2)``. + """ + + num_hidden_layers: int + hidden_size: int + intermediate_size: int + num_attention_heads: int + num_key_value_heads: int + head_dim: int + patch_size: int + pooling_kernel_size: int + position_embedding_size: int + rms_norm_eps: float + rope_theta: float + standardize: bool + use_clipped_linears: bool + hidden_activation: str + in_channels: int = 3 + initializer_range: float = 0.02 + + +def get_vision_config(model_family, model_name): + del model_family, model_name + return VisionConfig( + num_hidden_layers=27, + hidden_size=1152, + intermediate_size=4304, + num_attention_heads=16, + num_key_value_heads=16, + head_dim=72, + patch_size=16, + pooling_kernel_size=3, + position_embedding_size=10240, + rms_norm_eps=1e-06, + rope_theta=100.0, + standardize=True, + use_clipped_linears=False, + hidden_activation="gelu_pytorch_tanh", + in_channels=3, + initializer_range=0.02, + ) + + +def get_adapeter_config(model_family, model_name=None): + """Adapter config for Gemma4-VL: parameterless RMSNorm + bias-free Linear. + + The Gemma4 adapter is intentionally minimal — a single trainable tensor + (``embedding_projection.weight``) plus a parameterless input RMSNorm + (``with_scale=False``, no ``weight``). Mirrors HF ``Gemma4MultimodalEmbedder`` + (lines 2023-2047 of ``modeling_gemma4``); the LLM-side hidden_size used for + the projection out_features is derived at construction time from the LLM + config rather than carried here. + """ + del model_family, model_name + return AdapterConfig( + normalization="RMSNorm", + add_bias_linear=False, + layernorm_epsilon=1e-06, + ) diff --git a/aiak_training_llm/models/gemma4_vl/gemma4_vl_layer_spec.py b/aiak_training_llm/models/gemma4_vl/gemma4_vl_layer_spec.py new file mode 100644 index 00000000..4b7949a3 --- /dev/null +++ b/aiak_training_llm/models/gemma4_vl/gemma4_vl_layer_spec.py @@ -0,0 +1,170 @@ +"""Gemma4-VL layer specs. + +Three exports: + - get_vision_layer_with_spec: returns ``None`` — Gemma4 vision tower is a + plain ``nn.Module`` stack (HF-style verbatim port), not a Megatron + ``ModuleSpec``-driven ``TransformerBlock``. Kept as a no-op so + ``gemma4_vl_provider`` can call it uniformly with the LlavaOnevision2 + provider; the returned ``None`` is forwarded to ``Gemma4VL.__init__`` and + ignored there. + - get_adapter_layer_with_spec: same story — Gemma4 adapter is a plain + ``nn.Module`` (RMSNorm + Linear, see ``gemma4_adapter.py``). + - get_gemma4_layer_with_te_spec: Gemma4 LLM layer wiring + Gemma4TransformerLayer + Gemma4SelfAttention + Gemma4ParallelDenseMoE +""" + +from megatron.core.transformer.attention import SelfAttentionSubmodules +from megatron.core.transformer.enums import AttnMaskType +from megatron.core.transformer.identity_op import IdentityOp +from megatron.core.transformer.mlp import MLP, MLPSubmodules +from megatron.core.transformer.moe.experts import SequentialMLP, TEGroupedMLP +from megatron.core.transformer.spec_utils import ModuleSpec +from megatron.core.transformer.transformer_config import TransformerConfig + +from aiak_training_llm.models.dispatch import multiacc_modules +from aiak_training_llm.models.gemma4_vl.gemma4_attention import Gemma4SelfAttention +from aiak_training_llm.models.gemma4_vl.gemma4_mlp import ( + Gemma4ParallelDenseMoE, + Gemma4ParallelDenseMoESubmodules, +) +from aiak_training_llm.models.gemma4_vl.gemma4_moe_layer import ( + Gemma4MoELayer, + Gemma4MoESubmodules, +) +from aiak_training_llm.models.gemma4_vl.gemma4_transformer_layer import ( + Gemma4TransformerLayer, + Gemma4TransformerLayerSubmodules, +) +from aiak_training_llm.utils import is_te_min_version + + +__all__ = [ + "get_vision_layer_with_spec", + "get_adapter_layer_with_spec", + "get_gemma4_layer_with_te_spec", +] + + +def get_vision_layer_with_spec(): + """Gemma4 vision tower is not Megatron-spec-driven; return ``None``. + + The provider calls this for API parity with the LlavaOnevision2 provider + and forwards the result to ``Gemma4VL.__init__``, which builds the vision + tower via ``Gemma4VisionTower(...)`` directly and ignores this argument. + """ + return None + + +def get_adapter_layer_with_spec(): + """Gemma4 adapter is not Megatron-spec-driven; return ``None``. + + Same rationale as ``get_vision_layer_with_spec``. The Gemma4 adapter is + constructed as a plain ``nn.Module`` (see ``Gemma4Adapter``) inside + ``Gemma4VL.__init__``. + """ + return None + + +def _gemma4_norm_module() -> type: + return multiacc_modules.TENorm if is_te_min_version("1.9.0") else multiacc_modules.LocalNorm + + +def _gemma4_dense_mlp_spec() -> ModuleSpec: + return ModuleSpec( + module=MLP, + submodules=MLPSubmodules( + linear_fc1=multiacc_modules.TEColumnParallelLinear, + linear_fc2=multiacc_modules.TERowParallelLinear, + bias_activation_func_impl=multiacc_modules.bias_activation_func_impl, + ), + ) + + +def _gemma4_moe_spec(moe_grouped_gemm: bool) -> ModuleSpec: + if moe_grouped_gemm: + assert multiacc_modules.TEColumnParallelGroupedLinear is not None + expert_module = TEGroupedMLP + linear_fc1 = multiacc_modules.TEColumnParallelGroupedLinear + linear_fc2 = multiacc_modules.TERowParallelGroupedLinear + else: + expert_module = SequentialMLP + linear_fc1 = multiacc_modules.TEColumnParallelLinear + linear_fc2 = multiacc_modules.TERowParallelLinear + + experts_spec = ModuleSpec( + module=expert_module, + submodules=MLPSubmodules( + linear_fc1=linear_fc1, + linear_fc2=linear_fc2, + bias_activation_func_impl=multiacc_modules.bias_activation_func_impl, + ), + ) + # HF Gemma4 MoE has NO shared experts; the dense MLP is implemented as a + # peer branch in Gemma4ParallelDenseMoE (see plan v5 §531 R5). We must + # leave moe_shared_expert_intermediate_size unset (None) in the config — + # any other value, including 0, makes MoELayer's ``use_shared_expert`` + # flag True (it checks ``is not None``) and triggers a build of the + # shared-experts branch which we did not wire (would crash). + return ModuleSpec( + module=Gemma4MoELayer, + submodules=Gemma4MoESubmodules( + experts=experts_spec, + pre_feedforward_layernorm_2=_gemma4_norm_module(), + ), + ) + + +def _gemma4_parallel_dense_moe_spec(moe_grouped_gemm: bool) -> ModuleSpec: + norm = _gemma4_norm_module() + return ModuleSpec( + module=Gemma4ParallelDenseMoE, + submodules=Gemma4ParallelDenseMoESubmodules( + dense=_gemma4_dense_mlp_spec(), + moe=_gemma4_moe_spec(moe_grouped_gemm=moe_grouped_gemm), + pre_feedforward_layernorm=norm, + post_feedforward_layernorm_1=norm, + post_feedforward_layernorm_2=norm, + ), + ) + + +def _gemma4_self_attention_submodules(qk_norm) -> SelfAttentionSubmodules: + return SelfAttentionSubmodules( + linear_qkv=multiacc_modules.TELayerNormColumnParallelLinear, + core_attention=multiacc_modules.DotProductAttention, + linear_proj=multiacc_modules.TERowParallelLinear, + q_layernorm=qk_norm, + k_layernorm=qk_norm, + apply_rotary_fn=multiacc_modules.apply_rotary_pos_emb, + ) + + +def get_gemma4_layer_with_te_spec(config: TransformerConfig) -> ModuleSpec: + assert not config.multi_latent_attention, ( + "Gemma4-VL does not use multi-latent attention; got " + "config.multi_latent_attention=True" + ) + + qk_norm = _gemma4_norm_module() + mlp = _gemma4_parallel_dense_moe_spec(moe_grouped_gemm=config.moe_grouped_gemm) + + # ``pre_mlp_layernorm = IdentityOp`` because Gemma4ParallelDenseMoE + + # Gemma4MoELayer between them own all four pre-merge RMSNorms internally + # and must receive the raw post-attention residual (router math depends + # on this — see gemma4_mlp.py module docstring). + return ModuleSpec( + module=Gemma4TransformerLayer, + submodules=Gemma4TransformerLayerSubmodules( + input_layernorm=IdentityOp, + self_attention=ModuleSpec( + module=Gemma4SelfAttention, + params={"attn_mask_type": AttnMaskType.causal}, + submodules=_gemma4_self_attention_submodules(qk_norm=qk_norm), + ), + self_attn_bda=multiacc_modules.get_bias_dropout_add, + pre_mlp_layernorm=IdentityOp, + mlp=mlp, + mlp_bda=multiacc_modules.get_bias_dropout_add, + post_feedforward_layernorm=_gemma4_norm_module(), + ), + ) diff --git a/aiak_training_llm/models/gemma4_vl/gemma4_vl_provider.py b/aiak_training_llm/models/gemma4_vl/gemma4_vl_provider.py new file mode 100644 index 00000000..266bbe7b --- /dev/null +++ b/aiak_training_llm/models/gemma4_vl/gemma4_vl_provider.py @@ -0,0 +1,141 @@ +"""Provider for the Gemma4-VL model family. + +Mirrors :mod:`llava_onevision2_provider` but switches: + - ``register_model_provider`` family to ``VisionLanguageModelFamilies.GEMMA4_VL`` + - ``language_layer_spec`` builder to :func:`get_gemma4_layer_with_te_spec` + - top-level VLM class to :class:`Gemma4VL` + +All other knobs (TP/PP/SP gating, freeze gating, encoder PP) match the +LlavaOnevision2 provider verbatim so any base-provider improvements port +directly to this provider. +""" + +from copy import deepcopy +from dataclasses import asdict + +from megatron.core import mpu +from megatron.core.transformer.spec_utils import import_module + +from aiak_training_llm.models.factory import register_model_provider +from aiak_training_llm.models.gemma4_vl.gemma4_vl_config import ( + get_adapeter_config, + get_vision_config, +) +from aiak_training_llm.models.gemma4_vl.gemma4_vl_layer_spec import ( + get_adapter_layer_with_spec, + get_gemma4_layer_with_te_spec, + get_vision_layer_with_spec, +) +from aiak_training_llm.models.gemma4_vl.gemma4_vl_model import Gemma4VL +from aiak_training_llm.utils import build_transformer_config, get_args, print_rank_0 +from aiak_training_llm.utils.constants import VisionLanguageModelFamilies + + +@register_model_provider(model_family=[VisionLanguageModelFamilies.GEMMA4_VL]) +def gemma4_vl_model_provider( + pre_process: bool = True, + post_process: bool = True, + add_encoder: bool = True, + add_decoder: bool = True, + parallel_output: bool = True, +) -> Gemma4VL: + """Build a Gemma4-VL model instance. + + Args: + pre_process: include the embedding layer (pipeline-first stage). + post_process: include the output layer / loss (pipeline-last stage). + add_encoder / add_decoder: pipeline-encoder/decoder gating. + parallel_output: keep outputs split across TP ranks. + """ + args = get_args() + + print_rank_0(f"building {args.model_name} model ...") + + config = build_transformer_config(args) + + language_config = deepcopy(config) + vision_config = deepcopy(config) + adapter_config = deepcopy(config) + + from aiak_training_llm.models import get_model_family + + model_family = get_model_family(args.model_name) + for k, v in asdict(get_vision_config(model_family, args.model_name)).items(): + setattr(vision_config, k, v) + for k, v in asdict(get_adapeter_config(model_family, args.model_name)).items(): + setattr(adapter_config, k, v) + + setattr(language_config, "image_token_id", 258880) + setattr(language_config, "video_token_id", 258884) + + if args.encoder_pipeline_model_parallel_size in [0, None]: + vision_config.pipeline_model_parallel_size = 1 + vision_config.first_pipeline_num_layers = ( + vision_config.last_pipeline_num_layers + ) = None + vision_config.tp_comm_overlap = False + vision_config.context_parallel_size = 1 + vision_config.context_parallel_ulysses_degree = 1 + add_encoder = mpu.is_pipeline_first_stage() + add_decoder = True + else: + assert args.encoder_pipeline_model_parallel_size == 1, ( + "vision model and projection can only live on 1 pipeline stage." + ) + vision_config.pipeline_model_parallel_size = ( + args.encoder_pipeline_model_parallel_size + ) + if args.encoder_tensor_model_parallel_size > 0: + vision_config.tensor_model_parallel_size = ( + args.encoder_tensor_model_parallel_size + ) + vision_config.first_pipeline_num_layers = ( + vision_config.last_pipeline_num_layers + ) = None + vision_config.context_parallel_size = 1 + vision_config.tp_comm_overlap = False + + if args.use_legacy_models: + raise ValueError("Classic Megatron-LM models are not supported.") + + if args.spec is not None: + language_layer_spec = import_module(args.spec) + else: + adapter_layer_spec = get_adapter_layer_with_spec() + vision_layer_spec = get_vision_layer_with_spec() + language_layer_spec = get_gemma4_layer_with_te_spec(language_config) + + model = Gemma4VL( + language_config=language_config, + vision_config=vision_config, + adapter_config=adapter_config, + language_layer_spec=language_layer_spec, + vision_layer_spec=vision_layer_spec, + adapter_layer_spec=adapter_layer_spec, + language_vocab_size=args.padded_vocab_size, + language_max_sequence_length=args.max_position_embeddings, + pre_process=pre_process, + post_process=post_process, + add_encoder=add_encoder, + add_decoder=add_decoder, + fp16_lm_cross_entropy=args.fp16_lm_cross_entropy, + parallel_output=parallel_output, + share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights, + language_position_embedding_type=args.position_embedding_type, + language_rotary_percent=args.rotary_percent, + language_rotary_base=args.rotary_base, + language_rotary_base_sliding=getattr(args, "rotary_base_sliding", None), + seq_len_interpolation_factor=args.rotary_seq_len_interpolation_factor, + ) + + if args.trainable_modules != ["all"]: + train_language_model = "language_model" in args.trainable_modules + train_vision_model = "vision_model" in args.trainable_modules + train_adapter = "adapter" in args.trainable_modules + model.freeze( + freeze_language_model=not train_language_model, + freeze_vision_model=not train_vision_model, + freeze_adapter=not train_adapter, + ) + + return model From a02cb1523773e2c09bd234ecedade648428d38c1 Mon Sep 17 00:00:00 2001 From: XiangAn Date: Sun, 17 May 2026 23:35:25 +0800 Subject: [PATCH 3/6] feat: add Gemma4 VL vision stack Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- .../models/gemma4_vl/gemma4_adapter.py | 47 +++ .../gemma4_vl/gemma4_vision_attention.py | 150 ++++++++ .../models/gemma4_vl/gemma4_vision_common.py | 79 +++++ .../models/gemma4_vl/gemma4_vision_layer.py | 101 ++++++ .../models/gemma4_vl/gemma4_vision_mlp.py | 67 ++++ .../gemma4_vl/gemma4_vision_patch_embed.py | 65 ++++ .../models/gemma4_vl/gemma4_vision_pooler.py | 94 +++++ .../models/gemma4_vl/gemma4_vision_rotary.py | 181 ++++++++++ .../models/gemma4_vl/gemma4_vision_tower.py | 164 +++++++++ .../gemma4_vl/gemma4_vl_attention_mask.py | 126 +++++++ .../models/gemma4_vl/gemma4_vl_model.py | 331 ++++++++++++++++++ 11 files changed, 1405 insertions(+) create mode 100644 aiak_training_llm/models/gemma4_vl/gemma4_adapter.py create mode 100644 aiak_training_llm/models/gemma4_vl/gemma4_vision_attention.py create mode 100644 aiak_training_llm/models/gemma4_vl/gemma4_vision_common.py create mode 100644 aiak_training_llm/models/gemma4_vl/gemma4_vision_layer.py create mode 100644 aiak_training_llm/models/gemma4_vl/gemma4_vision_mlp.py create mode 100644 aiak_training_llm/models/gemma4_vl/gemma4_vision_patch_embed.py create mode 100644 aiak_training_llm/models/gemma4_vl/gemma4_vision_pooler.py create mode 100644 aiak_training_llm/models/gemma4_vl/gemma4_vision_rotary.py create mode 100644 aiak_training_llm/models/gemma4_vl/gemma4_vision_tower.py create mode 100644 aiak_training_llm/models/gemma4_vl/gemma4_vl_attention_mask.py create mode 100644 aiak_training_llm/models/gemma4_vl/gemma4_vl_model.py diff --git a/aiak_training_llm/models/gemma4_vl/gemma4_adapter.py b/aiak_training_llm/models/gemma4_vl/gemma4_adapter.py new file mode 100644 index 00000000..6d6a9e1d --- /dev/null +++ b/aiak_training_llm/models/gemma4_vl/gemma4_adapter.py @@ -0,0 +1,47 @@ +"""Gemma4 vision-to-LM adapter. + +Verbatim port of HF ``Gemma4MultimodalEmbedder`` (modeling_gemma4 lines 2023-2047). + +In the HF state_dict the adapter is stored under +``model.embed_vision.{embedding_pre_projection_norm,embedding_projection}.*``. +For Gemma4-26B-A4B-it specifically: +- ``embedding_pre_projection_norm`` is a ``Gemma4RMSNorm`` with + ``with_scale=False`` (no learnable weight, ckpt has no entry for it). +- ``embedding_projection`` is ``nn.Linear(vision.hidden_size=1152, + text.hidden_size=2816, bias=False)``; the only ckpt entry is + ``model.embed_vision.embedding_projection.weight``. + +So the entire adapter holds exactly ONE trainable tensor. Do not promote the +RMSNorm to ``with_scale=True`` even though it would seem more "complete" — HF +ckpt loading will fail on an unexpected ``embedding_pre_projection_norm.weight`` +key. The naming uses HF attribute names so converter mapping stays trivial. +""" + +from __future__ import annotations + +import torch +from torch import nn + +from .gemma4_vision_common import Gemma4RMSNorm + + +class Gemma4Adapter(nn.Module): + def __init__( + self, + vision_hidden_size: int, + text_hidden_size: int, + rms_norm_eps: float = 1e-6, + ) -> None: + super().__init__() + self.vision_hidden_size = vision_hidden_size + self.text_hidden_size = text_hidden_size + self.embedding_pre_projection_norm = Gemma4RMSNorm( + vision_hidden_size, eps=rms_norm_eps, with_scale=False + ) + self.embedding_projection = nn.Linear( + vision_hidden_size, text_hidden_size, bias=False + ) + + def forward(self, inputs_embeds: torch.Tensor) -> torch.Tensor: + embs_normed = self.embedding_pre_projection_norm(inputs_embeds) + return self.embedding_projection(embs_normed) diff --git a/aiak_training_llm/models/gemma4_vl/gemma4_vision_attention.py b/aiak_training_llm/models/gemma4_vl/gemma4_vision_attention.py new file mode 100644 index 00000000..32044911 --- /dev/null +++ b/aiak_training_llm/models/gemma4_vl/gemma4_vision_attention.py @@ -0,0 +1,150 @@ +"""Gemma4 vision multi-head attention. + +Verbatim port of HF ``Gemma4VisionAttention`` and ``eager_attention_forward`` +(modeling_gemma4 lines 779-936). + +Three non-obvious design choices that MUST be preserved bit-for-bit with HF: + +1. ``self.scaling = 1.0`` (NOT ``head_dim ** -0.5``). The standard + 1/sqrt(d) scale is *omitted* because ``q_norm`` and ``k_norm`` already + normalize Q and K to unit-RMS. Adding the scale here would double-scale + and break HF parity. + +2. ``v_norm`` is a parameterless RMSNorm (``with_scale=False``) applied to + value states *after* the V projection but *before* head transpose. This + is unusual (most attention impls don't normalize V); do not remove it. + +3. RoPE is applied via ``apply_multidimensional_rope`` to Q and K *before* + the head transpose, with ``unsqueeze_dim=2`` to match the + ``[B, N, H, D]`` layout. After RoPE, Q/K/V are transposed to + ``[B, H, N, D]`` for attention. + +The attention math is fp32 softmax (upcast inside ``eager_attention_forward``) +and uses standard GQA via ``repeat_kv``. We deliberately do NOT route this +through Megatron's fused TE attention because (a) vision sequences are short +(<=2240 patches) so the kernel speedup is marginal, and (b) TE's softmax +fp16/bf16 path diverges from HF's fp32 path at the 4th decimal which would +fail downstream consistency checks against the HF reference. +""" + +from __future__ import annotations + +import torch +from torch import nn + +from .gemma4_vision_common import Gemma4ClippableLinear, Gemma4RMSNorm +from .gemma4_vision_rotary import apply_multidimensional_rope + + +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand( + batch, num_key_value_heads, n_rep, slen, head_dim + ) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +def eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor | None, + dropout: float = 0.0, + scaling: float | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + if scaling is None: + scaling = module.head_dim**-0.5 + + key_states = repeat_kv(key, module.num_key_value_groups) + value_states = repeat_kv(value, module.num_key_value_groups) + + attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling + if attention_mask is not None: + attn_weights = attn_weights + attention_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) + attn_output = torch.matmul(attn_weights, value_states) + attn_output = attn_output.transpose(1, 2).contiguous() + return attn_output, attn_weights + + +class Gemma4VisionAttention(nn.Module): + def __init__( + self, + hidden_size: int, + num_attention_heads: int, + num_key_value_heads: int, + head_dim: int, + rms_norm_eps: float = 1e-6, + attention_dropout: float = 0.0, + use_clipped_linears: bool = False, + ) -> None: + super().__init__() + self.hidden_size = hidden_size + self.num_attention_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + self.head_dim = head_dim + self.num_key_value_groups = num_attention_heads // num_key_value_heads + self.scaling = 1.0 + self.attention_dropout = attention_dropout + self.is_causal = False + + self.q_proj = Gemma4ClippableLinear( + hidden_size, num_attention_heads * head_dim, use_clipped_linears=use_clipped_linears + ) + self.k_proj = Gemma4ClippableLinear( + hidden_size, num_key_value_heads * head_dim, use_clipped_linears=use_clipped_linears + ) + self.v_proj = Gemma4ClippableLinear( + hidden_size, num_key_value_heads * head_dim, use_clipped_linears=use_clipped_linears + ) + self.o_proj = Gemma4ClippableLinear( + num_attention_heads * head_dim, hidden_size, use_clipped_linears=use_clipped_linears + ) + + self.q_norm = Gemma4RMSNorm(dim=head_dim, eps=rms_norm_eps) + self.k_norm = Gemma4RMSNorm(dim=head_dim, eps=rms_norm_eps) + self.v_norm = Gemma4RMSNorm(dim=head_dim, eps=rms_norm_eps, with_scale=False) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + position_ids: torch.LongTensor, + attention_mask: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + cos, sin = position_embeddings + + query_states = self.q_proj(hidden_states).view(hidden_shape) + query_states = self.q_norm(query_states) + query_states = apply_multidimensional_rope(query_states, cos, sin, position_ids) + query_states = query_states.transpose(1, 2) + + key_states = self.k_proj(hidden_states).view(hidden_shape) + key_states = self.k_norm(key_states) + key_states = apply_multidimensional_rope(key_states, cos, sin, position_ids) + key_states = key_states.transpose(1, 2) + + value_states = self.v_proj(hidden_states).view(hidden_shape) + value_states = self.v_norm(value_states) + value_states = value_states.transpose(1, 2) + + attn_output, attn_weights = eager_attention_forward( + module=self, + query=query_states, + key=key_states, + value=value_states, + attention_mask=attention_mask, + dropout=self.attention_dropout if self.training else 0.0, + scaling=self.scaling, + ) + + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights diff --git a/aiak_training_llm/models/gemma4_vl/gemma4_vision_common.py b/aiak_training_llm/models/gemma4_vl/gemma4_vision_common.py new file mode 100644 index 00000000..27fc6b3f --- /dev/null +++ b/aiak_training_llm/models/gemma4_vl/gemma4_vision_common.py @@ -0,0 +1,79 @@ +"""Shared Gemma4 vision/adapter primitives: RMSNorm and ClippableLinear. + +Verbatim port of HF ``Gemma4RMSNorm`` (lines 168-187) and ``Gemma4ClippableLinear`` +(lines 139-167) from transformers ``modeling_gemma4``. + +Two non-obvious facts that determine HF checkpoint compatibility and MUST NOT +change: + +1. ``Gemma4ClippableLinear`` always wraps the real ``nn.Linear`` under attribute + name ``.linear``. HF checkpoints store keys as ``*..linear.weight`` + even when ``use_clipped_linears=False`` and the clamp branch is dead code. + Renaming this attribute (e.g. inlining the linear) breaks every HF→Megatron + converter mapping. + +2. ``Gemma4RMSNorm.with_scale=False`` produces a parameterless module (no + ``weight`` attribute at all). Both ``v_norm`` in vision attention and the + adapter's RMSNorm rely on this; do NOT add a default ``weight`` for "API + uniformity" because state_dict loading from HF will then fail with an + unexpected key. + +The HF clamp branch (``use_clipped_linears=True``) registers four buffers +(input/output min/max). Gemma4-26B-A4B-it ships with ``use_clipped_linears=False`` +so we keep the branch for fidelity but it is exercised only by post-training +quantization configs. +""" + +from __future__ import annotations + +import torch +from torch import nn + + +class Gemma4RMSNorm(nn.Module): + def __init__(self, dim: int, eps: float = 1e-6, with_scale: bool = True) -> None: + super().__init__() + self.eps = eps + self.with_scale = with_scale + if self.with_scale: + self.weight = nn.Parameter(torch.ones(dim), requires_grad=True) + + def _norm(self, hidden_states: torch.Tensor) -> torch.Tensor: + mean_squared = hidden_states.pow(2).mean(-1, keepdim=True) + self.eps + return hidden_states * torch.rsqrt(mean_squared) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + output = self._norm(hidden_states.float()) + if self.with_scale: + output = output * self.weight.float() + return output.type_as(hidden_states) + + def extra_repr(self) -> str: + weight_shape = tuple(self.weight.shape) if self.with_scale else None + return f"weight={weight_shape}, eps={self.eps}, with_scale={self.with_scale}" + + +class Gemma4ClippableLinear(nn.Module): + def __init__( + self, + in_features: int, + out_features: int, + use_clipped_linears: bool = False, + ) -> None: + super().__init__() + self.use_clipped_linears = use_clipped_linears + self.linear = nn.Linear(in_features, out_features, bias=False) + + if self.use_clipped_linears: + self.register_buffer("input_min", torch.tensor(-float("inf"))) + self.register_buffer("input_max", torch.tensor(float("inf"))) + self.register_buffer("output_min", torch.tensor(-float("inf"))) + self.register_buffer("output_max", torch.tensor(float("inf"))) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + if self.use_clipped_linears: + hidden_states = torch.clamp(hidden_states, self.input_min, self.input_max) + hidden_states = self.linear(hidden_states) + if self.use_clipped_linears: + hidden_states = torch.clamp(hidden_states, self.output_min, self.output_max) + return hidden_states diff --git a/aiak_training_llm/models/gemma4_vl/gemma4_vision_layer.py b/aiak_training_llm/models/gemma4_vl/gemma4_vision_layer.py new file mode 100644 index 00000000..b73a8605 --- /dev/null +++ b/aiak_training_llm/models/gemma4_vl/gemma4_vision_layer.py @@ -0,0 +1,101 @@ +"""Gemma4 vision encoder layer. + +Verbatim port of HF ``Gemma4VisionEncoderLayer`` (modeling_gemma4 lines 939-980). + +Sandwich normalization layout — this is the source of the layer's 4 RMSNorms +and the order is NOT the conventional pre-norm or post-norm transformer: + + residual = h + h = input_layernorm(h) # pre-attn norm + h, _ = self_attn(h, ...) + h = post_attention_layernorm(h) # post-attn norm, BEFORE residual add + h = residual + h + + residual = h + h = pre_feedforward_layernorm(h) # pre-mlp norm + h = mlp(h) + h = post_feedforward_layernorm(h) # post-mlp norm, BEFORE residual add + h = residual + h + +The two "post_*_layernorm" calls happen *between* the sublayer output and the +residual sum, not after. Reordering to standard pre-norm +(``residual + sublayer(norm(h))``) breaks HF parity because the sandwich +construction normalizes the sublayer output before the skip-add, which changes +the variance arithmetic of every subsequent layer. This pattern is shared with +the LLM hybrid layers (see ``gemma4_transformer_layer.py``); see plan v5 for +why we do not collapse the four norms into two. +""" + +from __future__ import annotations + +import torch +from torch import nn + +from .gemma4_vision_attention import Gemma4VisionAttention +from .gemma4_vision_common import Gemma4RMSNorm +from .gemma4_vision_mlp import Gemma4VisionMLP + + +class Gemma4VisionEncoderLayer(nn.Module): + def __init__( + self, + hidden_size: int, + num_attention_heads: int, + num_key_value_heads: int, + head_dim: int, + intermediate_size: int, + rms_norm_eps: float = 1e-6, + attention_dropout: float = 0.0, + hidden_activation: str = "gelu_pytorch_tanh", + use_clipped_linears: bool = False, + layer_idx: int | None = None, + ) -> None: + super().__init__() + self.hidden_size = hidden_size + self.layer_idx = layer_idx + + self.self_attn = Gemma4VisionAttention( + hidden_size=hidden_size, + num_attention_heads=num_attention_heads, + num_key_value_heads=num_key_value_heads, + head_dim=head_dim, + rms_norm_eps=rms_norm_eps, + attention_dropout=attention_dropout, + use_clipped_linears=use_clipped_linears, + ) + self.mlp = Gemma4VisionMLP( + hidden_size=hidden_size, + intermediate_size=intermediate_size, + hidden_activation=hidden_activation, + use_clipped_linears=use_clipped_linears, + ) + self.input_layernorm = Gemma4RMSNorm(hidden_size, eps=rms_norm_eps) + self.post_attention_layernorm = Gemma4RMSNorm(hidden_size, eps=rms_norm_eps) + self.pre_feedforward_layernorm = Gemma4RMSNorm(hidden_size, eps=rms_norm_eps) + self.post_feedforward_layernorm = Gemma4RMSNorm(hidden_size, eps=rms_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + position_ids: torch.LongTensor, + attention_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states, _ = self.self_attn( + hidden_states=hidden_states, + position_embeddings=position_embeddings, + position_ids=position_ids, + attention_mask=attention_mask, + ) + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.pre_feedforward_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = self.post_feedforward_layernorm(hidden_states) + hidden_states = residual + hidden_states + + return hidden_states diff --git a/aiak_training_llm/models/gemma4_vl/gemma4_vision_mlp.py b/aiak_training_llm/models/gemma4_vl/gemma4_vision_mlp.py new file mode 100644 index 00000000..acc6023e --- /dev/null +++ b/aiak_training_llm/models/gemma4_vl/gemma4_vision_mlp.py @@ -0,0 +1,67 @@ +"""Gemma4 vision gated MLP. + +Verbatim port of HF ``Gemma4VisionMLP`` (modeling_gemma4 lines 643-657). + +The MLP shape is the SwiGLU-style ``down(act(gate) * up)`` (NOT +``down(gate * act(up))``); the activation wraps the gate projection only. +This matches Gemma's LLM MLP convention and is the order that HF checkpoints +were trained with — swapping the wrap target produces silently wrong outputs +because ``gate_proj`` and ``up_proj`` share the same in/out shapes. + +``hidden_activation`` is ``gelu_pytorch_tanh`` for Gemma4 vision (config +default), which maps to ``torch.nn.functional.gelu(x, approximate="tanh")``. +The standard exact GELU diverges from HF outputs at the 5th decimal place, +which compounds across 27 layers and breaks consistency checks; do NOT +substitute the exact form. +""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F +from torch import nn + +from .gemma4_vision_common import Gemma4ClippableLinear + + +def _gelu_pytorch_tanh(x: torch.Tensor) -> torch.Tensor: + return F.gelu(x, approximate="tanh") + + +_ACTIVATIONS = { + "gelu_pytorch_tanh": _gelu_pytorch_tanh, + "gelu": F.gelu, + "silu": F.silu, + "relu": F.relu, +} + + +class Gemma4VisionMLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_activation: str = "gelu_pytorch_tanh", + use_clipped_linears: bool = False, + ) -> None: + super().__init__() + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.gate_proj = Gemma4ClippableLinear( + hidden_size, intermediate_size, use_clipped_linears=use_clipped_linears + ) + self.up_proj = Gemma4ClippableLinear( + hidden_size, intermediate_size, use_clipped_linears=use_clipped_linears + ) + self.down_proj = Gemma4ClippableLinear( + intermediate_size, hidden_size, use_clipped_linears=use_clipped_linears + ) + if hidden_activation not in _ACTIVATIONS: + raise ValueError( + f"Unknown hidden_activation={hidden_activation!r}; " + f"expected one of {sorted(_ACTIVATIONS)}." + ) + self.act_fn = _ACTIVATIONS[hidden_activation] + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) diff --git a/aiak_training_llm/models/gemma4_vl/gemma4_vision_patch_embed.py b/aiak_training_llm/models/gemma4_vl/gemma4_vision_patch_embed.py new file mode 100644 index 00000000..7712596c --- /dev/null +++ b/aiak_training_llm/models/gemma4_vl/gemma4_vision_patch_embed.py @@ -0,0 +1,65 @@ +"""Gemma4 vision patch embedder. + +Verbatim port of HF ``Gemma4VisionPatchEmbedder`` (modeling_gemma4 lines 550-582). + +Key non-obvious facts that callers MUST honor: +- Pre-projection scaling is ``2 * (pixel_values - 0.5)``, NOT a learned standardize + pair (``std_bias``/``std_scale``). The standardize buffers live on + ``Gemma4VisionTower`` and are applied *after* pooling, not before patch embed. +- ``position_embedding_table`` has shape ``[2, position_embedding_size, hidden_size]``: + the leading ``2`` indexes the spatial axis (x first, y second). Two PE lookups + per patch are summed; padding patches (position == -1) get zeroed out. +- Position embedding uses a one-hot @ matmul rather than ``nn.Embedding`` lookup + so it stays within the autograd graph for fused backward when patches are + shared across multiple soft-token outputs. +""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F +from torch import nn + + +class Gemma4VisionPatchEmbedder(nn.Module): + def __init__( + self, + patch_size: int, + hidden_size: int, + position_embedding_size: int, + in_channels: int = 3, + ) -> None: + super().__init__() + self.patch_size = patch_size + self.hidden_size = hidden_size + self.position_embedding_size = position_embedding_size + self.in_channels = in_channels + + self.input_proj = nn.Linear(in_channels * patch_size * patch_size, hidden_size, bias=False) + self.position_embedding_table = nn.Parameter( + torch.ones(2, position_embedding_size, hidden_size) + ) + + def _position_embeddings( + self, + pixel_position_ids: torch.Tensor, + padding_positions: torch.Tensor, + ) -> torch.Tensor: + clamped_positions = pixel_position_ids.clamp(min=0) + one_hot = F.one_hot(clamped_positions, num_classes=self.position_embedding_size) + one_hot = one_hot.permute(0, 2, 1, 3).to(self.position_embedding_table) + position_embeddings = one_hot @ self.position_embedding_table + position_embeddings = position_embeddings.sum(dim=1) + position_embeddings = torch.where(padding_positions.unsqueeze(-1), 0.0, position_embeddings) + return position_embeddings + + def forward( + self, + pixel_values: torch.Tensor, + pixel_position_ids: torch.Tensor, + padding_positions: torch.Tensor, + ) -> torch.Tensor: + pixel_values = 2 * (pixel_values - 0.5) + hidden_states = self.input_proj(pixel_values.to(self.input_proj.weight.dtype)) + position_embeddings = self._position_embeddings(pixel_position_ids, padding_positions) + return hidden_states + position_embeddings diff --git a/aiak_training_llm/models/gemma4_vl/gemma4_vision_pooler.py b/aiak_training_llm/models/gemma4_vl/gemma4_vision_pooler.py new file mode 100644 index 00000000..eabc2168 --- /dev/null +++ b/aiak_training_llm/models/gemma4_vl/gemma4_vision_pooler.py @@ -0,0 +1,94 @@ +"""Gemma4 vision spatial pooler. + +Verbatim port of HF ``Gemma4VisionPooler`` (modeling_gemma4 lines 584-640). + +Three subtleties that determine numerical parity with HF: + +1. The pooling kernel size ``k`` is derived dynamically from the *runtime* + ratio between input patch count and target soft-token count: + ``k = int(sqrt(input_seq_len // output_length))``. There is no + ``pooling_kernel_size`` field used at runtime — the config field of that + name is only consulted at processor / config-validation time to compute + ``num_soft_tokens``. So this module must accept ``output_length`` per call + and infer ``k`` from shapes; do NOT hard-code ``k=3`` even though that is + the value used by Gemma4 vision under default config. + +2. Pooling is done as ``(one_hot(kernel_idx) / k^2)^T @ hidden_states``: this + is mathematically equivalent to a 2D avg-pool by spatial-bucket index but + stays in the autograd graph as plain matmul, which is what the HF + reference uses to keep the operation TPU/XLA-friendly. Replacing it with + ``F.avg_pool2d`` is *not* equivalent because patches arrive in raster / + packed order, not as a dense [B, C, H, W] grid. + +3. The final ``hidden_states *= sqrt(hidden_size)`` scaling is applied + ALWAYS, even on the no-pool branch (when input length already equals + output length). This compensates for the post-LN variance reduction that + the soft-token consumer (Gemma4MultimodalEmbedder) was trained against. + Removing or moving this scale silently changes the LM input statistics + and breaks downstream consistency. + +The output ``mask`` is a *new* padding mask produced from the pooling weight +matrix: a pooled bucket is considered padding iff every input patch that +maps to it was already padding. The caller MUST use this returned mask, not +the input ``padding_positions``. +""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F +from torch import nn + + +class Gemma4VisionPooler(nn.Module): + def __init__(self, hidden_size: int) -> None: + super().__init__() + self.hidden_size = hidden_size + self.root_hidden_size = hidden_size**0.5 + + def _avg_pool_by_positions( + self, + hidden_states: torch.Tensor, + pixel_position_ids: torch.Tensor, + length: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + input_seq_len = hidden_states.shape[1] + k = int((input_seq_len // length) ** 0.5) + k_squared = k * k + if k_squared * length != input_seq_len: + raise ValueError( + f"Cannot pool {tuple(hidden_states.shape)} to {length}: " + f"k={k}^2 times length={length} must equal input_seq_len={input_seq_len}." + ) + + clamped_positions = pixel_position_ids.clamp(min=0) + max_x = clamped_positions[..., 0].max(dim=-1, keepdim=True)[0] + 1 + kernel_idxs = torch.div(clamped_positions, k, rounding_mode="floor") + kernel_idxs = kernel_idxs[..., 0] + (max_x // k) * kernel_idxs[..., 1] + weights = F.one_hot(kernel_idxs.long(), length).float() / k_squared + output = weights.transpose(1, 2) @ hidden_states.float() + mask = torch.logical_not((weights == 0).all(dim=1)) + return output.to(hidden_states.dtype), mask + + def forward( + self, + hidden_states: torch.Tensor, + pixel_position_ids: torch.Tensor, + padding_positions: torch.Tensor, + output_length: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + if output_length > hidden_states.shape[1]: + raise ValueError( + f"Cannot output more soft tokens (requested {output_length}) than there are " + f"patches ({hidden_states.shape[1]}). Change `num_soft_tokens` upstream." + ) + + hidden_states = hidden_states.masked_fill(padding_positions.unsqueeze(-1), 0.0) + + if hidden_states.shape[1] != output_length: + hidden_states, padding_positions = self._avg_pool_by_positions( + hidden_states, pixel_position_ids, output_length + ) + + hidden_states = hidden_states * self.root_hidden_size + return hidden_states, padding_positions diff --git a/aiak_training_llm/models/gemma4_vl/gemma4_vision_rotary.py b/aiak_training_llm/models/gemma4_vl/gemma4_vision_rotary.py new file mode 100644 index 00000000..90581eb6 --- /dev/null +++ b/aiak_training_llm/models/gemma4_vl/gemma4_vision_rotary.py @@ -0,0 +1,181 @@ +"""Gemma4 vision 2D rotary position embedding. + +Verbatim port of HF ``Gemma4VisionRotaryEmbedding`` and ``apply_multidimensional_rope`` +from transformers ``modeling_gemma4`` (lines 659-866). Kept as a standalone module so +that the Megatron port stays bit-identical to HF; do NOT route this through the +shared Megatron RoPE plumbing because Megatron's RoPE assumes a single 1-D position +sequence whereas Gemma4 vision splits the head dimension across two spatial axes +(x, y) and applies independent rope to each half. + +Math contract: +- ``head_dim`` is even and divisible by ``2 * ndim`` (ndim=2 here). +- ``spatial_dim = head_dim // ndim`` (=36 for head_dim=72): each axis gets its own + ``inv_freq`` table of length ``spatial_dim // 2`` (=18 sin/cos pairs per axis). +- ``forward(x, position_ids)`` returns ``(cos, sin)`` each of shape + ``[B, N, head_dim]`` where the last dim is the concatenation of the per-axis + cos/sin tables (axis-x first ``spatial_dim`` channels, axis-y last ``spatial_dim``). +- ``apply_multidimensional_rope`` splits ``x``, ``cos``, ``sin`` along the last axis + into ``ndim`` equal parts and applies standard rope independently to each, then + concatenates back. This matches the HF reference and is the only way to reproduce + ``model.embed_vision.embedding_projection`` numerics. +""" + +from __future__ import annotations + +import torch +from torch import nn + + +def rotate_half(x: torch.Tensor) -> torch.Tensor: + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def apply_rotary_pos_emb( + x: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + unsqueeze_dim: int = 1, +) -> torch.Tensor: + """Apply standard 1-D rotary position embedding to ``x``. + + ``cos`` / ``sin`` have shape ``[B, N, D]`` and are unsqueezed at + ``unsqueeze_dim`` to broadcast against ``x`` whose layout is either + ``[B, H, N, D]`` (unsqueeze_dim=1, post-transpose) or + ``[B, N, H, D]`` (unsqueeze_dim=2, pre-transpose, used by Gemma4 vision). + """ + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + return (x * cos) + (rotate_half(x) * sin) + + +def apply_multidimensional_rope( + x: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + position_ids: torch.Tensor, + unsqueeze_dim: int = 2, +) -> torch.Tensor: + """Apply multi-dimensional RoPE to ``x``. + + Splits ``x`` along the last (channel) dimension into ``ndim`` parts + (where ``ndim = position_ids.shape[-1]``), applies standard rotary embedding + independently to each part using its own slice of ``cos`` / ``sin``, then + concatenates the parts back. For Gemma4 vision ``ndim=2`` (x and y axes). + + Args: + x: Tensor of shape ``[B, N, H, D]`` (queries / keys before transpose). + cos: Tensor of shape ``[B, N, D]`` produced by ``Gemma4VisionRotaryEmbedding``. + sin: Tensor of shape ``[B, N, D]`` produced by ``Gemma4VisionRotaryEmbedding``. + position_ids: Tensor of shape ``[B, N, ndim]`` with per-axis patch coordinates. + unsqueeze_dim: Broadcast axis when applying rope (defaults to 2 to match + Gemma4VisionAttention call site, where ``x`` is pre-transpose ``[B,N,H,D]``). + + Returns: + Tensor of shape ``[B, N, H, D]`` with rope applied. + """ + ndim = position_ids.shape[-1] + num_input_channels = x.shape[-1] + num_rotated_channels_per_dim = 2 * (num_input_channels // (2 * ndim)) + + if num_rotated_channels_per_dim <= 0: + raise ValueError( + "Invalid configuration: num_rotated_channels_per_dim must be > 0, got " + f"{num_rotated_channels_per_dim} (num_input_channels={num_input_channels}, " + f"ndim={ndim})" + ) + + split_sizes = [num_rotated_channels_per_dim] * ndim + x_parts = torch.split(x, split_sizes, dim=-1) + cos_parts = torch.split(cos, split_sizes, dim=-1) + sin_parts = torch.split(sin, split_sizes, dim=-1) + y_parts = [ + apply_rotary_pos_emb( + x=x_parts[k], + cos=cos_parts[k], + sin=sin_parts[k], + unsqueeze_dim=unsqueeze_dim, + ) + for k in range(ndim) + ] + return torch.cat(y_parts, dim=-1) + + +class Gemma4VisionRotaryEmbedding(nn.Module): + """Gemma4 vision 2-D rotary embedding. + + Computes per-axis ``cos`` / ``sin`` tables for x and y patch coordinates and + concatenates them along the channel dim. The result is a single + ``[B, N, head_dim]`` cos/sin pair that ``apply_multidimensional_rope`` knows + how to split back into per-axis pieces. + + The frequency table is computed once per axis with + ``spatial_dim = head_dim // ndim`` and ``base = rope_theta``; ``inv_freq`` is + a non-persistent buffer so it is re-derived from config on load and never + stored in checkpoints. + """ + + inv_freq: torch.Tensor + + def __init__( + self, + head_dim: int, + rope_theta: float = 100.0, + ndim: int = 2, + device: torch.device | None = None, + ) -> None: + super().__init__() + if head_dim % (2 * ndim) != 0: + raise ValueError( + f"head_dim={head_dim} must be divisible by 2*ndim={2 * ndim} " + "to support multidimensional rotary embedding." + ) + self.head_dim = head_dim + self.ndim = ndim + self.rope_theta = rope_theta + self.attention_scaling = 1.0 + + spatial_dim = head_dim // ndim + inv_freq = 1.0 / ( + rope_theta + ** ( + torch.arange(0, spatial_dim, 2, dtype=torch.int64).to( + device=device, dtype=torch.float + ) + / spatial_dim + ) + ) + self.register_buffer("inv_freq", inv_freq, persistent=False) + self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False) + + @torch.no_grad() + def forward(self, x: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + inv_freq_expanded = ( + self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device) + ) + device_type = ( + x.device.type + if isinstance(x.device.type, str) and x.device.type != "mps" + else "cpu" + ) + + all_cos: list[torch.Tensor] = [] + all_sin: list[torch.Tensor] = [] + for i in range(self.ndim): + dim_position_ids = position_ids[:, :, i] + dim_position_ids_expanded = dim_position_ids[:, None, :].float() + + with torch.autocast(device_type=device_type, enabled=False): + freqs = ( + inv_freq_expanded.float() @ dim_position_ids_expanded.float() + ).transpose(1, 2) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() * self.attention_scaling + sin = emb.sin() * self.attention_scaling + all_cos.append(cos) + all_sin.append(sin) + + cos = torch.cat(all_cos, dim=-1).to(dtype=x.dtype) + sin = torch.cat(all_sin, dim=-1).to(dtype=x.dtype) + return cos, sin diff --git a/aiak_training_llm/models/gemma4_vl/gemma4_vision_tower.py b/aiak_training_llm/models/gemma4_vl/gemma4_vision_tower.py new file mode 100644 index 00000000..86b6fe64 --- /dev/null +++ b/aiak_training_llm/models/gemma4_vl/gemma4_vision_tower.py @@ -0,0 +1,164 @@ +"""Gemma4 vision tower. + +Verbatim port of HF ``Gemma4VisionEncoder`` + ``Gemma4VisionModel`` +(modeling_gemma4 lines 983-2020) bundled into a single Megatron-side +``VisionModule`` subclass. + +Forward signature matches HF exactly: + forward(pixel_values, pixel_position_ids) -> hidden_states +where: + pixel_values: [B, num_patches, in_channels * patch_size**2] (flattened patches) + pixel_position_ids: [B, num_patches, 2] with (-1, -1) marking padding patches. + +The output is a *flat* tensor ``[total_valid_soft_tokens, hidden_size]`` where +padding has been stripped via ``hidden_states[pooler_mask]``. This matches HF +and is what ``Gemma4Adapter`` (= ``Gemma4MultimodalEmbedder``) expects. The +LLM-side scatter logic in Gemma4VL.forward must use the per-image patch counts +to slice this flat tensor back into per-sample chunks before injecting into +the LLM input embeddings. + +Two design points worth flagging for the converter (Phase C): + +1. ``std_bias`` and ``std_scale`` are registered ONLY when + ``config.standardize=True``. Gemma4-26B-A4B-it has ``standardize=true``, + so both buffers exist as real ckpt entries. The Megatron checkpoint MUST + include them as buffer tensors (non-persistent=False), and the converter + must round-trip them. They live on the VisionModel itself, NOT inside + ``patch_embedder`` (a common misreading from earlier OV2 ports). + +2. The encoder runs a non-causal additive attention mask derived from + ``~padding_positions``. We build it here as a ``[B, 1, 1, N]`` broadcast + mask (valid=0, pad=finfo.min), avoiding the full ``[B, 1, N, N]`` materialize + that ``create_bidirectional_mask`` produces in HF. The arithmetic is + identical because the Gemma vision encoder applies the mask on the + key dimension only (no cross-sample masking, no per-q masking). +""" + +from __future__ import annotations + +import torch +from megatron.core.models.common.vision_module.vision_module import VisionModule +from megatron.core.transformer.transformer_config import TransformerConfig +from torch import nn + +from .gemma4_vision_layer import Gemma4VisionEncoderLayer +from .gemma4_vision_patch_embed import Gemma4VisionPatchEmbedder +from .gemma4_vision_pooler import Gemma4VisionPooler +from .gemma4_vision_rotary import Gemma4VisionRotaryEmbedding + + +class Gemma4VisionTower(VisionModule): + def __init__( + self, + transformer_config: TransformerConfig, + hidden_size: int, + num_hidden_layers: int, + num_attention_heads: int, + num_key_value_heads: int, + head_dim: int, + intermediate_size: int, + patch_size: int, + in_channels: int, + position_embedding_size: int, + pooling_kernel_size: int, + rope_theta: float = 100.0, + rms_norm_eps: float = 1e-6, + attention_dropout: float = 0.0, + hidden_activation: str = "gelu_pytorch_tanh", + use_clipped_linears: bool = False, + standardize: bool = True, + ) -> None: + super().__init__(config=transformer_config) + self.hidden_size = hidden_size + self.num_hidden_layers = num_hidden_layers + self.patch_size = patch_size + self.pooling_kernel_size = pooling_kernel_size + self.standardize = standardize + + self.patch_embedder = Gemma4VisionPatchEmbedder( + patch_size=patch_size, + hidden_size=hidden_size, + position_embedding_size=position_embedding_size, + in_channels=in_channels, + ) + self.rotary_emb = Gemma4VisionRotaryEmbedding( + head_dim=head_dim, + rope_theta=rope_theta, + ndim=2, + ) + self.layers = nn.ModuleList( + [ + Gemma4VisionEncoderLayer( + hidden_size=hidden_size, + num_attention_heads=num_attention_heads, + num_key_value_heads=num_key_value_heads, + head_dim=head_dim, + intermediate_size=intermediate_size, + rms_norm_eps=rms_norm_eps, + attention_dropout=attention_dropout, + hidden_activation=hidden_activation, + use_clipped_linears=use_clipped_linears, + layer_idx=i, + ) + for i in range(num_hidden_layers) + ] + ) + self.pooler = Gemma4VisionPooler(hidden_size=hidden_size) + + if self.standardize: + self.register_buffer( + "std_bias", torch.empty(hidden_size), persistent=True + ) + self.register_buffer( + "std_scale", torch.empty(hidden_size), persistent=True + ) + + @staticmethod + def _build_additive_mask( + valid_mask: torch.Tensor, dtype: torch.dtype + ) -> torch.Tensor: + neg_inf = torch.finfo(dtype).min + additive = torch.zeros_like(valid_mask, dtype=dtype) + additive = additive.masked_fill(~valid_mask, neg_inf) + return additive[:, None, None, :] + + def forward( + self, + pixel_values: torch.Tensor, + pixel_position_ids: torch.Tensor, + ) -> torch.Tensor: + output_length = pixel_values.shape[-2] // ( + self.pooling_kernel_size * self.pooling_kernel_size + ) + + padding_positions = (pixel_position_ids == -1).all(dim=-1) + valid_positions = ~padding_positions + + inputs_embeds = self.patch_embedder( + pixel_values, pixel_position_ids, padding_positions + ) + attention_mask = self._build_additive_mask(valid_positions, inputs_embeds.dtype) + position_embeddings = self.rotary_emb(inputs_embeds, pixel_position_ids) + + hidden_states = inputs_embeds + for layer in self.layers: + hidden_states = layer( + hidden_states=hidden_states, + position_embeddings=position_embeddings, + position_ids=pixel_position_ids, + attention_mask=attention_mask, + ) + + hidden_states, pooler_mask = self.pooler( + hidden_states=hidden_states, + pixel_position_ids=pixel_position_ids, + padding_positions=padding_positions, + output_length=output_length, + ) + + hidden_states = hidden_states[pooler_mask] + + if self.standardize: + hidden_states = (hidden_states - self.std_bias) * self.std_scale + + return hidden_states diff --git a/aiak_training_llm/models/gemma4_vl/gemma4_vl_attention_mask.py b/aiak_training_llm/models/gemma4_vl/gemma4_vl_attention_mask.py new file mode 100644 index 00000000..71d3c3e5 --- /dev/null +++ b/aiak_training_llm/models/gemma4_vl/gemma4_vl_attention_mask.py @@ -0,0 +1,126 @@ +"""Gemma4-VL multimodal attention mask with bidirectional vision-span overlay. + +This module implements HuggingFace ``create_causal_mask_mapping`` semantics for +Megatron's ``TransformerBlock.forward(attention_mask=...)`` contract: + + 1. Base causal mask (lower-triangular True-is-VISIBLE inside the helper, then + inverted to Megatron True-is-MASKED at return). + + 2. Vision overlay: tokens belonging to the same vision span (contiguous run of + ``mm_token_type_ids in {1, 2}``) attend bidirectionally to one another. + Mathematically ``mask = causal | (q_group == kv_group) & (q_group >= 0)``, + mirroring HF ``token_type_ids_mask_function`` and ``or_masks`` composition + in ``transformers.models.gemma4.modeling_gemma4``. + +The helper returns a SINGLE BoolTensor ``[B, 1, S, S]`` (True = MASKED), shared +by every ``TransformerLayer``. Sliding layers further restrict this base mask at +runtime via TE's ``window_size=(sliding_window-1, 0)`` parameter, which is wired +per-layer by ``gemma4_attention.py:_resolve_per_layer_overrides``. Because each +Gemma4 vision span is bounded by the ViT grid (16x16 = 256 patches) and the +sliding window is 1024, the OR overlay is always strictly contained inside the +sliding window — TE's hard window cutoff therefore composes losslessly with the +overlay, recovering exact HF semantics WITHOUT a per-layer mask dict. + +A runtime guard (``_assert_vision_spans_within_window``) enforces the +``max_vision_span <= sliding_window`` invariant and fails fast if a future video +packing scheme produces a span that would break the equivalence. + +External callers: ``Gemma4VLModel.forward`` (P5.5) — invoked when the batch +contains vision tokens. Pure-text batches skip this helper and pass the stock +causal mask produced by the data path. +""" + +from __future__ import annotations + +import torch + + +VISION_TOKEN_TYPE_IDS: tuple[int, ...] = (1, 2) + + +def _compute_vision_group_ids(mm_token_type_ids: torch.Tensor) -> torch.Tensor: + """Assign each token a vision-span group id (text tokens get -1). + + Mirrors HF ``get_block_sequence_ids_for_mask`` (modeling_gemma4.py:2078-2087): + a new group starts whenever a vision token follows a non-vision token. + """ + is_vision = torch.zeros_like(mm_token_type_ids, dtype=torch.bool) + for token_id in VISION_TOKEN_TYPE_IDS: + is_vision = is_vision | (mm_token_type_ids == token_id) + + is_prev_vision = torch.roll(is_vision, shifts=1, dims=-1) + is_prev_vision[..., 0] = False + new_starts = is_vision & ~is_prev_vision + vision_group_ids = torch.cumsum(new_starts.to(torch.long), dim=-1) - 1 + return torch.where(is_vision, vision_group_ids, torch.full_like(vision_group_ids, -1)) + + +def _assert_vision_spans_within_window( + vision_group_ids: torch.Tensor, + sliding_window: int, +) -> None: + """Fail fast if any vision span exceeds the sliding window. + + Gemma4 single-image spans are physically bounded by ViT grid (<=256 patches), + so this guard only triggers on hypothetical future video packing. Without it, + TE's hard sliding cutoff would silently truncate the OR overlay, producing + NaN-equivalent silent divergence from HF rather than a fail-fast error. + """ + valid_mask = vision_group_ids >= 0 + if not valid_mask.any(): + return + flat_groups = vision_group_ids[valid_mask] + span_lengths = torch.bincount(flat_groups) + max_span = int(span_lengths.max().item()) + assert max_span <= sliding_window, ( + f"Gemma4-VL vision span length {max_span} exceeds sliding_window {sliding_window}; " + "the OR-overlay no longer composes losslessly with TE per-layer window cutoff. " + "Either disable sliding window for vision batches or shorten the span." + ) + + +def build_gemma4_mm_attention_mask( + input_ids: torch.Tensor, + mm_token_type_ids: torch.Tensor, + sliding_window: int, +) -> torch.Tensor: + """Build a Megatron-format multimodal attention mask for Gemma4-VL. + + Args: + input_ids: ``[B, S]`` token ids (only shape and device are used). + mm_token_type_ids: ``[B, S]`` int tensor with values in + ``{0=text, 1=image, 2=video}``. + sliding_window: Sliding window size from ``Gemma4VLConfig.sliding_window`` + (1024 for 26B-A4B). Used only by the runtime guard; sliding layers + apply the cutoff themselves via TE ``window_size``. + + Returns: + BoolTensor ``[B, 1, S, S]`` with ``True = MASKED`` (Megatron convention). + Layers select between the full and sliding behavior via TE's per-layer + ``window_size``; this single mask serves both. + """ + if input_ids.shape != mm_token_type_ids.shape: + raise ValueError( + f"input_ids shape {tuple(input_ids.shape)} does not match " + f"mm_token_type_ids shape {tuple(mm_token_type_ids.shape)}" + ) + if input_ids.dim() != 2: + raise ValueError(f"input_ids must be 2D [B, S]; got shape {tuple(input_ids.shape)}") + + batch_size, seq_len = input_ids.shape + device = input_ids.device + + causal_visible = torch.tril( + torch.ones(seq_len, seq_len, dtype=torch.bool, device=device) + ).unsqueeze(0).expand(batch_size, -1, -1) + + vision_group_ids = _compute_vision_group_ids(mm_token_type_ids.to(device)) + _assert_vision_spans_within_window(vision_group_ids, sliding_window) + + q_groups = vision_group_ids.unsqueeze(-1) + kv_groups = vision_group_ids.unsqueeze(-2) + vision_overlay_visible = (q_groups == kv_groups) & (q_groups >= 0) + + visible = causal_visible | vision_overlay_visible + masked = ~visible + return masked.unsqueeze(1) diff --git a/aiak_training_llm/models/gemma4_vl/gemma4_vl_model.py b/aiak_training_llm/models/gemma4_vl/gemma4_vl_model.py new file mode 100644 index 00000000..4a1c0223 --- /dev/null +++ b/aiak_training_llm/models/gemma4_vl/gemma4_vl_model.py @@ -0,0 +1,331 @@ +"""Top-level Gemma4-VL VLM (Gemma4 vision tower + Gemma4 adapter + Gemma4 LLM). + +Subclasses :class:`LlavaOnevision2` so the provider/freeze/set_input_tensor +plumbing is reused, but overrides ``__init__`` and ``forward`` to swap in +the HF-verbatim Gemma4 vision tower / adapter and the Gemma4 MoE LLM. + +``vision_layer_spec`` and ``adapter_layer_spec`` arguments are accepted for +provider-call uniformity but are IGNORED — the Gemma4 vision tower / adapter +are plain ``nn.Module`` stacks (see ``gemma4_vision_tower.py`` / +``gemma4_adapter.py``), not Megatron ``ModuleSpec``-driven blocks. + +The vision data path (``images is not None``) currently raises +``NotImplementedError``: the production Gemma4 image processor / dataloader +that produces ``pixel_values`` + ``pixel_position_ids`` in the HF layout the +tower expects has not yet been wired into the OV2 dataloader (P3.x). LM-only +forward (``images is None``, P1 smoke + pure-text training) is fully +functional. +""" + +import logging +from functools import partial +from typing import Optional + +import torch +from megatron.core import InferenceParams, parallel_state, tensor_parallel +from megatron.core.packed_seq_params import PackedSeqParams +from megatron.core.transformer.enums import AttnMaskType +from megatron.core.transformer.spec_utils import ModuleSpec + +from aiak_training_llm.models.gemma4_vl.gemma4_adapter import Gemma4Adapter +from aiak_training_llm.models.gemma4_vl.gemma4_model import Gemma4Model +from aiak_training_llm.models.gemma4_vl.gemma4_vision_tower import Gemma4VisionTower +from aiak_training_llm.models.llava_onevision2.llava_onevision2_model import ( + LlavaOnevision2, + _load_state_dict_hook_ignore_param_names, +) + + +class Gemma4VL(LlavaOnevision2): + def __init__( + self, + language_config, + vision_config, + adapter_config, + language_layer_spec: ModuleSpec, + vision_layer_spec: ModuleSpec, + adapter_layer_spec: ModuleSpec, + language_vocab_size: int, + language_max_sequence_length: int, + allow_missing_adapter_checkpoint: bool = False, + parallel_output: bool = True, + language_position_embedding_type: str = "rope", + language_rotary_percent: float = 1.0, + pre_process: bool = True, + post_process: bool = True, + add_encoder: bool = True, + add_decoder: bool = True, + language_rotary_base: int = 10000, + language_rotary_base_sliding: Optional[int] = None, + fp16_lm_cross_entropy: bool = False, + share_embeddings_and_output_weights: bool = True, + seq_len_interpolation_factor: Optional[float] = None, + ) -> None: + del vision_layer_spec, adapter_layer_spec + + from megatron.core.transformer import MegatronModule + + MegatronModule.__init__(self, config=language_config) + + self.pre_process = pre_process + self.post_process = post_process + self.add_encoder = add_encoder + self.add_decoder = add_decoder + + self.encoder_hidden_state = None + self.vision_model = None + self.adapter = None + self.language_model = None + + if self.add_encoder: + self.vision_model = Gemma4VisionTower( + transformer_config=vision_config, + hidden_size=vision_config.hidden_size, + num_hidden_layers=vision_config.num_hidden_layers, + num_attention_heads=vision_config.num_attention_heads, + num_key_value_heads=vision_config.num_key_value_heads, + head_dim=vision_config.head_dim, + intermediate_size=vision_config.intermediate_size, + patch_size=vision_config.patch_size, + in_channels=vision_config.in_channels, + position_embedding_size=vision_config.position_embedding_size, + pooling_kernel_size=vision_config.pooling_kernel_size, + rope_theta=vision_config.rope_theta, + rms_norm_eps=vision_config.rms_norm_eps, + hidden_activation=vision_config.hidden_activation, + use_clipped_linears=vision_config.use_clipped_linears, + standardize=vision_config.standardize, + ) + self.adapter = Gemma4Adapter( + vision_hidden_size=vision_config.hidden_size, + text_hidden_size=language_config.hidden_size, + rms_norm_eps=adapter_config.layernorm_epsilon, + ) + if allow_missing_adapter_checkpoint: + adapter_param_names = [ + f"adapter.{name}" for name in self.adapter.state_dict().keys() + ] + self.adapter.register_load_state_dict_post_hook( + partial(_load_state_dict_hook_ignore_param_names, adapter_param_names) + ) + + if self.add_decoder: + self.language_model = Gemma4Model( + config=language_config, + transformer_layer_spec=language_layer_spec, + vocab_size=language_vocab_size, + max_sequence_length=language_max_sequence_length, + parallel_output=parallel_output, + position_embedding_type=language_position_embedding_type, + rotary_percent=language_rotary_percent, + pre_process=self.pre_process, + post_process=self.post_process, + rotary_base=language_rotary_base, + rotary_base_sliding=language_rotary_base_sliding, + share_embeddings_and_output_weights=share_embeddings_and_output_weights, + scatter_embedding_sequence_parallel=False, + fp16_lm_cross_entropy=fp16_lm_cross_entropy, + seq_len_interpolation_factor=seq_len_interpolation_factor, + ) + self.share_embeddings_and_output_weights = ( + self.language_model.share_embeddings_and_output_weights + ) + + def set_input_tensor(self, input_tensor) -> None: + if not isinstance(input_tensor, list): + input_tensor = [input_tensor] + assert len(input_tensor) == 1, ( + "input_tensor should only be length 1 for Gemma4-VL" + ) + + if self.add_encoder and self.add_decoder: + raise NotImplementedError( + "Gemma4-VL encoder-stage set_input_tensor requires Gemma4VisionTower " + "to expose a Megatron-style set_input_tensor; deferred to P3.x along " + "with the vision data path." + ) + if self.add_encoder: + raise NotImplementedError( + "Gemma4-VL encoder-only pipeline stage requires Gemma4VisionTower " + "to expose a Megatron-style set_input_tensor; deferred to P3.x." + ) + if self.pre_process: + self.encoder_hidden_state = input_tensor[0] + else: + self.language_model.set_input_tensor(input_tensor[0]) + + def forward_debug(self, *args, **kwargs): + raise NotImplementedError( + "Gemma4-VL forward_debug is deferred to P3.x; the inherited " + "LlavaOnevision2.forward_debug calls Gemma4VisionTower.forward_debug " + "which does not exist (the Gemma4 vision tower is a plain nn.Module, " + "not a Megatron block). The HF↔Megatron consistency check in Phase C " + "must call Gemma4VisionTower.forward / Gemma4Adapter.forward directly." + ) + + @staticmethod + def _prepare_vision_inputs( + images: torch.Tensor, + image_grid_thw: torch.Tensor, + patch_positions: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + patch_counts = image_grid_thw.prod(dim=-1).to(torch.long) + if int(patch_counts.sum().item()) != images.shape[0]: + raise ValueError( + f"Gemma4-VL image patch count mismatch: images has {images.shape[0]} " + f"patches but image_grid_thw sums to {int(patch_counts.sum().item())}." + ) + + max_patches = int(patch_counts.max().item()) + batch_size = int(patch_counts.numel()) + pixel_values = images.new_zeros((batch_size, max_patches, images.shape[-1])) + pixel_position_ids = torch.full( + (batch_size, max_patches, 2), + -1, + dtype=torch.long, + device=images.device, + ) + + offset = 0 + for batch_idx, patch_count_tensor in enumerate(patch_counts): + patch_count = int(patch_count_tensor.item()) + next_offset = offset + patch_count + pixel_values[batch_idx, :patch_count] = images[offset:next_offset] + + if patch_positions is None: + _t, h, w = image_grid_thw[batch_idx].tolist() + h_coords = torch.arange(h, dtype=torch.long, device=images.device).repeat_interleave(w) + w_coords = torch.arange(w, dtype=torch.long, device=images.device).repeat(h) + coords = torch.stack((w_coords, h_coords), dim=-1) + if coords.shape[0] != patch_count: + raise ValueError( + f"Default Gemma4-VL pixel positions only support single-frame images; " + f"got image_grid_thw={image_grid_thw[batch_idx].tolist()}." + ) + else: + coords = patch_positions[offset:next_offset, -2:].to(device=images.device, dtype=torch.long) + coords = torch.stack((coords[:, 1], coords[:, 0]), dim=-1) + + pixel_position_ids[batch_idx, :patch_count] = coords + offset = next_offset + + return pixel_values, pixel_position_ids + + def forward( + self, + images: torch.Tensor, + image_grid_thw: torch.Tensor, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + attention_mask: torch.Tensor, + attn_mask_type: AttnMaskType | None = None, + labels: torch.Tensor = None, + packed_seq_params: PackedSeqParams = None, + inference_params: InferenceParams = None, + pixel_values_videos: torch.Tensor = None, + video_grid_thw: torch.Tensor = None, + patch_positions: list[torch.Tensor] | None = None, + ) -> torch.Tensor: + del position_ids + + if pixel_values_videos is not None or video_grid_thw is not None: + raise NotImplementedError( + "Gemma4-VL video path not implemented; pixel_values_videos / " + "video_grid_thw must be None." + ) + + use_inference_kv_cache = ( + inference_params is not None + and "image_tokens_count" in inference_params.key_value_memory_dict + ) + + if images is not None and self.add_encoder and not use_inference_kv_cache: + if image_grid_thw is None: + raise ValueError("Gemma4-VL image_grid_thw is required when images are provided.") + pixel_values, pixel_position_ids = self._prepare_vision_inputs( + images, image_grid_thw, patch_positions + ) + image_embeddings = self.vision_model(pixel_values, pixel_position_ids) + image_embeddings = self.adapter(image_embeddings) + + n_image_tokens = (input_ids == self.config.image_token_id).sum().item() + n_image_features = image_embeddings.shape[0] + if n_image_features > n_image_tokens: + logging.getLogger(__name__).warning( + "Trimming %d extra Gemma4 image embedding(s) " + "(n_image_features=%d, n_image_tokens=%d).", + n_image_features - n_image_tokens, + n_image_features, + n_image_tokens, + ) + image_embeddings = image_embeddings[:n_image_tokens] + elif n_image_features < n_image_tokens: + raise ValueError( + f"Gemma4 image features {n_image_features} < image tokens {n_image_tokens}" + ) + + if inference_params is not None: + inference_params.key_value_memory_dict["image_tokens_count"] = image_embeddings.shape[0] + else: + image_embeddings = None + + if not self.add_decoder: + raise NotImplementedError( + "Gemma4-VL encoder-only pipeline stage requires the vision data " + "path which is not yet implemented; see images-not-None branch." + ) + + if self.pre_process: + language_embeddings = self.language_model.embedding( + input_ids=input_ids, position_ids=None + ) + if use_inference_kv_cache or images is None: + combined_embeddings = language_embeddings + elif (input_ids == self.config.image_token_id).any().item(): + images_mask = ( + (input_ids == self.config.image_token_id) + .transpose(0, 1) + .unsqueeze(-1) + .expand_as(language_embeddings) + .to(language_embeddings.device) + ) + image_embeddings = image_embeddings.to( + language_embeddings.device, language_embeddings.dtype + ) + combined_embeddings = language_embeddings.masked_scatter( + images_mask, image_embeddings + ) + else: + combined_embeddings = language_embeddings + + if self.config.sequence_parallel: + seq_len = combined_embeddings.size(0) + tp_world_size = parallel_state.get_tensor_model_parallel_world_size() + remainder = seq_len % tp_world_size + if remainder != 0: + pad = tp_world_size - remainder + pad_shape = (pad,) + combined_embeddings.shape[1:] + pad_tensor = combined_embeddings.new_zeros(pad_shape) + combined_embeddings = torch.cat( + (combined_embeddings, pad_tensor), dim=0 + ) + combined_embeddings = ( + tensor_parallel.scatter_to_sequence_parallel_region( + combined_embeddings + ) + ) + else: + combined_embeddings = None + + return self.language_model( + input_ids=None, + position_ids=None, + attention_mask=attention_mask, + attn_mask_type=attn_mask_type, + decoder_input=combined_embeddings, + labels=labels, + rotary_pos_emb=None, + inference_params=inference_params, + packed_seq_params=packed_seq_params, + extra_block_kwargs={}, + ) From 2e0bae331d651ca59ec6c7b010dec7868c1374ba Mon Sep 17 00:00:00 2001 From: XiangAn Date: Sun, 17 May 2026 23:35:50 +0800 Subject: [PATCH 4/6] feat: add Gemma4 VL data pipeline Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- aiak_training_llm/data/chat_templete.py | 17 +- aiak_training_llm/data/mm_plugin.py | 177 ++++ .../data/multimodal/gemma4_vl_task_encoder.py | 934 ++++++++++++++++++ .../train/pretrain/pretrain_gemma4_vl.py | 406 ++++++++ aiak_training_llm/train/sft/sft_gemma4_vl.py | 25 + 5 files changed, 1557 insertions(+), 2 deletions(-) create mode 100644 aiak_training_llm/data/multimodal/gemma4_vl_task_encoder.py create mode 100644 aiak_training_llm/train/pretrain/pretrain_gemma4_vl.py create mode 100644 aiak_training_llm/train/sft/sft_gemma4_vl.py diff --git a/aiak_training_llm/data/chat_templete.py b/aiak_training_llm/data/chat_templete.py index 82262cc4..8eebdefe 100644 --- a/aiak_training_llm/data/chat_templete.py +++ b/aiak_training_llm/data/chat_templete.py @@ -13,10 +13,11 @@ import re from abc import ABC, abstractmethod from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Type, Dict, List, Optional, Sequence, Set, Tuple, Union +from typing import TYPE_CHECKING, Dict, List, Optional, Sequence, Set, Tuple, Type, Union from aiak_training_llm.utils.constants import DataRoles -from .mm_plugin import MMPlugin, Qwen2VLPlugin + +from .mm_plugin import Gemma4VLPlugin, MMPlugin, Qwen2VLPlugin if TYPE_CHECKING: @@ -459,3 +460,15 @@ def get_support_templates() -> List[str]: format_user=StringFormatter(slots=["<｜User｜>{{content}}<｜Assistant｜>"]), format_prefix=EmptyFormatter(slots=[{"bos_token"}]), ) + +_register_chat_template( + name="gemma4", + format_user=StringFormatter( + slots=["<|turn>user\n{{content}}\n<|turn>model\n"] + ), + format_assistant=StringFormatter(slots=["{{content}}\n"]), + format_system=StringFormatter(slots=["<|turn>system\n{{content}}\n"]), + format_separator=EmptyFormatter(slots=[""]), + format_prefix=EmptyFormatter(slots=[{"bos_token"}]), + mm_plugin=Gemma4VLPlugin(image_token="<|image|>", video_token=None), +) diff --git a/aiak_training_llm/data/mm_plugin.py b/aiak_training_llm/data/mm_plugin.py index 6438e282..09a7d57d 100644 --- a/aiak_training_llm/data/mm_plugin.py +++ b/aiak_training_llm/data/mm_plugin.py @@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Dict, List, Optional, Sequence, Tuple, Type, TypedDict, Union import numpy as np +import torch from PIL import Image from PIL.Image import Image as ImageObject from typing_extensions import override @@ -16,6 +17,7 @@ import torch from transformers.image_processing_utils import BaseImageProcessor + from transformers.processing_utils import ProcessorMixin class EncodedImage(TypedDict): """Encoded image type.""" @@ -292,3 +294,178 @@ def get_mm_inputs( ) -> Dict[str, Union[List[int], "torch.Tensor"]]: self._validate_input(images, videos) return self._get_mm_inputs(images, videos, processor) + + +class Gemma4VLPlugin(MMPlugin): + """Gemma4-VL passthrough plugin: process_messages returns (messages, mm_inputs) + where mm_inputs follows the existing OV2 flattened-patch contract. + + Cross-module contract (downstream consumers depend on these invariants): + - ``pixel_values`` shape ``[total_imgs_in_batch, P, D]`` — FLATTENED, NOT ``[B, ...]``. + Indexing by batch index will silently return wrong tensor. + - ``image_grid_thw`` is synthesized from HF ``image_position_ids`` as + ``[num_images, 3]`` rows of ``[1, H_p, W_p]``. + - Text-only batches OMIT multimodal keys entirely (no zero-shape sentinel). + All downstream consumers MUST guard with ``in mm_inputs``. + """ + + @staticmethod + def _flatten_gemma4_image_outputs( + image_outputs: dict[str, "torch.Tensor"], + ) -> dict[str, Union["torch.Tensor", list["torch.Tensor"]]]: + pixel_values = image_outputs["pixel_values"] + image_position_ids = image_outputs["image_position_ids"] + num_soft_tokens_per_image = image_outputs["num_soft_tokens_per_image"] + + valid_mask = (image_position_ids != -1).all(dim=-1) + flat_pixel_values = pixel_values[valid_mask] + + image_grid_rows: list[list[int]] = [] + patch_positions: list[torch.Tensor] = [] + for image_idx in range(image_position_ids.shape[0]): + valid_positions = image_position_ids[image_idx][valid_mask[image_idx]].to(dtype=torch.int64) + if valid_positions.numel() == 0: + raise ValueError(f"Gemma4 image {image_idx} has no valid patch positions.") + + width = int(valid_positions[:, 0].max().item()) + 1 + height = int(valid_positions[:, 1].max().item()) + 1 + patch_count = int(valid_positions.shape[0]) + if patch_count != height * width: + raise ValueError( + "Gemma4 image patch positions are not a dense single-frame grid: " + f"image_idx={image_idx}, patch_count={patch_count}, height={height}, width={width}." + ) + + image_grid_rows.append([1, height, width]) + patch_positions.append( + torch.stack( + ( + torch.zeros(patch_count, dtype=torch.int64, device=valid_positions.device), + valid_positions[:, 1], + valid_positions[:, 0], + ), + dim=-1, + ) + ) + + image_grid_thw = torch.tensor( + image_grid_rows, + dtype=torch.int32, + device=image_position_ids.device, + ) + + return { + "pixel_values": flat_pixel_values, + "image_grid_thw": image_grid_thw, + "patch_positions": patch_positions, + "image_position_ids": image_position_ids, + "num_soft_tokens_per_image": num_soft_tokens_per_image, + } + + def _build_gemma4_mm_inputs( + self, + images: Sequence["ImageInput"], + processor: Optional["ProcessorMixin"], + ) -> tuple[Optional[list["ImageObject"]], dict[str, Union["torch.Tensor", list["torch.Tensor"]]]]: + regularized_images = self._regularize_images(images) if len(images) != 0 else None + if regularized_images is None: + return None, {} + + image_outputs = processor.image_processor(regularized_images, return_tensors="pt") + return regularized_images, self._flatten_gemma4_image_outputs(dict(image_outputs)) + + def _expand_image_placeholders( + self, + messages: Sequence[dict[str, str]], + num_soft_tokens_per_image: Sequence[int], + processor: Optional["ProcessorMixin"], + ) -> list[dict[str, str]]: + messages = deepcopy(messages) + actual_num_images = len(num_soft_tokens_per_image) + + image_placeholder_count = sum(message["content"].count(Placeholder.IMAGE) for message in messages) + if actual_num_images > 0 and image_placeholder_count != actual_num_images: + for message in messages: + message["content"] = message["content"].replace(Placeholder.IMAGE, "") + + first_user_msg = None + for message in messages: + if message.get("role") == "user": + first_user_msg = message + break + + if first_user_msg is None: + raise ValueError("Cannot rebuild Gemma4 image placeholders: no user message found.") + + image_placeholders = "\n".join([Placeholder.IMAGE] * actual_num_images) + user_content = first_user_msg["content"].lstrip("\n") + first_user_msg["content"] = f"{image_placeholders}\n{user_content}" + + image_idx = 0 + for message in messages: + content = message["content"] + while Placeholder.IMAGE in content: + if image_idx >= actual_num_images: + raise ValueError( + f"The number of {Placeholder.IMAGE} tokens is greater than available images." + ) + + n_soft_tokens = int(num_soft_tokens_per_image[image_idx]) + replacement = ( + f"{processor.boi_token}{self.image_token * n_soft_tokens}{processor.eoi_token}" + ) + content = content.replace(Placeholder.IMAGE, replacement, 1) + image_idx += 1 + + if Placeholder.VIDEO in content: + raise ValueError("Gemma4-VL video placeholders are not supported in this OV2 path yet.") + + message["content"] = content + + if image_idx != actual_num_images: + raise ValueError( + f"The number of images ({actual_num_images}) does not match expanded placeholders ({image_idx})." + ) + + return messages + + @override + def _preprocess_image(self, image: "ImageObject", **kwargs) -> "ImageObject": + return super()._preprocess_image(image, **kwargs) + + @override + def process_messages( + self, + messages: Sequence[dict[str, str]], + images: Sequence["ImageInput"], + videos: Sequence["VideoInput"], + processor: Optional["ProcessorMixin"], + ) -> tuple[list[dict[str, str]], dict[str, "torch.Tensor"]]: + self._validate_input(images, videos) + _regularized_images, mm_inputs = self._build_gemma4_mm_inputs(images, processor) + + if "num_soft_tokens_per_image" in mm_inputs: + messages = self._expand_image_placeholders( + messages, + mm_inputs["num_soft_tokens_per_image"], + processor, + ) + else: + messages = list(messages) + + return messages, dict(mm_inputs) + + @override + def get_mm_inputs( + self, + images: Sequence["ImageInput"], + videos: Sequence["VideoInput"], + imglens: Sequence[int], + vidlens: Sequence[int], + seqlens: Sequence[int], + processor: Optional["ProcessorMixin"], + ) -> dict[str, Union[list[int], "torch.Tensor"]]: + self._validate_input(images, videos) + del imglens, vidlens, seqlens + _regularized_images, mm_inputs = self._build_gemma4_mm_inputs(images, processor) + return dict(mm_inputs) diff --git a/aiak_training_llm/data/multimodal/gemma4_vl_task_encoder.py b/aiak_training_llm/data/multimodal/gemma4_vl_task_encoder.py new file mode 100644 index 00000000..0f47fa13 --- /dev/null +++ b/aiak_training_llm/data/multimodal/gemma4_vl_task_encoder.py @@ -0,0 +1,934 @@ +"""Gemma4VLTaskEncoder class. + +Verbatim mirror of :mod:`qwen2vl_task_encoder` for the Gemma4-VL family. +The structural pipeline (mm_plugin → encode_multiturn → packing → batching) is +identical; only Gemma4-specific surface differs: + +- Special-token strings (Gemma4 vocab is unrelated to Qwen): + * boi (vision_start) -> ``<|image>`` (id=255999) + * image pad -> ``<|image|>`` (id=258880) + * eoi (vision_end) -> ```` (id=258882) + * video -> ``<|video|>`` (id=258884) +- Spatial merge: Gemma4 uses ``image_processor.pooling_kernel_size=3`` instead of + Qwen's ``merge_size=2``; the existing ``getattr(..., 'merge_size', 2)`` lookup + reads the wrong attribute on Gemma4Processor so we resolve it via a small + helper (``_gemma4_spatial_merge_size``) that prefers ``pooling_kernel_size``. +- ``smart_resize`` factor: Qwen uses 28 (= patch_size 14 × merge 2); Gemma4 uses + patch_size 16 × pooling 3 = 48. We pass ``size_factor=48`` through. + +We keep this as a fork rather than a subclass per user direction (avoids leaky +class-attribute overrides on the parent Qwen module-level constants). + +CAVEAT — image placeholder wrapping mismatch (deferred fix, ack'd in P4.9): + Qwen2VL inline format is ``<|vision_start|><|image_pad|><|vision_end|>`` (3 + tokens, fence + pad + fence). Gemma4 chat-template instead emits a single + ``<|image|>`` (no boi/eoi wrap). Our verbatim mirror keeps the 3-token wrap + using Gemma4 ids — correct for fence semantics but produces 1 extra image + placeholder vs the reference Gemma4 jinja path. Forward will run; lm_loss + numerics will be off until a future revision rewrites the + ````-substitution path to call ``processor.apply_chat_template`` + directly. Tracked as a Phase-P4.9 known issue. +""" + +import re +import sys +from collections.abc import Callable +from dataclasses import dataclass +from typing import Optional, TypeVar, Union + +import numpy as np +import torch +from megatron.energon import CaptioningSample, SkipSample, VQASample +from megatron.energon.flavors.base_dataset import ( + BaseCoreDatasetFactory, + SavableDataset, +) +from megatron.energon.flavors.crude import CrudeWebdataset +from megatron.energon.flavors.webdataset import VideoData +from megatron.energon.metadataset.loader_interface import DatasetBlendMode +from megatron.energon.task_encoder.base import stateless +from megatron.energon.worker import WorkerConfig +from megatron.energon.wrappers import ( + BlendDataset, + EpochizeDataset, + LogSampleDataset, + ShuffleBufferDataset, +) +from megatron.energon.wrappers.repeat_dataset import RepeatDataset +from qwen_vl_utils.vision_process import smart_nframes, smart_resize +from torchvision import transforms +from torchvision.transforms import InterpolationMode +from typing_extensions import override + +from aiak_training_llm.data.multimodal import MultiMixQASample +from aiak_training_llm.data.multimodal.length_sort_dataset import LengthPoolSortDataset +from aiak_training_llm.data.multimodal.packed_sort_dataset import PackedSeparateSortDataset +from aiak_training_llm.utils import constants, get_chat_template +from transformers import AutoProcessor + +from .task_encoder import ImageTaskBatchPacked, ImageTaskSample, ImageTaskSamplePacked, TaskEncoder + + +T = TypeVar("T") +V = TypeVar("V") +T_sample = TypeVar("T_sample") +T_encoded_sample = TypeVar("T_encoded_sample") +T_raw_batch = TypeVar("T_raw_batch") +T_batch = TypeVar("T_batch") + + +IGNORE_INDEX = -100 # ID for labels that should be ignored. +IMAGE_TOKEN = "<|image|>" +VIDEO_TOKEN = "<|video|>" +VISION_TAGS = ["<|image>", ""] +IMAGE_TOKEN_WITH_TAGS = VISION_TAGS[0] + IMAGE_TOKEN + VISION_TAGS[1] +VIDEO_TOKEN_WITH_TAGS = VISION_TAGS[0] + VIDEO_TOKEN + VISION_TAGS[1] +SKIP_LOG_LIMIT = 5 +_SKIP_LOG_COUNTS: dict[str, int] = {} + + +def _gemma4_spatial_merge_size(image_processor) -> int: + """Resolve Gemma4's spatial merge factor. + + Gemma4Processor exposes ``pooling_kernel_size`` (default 3); the parent + Qwen lookup ``merge_size`` (default 2) returns the wrong fallback on Gemma4 + and silently corrupts patch_position block layout. Prefer the Gemma4 name + when present. + """ + val = getattr(image_processor, "pooling_kernel_size", None) + if val is not None: + return int(val) + return int(getattr(image_processor, "merge_size", 2)) + + +def skip_malformed_multimodal_sample(sample_key: str, signature: str, detail: str) -> None: + count = _SKIP_LOG_COUNTS.get(signature, 0) + 1 + _SKIP_LOG_COUNTS[signature] = count + + if count <= SKIP_LOG_LIMIT: + print(f"Skipping malformed multimodal sample {sample_key}: {detail}", file=sys.stderr) + if count == SKIP_LOG_LIMIT: + print( + f"Further '{signature}' skip logs will be suppressed for this worker.", + file=sys.stderr, + ) + + raise SkipSample(f"{sample_key}: {detail}") + + +def get_stateless(fn: Callable[..., T_sample]) -> bool: + """Get whether a function is stateless.""" + return getattr(fn, "__stateless__", False) + + +def convert_positions_to_block_layout( + positions: torch.Tensor, t: int, h: int, w: int, spatial_merge_size: int = 2 +) -> torch.Tensor: + """ + Convert patch positions from row-major order to 2x2 block layout. + + This function reorders patch positions to match the 2x2 block arrangement + used by the image processor. Uses index-based reordering instead of reshape. + + Args: + positions: Patch positions in row-major order, shape [t*h*w, 3] + t: temporal dimension + h: height (unmerged patch count) + w: width (unmerged patch count) + spatial_merge_size: size of spatial merge blocks (default: 2) + + Returns: + torch.Tensor: Patch positions in 2x2 block order, same shape [t*h*w, 3] + """ + sms = spatial_merge_size + if sms == 1: + return positions + + device = positions.device + total_patches = t * h * w + + # Generate row-major indices: [0, 1, 2, ..., t*h*w-1] + # Reshape to [t, h, w] + indices = torch.arange(total_patches, device=device).view(t, h, w) + + # Calculate merged dimensions + h_merged = h // sms + w_merged = w // sms + + # Reshape to [t, h_merged, sms, w_merged, sms] + indices = indices.view(t, h_merged, sms, w_merged, sms) + + # Permute to [t, h_merged, w_merged, sms_h, sms_w] - 2x2 block order + indices = indices.permute(0, 1, 3, 2, 4).contiguous() + + # Flatten to get the reordering indices + indices = indices.view(total_patches) + + # Apply the reordering to positions + return positions[indices] + + +@dataclass +class Gemma4VLImageTaskSample(ImageTaskSample): + """An image task sample with a grid of tokens and their corresponding pixel values.""" + + image_grid_thw: torch.Tensor = None + video_grid_thw: torch.Tensor = None + + def __init__(self, image_grid_thw: str, video_grid_thw=None, **kwargs): + super().__init__(**kwargs) + self.image_grid_thw = image_grid_thw + self.video_grid_thw = video_grid_thw + + +@dataclass +class Gemma4VLImageTaskSamplePacked(ImageTaskSamplePacked): + """An image task sample with a grid of tokens and their corresponding pixel values.""" + + image_grid_thw: torch.Tensor = None + video_grid_thw: torch.Tensor = None + + def __init__(self, sample: ImageTaskSample, image_grid_thw: str, video_grid_thw=None): + super().__init__(**vars(sample)) + self.image_grid_thw = image_grid_thw + self.video_grid_thw = video_grid_thw + + +@dataclass +class Gemma4VLImageTaskBatchPacked(ImageTaskBatchPacked): + """An image task sample with a grid of tokens and their corresponding pixel values.""" + + image_grid_thw: torch.Tensor = None + video_grid_thw: torch.Tensor = None + + def __init__(self, sample: ImageTaskSample, image_grid_thw: str, video_grid_thw=None): + super().__init__(**vars(sample)) + self.image_grid_thw = image_grid_thw + self.video_grid_thw = video_grid_thw + + +class Gemma4VLTaskEncoder(TaskEncoder): + """A simple task encoder for VLMs.""" + + def __init__(self, args): + super().__init__() + if args.training_phase in ["sft"]: + self.chat_template = get_chat_template() + self.processor = AutoProcessor.from_pretrained(self.args.hf_tokenizer_path, trust_remote_code=True) + + # image + self.min_pixels = args.min_pixels + self.max_pixels = args.max_pixels + + def _normalize_image_backed_video_placeholders( + self, + messages: list[dict[str, str]], + image: list | None, + video: list | None, + ) -> list[dict[str, str]]: + """Rewrite image-backed