From fdb5153dae2e8d09b4a9b0be423a1202069787fe Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Thu, 20 Aug 2026 03:45:13 +0000 Subject: [PATCH] feat: restore native Gemma4 support Restore Gemma4 with the native model-plugin and direct HF conversion architecture used by Slime before Gemma4 was removed. Signed-off-by: aoshen02 --- scripts/models/gemma4-26B-A4B.sh | 28 + tests/test_hf_to_megatron.py | 1 + tests/utils/test_gemma4.py | 234 ++++ .../megatron_utils/hf_to_megatron/__init__.py | 5 + .../megatron_utils/hf_to_megatron/gemma4.py | 91 ++ .../megatron_utils/megatron_to_hf/__init__.py | 5 +- .../megatron_utils/megatron_to_hf/gemma4.py | 177 +++ vime_plugins/models/gemma4.py | 1188 +++++++++++++++++ vime_plugins/models/gemma4_provider.py | 344 +++++ 9 files changed, 2072 insertions(+), 1 deletion(-) create mode 100644 scripts/models/gemma4-26B-A4B.sh create mode 100644 tests/utils/test_gemma4.py create mode 100644 vime/backends/megatron_utils/hf_to_megatron/gemma4.py create mode 100644 vime/backends/megatron_utils/megatron_to_hf/gemma4.py create mode 100644 vime_plugins/models/gemma4.py create mode 100644 vime_plugins/models/gemma4_provider.py diff --git a/scripts/models/gemma4-26B-A4B.sh b/scripts/models/gemma4-26B-A4B.sh new file mode 100644 index 000000000..9601e4009 --- /dev/null +++ b/scripts/models/gemma4-26B-A4B.sh @@ -0,0 +1,28 @@ +MODEL_ARGS=( + --spec "vime_plugins.models.gemma4" "get_gemma4_spec" + --custom-model-provider-path "vime_plugins.models.gemma4_provider.model_provider" + --num-layers 30 + --hidden-size 2816 + --ffn-hidden-size 2112 + --num-attention-heads 16 + --group-query-attention + --num-query-groups 8 + --kv-channels 256 + --use-rotary-position-embeddings + --disable-bias-linear + --normalization "RMSNorm" + --norm-epsilon 1e-6 + --rotary-base 10000 + --rotary-percent 1.0 + --vocab-size 262144 + --qk-layernorm + --num-experts 128 + --moe-ffn-hidden-size 704 + --moe-router-topk 8 + --moe-router-dtype fp32 + --moe-router-score-function softmax + --moe-router-load-balancing-type none + --moe-aux-loss-coeff 0.0 + --moe-token-dispatcher-type alltoall + --moe-grouped-gemm +) diff --git a/tests/test_hf_to_megatron.py b/tests/test_hf_to_megatron.py index 8e180d946..31b36a420 100644 --- a/tests/test_hf_to_megatron.py +++ b/tests/test_hf_to_megatron.py @@ -350,6 +350,7 @@ def test_loader_scope_stays_explicit(): assert set(_LOADERS) == { "deepseek_v3", "deepseek_v32", + "gemma4_text", "glm4", "glm4_moe", "glm4_moe_lite", diff --git a/tests/utils/test_gemma4.py b/tests/utils/test_gemma4.py new file mode 100644 index 000000000..69b57fa60 --- /dev/null +++ b/tests/utils/test_gemma4.py @@ -0,0 +1,234 @@ +import importlib.util +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch + +from vime.backends.megatron_utils.hf_to_megatron import _LOADERS +from vime.backends.megatron_utils.hf_to_megatron.gemma4 import gemma4_hf_tensor +from vime.backends.megatron_utils.megatron_to_hf.gemma4 import _config_cache, _expert_buffers, convert_gemma4_to_hf + +NUM_GPUS = 0 + +try: + _has_megatron = importlib.util.find_spec("megatron.core") is not None +except ModuleNotFoundError: + _has_megatron = False + +requires_megatron = pytest.mark.skipif(not _has_megatron, reason="requires the Megatron runtime") + + +class Reader(dict): + def get_tensor(self, name): + return self[name] + + +def conversion_config(config): + return { + "global_attn_layers": { + index for index, layer_type in enumerate(config.layer_types) if layer_type == "full_attention" + }, + "local_head_dim": config.head_dim, + "global_head_dim": config.global_head_dim, + "num_attention_heads": config.num_attention_heads, + "local_num_kv_heads": config.num_key_value_heads, + "global_num_kv_heads": config.num_global_key_value_heads, + "hidden_size": config.hidden_size, + "num_experts": config.num_experts, + } + + +@pytest.fixture +def config(): + return SimpleNamespace( + model_type="gemma4_text", + enable_moe_block=True, + layer_types=["sliding_attention", "full_attention"], + num_attention_heads=4, + num_key_value_heads=2, + num_global_key_value_heads=1, + head_dim=2, + global_head_dim=4, + hidden_size=3, + num_experts=2, + ) + + +@pytest.mark.unit +def test_registers_native_hf_loader(): + assert _LOADERS["gemma4_text"] is gemma4_hf_tensor + + +@pytest.mark.unit +def test_native_model_does_not_import_bridge(): + source = (Path(__file__).resolve().parents[2] / "vime_plugins/models/gemma4.py").read_text() + + assert "megatron.bridge" not in source + assert "mbridge" not in source + + +@pytest.mark.unit +@requires_megatron +def test_native_model_extends_transformer_config(): + from megatron.core.transformer.transformer_config import TransformerConfig + + from vime_plugins.models.gemma4 import Gemma4TransformerConfig + + assert issubclass(Gemma4TransformerConfig, TransformerConfig) + + +@pytest.mark.unit +@requires_megatron +def test_marks_experts_for_direct_hf_loading(): + from vime_plugins.models.gemma4_provider import _mark_expert_weights_for_direct_hf_loading + + fc1 = torch.nn.Parameter(torch.zeros(1)) + fc2 = torch.nn.Parameter(torch.zeros(1)) + dense = torch.nn.Parameter(torch.zeros(1)) + model = SimpleNamespace( + named_parameters=lambda: iter( + [ + ("decoder.layers.0.mlp.experts.linear_fc1.weight0", fc1), + ("decoder.layers.0.mlp.experts.linear_fc2.weight0", fc2), + ("decoder.layers.0.self_attention.linear_proj.weight", dense), + ] + ) + ) + + _mark_expert_weights_for_direct_hf_loading(model) + + assert (fc1.tensor_model_parallel, fc1.partition_dim, fc1.partition_stride) == (True, 0, 1) + assert (fc2.tensor_model_parallel, fc2.partition_dim, fc2.partition_stride) == (True, 1, 1) + assert not hasattr(dense, "tensor_model_parallel") + + +@pytest.mark.unit +@requires_megatron +def test_promotes_layer_scalars_for_direct_hf_loading(): + from vime_plugins.models.gemma4_provider import _promote_layer_scalars + + model = torch.nn.Module() + model.layer = torch.nn.Module() + model.layer.register_buffer("layer_scalar", torch.tensor([0.5])) + + _promote_layer_scalars(model) + + assert torch.equal(dict(model.named_parameters())["layer.layer_scalar"], torch.tensor([0.5])) + assert not model.layer.layer_scalar.requires_grad + + +@pytest.mark.unit +def test_tied_output_weight_is_not_exported(config): + args = SimpleNamespace(hf_checkpoint="checkpoint") + _config_cache[args.hf_checkpoint] = conversion_config(config) + + assert convert_gemma4_to_hf(args, "module.module.output_layer.weight", torch.empty(1)) == [] + + +@pytest.mark.unit +@pytest.mark.parametrize("layer", [0, 1]) +def test_qkv_round_trip(config, layer): + attention_config = ( + config + if layer == 0 + else SimpleNamespace( + num_attention_heads=config.num_attention_heads, + num_key_value_heads=config.num_global_key_value_heads, + head_dim=config.global_head_dim, + hidden_size=config.hidden_size, + ) + ) + q = torch.arange( + attention_config.num_attention_heads * attention_config.head_dim * config.hidden_size, + dtype=torch.float32, + ).reshape(-1, config.hidden_size) + k = torch.arange( + attention_config.num_key_value_heads * attention_config.head_dim * config.hidden_size, + dtype=torch.float32, + ).reshape(-1, config.hidden_size) + v = k if layer == 1 else k + 1000 + prefix = f"model.language_model.layers.{layer}.self_attn" + reader = Reader( + { + f"{prefix}.q_proj.weight": q, + f"{prefix}.k_proj.weight": k, + **({} if layer == 1 else {f"{prefix}.v_proj.weight": v}), + } + ) + + megatron = gemma4_hf_tensor( + f"decoder.layers.{layer}.self_attention.linear_qkv.weight", + reader, + config, + ) + args = SimpleNamespace(hf_checkpoint="checkpoint") + _config_cache[args.hf_checkpoint] = conversion_config(config) + converted = dict( + convert_gemma4_to_hf( + args, + f"module.module.decoder.layers.{layer}.self_attention.linear_qkv.weight", + megatron, + ) + ) + + assert torch.equal(converted[f"{prefix}.q_proj.weight"], q) + assert torch.equal(converted[f"{prefix}.k_proj.weight"], k) + if layer == 0: + assert torch.equal(converted[f"{prefix}.v_proj.weight"], v) + else: + assert f"{prefix}.v_proj.weight" not in converted + + +@pytest.mark.unit +def test_expert_weights_round_trip_as_packed_hf_tensor(config): + args = SimpleNamespace(hf_checkpoint="checkpoint") + _config_cache[args.hf_checkpoint] = conversion_config(config) + _expert_buffers.clear() + packed = torch.arange(2 * 6 * 3, dtype=torch.float32).reshape(2, 6, 3) + reader = Reader({"model.language_model.layers.0.experts.gate_up_proj": packed}) + + first = gemma4_hf_tensor("decoder.layers.0.mlp.experts.linear_fc1.weight0", reader, config) + second = gemma4_hf_tensor("decoder.layers.0.mlp.experts.linear_fc1.weight1", reader, config) + assert ( + convert_gemma4_to_hf( + args, + "module.module.decoder.layers.0.mlp.experts.linear_fc1.weight0", + first, + ) + == [] + ) + converted = convert_gemma4_to_hf( + args, + "module.module.decoder.layers.0.mlp.experts.linear_fc1.weight1", + second, + ) + + assert converted[0][0] == "model.language_model.layers.0.experts.gate_up_proj" + assert torch.equal(converted[0][1], packed) + assert not _expert_buffers + + +@pytest.mark.unit +def test_dense_mlp_and_router_scale(config): + prefix = "model.language_model.layers.0" + gate = torch.arange(18, dtype=torch.float32).reshape(6, 3) + up = gate + 100 + router_scale = torch.arange(3, dtype=torch.float32) + reader = Reader( + { + f"{prefix}.mlp.gate_proj.weight": gate, + f"{prefix}.mlp.up_proj.weight": up, + f"{prefix}.router.scale": router_scale, + } + ) + + assert torch.equal( + gemma4_hf_tensor("decoder.layers.0.dense_mlp.linear_fc1.weight", reader, config), + torch.cat((gate, up)), + ) + assert torch.equal(gemma4_hf_tensor("decoder.layers.0.mlp.router.scale", reader, config), router_scale) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__])) diff --git a/vime/backends/megatron_utils/hf_to_megatron/__init__.py b/vime/backends/megatron_utils/hf_to_megatron/__init__.py index 234a410b9..5899ee7a5 100644 --- a/vime/backends/megatron_utils/hf_to_megatron/__init__.py +++ b/vime/backends/megatron_utils/hf_to_megatron/__init__.py @@ -4,6 +4,7 @@ from .common import load_model_hf_weights from .deepseek import deepseek_hf_tensor +from .gemma4 import gemma4_hf_tensor from .glm import glm4_hf_tensor, glm4_moe_hf_tensor from .qwen import mimo_hf_tensor, minimax_m2_hf_tensor, qwen_hf_tensor, qwen_moe_hf_tensor from .qwen3_5 import qwen3_5_hf_tensor @@ -12,6 +13,7 @@ _LOADERS = { "deepseek_v3": deepseek_hf_tensor, "deepseek_v32": deepseek_hf_tensor, + "gemma4_text": gemma4_hf_tensor, "glm4": glm4_hf_tensor, "glm4_moe": glm4_moe_hf_tensor, "glm4_moe_lite": deepseek_hf_tensor, @@ -32,11 +34,14 @@ def supports_hf_weight_loading(path: str | Path) -> bool: config = AutoConfig.from_pretrained(path, trust_remote_code=True) + config = getattr(config, "text_config", config) return config.model_type in _LOADERS def load_hf_weights(args, model, path: str | Path) -> None: config = AutoConfig.from_pretrained(path, trust_remote_code=True) + config = getattr(config, "text_config", config) + try: get_hf_tensor = _LOADERS[config.model_type] except KeyError as exc: diff --git a/vime/backends/megatron_utils/hf_to_megatron/gemma4.py b/vime/backends/megatron_utils/hf_to_megatron/gemma4.py new file mode 100644 index 000000000..b85b2e8ed --- /dev/null +++ b/vime/backends/megatron_utils/hf_to_megatron/gemma4.py @@ -0,0 +1,91 @@ +import re +from types import SimpleNamespace + +import torch + +from .common import SafetensorReader, merge_gate_up, merge_qkv, strip_mcore_wrappers, text_config + + +def _layer(name: str) -> tuple[int, str]: + match = re.fullmatch(r"decoder\.layers\.(\d+)\.(.+)", name) + if not match: + raise KeyError(f"Unsupported Gemma4 Megatron parameter {name!r}") + return int(match.group(1)), match.group(2) + + +def _attention_config(config, layer: int): + config = text_config(config) + if hasattr(config, "per_layer_config"): + config = config.per_layer_config[layer] + return SimpleNamespace( + num_attention_heads=config.num_attention_heads, + num_key_value_heads=config.num_key_value_heads, + head_dim=config.head_dim, + hidden_size=config.hidden_size, + ) + if config.layer_types[layer] != "full_attention": + return config + return SimpleNamespace( + num_attention_heads=config.num_attention_heads, + num_key_value_heads=config.num_global_key_value_heads, + head_dim=config.global_head_dim, + hidden_size=config.hidden_size, + ) + + +def gemma4_hf_tensor(name: str, reader: SafetensorReader, config) -> torch.Tensor: + name = strip_mcore_wrappers(name) + direct = { + "embedding.word_embeddings.weight": "model.language_model.embed_tokens.weight", + "decoder.final_layernorm.weight": "model.language_model.norm.weight", + "output_layer.weight": "model.language_model.embed_tokens.weight", + } + if name in direct: + return reader.get_tensor(direct[name]) + + layer, rest = _layer(name) + prefix = f"model.language_model.layers.{layer}" + mapping = { + "self_attention.linear_qkv.layer_norm_weight": "input_layernorm.weight", + "self_attention.q_layernorm.weight": "self_attn.q_norm.weight", + "self_attention.k_layernorm.weight": "self_attn.k_norm.weight", + "self_attention.linear_proj.weight": "self_attn.o_proj.weight", + "pre_mlp_layernorm.weight": "pre_feedforward_layernorm.weight", + "dense_mlp.linear_fc1.layer_norm_weight": "pre_feedforward_layernorm.weight", + "post_attention_layernorm.weight": "post_attention_layernorm.weight", + "post_feedforward_layernorm.weight": "post_feedforward_layernorm.weight", + "mlp.pre_feedforward_layernorm_2.weight": "pre_feedforward_layernorm_2.weight", + "pre_feedforward_layernorm_2.weight": "pre_feedforward_layernorm_2.weight", + "post_feedforward_layernorm_1.weight": "post_feedforward_layernorm_1.weight", + "post_feedforward_layernorm_2.weight": "post_feedforward_layernorm_2.weight", + "mlp.router.proj.weight": "router.proj.weight", + "mlp.router.scale": "router.scale", + "mlp.router.per_expert_scale": "router.per_expert_scale", + "layer_scalar": "layer_scalar", + } + if rest in mapping: + return reader.get_tensor(f"{prefix}.{mapping[rest]}") + + if rest == "self_attention.linear_qkv.weight": + attention_config = _attention_config(config, layer) + q = reader.get_tensor(f"{prefix}.self_attn.q_proj.weight") + k = reader.get_tensor(f"{prefix}.self_attn.k_proj.weight") + v_name = f"{prefix}.self_attn.v_proj.weight" + v = reader.get_tensor(v_name) if v_name in reader else k + return merge_qkv(q, k, v, attention_config) + + if rest in {"mlp.linear_fc1.weight", "dense_mlp.linear_fc1.weight"}: + return merge_gate_up( + reader.get_tensor(f"{prefix}.mlp.gate_proj.weight"), + reader.get_tensor(f"{prefix}.mlp.up_proj.weight"), + ) + if rest in {"mlp.linear_fc2.weight", "dense_mlp.linear_fc2.weight"}: + return reader.get_tensor(f"{prefix}.mlp.down_proj.weight") + + match = re.fullmatch(r"mlp\.experts\.linear_fc([12])\.weight(\d+)", rest) + if match: + projection, expert = map(int, match.groups()) + packed_name = "gate_up_proj" if projection == 1 else "down_proj" + return reader.get_tensor(f"{prefix}.experts.{packed_name}")[expert].contiguous() + + raise KeyError(f"Unsupported Gemma4 Megatron parameter {name!r}") diff --git a/vime/backends/megatron_utils/megatron_to_hf/__init__.py b/vime/backends/megatron_utils/megatron_to_hf/__init__.py index 33bb45eb8..b10d188b9 100644 --- a/vime/backends/megatron_utils/megatron_to_hf/__init__.py +++ b/vime/backends/megatron_utils/megatron_to_hf/__init__.py @@ -1,4 +1,5 @@ from .deepseekv3 import convert_deepseekv3_to_hf +from .gemma4 import convert_gemma4_to_hf from .glm4 import convert_glm4_to_hf from .glm4moe import convert_glm4moe_to_hf from .llama import convert_llama_to_hf @@ -36,7 +37,9 @@ def convert_to_hf(args, model_name, name, param, quantization_config=None, trans # TODO optimize code details def _convert_to_hf_core(args, model_name, name, param): model_name = model_name.lower().replace("_", "").replace("-", "") - if "minimaxm2" in model_name: + if "gemma4" in model_name: + converted_named_tensors = convert_gemma4_to_hf(args, name, param) + elif "minimaxm2" in model_name: converted_named_tensors = convert_minimax_m2_to_hf(args, name, param) elif any(family in model_name for family in ("glm4moelite", "deepseekv3", "deepseekv32", "glmmoedsa", "kimi")): converted_named_tensors = convert_deepseekv3_to_hf(args, name, param) diff --git a/vime/backends/megatron_utils/megatron_to_hf/gemma4.py b/vime/backends/megatron_utils/megatron_to_hf/gemma4.py new file mode 100644 index 000000000..a7e326b90 --- /dev/null +++ b/vime/backends/megatron_utils/megatron_to_hf/gemma4.py @@ -0,0 +1,177 @@ +import re +import torch + +_config_cache: dict[str, dict] = {} + +# Per-layer buffers for stacked expert tensors. vllm's Gemma4 loader expects +# `experts.gate_up_proj` as a single 3D tensor of shape [E, 2I, H] and +# `experts.down_proj` as [E, H, I] - it walks all experts inside the loader +# and would silently drop per-expert 2D inputs. We accumulate expert tensors +# as they stream through and emit the stacked form once all num_experts arrive. +_expert_buffers: dict = {} + + +def reset_expert_buffers() -> None: + """Drop any partial expert buckets. Callers that drive the converter from a + long-lived process (tests, repeated conversions) should invoke this between + runs so an interrupted prior conversion doesn't leak its partial state.""" + _expert_buffers.clear() + + +def _get_config(args): + checkpoint = args.hf_checkpoint + if checkpoint not in _config_cache: + from transformers import AutoConfig + + hf_config = AutoConfig.from_pretrained(checkpoint, trust_remote_code=True) + hf_text = hf_config.text_config if hasattr(hf_config, "text_config") else hf_config + if hasattr(hf_text, "per_layer_config"): + local_config = next( + layer_config + for layer_type, layer_config in zip(hf_text.layer_types, hf_text.per_layer_config, strict=True) + if layer_type != "full_attention" + ) + global_config = next( + layer_config + for layer_type, layer_config in zip(hf_text.layer_types, hf_text.per_layer_config, strict=True) + if layer_type == "full_attention" + ) + else: + local_config = hf_text + global_config = hf_text + _config_cache[checkpoint] = { + "global_attn_layers": {i for i, t in enumerate(hf_text.layer_types) if t == "full_attention"}, + "local_head_dim": local_config.head_dim, + "global_head_dim": global_config.head_dim, + "num_attention_heads": local_config.num_attention_heads, + "local_num_kv_heads": local_config.num_key_value_heads, + "global_num_kv_heads": global_config.num_key_value_heads, + "hidden_size": local_config.hidden_size, + "num_experts": getattr(hf_text, "num_experts", 0), + } + return _config_cache[checkpoint] + + +def convert_gemma4_to_hf(args, name, param): + cfg = _get_config(args) + prefix = "model.language_model." + + if name == "module.module.embedding.word_embeddings.weight": + return [(f"{prefix}embed_tokens.weight", param)] + if name == "module.module.output_layer.weight": + return [] # tied to embedding weight + if name == "module.module.decoder.final_layernorm.weight": + return [(f"{prefix}norm.weight", param)] + + match = re.match(r"module\.module\.decoder\.layers\.(\d+)\.(.+)", name) + if match: + layer_idx = int(match.group(1)) + rest = match.group(2) + L = f"{prefix}layers.{layer_idx}" + is_global = layer_idx in cfg["global_attn_layers"] + + if rest == "self_attention.linear_proj.weight": + return [(f"{L}.self_attn.o_proj.weight", param)] + elif rest == "self_attention.linear_qkv.weight": + if is_global: + head_dim = cfg["global_head_dim"] + num_kv_heads = cfg["global_num_kv_heads"] + else: + head_dim = cfg["local_head_dim"] + num_kv_heads = cfg["local_num_kv_heads"] + + q_heads_per_kv = cfg["num_attention_heads"] // num_kv_heads + hidden_size = cfg["hidden_size"] + param = param.view(num_kv_heads, (q_heads_per_kv + 2) * head_dim, hidden_size) + q_dim = q_heads_per_kv * head_dim + q_param = param[:, :q_dim, :].reshape(-1, hidden_size) + k_param = param[:, q_dim : q_dim + head_dim, :].reshape(-1, hidden_size) + + if is_global: + return [ + (f"{L}.self_attn.q_proj.weight", q_param), + (f"{L}.self_attn.k_proj.weight", k_param), + ] + else: + v_param = param[:, q_dim + head_dim :, :].reshape(-1, hidden_size) + return [ + (f"{L}.self_attn.q_proj.weight", q_param), + (f"{L}.self_attn.k_proj.weight", k_param), + (f"{L}.self_attn.v_proj.weight", v_param), + ] + elif rest == "self_attention.linear_qkv.layer_norm_weight": + return [(f"{L}.input_layernorm.weight", param)] + elif rest == "self_attention.q_layernorm.weight": + return [(f"{L}.self_attn.q_norm.weight", param)] + elif rest == "self_attention.k_layernorm.weight": + return [(f"{L}.self_attn.k_norm.weight", param)] + elif rest in ("mlp.linear_fc1.weight", "dense_mlp.linear_fc1.weight"): + gate_weight, up_weight = param.chunk(2, dim=0) + return [ + (f"{L}.mlp.gate_proj.weight", gate_weight), + (f"{L}.mlp.up_proj.weight", up_weight), + ] + elif rest in ("mlp.linear_fc2.weight", "dense_mlp.linear_fc2.weight"): + return [(f"{L}.mlp.down_proj.weight", param)] + elif rest in ("mlp.linear_fc1.layer_norm_weight", "dense_mlp.linear_fc1.layer_norm_weight"): + return [(f"{L}.pre_feedforward_layernorm.weight", param)] + elif rest == "pre_mlp_layernorm.weight": + return [(f"{L}.pre_feedforward_layernorm.weight", param)] + elif rest == "post_attention_layernorm.weight": + return [(f"{L}.post_attention_layernorm.weight", param)] + elif rest == "post_feedforward_layernorm.weight": + return [(f"{L}.post_feedforward_layernorm.weight", param)] + elif rest == "layer_scalar": + return [(f"{L}.layer_scalar", param)] + elif rest == "mlp.router.proj.weight": + return [(f"{L}.router.proj.weight", param)] + elif rest == "mlp.router.scale": + return [(f"{L}.router.scale", param)] + elif rest == "mlp.router.per_expert_scale": + return [(f"{L}.router.per_expert_scale", param)] + else: + expert_match = re.match(r"mlp\.experts\.linear_fc([12])\.weight(\d+)", rest) + if expert_match: + fc, expert_idx = expert_match.group(1), int(expert_match.group(2)) + return _buffer_expert_and_maybe_flush( + layer_idx, + fc, + expert_idx, + param, + L, + num_experts=cfg["num_experts"], + ) + + if rest == "pre_feedforward_layernorm_2.weight": + return [(f"{L}.pre_feedforward_layernorm_2.weight", param)] + elif rest == "mlp.pre_feedforward_layernorm_2.weight": + return [(f"{L}.pre_feedforward_layernorm_2.weight", param)] + elif rest == "post_feedforward_layernorm_2.weight": + return [(f"{L}.post_feedforward_layernorm_2.weight", param)] + elif rest == "post_feedforward_layernorm_1.weight": + return [(f"{L}.post_feedforward_layernorm_1.weight", param)] + + raise ValueError(f"Unknown Gemma4 parameter name: {name}") + + +def _buffer_expert_and_maybe_flush(layer_idx, fc, expert_idx, param, L_prefix, num_experts): + """Buffer per-expert tensor; emit stacked 3D `experts.gate_up_proj` / `experts.down_proj` + once the bucket for (layer, fc) has all `num_experts` experts.""" + assert ( + num_experts and num_experts > 0 + ), f"num_experts must be known for MoE layer expert conversion, got {num_experts}" + key = (layer_idx, fc) + bucket = _expert_buffers.setdefault(key, {}) + bucket[expert_idx] = param + + if len(bucket) < num_experts: + return [] + + ordered = [bucket[i] for i in range(num_experts)] + stacked = torch.stack(ordered, dim=0).contiguous() + del _expert_buffers[key] + + if fc == "1": + return [(f"{L_prefix}.experts.gate_up_proj", stacked)] + else: + return [(f"{L_prefix}.experts.down_proj", stacked)] diff --git a/vime_plugins/models/gemma4.py b/vime_plugins/models/gemma4.py new file mode 100644 index 000000000..89ebca5cc --- /dev/null +++ b/vime_plugins/models/gemma4.py @@ -0,0 +1,1188 @@ +"""Native Megatron Gemma4 transformer layer and config. + +Extends the native Gemma3 implementation with Gemma4-specific features: +- Heterogeneous attention: global layers use head_dim=512, num_kv_heads=4; + sliding layers use head_dim=256, num_kv_heads=16. +- attention_k_eq_v: global layers reuse K output as V (no v_proj). +- v_norm: RMSNorm without learnable scale applied to V states. +- layer_scalar: buffer multiplied after residual (not learned). +- final_logit_softcapping: applied to output logits in the model wrapper. +- MoE block (26B-A4B): Gemma4's custom router (with per-expert scale) plugged + into Megatron's MoE infrastructure for proper expert-parallel sharding. + The router is still custom (see Gemma4Router); dispatching + grouped-GEMM + come from Megatron's MoELayer + TEGroupedMLP. +""" + +import functools +import logging +from dataclasses import dataclass +from dataclasses import replace as dc_replace + +import torch +import torch.nn as nn +import torch.nn.functional as F +from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add +from megatron.core.transformer.attention import SelfAttention, 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.moe_layer import BaseMoELayer, MoELayer +from megatron.core.transformer.spec_utils import ModuleSpec, build_module +from megatron.core.transformer.transformer_layer import TransformerLayer, TransformerLayerSubmodules +from megatron.core.utils import make_viewless_tensor + +try: + from megatron.core.extensions.transformer_engine import ( + TEColumnParallelLinear, + TEDotProductAttention, + TELayerNormColumnParallelLinear, + TENorm, + TERowParallelLinear, + ) + + HAVE_TE = True +except ImportError: + HAVE_TE = False + +from megatron.core.transformer.transformer_config import TransformerConfig + +# Gemma uses GeGLU, not SwiGLU. +_gelu_tanh = functools.partial(F.gelu, approximate="tanh") + + +@dataclass +class Gemma4TransformerConfig(TransformerConfig): + """Gemma4-specific config extending Gemma3.""" + + rope_local_base_freq: float = 10000.0 + sliding_window_pattern: int = 6 + sliding_window: int = 1024 + global_kv_channels: int = 512 + global_num_query_groups: int = 4 + global_partial_rotary_factor: float = 0.25 # fraction of global head_dim that gets RoPE + attention_k_eq_v: bool = True # global layers: V = K (no v_proj) + enable_moe_block: bool = False # 26B-A4B MoE variant + + +class VNorm(nn.Module): + """RMSNorm without learnable scale, matching Gemma4's v_norm.""" + + def __init__(self, dim: int, eps: float = 1e-6): + super().__init__() + self.eps = eps + self.dim = dim + + def forward(self, x: torch.Tensor) -> torch.Tensor: + dtype = x.dtype + x = x.float() + return (x * torch.pow(x.pow(2).mean(-1, keepdim=True) + self.eps, -0.5)).to(dtype) + + +@dataclass +class Gemma4TransformerLayerSubmodules(TransformerLayerSubmodules): + post_attention_layernorm: ModuleSpec | type = IdentityOp + post_feedforward_layernorm: ModuleSpec | type = IdentityOp + # For MoE-enabled variants (26B-A4B), the primary `mlp` submodule is swapped + # to a Gemma4MoELayer and the original dense MLP moves to `dense_mlp`. This + # keeps the `.mlp.experts.linear_fc...` naming that the weight-conversion path + # expects while preserving Gemma4's dense+MoE-in-parallel structure. + dense_mlp: ModuleSpec | type = IdentityOp + + +class Gemma4Router(nn.Module): + """Gemma4 MoE router. + + The router equation (mirroring HF ``Gemma4TextTopkRouter``) is: + + h_norm = RMSNorm_no_scale(h) # VNorm: no learnable scale + h_scaled = h_norm * scale / sqrt(H) # learnable per-hidden scale + logits = proj(h_scaled) # [T, E] + probs = softmax(logits, dim=-1) + top_w, top_i = topk(probs, k=top_k) + top_w = top_w / top_w.sum(dim=-1, keepdim=True) # renormalize + top_w = top_w * per_expert_scale[top_i] # per-expert scale + + The renormalise-then-scale order is load-bearing and must match HF: it + produces ``top_w.sum() == per_expert_scale.mean_over_selected`` rather + than a renormalised-back-to-1 distribution. Reversing the order (scale + first, then renormalise) would cancel ``per_expert_scale``. + ``test_router_matches_hf_reference_equation`` guards this. + """ + + def __init__(self, config): + super().__init__() + self.hidden_size = config.hidden_size + self.num_experts = config.num_moe_experts + self.top_k = config.moe_router_topk + self.scalar_root_size = self.hidden_size**-0.5 + self.norm = VNorm(self.hidden_size, eps=config.layernorm_epsilon) + self.proj = nn.Linear(self.hidden_size, self.num_experts, bias=False) + self.scale = nn.Parameter(torch.ones(self.hidden_size)) + self.per_expert_scale = nn.Parameter(torch.ones(self.num_experts)) + + def forward(self, hidden_states): + h = self.norm(hidden_states) + h = h * self.scale * self.scalar_root_size + logits = self.proj(h) + probs = torch.softmax(logits, dim=-1) + top_k_weights, top_k_index = torch.topk(probs, k=self.top_k, dim=-1) + top_k_weights = top_k_weights / top_k_weights.sum(dim=-1, keepdim=True) + top_k_weights = top_k_weights * self.per_expert_scale[top_k_index] + return top_k_weights, top_k_index + + def set_layer_number(self, layer_number): + pass + + +class Gemma4MoELayer(MoELayer): + """Gemma4 MoE block: Megatron's MoELayer with Gemma4's custom router. + + Megatron's MoELayer hardcodes its own ``TopKRouter`` which uses a + softmax-with-expert-bias scheme. Gemma4 has its own router semantics + (no-scale RMSNorm -> learnable per-hidden scale -> proj -> softmax -> topk -> + per-expert scale multiplier). We reuse all of Megatron's infrastructure + for dispatching (alltoall), expert parallelism, and grouped-GEMM expert + computation - but swap in our ``Gemma4Router`` and convert its compact + (top_k_weights [T, K], top_k_index [T, K]) output into Megatron's + expected (probs [T, E], routing_map [T, E]) format inside ``route()``. + """ + + def __init__(self, config, submodules=None, layer_number=None, pg_collection=None): + # Fall back to Megatron's global parallel_state when pg_collection isn't + # explicitly passed. TransformerLayer only forwards pg_collection when + # submodules.mlp.module is *exactly* one of + # (MoELayer, GroupedMLP, TEGroupedMLP, SequentialMLP) - an identity check + # via `in`, so Gemma4MoELayer (a MoELayer subclass) slips through and + # receives None. BaseMoELayer.__init__ then crashes on `pg_collection.ep`. + # Same fallback MoELayer.__init__ uses when invoked directly. + if pg_collection is None: + from megatron.core.transformer.moe.moe_utils import get_default_pg_collection + + pg_collection = get_default_pg_collection() + BaseMoELayer.__init__(self, config=config, layer_number=layer_number, pg_collection=pg_collection) + self.moe_layer_recompute = False + self.shared_experts_recompute = False + self.submodules = submodules + + self.router = Gemma4Router(config) + + from megatron.core.transformer.moe.token_dispatcher import ( + MoEAllGatherTokenDispatcher, + MoEAlltoAllTokenDispatcher, + MoEFlexTokenDispatcher, + ) + + if config.moe_token_dispatcher_type == "allgather": + self.token_dispatcher = MoEAllGatherTokenDispatcher( + self.num_local_experts, + self.local_expert_indices, + config=self.config, + pg_collection=pg_collection, + ) + elif config.moe_token_dispatcher_type == "alltoall": + self.token_dispatcher = MoEAlltoAllTokenDispatcher( + self.num_local_experts, + self.local_expert_indices, + config=self.config, + pg_collection=pg_collection, + ) + elif config.moe_token_dispatcher_type == "flex": + self.token_dispatcher = MoEFlexTokenDispatcher( + self.num_local_experts, + self.local_expert_indices, + config=self.config, + pg_collection=pg_collection, + ) + else: + raise ValueError(f"Unsupported token dispatcher type: {config.moe_token_dispatcher_type}") + + self.experts = build_module( + self.submodules.experts, + self.num_local_experts, + self.config, + pg_collection=pg_collection, + ) + + self.shared_experts = None + + from megatron.core.transformer.moe.moe_utils import MoECudaGraphTensorStore + + self.cudagraph_tensor_store = MoECudaGraphTensorStore() + + # pre_feedforward_layernorm_2: applied to experts' input ONLY (router + # input stays un-normed). Matches HF Gemma4TextDecoderLayer: + # hidden_states_flat = residual # router input (un-normed) + # hidden_states_2 = pre_feedforward_layernorm_2(hidden_states_flat) + # hidden_states_2 = experts(hidden_states_2, top_k_index, top_k_weights) + self.pre_feedforward_layernorm_2 = TENorm( + config=config, + hidden_size=config.hidden_size, + eps=config.layernorm_epsilon, + ) + + def route(self, hidden_states: torch.Tensor): + """Call ``Gemma4Router`` and pack its output into Megatron's + ``(probs, routing_map)`` format. + + ``Gemma4Router`` emits compact top-k tensors: + top_k_weights: [T, K] - routing weights (already scaled by per_expert_scale) + top_k_index: [T, K] - which experts each token routes to + Megatron's dispatcher wants: + probs: [T, E] - weight per (token, expert), 0 where not routed + routing_map: [T, E] - boolean mask + """ + flat = hidden_states.reshape(-1, hidden_states.shape[-1]) + top_k_weights, top_k_index = self.router(flat) + + num_tokens = flat.shape[0] + num_experts = self.config.num_moe_experts + probs = torch.zeros( + num_tokens, + num_experts, + dtype=top_k_weights.dtype, + device=top_k_weights.device, + ) + probs.scatter_(1, top_k_index, top_k_weights) + routing_map = probs != 0 + return probs, routing_map + + def forward( + self, + hidden_states: torch.Tensor, + router_input: torch.Tensor | None = None, + ): + """Gemma4 MoE forward with split router / experts inputs. + + HF's ``Gemma4TextDecoderLayer`` routes based on the *un-normed* residual + but feeds the experts the *pre-ff-norm-2'd* residual: + + hidden_states_flat = residual # un-normed + _, tk_w, tk_i = self.router(hidden_states_flat) + experts_input = self.pre_feedforward_layernorm_2(hidden_states_flat) + output = self.experts(experts_input, tk_i, tk_w) + + We take the un-normed residual in ``hidden_states`` and apply + ``pre_feedforward_layernorm_2`` internally to obtain the experts + input. The router path uses the un-normed residual directly. Callers + may pass a different ``router_input`` for tests or ablations; when + ``router_input is None`` (the normal case) the router sees the same + un-normed residual the layer was called with. + + We inline the Megatron parent's ``forward`` body here - rather than + calling ``super().forward`` with a side-channel stash - so the + router input is passed explicitly end-to-end and the code is safe + under activation checkpointing / recomputation. + """ + if self.training and self.attn_tp_group.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." + ) + + router_in = router_input if router_input is not None else hidden_states + experts_in = self.pre_feedforward_layernorm_2(hidden_states) + + def custom_forward(experts_in, router_in): + # Gemma4 has no shared experts; shared_experts_compute returns None. + shared_expert_output = self.shared_experts_compute(experts_in) + probs, routing_map = self.route(router_in) + experts_in2, probs = self.preprocess(experts_in, probs, routing_map) + dispatched_input, probs = self.dispatch(experts_in2, probs) + output, mlp_bias = self.routed_experts_compute(dispatched_input, probs) + output = self.combine(output) + output = self.postprocess(output, shared_expert_output) + return output, mlp_bias + + # moe_layer_recompute is forced to False in __init__; call directly. + return custom_forward(experts_in, router_in) + + +class Gemma4TransformerLayer(TransformerLayer): + """Gemma4 transformer layer with heterogeneous attention and layer_scalar.""" + + def __init__( + self, + config: Gemma4TransformerConfig, + submodules: Gemma4TransformerLayerSubmodules, + layer_number: int = 1, + hidden_dropout: float = None, + **kwargs, + ): + from megatron.core.transformer.transformer_layer import get_transformer_layer_offset + + global_layer_number = layer_number + get_transformer_layer_offset(config) + # Megatron passes `layer_number` as 1-indexed (default 1), so in 0-indexed + # HF space a global layer is `(i+1) % pattern == 0` -> `i % pattern == pattern-1`. + # Equivalently: `is_sliding` when `global_layer_number % pattern != 0`. + self.is_sliding = bool(global_layer_number % config.sliding_window_pattern) + self._is_global = not self.is_sliding + + # Global layers have different head_dim (kv_channels) and num_kv_heads + # (num_query_groups). Build the layer against a *cloned* config with + # those overrides so we never mutate the shared transformer config. + # Mutation would be reentrant-unsafe under concurrent layer + # construction and leak global-layer shapes into sibling sliding + # layers if an exception were raised during super().__init__. + layer_config = ( + dc_replace( + config, + kv_channels=config.global_kv_channels, + num_query_groups=config.global_num_query_groups, + ) + if self._is_global + else config + ) + super().__init__( + config=layer_config, + submodules=submodules, + layer_number=layer_number, + hidden_dropout=hidden_dropout, + **kwargs, + ) + + self.self_attention._is_global = self._is_global + + # Global layers require this because head_dim=512 exceeds flash attention's limit (256). + # Local layers also use SDPA for consistency. + self.self_attention.core_attention = SDPACoreAttention( + config=config, + layer_number=self.layer_number, + attn_mask_type=AttnMaskType.causal, + softmax_scale=config.softmax_scale, + ) + self.self_attention.core_attention._is_sliding = self.is_sliding + + self.post_attention_layernorm = build_module( + submodules.post_attention_layernorm, + config=self.config, + hidden_size=self.config.hidden_size, + eps=self.config.layernorm_epsilon, + ) + self.post_feedforward_layernorm = build_module( + submodules.post_feedforward_layernorm, + config=self.config, + hidden_size=self.config.hidden_size, + eps=self.config.layernorm_epsilon, + ) + + # Layer scalar (buffer, not learned). Kept in fp32 intentionally - + # HF stores this scalar in fp32 and relies on the implicit upcast of + # ``bf16_hidden * fp32_scalar`` at multiply time (see HF Gemma4 + # ``Gemma4TextDecoderLayer.__init__`` at modeling_gemma4.py:1331). + # Don't switch to ``dtype=self.config.params_dtype``; that would + # silently change the arithmetic. + self.register_buffer("layer_scalar", torch.ones(1)) + + # MoE block (26B-A4B): super().__init__ already built self.mlp from the + # layer spec, which when enable_moe_block=True is a Gemma4MoELayer (not + # a dense MLP). We also build a parallel `dense_mlp` for Gemma4's + # dense + MoE combined-FFN pattern. The two outputs are summed in + # forward(). + self.enable_moe_block = getattr(config, "enable_moe_block", False) + if self.enable_moe_block: + self.dense_mlp = build_module( + submodules.dense_mlp, + config=config, + ) + self.post_feedforward_layernorm_1 = TENorm( + config=config, + hidden_size=config.hidden_size, + eps=config.layernorm_epsilon, + ) + # pre_feedforward_layernorm_2 now lives INSIDE Gemma4MoELayer + # (matching HF Gemma4TextDecoderLayer semantics: router sees un-normed + # residual, experts see pre_feedforward_layernorm_2(residual)). This + # attribute is kept on the MoE block so weight-conversion paths + # don't change. + self.post_feedforward_layernorm_2 = TENorm( + config=config, + hidden_size=config.hidden_size, + eps=config.layernorm_epsilon, + ) + + def _forward_dense_ffn(self, pre_mlp_ln): + """Run the dense MLP. ``self.mlp`` is the dense MLP directly for the + 31B variant.""" + out, bias = self.mlp(pre_mlp_ln) + return out + bias if bias is not None else out + + def _forward_moe_ffn(self, residual, pre_mlp_ln): + """Run dense + MoE in parallel and sum (26B-A4B variant). + + Mirrors HF ``Gemma4TextDecoderLayer.forward`` (transformers + modeling_gemma4.py:1376-1391): dense branch goes through + ``post_feedforward_layernorm_1``, MoE branch through + ``post_feedforward_layernorm_2``, the two are summed, and the outer + ``Gemma4TransformerLayer.forward`` applies ``post_feedforward_layernorm`` + to the sum - 3 post-FFN LNs total for MoE layers is correct. + + HF routes on the un-normed residual but feeds experts the + ``pre_feedforward_layernorm_2``'d residual; Gemma4MoELayer applies + that norm internally, so we pass the un-normed residual directly. + """ + dense_out, dense_bias = self.dense_mlp(pre_mlp_ln) + if dense_bias is not None: + dense_out = dense_out + dense_bias + mlp_output = self.post_feedforward_layernorm_1(dense_out) + + moe_output, _ = self.mlp(residual) + moe_output = self.post_feedforward_layernorm_2(moe_output) + + return mlp_output + moe_output + + def forward( + self, + hidden_states, + attention_mask=None, + context=None, + context_mask=None, + rotary_pos_emb=None, + rotary_pos_cos=None, + rotary_pos_sin=None, + attention_bias=None, + inference_context=None, + inference_params=None, + packed_seq_params=None, + sequence_len_offset=None, + **kwargs, + ): + if rotary_pos_emb is not None and not isinstance(rotary_pos_emb, tuple): + global_dim = getattr(self.config, "dual_rope_global_dim", 0) + if global_dim > 0 and rotary_pos_emb.shape[-1] > global_dim: + if self.is_sliding: + rotary_pos_emb = rotary_pos_emb[..., global_dim:] + else: + rotary_pos_emb = rotary_pos_emb[..., :global_dim] + elif isinstance(rotary_pos_emb, tuple): + rotary_pos_emb = rotary_pos_emb[1] if self.is_sliding else rotary_pos_emb[0] + if isinstance(attention_mask, tuple): + attention_mask = attention_mask[1] if self.is_sliding else attention_mask[0] + + # Global layers use partial RoPE (25% of head_dim=512 = 128 dims) + # Local layers use full RoPE (100% of head_dim=256 = 256 dims) + # With DualRotaryEmbedding, global RoPE is full-size (512 dims) with zero-padded + # non-rotated dims, so no truncation needed. + # With single RoPE (local only, 256 dims), truncate for global layers. + if not self.is_sliding and rotary_pos_emb is not None: + global_rope_dim = int(self.config.global_kv_channels * self.config.global_partial_rotary_factor) + if ( + rotary_pos_emb.shape[-1] != self.config.global_kv_channels + and rotary_pos_emb.shape[-1] > global_rope_dim + ): + rotary_pos_emb = rotary_pos_emb[..., :global_rope_dim] + + residual = hidden_states + + extra_kwargs = {} + if inference_context is not None: + extra_kwargs["inference_context"] = inference_context + elif inference_params is not None: + extra_kwargs["inference_params"] = inference_params + + input_layernorm_output = self.input_layernorm(hidden_states) + + hidden_states, hidden_states_bias = self.self_attention( + input_layernorm_output, + attention_mask=attention_mask, + 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, + **extra_kwargs, + ) + + if hidden_states_bias is not None: + hidden_states = hidden_states + hidden_states_bias + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = residual + hidden_states + + residual = hidden_states + pre_mlp_layernorm_output = self.pre_mlp_layernorm(hidden_states) + if self.enable_moe_block: + hidden_states = self._forward_moe_ffn(residual, pre_mlp_layernorm_output) + else: + hidden_states = self._forward_dense_ffn(pre_mlp_layernorm_output) + hidden_states = self.post_feedforward_layernorm(hidden_states) + hidden_states = residual + hidden_states + + hidden_states = hidden_states * self.layer_scalar + + output = make_viewless_tensor( + inp=hidden_states, + requires_grad=hidden_states.requires_grad, + keep_graph=True, + ) + + if self.config.external_cuda_graph and self.training: + return output + return output, context + + +class SDPACoreAttention(nn.Module): + """Gemma4 core attention. + + Replaces TE's DotProductAttention because: + - Global layers have head_dim=512, which flash-attn 2.x doesn't support. + - Sliding-window layers need an explicit left-window mask (HF behavior). + - Context-parallelism on the global layers needs an all-gather+full-attn + path with a differentiable K/V gather. + + Dispatch at call time (packed / thd shape): + - CP > 1 (any layer) : all-gather K/V, apply causal + optional + sliding-window mask computed from vime zig-zag global indices. + - global + CP == 1 : sub-sequence causal SDPA (no O(T^2) mask alloc). + - sliding + CP == 1 : flash_attn_varlen_func with (sw-1, 0) window. + """ + + def __init__( + self, + config, + layer_number, + attn_mask_type, + attention_type="self", + attention_dropout=None, + softmax_scale=None, + **kwargs, + ): + super().__init__() + # Megatron's SelfAttention.__init__ passes a few kwargs (e.g. cp_comm_type, + # model_comm_pgs) intended for TE's DotProductAttention. We accept-and-ignore + # by name rather than asserting empty; a strict assert breaks whenever + # Megatron/TE add a new kwarg. If a kwarg shows up here that we *should* + # honor (e.g. a new softmax dtype), it will surface as a behavioral bug + # in parity, which is what the test suite covers. + del kwargs + self.config = config + self.softmax_scale = softmax_scale + self.dropout_p = config.attention_dropout if attention_dropout is None else attention_dropout + self._is_sliding = False # set by Gemma4TransformerLayer + + def _resolve_scale(self, hn: int) -> float: + return self.softmax_scale if self.softmax_scale is not None else (hn**-0.5) + + @staticmethod + def _zigzag_global_indices(local_len, cp_rank, cp_size, device): + """Global positions of this rank's local Q tokens under vime's + zig-zag CP layout (matches cp_utils.slice_with_cp). + + Local tokens on rank r occupy two global sub-ranges: + [r*cs, (r+1)*cs) and [(2*cp-r-1)*cs, (2*cp-r)*cs) + where cs = local_len / 2 = seq_len / (2*cp_size). + """ + cs = local_len // 2 + first = torch.arange(cp_rank * cs, (cp_rank + 1) * cs, device=device) + second = torch.arange( + (2 * cp_size - cp_rank - 1) * cs, + (2 * cp_size - cp_rank) * cs, + device=device, + ) + return torch.cat([first, second]) + + @staticmethod + def _cp_unzigzag_permutation(cu_seqlens_list, cp_size, device): + """Map rank-major CP-gathered K/V tokens back to packed global order.""" + total_local_len = sum( + (cu_seqlens_list[i + 1] - cu_seqlens_list[i]) // cp_size for i in range(len(cu_seqlens_list) - 1) + ) + local_prefix = 0 + perm_parts = [] + for s_idx in range(len(cu_seqlens_list) - 1): + seq_len_global = cu_seqlens_list[s_idx + 1] - cu_seqlens_list[s_idx] + cs = seq_len_global // (2 * cp_size) + g = torch.arange(seq_len_global, device=device) + chunk = g // cs + owner = torch.where(chunk < cp_size, chunk, 2 * cp_size - 1 - chunk) + local_in_rank = torch.where( + chunk < cp_size, + g - owner * cs, + cs + (g - (2 * cp_size - 1 - owner) * cs), + ) + perm_parts.append(owner * total_local_len + local_prefix + local_in_rank) + local_prefix += seq_len_global // cp_size + return torch.cat(perm_parts) + + def _forward_cp_subseq_mask(self, query, key, value, packed_seq_params, sliding_window=None): + """CP>1 path for any layer: all-gather K/V, then loop over sub-seqs + and apply a per-sub-seq attention mask built from zig-zag global + positions. Supports causal-only (global layers) and causal + + sliding-window (sliding layers). + + Under vime's CP convention, ``packed_seq_params.cu_seqlens_q`` holds + GLOBAL boundaries: each packed sub-sequence on this rank represents + ``(cu[i+1] - cu[i])`` tokens globally but only ``(cu[i+1] - cu[i]) // + cp_size`` tokens locally (the zig-zag slice of this rank's two + chunks, concatenated as [first, second]). + """ + from megatron.core import parallel_state + from megatron.core.tensor_parallel.mappings import gather_from_sequence_parallel_region + + cp_group = parallel_state.get_context_parallel_group() + cp_size = parallel_state.get_context_parallel_world_size() + cp_rank = parallel_state.get_context_parallel_rank() + + t_local = query.shape[0] + np_q, hn = query.shape[1], query.shape[2] + nk = key.shape[1] + scale = self._resolve_scale(hn) + + # Differentiable all-gather along the token dim. forward: AG, + # backward: RS - so K/V grads on non-owning ranks flow back to the + # originating rank. The raw `dist.all_gather_into_tensor` has no + # autograd rule and PyTorch prints a "silently incorrect behavior" + # warning + drops those grads. + k_full = gather_from_sequence_parallel_region(key.contiguous(), group=cp_group) + v_full = gather_from_sequence_parallel_region(value.contiguous(), group=cp_group) + # gather_from_sequence_parallel_region stacks each rank's chunk + # consecutively in rank order. Under zig-zag, each rank's [2*cs] + # local tokens are [chunk_r_first, chunk_r_second]. So the gathered + # tensor layout is [r0_first, r0_second, r1_first, r1_second, ...]. + # We need to un-zig-zag into pure global order so mask indices line + # up. Build a permutation that maps gathered index -> global index. + device = query.device + dtype = query.dtype + cu_seqlens = packed_seq_params.cu_seqlens_q if packed_seq_params is not None else None + + # Sanity: for each packed sub-seq, the GLOBAL length must be + # divisible by 2*cp_size so chunk_size is integer. With cp_size=1 this + # reduces to even-length, which the CP=1 parity-test harness may + # violate (no zig-zag pre-slicing). Skip the check there; permutation + # is identity under cp_size=1 so odd length is harmless. + if cu_seqlens is not None and cp_size > 1: + expected_t_local = 0 + for s_idx in range(len(cu_seqlens) - 1): + s_len = (cu_seqlens[s_idx + 1] - cu_seqlens[s_idx]).item() + assert s_len % (2 * cp_size) == 0, ( + f"sub-sequence {s_idx} global length ({s_len}) is not " + f"divisible by 2*cp_size ({2 * cp_size}); `slice_with_cp` " + "should pad before packing" + ) + expected_t_local += s_len // cp_size + assert expected_t_local == t_local, ( + f"packed-seq local length mismatch: sum(seq_len // cp_size) = " + f"{expected_t_local}, but query.shape[0] = {t_local}" + ) + + if cu_seqlens is None: + t_full_total = k_full.shape[0] + cu_seqlens_list = [0, t_full_total] + else: + cu_seqlens_list = cu_seqlens.tolist() + + # With cp_size=1 the zigzag degenerates to identity and all-gather is + # a no-op; skip the permutation (and the floor-div that would drop the + # trailing odd token for seq_len_global % 2 == 1). + if cp_size > 1: + perm = self._cp_unzigzag_permutation(cu_seqlens_list, cp_size, device) + k_full = k_full.index_select(0, perm) + v_full = v_full.index_select(0, perm) + + out = torch.empty(t_local, np_q * hn, dtype=dtype, device=device) + + local_offset = 0 + for s_idx in range(len(cu_seqlens_list) - 1): + seq_start = cu_seqlens_list[s_idx] + seq_len_global = cu_seqlens_list[s_idx + 1] - seq_start + local_len = seq_len_global // cp_size # this sub-seq's local Q count + + q_seq = query[local_offset : local_offset + local_len] + k_seq = k_full[seq_start : seq_start + seq_len_global] + v_seq = v_full[seq_start : seq_start + seq_len_global] + + q4 = q_seq.unsqueeze(0).transpose(1, 2) # [1, np, local_len, hn] + k4 = k_seq.unsqueeze(0).transpose(1, 2) # [1, nk, seq_len, hn] + v4 = v_seq.unsqueeze(0).transpose(1, 2) + + # Global positions of local Q tokens. cp_size=1 degenerates to + # identity; use arange to preserve odd-length seqs (zigzag helper + # floor-divides, dropping the trailing token). + if cp_size > 1: + row_idx = self._zigzag_global_indices(local_len, cp_rank, cp_size, device) + else: + row_idx = torch.arange(local_len, device=device) + col_idx = torch.arange(seq_len_global, device=device) + forbid_future = col_idx[None, :] > row_idx[:, None] + if sliding_window is not None and sliding_window > 0: + forbid_past = col_idx[None, :] < (row_idx[:, None] - (sliding_window - 1)) + forbid = forbid_future | forbid_past + else: + forbid = forbid_future + mask = torch.where( + forbid, + torch.finfo(dtype).min, + 0.0, + ).to(dtype=dtype) + + o = F.scaled_dot_product_attention( + q4, + k4, + v4, + attn_mask=mask[None, None, :, :], + dropout_p=self.dropout_p if self.training else 0.0, + scale=scale, + enable_gqa=(np_q != nk), + ) + out[local_offset : local_offset + local_len] = o.transpose(1, 2).reshape(local_len, -1) + local_offset += local_len + + return out + + def _forward_thd_flash(self, query, key, value, cu_seqlens): + """Sliding-window or head_dim<=256 path via flash_attn_varlen_func. + + CP==1 only. For CP>1, `_forward_cp_subseq_mask` handles zig-zag. + + Sliding-window layers must pass `window_size=(sliding_window-1, 0)` so + only tokens within `sliding_window` positions back are attended to - + this matches HF's `sliding_window_mask_function`. Global layers and + dense-attention sliding layers use the default full-causal window. + """ + from flash_attn import flash_attn_varlen_func + + window_size = (-1, -1) # full causal when causal=True + if self._is_sliding: + sw = getattr(self.config, "sliding_window", None) + if sw and sw > 0: + window_size = (int(sw) - 1, 0) + + cu = cu_seqlens.to(torch.int32) + max_seqlen = (cu[1:] - cu[:-1]).max().item() + out = flash_attn_varlen_func( + query.contiguous(), + key.contiguous(), + value.contiguous(), + cu_seqlens_q=cu, + cu_seqlens_k=cu, + max_seqlen_q=max_seqlen, + max_seqlen_k=max_seqlen, + dropout_p=self.dropout_p if self.training else 0.0, + softmax_scale=self._resolve_scale(query.shape[2]), + causal=True, + window_size=window_size, + ) + return out.reshape(query.shape[0], -1) + + def _forward_thd_sdpa_per_subseq(self, query, key, value, cu_seqlens): + """Per-sub-sequence causal SDPA - used when flash-attn can't handle + head_dim (global layer w/o CP). Avoids materializing a [T, T] mask. + """ + np_q, hn = query.shape[1], query.shape[2] + nk = key.shape[1] + scale = self._resolve_scale(hn) + out = torch.empty(query.shape[0], np_q * hn, dtype=query.dtype, device=query.device) + for i in range(len(cu_seqlens) - 1): + s = cu_seqlens[i].item() + e = cu_seqlens[i + 1].item() + q4 = query[s:e].unsqueeze(0).transpose(1, 2) # [1, np, L, hn] + k4 = key[s:e].unsqueeze(0).transpose(1, 2) + v4 = value[s:e].unsqueeze(0).transpose(1, 2) + o = F.scaled_dot_product_attention( + q4, + k4, + v4, + dropout_p=self.dropout_p if self.training else 0.0, + scale=scale, + is_causal=True, + enable_gqa=(np_q != nk), + ) + out[s:e] = o.transpose(1, 2).reshape(e - s, -1) + return out + + def forward(self, query, key, value, attention_mask=None, attn_mask_type=None, packed_seq_params=None, **kwargs): + cp_size = getattr(self.config, "context_parallel_size", 1) or 1 + is_thd = query.dim() == 3 + + force_cp_path = getattr(self.config, "force_cp_subseq_mask", False) + + if is_thd: + if cp_size > 1 or force_cp_path: + sw = None + if self._is_sliding: + sw_cfg = getattr(self.config, "sliding_window", None) + if sw_cfg and sw_cfg > 0: + sw = int(sw_cfg) + return self._forward_cp_subseq_mask( + query, + key, + value, + packed_seq_params, + sliding_window=sw, + ) + + cu_seqlens = None + if packed_seq_params is not None: + cu_seqlens = packed_seq_params.cu_seqlens_q + + hn = query.shape[2] + if cu_seqlens is not None: + if hn <= 256: + return self._forward_thd_flash(query, key, value, cu_seqlens) + return self._forward_thd_sdpa_per_subseq(query, key, value, cu_seqlens) + + q = query.unsqueeze(0).transpose(1, 2) + k = key.unsqueeze(0).transpose(1, 2) + v = value.unsqueeze(0).transpose(1, 2) + nq, nk = q.shape[1], k.shape[1] + out = F.scaled_dot_product_attention( + q, + k, + v, + dropout_p=self.dropout_p if self.training else 0.0, + scale=self._resolve_scale(hn), + is_causal=True, + enable_gqa=(nq != nk), + ) + return out.transpose(1, 2).reshape(query.shape[0], -1) + + q = query.permute(1, 2, 0, 3) + k = key.permute(1, 2, 0, 3) + v = value.permute(1, 2, 0, 3) + nq, nk = q.shape[1], k.shape[1] + out = F.scaled_dot_product_attention( + q, + k, + v, + dropout_p=self.dropout_p if self.training else 0.0, + scale=self._resolve_scale(query.shape[3]), + is_causal=True, + enable_gqa=(nq != nk), + ) + return out.permute(2, 0, 1, 3).reshape(out.size(2), out.size(0), -1) + + +class Gemma4SelfAttention(SelfAttention): + """SelfAttention with Gemma4-specific modifications: + - v_norm: RMSNorm without learnable scale applied to value states. + - attention_k_eq_v: on global layers the linear_qkv projection emits + ``[q, k]`` only (no v_proj) and V is derived from K - specifically + ``V = v_norm(raw_k)`` while ``K = k_norm(raw_k)``. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._is_global = False # set by Gemma4TransformerLayer after construction + self.v_norm = VNorm(self.hidden_size_per_attention_head, eps=self.config.layernorm_epsilon) + + def _split_qkv_global_k_eq_v(self, hidden_states): + """Split linear_qkv output for global K=V layers. + + The Mcore linear_qkv weight for a K=V global layer is built with + ``v_proj_weight == k_proj_weight`` (see Gemma4Bridge + convert_gemma4_to_hf), + so ``linear_qkv(h)`` emits Q/K/V with ``raw_k == raw_v``. Gemma4's + per-head norms then apply as ``key = k_norm(raw_k)`` and + ``value = v_norm(raw_k)`` - *not* ``v_norm(k_norm(raw_k))``. We + reimplement the split here rather than calling the parent so we + don't have to mutate ``self.k_layernorm`` mid-forward. + + Returns (query[sq,b,np,hn], key[sq,b,ng,hn], value[sq,b,ng,hn]). + """ + mixed_qkv, _ = self.linear_qkv(hidden_states) + num_query_heads_per_group = self.num_attention_heads_per_partition // self.num_query_groups_per_partition + new_shape = mixed_qkv.size()[:-1] + ( + self.num_query_groups_per_partition, + (num_query_heads_per_group + 2) * self.hidden_size_per_attention_head, + ) + mixed_qkv = mixed_qkv.view(*new_shape) + + q_width = num_query_heads_per_group * self.hidden_size_per_attention_head + hn = self.hidden_size_per_attention_head + query, raw_key, _raw_value = torch.split(mixed_qkv, [q_width, hn, hn], dim=3) + query = query.reshape(query.size(0), query.size(1), -1, hn) + + if self.q_layernorm is not None: + query = self.q_layernorm(query) + + value = self.v_norm(raw_key) + key = self.k_layernorm(raw_key) if self.k_layernorm is not None else raw_key + return query, key, value + + def get_query_key_value_tensors(self, hidden_states, key_value_states=None, output_gate=False, split_qkv=True): + if self._is_global and self.config.attention_k_eq_v and split_qkv: + if output_gate: + raise NotImplementedError("output_gate is not supported together with attention_k_eq_v") + return self._split_qkv_global_k_eq_v(hidden_states) + + result = super().get_query_key_value_tensors( + hidden_states, key_value_states, output_gate=output_gate, split_qkv=split_qkv + ) + if not split_qkv: + return result + + if output_gate: + query, key, value, gate = result + value = self.v_norm(value) + return query, key, value, gate + + query, key, value = result + value = self.v_norm(value) + return query, key, value + + +def _build_moe_submodule_spec(config): + """Build the MoE submodule spec (Gemma4MoELayer + TE GroupedMLP experts).""" + from megatron.core.extensions.transformer_engine_spec_provider import TESpecProvider + from megatron.core.models.gpt.moe_module_specs import get_moe_module_spec_for_backend + + base_spec = get_moe_module_spec_for_backend( + backend=TESpecProvider(), + num_experts=config.num_moe_experts, + moe_grouped_gemm=config.moe_grouped_gemm, + use_te_activation_func=False, # use plain F.gelu(approximate='tanh') from config.activation_func + ) + return ModuleSpec( + module=Gemma4MoELayer, + submodules=base_spec.submodules, + metainfo=base_spec.metainfo, + ) + + +def get_gemma4_layer_spec_te(config=None) -> ModuleSpec: + """Layer spec for Gemma4 using native Megatron attention with TE. + + If ``config.enable_moe_block`` is set, the main ``mlp`` submodule is a + :class:`Gemma4MoELayer` (so that the state-dict path + ``.mlp.experts.linear_fc*.weight*`` matches the weight-conversion path), + and the original dense MLP moves to a sibling ``dense_mlp`` submodule that + the layer forward sums with the MoE output. For the 31B dense variant, + ``enable_moe_block=False`` and ``mlp`` stays as the normal Megatron MLP. + """ + # dense_mlp: use a plain (non-fused-layernorm) linear_fc1 so our explicit + # `pre_mlp_layernorm` in the layer forward is the sole norm applied to the + # MLP input. Using TELayerNormColumnParallelLinear here would apply a + # SECOND layernorm inside fc1, resulting in double-normalization and + # ~8x inflated MLP outputs. + dense_mlp_spec = ModuleSpec( + module=MLP, + submodules=MLPSubmodules( + linear_fc1=TEColumnParallelLinear, + linear_fc2=TERowParallelLinear, + ), + ) + if config is not None and getattr(config, "enable_moe_block", False): + mlp_spec = _build_moe_submodule_spec(config) + dense_spec = dense_mlp_spec + else: + mlp_spec = dense_mlp_spec + dense_spec = IdentityOp + + submods = Gemma4TransformerLayerSubmodules( + self_attention=ModuleSpec( + module=Gemma4SelfAttention, + params={"attn_mask_type": AttnMaskType.causal}, + submodules=SelfAttentionSubmodules( + linear_qkv=TELayerNormColumnParallelLinear, + core_attention=TEDotProductAttention, + linear_proj=TERowParallelLinear, + q_layernorm=TENorm, + k_layernorm=TENorm, + ), + ), + self_attn_bda=get_bias_dropout_add, + pre_mlp_layernorm=IdentityOp, + mlp=mlp_spec, + mlp_bda=get_bias_dropout_add, + post_attention_layernorm=TENorm, + post_feedforward_layernorm=TENorm, + dense_mlp=dense_spec, + ) + return ModuleSpec(module=Gemma4TransformerLayer, submodules=submods) + + +@functools.lru_cache(maxsize=4) +def _load_hf_text_config(hf_checkpoint): + """Load HF config and unwrap `text_config` if it's a multimodal wrapper. + + Cached via lru_cache so the model provider and weight converters share the + same parsed object. + """ + from transformers import AutoConfig + + cfg = AutoConfig.from_pretrained(hf_checkpoint, trust_remote_code=True) + return cfg.text_config if hasattr(cfg, "text_config") else cfg + + +class _Gemma4MoELayerWarningFilter(logging.Filter): + """Silence the once-per-layer Megatron warning: + 'Unknown MLP type: . Using default kwargs.' + Megatron's TransformerLayer.__init__ recognizes a hardcoded tuple of MLP + classes via `==` (not issubclass), so Gemma4MoELayer (a MoELayer subclass) + falls through to the default-kwargs branch. That branch is correct for us + - Gemma4MoELayer.__init__ fetches its own pg_collection via + get_default_pg_collection - but the warning spams 30 lines per layer at + init and confuses log readers. See gemma4_provider.py install hook. + """ + + def filter(self, record: logging.LogRecord) -> bool: + msg = record.getMessage() + return not ("Unknown MLP type" in msg and "Gemma4MoELayer" in msg) + + +def _install_moe_warning_filter(): + """Silence the per-layer "Unknown MLP type: Gemma4MoELayer" warning. + + Megatron's TransformerLayer compares MLP class identity via ``==``, so + MoELayer subclasses hit the default-kwargs branch and log a warning. + The default-kwargs branch is correct for us (Gemma4MoELayer fetches + pg_collection itself); filter the noise. + """ + tl_logger = logging.getLogger("megatron.core.transformer.transformer_layer") + if getattr(tl_logger, "_gemma4_moe_filter_installed", False): + return + tl_logger.addFilter(_Gemma4MoELayerWarningFilter()) + tl_logger._gemma4_moe_filter_installed = True + + +def _assert_hf_features_supported(hf_text): + """Fail loudly on Gemma4 HF features this plugin doesn't implement.""" + if getattr(hf_text, "hidden_size_per_layer_input", 0): + raise NotImplementedError( + "Gemma4 per-layer input mechanism " + f"(hidden_size_per_layer_input={hf_text.hidden_size_per_layer_input}) " + "is not implemented. See Gemma4TextDecoderLayer.per_layer_input_gate in HF." + ) + if getattr(hf_text, "num_kv_shared_layers", 0): + raise NotImplementedError( + "Gemma4 KV-sharing across the last N layers " + f"(num_kv_shared_layers={hf_text.num_kv_shared_layers}) is not implemented." + ) + if getattr(hf_text, "use_double_wide_mlp", False): + raise NotImplementedError("Gemma4 use_double_wide_mlp is not implemented.") + # Text-only training assumes causal attention; HF's "all" mode disables it. + if getattr(hf_text, "use_bidirectional_attention", "vision") == "all": + raise NotImplementedError("Gemma4 use_bidirectional_attention='all' disables causal masking; not supported.") + + +def _apply_core_config(config, hf_text): + """Set Gemma4's non-MoE, non-RoPE config fields. + + Mutates ``config`` in place. Promotes its ``__class__`` to + ``Gemma4TransformerConfig`` so the new dataclass fields are reachable + from downstream Megatron code. + """ + # Gemma uses GeGLU (gated gelu-tanh), not SwiGLU. + config.gated_linear_unit = True + config.activation_func = _gelu_tanh + config.bias_activation_fusion = False + + # No MoE-vs-dense layer scheduling: every layer is our Gemma4TransformerLayer + # and the MoE block lives inside its forward. An all-zero list keeps + # transformer_block's non_homogeneous_layers=True branch active (correct for + # 26B's differing global vs sliding head_dim / num_kv_heads). + # Rationale for using moe_layer_freq as the flag: Megatron's + # TransformerBlock.__init__ sets ``non_homogeneous_layers = True`` iff + # ``config.moe_layer_freq is not None``. We only need that flag on - + # the actual dense/MoE dispatch happens inside + # Gemma4TransformerLayer.forward, so the list contents are never + # consulted by TransformerBlock itself. If a future Megatron refactor + # starts reading the list per-layer, we need a Gemma4-specific schedule + # instead. + config.moe_layer_freq = [0] * config.num_layers + + # Mirror Megatron's own misspelling (`hetereogenous_*`) - correcting it + # would silently no-op on Megatron's read path. + config.hetereogenous_dist_checkpoint = True + + config.__class__ = Gemma4TransformerConfig + if hasattr(hf_text, "per_layer_config"): + global_config = next( + layer_config + for layer_type, layer_config in zip(hf_text.layer_types, hf_text.per_layer_config, strict=True) + if layer_type == "full_attention" + ) + config.global_kv_channels = global_config.head_dim + config.global_num_query_groups = global_config.num_key_value_heads + else: + config.global_kv_channels = hf_text.global_head_dim + config.global_num_query_groups = hf_text.num_global_key_value_heads + config.attention_k_eq_v = getattr(hf_text, "attention_k_eq_v", True) + config.final_logit_softcapping = getattr(hf_text, "final_logit_softcapping", 30.0) + config.sliding_window = hf_text.sliding_window + + # `sliding_window_pattern` isn't in Gemma4 HF configs - infer from + # layer_types (first full_attention layer's 1-indexed position). + layer_types = list(getattr(hf_text, "layer_types", [])) + try: + config.sliding_window_pattern = layer_types.index("full_attention") + 1 + except ValueError: + config.sliding_window_pattern = 6 + + # Q/K norms handle softmax scaling; Megatron's default of 1/sqrt(hn) is wrong. + config.softmax_scale = 1.0 + # Fused RoPE ignores zeroed inv_freq tails; we need unfused for partial-rotary. + config.apply_rope_fusion = False + + +def _apply_moe_config(config, hf_text): + """Set MoE fields if this is a MoE variant (26B-A4B).""" + config.enable_moe_block = getattr(hf_text, "enable_moe_block", False) + if not config.enable_moe_block: + return + + config.num_moe_experts = hf_text.num_experts + config.moe_router_topk = hf_text.top_k_experts + config.moe_ffn_hidden_size = hf_text.moe_intermediate_size + # Megatron MoE infrastructure reads these even though our custom router + # bypasses its scoring logic; defaults mirror a working Qwen3.5-A3B config. + config.moe_token_dispatcher_type = getattr(config, "moe_token_dispatcher_type", None) or "alltoall" + config.moe_grouped_gemm = getattr(config, "moe_grouped_gemm", None) or True + config.moe_aux_loss_coeff = 0.0 # Gemma4 router has no aux loss + config.moe_router_load_balancing_type = getattr(config, "moe_router_load_balancing_type", None) or "none" + config.moe_router_score_function = getattr(config, "moe_router_score_function", None) or "softmax" + config.moe_router_topk_scaling_factor = getattr(config, "moe_router_topk_scaling_factor", None) or 1.0 + config.moe_router_pre_softmax = False + + +def get_rope_local_base_freq(hf_text) -> float: + """Extract sliding-attention RoPE theta from an HF Gemma4 text config. + + Single source of truth for both the model provider and model config path; + otherwise the 10000.0 default would drift between + call sites. + """ + return (getattr(hf_text, "rope_parameters", {}) or {}).get("sliding_attention", {}).get("rope_theta", 10000.0) + + +def _apply_rope_config(config, hf_text): + rope_params = getattr(hf_text, "rope_parameters", {}) or {} + config.rope_local_base_freq = get_rope_local_base_freq(hf_text) + config.global_partial_rotary_factor = rope_params.get("full_attention", {}).get("partial_rotary_factor", 0.25) + + +def _guard_cp_sliding_window(args, config): + """Fail if per-rank CP token cap is smaller than the sliding window. + + Strong signal of a miscounted CP sizing - we'd train on truncated + attention windows otherwise. + """ + cp_size = getattr(args, "context_parallel_size", 1) or 1 + if cp_size <= 1: + return + max_tokens = getattr(args, "max_tokens_per_gpu", None) + if max_tokens is not None and max_tokens < config.sliding_window: + raise ValueError( + f"context_parallel_size={cp_size} with max_tokens_per_gpu={max_tokens} " + f"< sliding_window={config.sliding_window}: per-rank CP chunk cap is " + "smaller than the sliding window. Reduce CP or raise max_tokens_per_gpu." + ) + + +def get_gemma4_spec(args, config, vp_stage): + """Return the native Gemma4 layer spec with proper config overrides.""" + hf_text = _load_hf_text_config(args.hf_checkpoint) + + _install_moe_warning_filter() + _assert_hf_features_supported(hf_text) + _apply_core_config(config, hf_text) + _apply_moe_config(config, hf_text) + _apply_rope_config(config, hf_text) + _guard_cp_sliding_window(args, config) + + spec = get_gemma4_layer_spec_te(config) + from megatron.core.extensions.transformer_engine_spec_provider import TESpecProvider + + if not getattr(config, "enable_moe_block", False): + spec.submodules.mlp.submodules.linear_fc1 = TEColumnParallelLinear + spec.submodules.mlp.metainfo = {"fuse_pre_mlp_layernorm": False} + spec.submodules.pre_mlp_layernorm = TESpecProvider().layer_norm() + return spec diff --git a/vime_plugins/models/gemma4_provider.py b/vime_plugins/models/gemma4_provider.py new file mode 100644 index 000000000..0411a9031 --- /dev/null +++ b/vime_plugins/models/gemma4_provider.py @@ -0,0 +1,344 @@ +"""Custom model provider for Gemma4. + +Installs Gemma4-specific behaviors that sit outside the transformer layer: +- embedding scaling (multiply embeddings by sqrt(hidden_size)) +- logit softcapping (`final_logit_softcapping`) +- dual-RoPE (different rope_theta + partial-rotary for global vs sliding layers) +- layer_scalar buffers loaded from the HF checkpoint +""" + +import json +import logging +import os + +import torch +from megatron.core.models.gpt.gpt_model import GPTModel +from megatron.core.transformer.spec_utils import import_module +from megatron.training import get_args +from megatron.training.arguments import core_transformer_config_from_args + +from vime_plugins.models.gemma4 import _load_hf_text_config + +logger = logging.getLogger(__name__) + + +def _is_rank_zero() -> bool: + if not torch.distributed.is_available() or not torch.distributed.is_initialized(): + return True + return torch.distributed.get_rank() == 0 + + +def model_provider(pre_process=True, post_process=True, vp_stage=None): + args = get_args() + config = core_transformer_config_from_args(args) + + transformer_layer_spec = import_module(args.spec) + if callable(transformer_layer_spec): + transformer_layer_spec = transformer_layer_spec(args, config, vp_stage) + + model = GPTModel( + config=config, + transformer_layer_spec=transformer_layer_spec, + vocab_size=args.padded_vocab_size, + max_sequence_length=args.max_position_embeddings, + pre_process=pre_process, + post_process=post_process, + fp16_lm_cross_entropy=args.fp16_lm_cross_entropy, + parallel_output=True, + share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights, + position_embedding_type=args.position_embedding_type, + rotary_percent=args.rotary_percent, + rotary_base=args.rotary_base, + rope_scaling=args.use_rope_scaling, + ) + + _mark_expert_weights_for_direct_hf_loading(model) + _install_hooks(model, args, config, pre_process, post_process) + _promote_layer_scalars(model) + return model + + +def _mark_expert_weights_for_direct_hf_loading(model): + from megatron.core.tensor_parallel import set_tensor_model_parallel_attributes + + for name, parameter in model.named_parameters(): + if ".mlp.experts.linear_fc1.weight" in name: + set_tensor_model_parallel_attributes(parameter, True, 0, 1) + elif ".mlp.experts.linear_fc2.weight" in name: + set_tensor_model_parallel_attributes(parameter, True, 1, 1) + + +def _promote_layer_scalars(model): + for module in model.modules(): + layer_scalar = module._buffers.pop("layer_scalar", None) + if layer_scalar is not None: + module.register_parameter("layer_scalar", torch.nn.Parameter(layer_scalar, requires_grad=False)) + + +class DualRotaryEmbedding(torch.nn.Module): + """Wraps a (global, local) pair of RotaryEmbedding modules and emits a + single concatenated tensor (global part first). ``Gemma4TransformerLayer`` + slices it per-layer based on ``is_sliding``. Concat (not tuple) because + Megatron's ``SelfAttention.forward`` reads a 2-tuple as + ``(self_attn, cross_attn)`` RoPE and would misread our pair. + """ + + def __init__(self, local_rope, global_rope, global_dim: int): + super().__init__() + self.local_rope = local_rope + self.global_rope = global_rope + self.global_dim = global_dim + + def get_rotary_seq_len(self, *args, **kwargs): + return self.local_rope.get_rotary_seq_len(*args, **kwargs) + + def forward(self, seq_len, **kwargs): + global_emb = self.global_rope(seq_len, **kwargs) + local_emb = self.local_rope(seq_len, **kwargs) + return torch.cat([global_emb, local_emb], dim=-1) + + +class _Gemma4LogitSoftcap(torch.autograd.Function): + """Apply Gemma4 final logit softcapping without allocating new logits.""" + + @staticmethod + def forward(ctx, logits: torch.Tensor, scale: float) -> torch.Tensor: + ctx.scale = scale + ctx.mark_dirty(logits) + logits.div_(scale) + logits.tanh_() + logits.mul_(scale) + ctx.save_for_backward(logits) + return logits + + @staticmethod + def backward(ctx, grad_output: torch.Tensor) -> tuple[torch.Tensor, None]: + (softcapped,) = ctx.saved_tensors + scale = ctx.scale + grad_logits = softcapped / scale + grad_logits.pow_(2) + grad_logits.neg_() + grad_logits.add_(1.0) + grad_logits.mul_(grad_output) + return grad_logits, None + + +def _logit_softcapping(logits: torch.Tensor, scale: float) -> torch.Tensor: + if scale <= 0: + return logits + return _Gemma4LogitSoftcap.apply(logits, float(scale)) + + +def _install_hooks(model, args, config, pre_process, post_process): + """Install Gemma4-specific pre/post-process hooks on a built GPTModel. + + We use ``register_forward_hook`` rather than subclassing GPTModel + because: + - Two independent behaviors (embed scale, softcap) on two different + submodules. Subclassing would require overriding + ``GPTModel.forward`` and branching on pp/vp stage. + - The hooks are shape- and dtype-preserving, so they compose cleanly + with PP (only first-stage runs embedding, only last-stage runs + output_layer) - we gate registration on ``pre_process`` / + ``post_process`` accordingly. + - Keeps the diff local to this plugin: we don't need to shadow any + Megatron-maintained class. + """ + hf_text = _load_hf_text_config(args.hf_checkpoint) + hidden_size = config.hidden_size + + inner = model.module if hasattr(model, "module") else model + + # Embedding scaling - HF applies this inside the embedding module. + # See ``Gemma4TextScaledWordEmbedding``: the scale is stored as an fp32 + # tensor and cast to the embedding weight's dtype at forward time, so + # the scale-as-applied depends on the current weight dtype (bf16 during + # training, fp32 during some eval paths). We match that behavior here. + if pre_process and hasattr(inner, "embedding"): + embed_scale = torch.tensor(hidden_size**0.5) # fp32 + + def _embed_hook(module, inp, output): + return output * embed_scale.to(output.dtype) + + inner.embedding.register_forward_hook(_embed_hook) + + # Final logit softcapping - HF applies tanh(logits / cap) * cap. + # Some Megatron output_layer variants (parallel_output paths) return + # ``(logits, bias)``; we pass the non-logit tail through unchanged. + softcap = getattr(hf_text, "final_logit_softcapping", None) + if post_process and softcap and hasattr(inner, "output_layer"): + + def _softcap_hook(module, inp, output): + if isinstance(output, tuple): + return (_logit_softcapping(output[0], softcap),) + output[1:] + return _logit_softcapping(output, softcap) + + inner.output_layer.register_forward_hook(_softcap_hook) + + # Dual RoPE: replace Megatron's single rotary_pos_emb with a wrapper that + # produces (global, local) RoPE side-by-side. Gemma4 uses partial-rotary + # on global layers (implemented here by zeroing the tail of inv_freq). + if hasattr(inner, "rotary_pos_emb"): + from megatron.core.models.common.embeddings.rotary_pos_embedding import RotaryEmbedding + + rope_params = getattr(hf_text, "rope_parameters", {}) or {} + full = rope_params.get("full_attention", {}) or {} + sliding = rope_params.get("sliding_attention", {}) or {} + global_theta = full.get("rope_theta", 1_000_000.0) + local_theta = sliding.get("rope_theta", 10_000.0) + global_head_dim = config.global_kv_channels + global_partial = full.get("partial_rotary_factor", 0.25) + + local_rope = inner.rotary_pos_emb # already built with args.rotary_base + + global_rope = RotaryEmbedding( + kv_channels=global_head_dim, + rotary_percent=1.0, + rotary_base=global_theta, + ) + # HF "proportional" RoPE: first (partial * head_dim // 2) inv_freq + # entries are live, the rest are zero (no rotation on those dims). + # Writing this to the existing buffer keeps device/dtype correct. + rope_angles = int(global_partial * global_head_dim // 2) + half = global_head_dim // 2 + # Guard the RoPE geometry: 0 means "no rotation" (nonsensical here); + # > half would produce nope<0 and a shape-mismatched copy_. Both + # should fail loudly rather than silently writing garbage. + assert 0 < rope_angles <= half, ( + f"global_partial_rotary_factor={global_partial} with " + f"global_head_dim={global_head_dim} produced rope_angles=" + f"{rope_angles}; must be in (0, {half}]." + ) + inv_freq_live = 1.0 / ( + global_theta ** (torch.arange(0, 2 * rope_angles, 2, dtype=torch.float) / global_head_dim) + ) + nope = half - rope_angles + inv_freq = torch.cat([inv_freq_live, torch.zeros(nope)]) if nope > 0 else inv_freq_live + assert inv_freq.shape == global_rope.inv_freq.shape, ( + f"inv_freq shape {tuple(inv_freq.shape)} doesn't match " + f"global_rope.inv_freq shape {tuple(global_rope.inv_freq.shape)}; " + "Megatron RotaryEmbedding layout may have changed." + ) + global_rope.inv_freq.copy_(inv_freq.to(global_rope.inv_freq.device)) + + inner.rotary_pos_emb = DualRotaryEmbedding(local_rope, global_rope, global_head_dim) + config.dual_rope_global_dim = global_head_dim + if _is_rank_zero(): + logger.info( + "DualRotaryEmbedding: local_theta=%s global_theta=%s " "global_dim=%s rope_angles=%d (nope=%d)", + local_theta, + global_theta, + global_head_dim, + rope_angles, + nope, + ) + + if hasattr(inner, "decoder") and args.hf_checkpoint: + _load_layer_scalars(inner, args.hf_checkpoint, config) + + +def _read_layer_scalars_from_safetensors(hf_checkpoint: str) -> dict[int, float] | None: + """Read all ``layer_scalar`` values from the HF safetensors checkpoint. + + Returns ``{global_layer_idx: scalar}`` or ``None`` if the checkpoint has + no safetensors index (older HF layouts) or no layer_scalar weights. Only + called on rank 0 - results are broadcast to the other ranks. + """ + index_path = os.path.join(hf_checkpoint, "model.safetensors.index.json") + if not os.path.exists(index_path): + logger.warning("No safetensors index at %s; skipping layer scalars", index_path) + return None + + from safetensors import safe_open + + with open(index_path) as f: + index = json.load(f) + + scalars: dict[int, float] = {} + for key, filename in index["weight_map"].items(): + if "layer_scalar" not in key: + continue + layer_idx = int(key.split(".layers.")[1].split(".")[0]) + with safe_open(os.path.join(hf_checkpoint, filename), framework="pt", device="cpu") as sf: + scalars[layer_idx] = sf.get_tensor(key).item() + + if not scalars: + logger.warning("No layer_scalar weights found in checkpoint %s", hf_checkpoint) + return None + return scalars + + +def _broadcast_layer_scalars(scalars: dict[int, float] | None) -> dict[int, float] | None: + """Broadcast the rank-0-read ``scalars`` dict to every rank. + + safetensors reads on every rank cause an O(world_size) fan-out of tiny + reads on the shared filesystem; the dict itself is a few kilobytes. If + ``torch.distributed`` isn't initialized (single-process run), we simply + return the input dict. + """ + if not torch.distributed.is_available() or not torch.distributed.is_initialized(): + return scalars + obj = [scalars] if torch.distributed.get_rank() == 0 else [None] + torch.distributed.broadcast_object_list(obj, src=0) + return obj[0] + + +def _load_layer_scalars(inner, hf_checkpoint, config): + # Wrong layer_scalars materially change activations vs HF (they're per- + # layer multiplicative gains on the residual stream, not decorative), so + # by default we fail hard if the load breaks. Set + # GEMMA4_ALLOW_MISSING_LAYER_SCALARS=1 to downgrade to a warning and + # train with the default value of 1.0 - only useful for debug runs + # against a checkpoint that genuinely lacks these buffers. + allow_missing = os.environ.get("GEMMA4_ALLOW_MISSING_LAYER_SCALARS") == "1" + try: + scalars = _read_layer_scalars_from_safetensors(hf_checkpoint) if _is_rank_zero() else None + scalars = _broadcast_layer_scalars(scalars) + if not scalars: + if allow_missing: + return + raise RuntimeError( + "No layer_scalar weights found in checkpoint; set " + "GEMMA4_ALLOW_MISSING_LAYER_SCALARS=1 to proceed with " + "default values (not numerically equivalent to HF)." + ) + + # Under pipeline-parallelism, inner.decoder.layers holds only this + # rank's local subset. Translate the local index back to the global + # (HF 0-indexed) layer index so we apply the right scalar per layer. + from megatron.core.transformer.transformer_layer import get_transformer_layer_offset + + pp_offset = get_transformer_layer_offset(config) + + loaded = 0 + for i, layer in enumerate(inner.decoder.layers): + if hasattr(layer, "layer_scalar"): + global_idx = i + pp_offset + if global_idx not in scalars: + if allow_missing: + logger.warning( + "layer_scalar for global layer %d missing; using default 1.0", + global_idx, + ) + else: + raise KeyError( + f"layer_scalar for global layer {global_idx} " + f"missing in checkpoint (have: {sorted(scalars)[:10]}...); " + "checkpoint may be truncated." + ) + layer.layer_scalar.fill_(scalars.get(global_idx, 1.0)) + loaded += 1 + if _is_rank_zero(): + logger.info( + "Applied %d/%d layer scalars (pp_offset=%d, range=%.4f..%.4f)", + loaded, + len(inner.decoder.layers), + pp_offset, + min(scalars.values()), + max(scalars.values()), + ) + except (FileNotFoundError, json.JSONDecodeError) as e: + if allow_missing: + logger.warning("layer scalars unavailable (%s: %s); using default 1.0", type(e).__name__, e) + return + raise