From dc614a2611289d11dcc04a53cd016d5a0b9d0879 Mon Sep 17 00:00:00 2001 From: Hollow Man Date: Thu, 10 Sep 2026 02:18:23 +0300 Subject: [PATCH 1/8] feat(model): add GLM-5.3-Flash support Signed-off-by: Hollow Man --- src/megatron/bridge/models/__init__.py | 8 + .../bridge/models/conversion/__init__.py | 4 + .../bridge/models/conversion/model_bridge.py | 29 + .../bridge/models/conversion/param_mapping.py | 81 ++ .../bridge/models/conversion/peft_bridge.py | 14 +- .../models/deepseek/deepseek_v4_bridge.py | 100 +- .../bridge/models/glm5next/__init__.py | 26 + .../models/glm5next/glm53flash_bridge.py | 929 ++++++++++++++++++ .../bridge/models/hybrid/hybrid_provider.py | 2 + .../bridge/models/hybrid_mla_provider.py | 148 +++ .../qwen_vl/modelling_qwen3_vl/model.py | 2 +- .../qwen_vl/modelling_qwen3_vl/utils.py | 109 -- .../bridge/training/utils/packed_seq_utils.py | 106 +- .../glm5next/test_glm53flash_mapping.py | 160 +++ .../qwen_vl/modelling_qwen3_vl/test_utils.py | 29 + .../models/test_model_bridge_lora.py | 20 + 16 files changed, 1556 insertions(+), 211 deletions(-) create mode 100644 src/megatron/bridge/models/glm5next/__init__.py create mode 100644 src/megatron/bridge/models/glm5next/glm53flash_bridge.py create mode 100644 src/megatron/bridge/models/hybrid_mla_provider.py create mode 100644 tests/unit_tests/models/glm5next/test_glm53flash_mapping.py diff --git a/src/megatron/bridge/models/__init__.py b/src/megatron/bridge/models/__init__.py index a542e55c73..3e7736a6da 100644 --- a/src/megatron/bridge/models/__init__.py +++ b/src/megatron/bridge/models/__init__.py @@ -93,6 +93,11 @@ GLM45Bridge, GLM47FlashBridge, ) +from megatron.bridge.models.glm5next import ( + GLM53FlashBridge, + GLM53FlashModel, + GLM53FlashModelProvider, +) from megatron.bridge.models.glm_moe_dsa import ( GLM5Bridge, ) @@ -258,6 +263,9 @@ "GLM45Bridge", "GLM47FlashBridge", "GLM5Bridge", + "GLM53FlashBridge", + "GLM53FlashModel", + "GLM53FlashModelProvider", "GLM45VBridge", "GLM45VModelProvider", "GPTModelProvider", diff --git a/src/megatron/bridge/models/conversion/__init__.py b/src/megatron/bridge/models/conversion/__init__.py index 7b4b493dec..d7c9d58d46 100644 --- a/src/megatron/bridge/models/conversion/__init__.py +++ b/src/megatron/bridge/models/conversion/__init__.py @@ -22,6 +22,8 @@ FusedExpertMapping, FusedGatedExpertMapping, GatedMLPMapping, + HCAlphaMapping, + InitOnlyMapping, MegatronParamMapping, QKVMapping, ReplicatedMapping, @@ -43,5 +45,7 @@ "ReplicatedMapping", "RowParallelMapping", "AutoMapping", + "HCAlphaMapping", + "InitOnlyMapping", "weights_verification_table", ] diff --git a/src/megatron/bridge/models/conversion/model_bridge.py b/src/megatron/bridge/models/conversion/model_bridge.py index feb9645102..1a24f2de85 100644 --- a/src/megatron/bridge/models/conversion/model_bridge.py +++ b/src/megatron/bridge/models/conversion/model_bridge.py @@ -1044,6 +1044,35 @@ def mapping_registry(self): """ raise NotImplementedError("Subclass must implement mapping_registry method") + @staticmethod + def _prefix_mapping_registry( + registry: MegatronMappingRegistry, megatron_prefix: str, *extra: MegatronParamMapping + ) -> MegatronMappingRegistry: + """Re-key a text-model mapping registry for a VLM wrapper. + + VLM bridges typically reuse a text bridge's ``mapping_registry`` but must + re-prefix every Megatron parameter (e.g. ``embedding.*`` → + ``language_model.embedding.*``) because the language model is a submodule of + the VLM wrapper. Concatenation/delegation mappings that hold a nested + ``_tp_mapping`` keep it in sync with the parent ``megatron_param``. + + Args: + registry: The text-model registry returned by ``super().mapping_registry()``. + megatron_prefix: Prefix prepended to each ``megatron_param``. + *extra: Additional mappings appended verbatim (e.g. replicated vision + weights); they are NOT prefixed. + + Returns: + A new ``MegatronMappingRegistry`` with prefixed text mappings + extras. + """ + mappings = list(registry.mappings) + for mapping in mappings: + mapping.megatron_param = megatron_prefix + mapping.megatron_param + nested = getattr(mapping, "_tp_mapping", None) + if nested is not None: + nested.megatron_param = mapping.megatron_param + return MegatronMappingRegistry(*mappings, *extra) + def _megatron_global_param_names_all_pp_ranks( self, megatron_model: Union[MegatronModel, List[MegatronModel]] ) -> List[str]: diff --git a/src/megatron/bridge/models/conversion/param_mapping.py b/src/megatron/bridge/models/conversion/param_mapping.py index 2eafdd682d..fa6e42ac42 100644 --- a/src/megatron/bridge/models/conversion/param_mapping.py +++ b/src/megatron/bridge/models/conversion/param_mapping.py @@ -4397,3 +4397,84 @@ def split_kv_weights(provider: TransformerConfig, kv: torch.Tensor) -> Tuple[tor k = kv_reshaped[k_slice].reshape(-1, hidden_size) v = kv_reshaped[v_slice].reshape(-1, hidden_size) return k, v + + +class HCAlphaMapping(MegatronParamMapping): + """Map a 3-element HF hyper-connection scale tensor to/from Megatron's + three separate ``alpha_pre`` / ``alpha_post`` / ``alpha_res`` parameters. + + The HF checkpoint stores the three scaling coefficients stacked along dim 0 + as ``.hc_{kind}_scale`` (shape ``[3]``); Megatron keeps them as three + independent ``nn.Parameter([1])`` tensors on a single module. + + One :class:`HCAlphaMapping` is registered **per alpha parameter** (index 0, 1, + or 2). Import extracts element ``[index:index+1]`` from the HF scale tensor. + Only the ``index == 0`` (``alpha_pre``) mapping exports: it gathers all three + alphas from the same module and concatenates them back into the ``[3]`` tensor. + The ``index != 0`` mappings return ``{}`` on export (the primary already wrote + all three), so they exist only to drive the import side and suppress export + "no mapping found" warnings. + + This serves both models with a single wrapped ``hyper_connection`` submodule + (e.g. GLM-5.3-Flash) and models with separate ``self_attention_hyper_connection`` / + ``mlp_hyper_connection`` submodules (e.g. DeepSeek-V4); the caller chooses the + ``megatron_param`` path accordingly. + """ + + def __init__(self, megatron_param: str, hf_param: str, index: int): + super().__init__(megatron_param=megatron_param, hf_param=hf_param) + self._index = index + # index != 0 imports from the shared HF scale tensor but never exports + # (the index-0 mapping writes all three alphas). Its HF key collides with + # the index-0 key, so skip the strict hf_keys presence check. + self.allow_hf_name_mismatch = index != 0 + + def resolve(self, captures): + resolved_mg, resolved_hf = self._resolve_names(captures) + return HCAlphaMapping(resolved_mg, resolved_hf, self._index) + + def hf_to_megatron(self, hf_weights, megatron_module): + attr = ("alpha_pre", "alpha_post", "alpha_res")[self._index] + return hf_weights[self._index : self._index + 1].to(getattr(megatron_module, attr).device) + + def megatron_to_hf(self, megatron_weights, megatron_module): + if self._index != 0: + return {} + parts = [] + for attr, weight in ( + ("alpha_pre", megatron_weights), + ("alpha_post", getattr(megatron_module, "alpha_post", None)), + ("alpha_res", getattr(megatron_module, "alpha_res", None)), + ): + parts.append(self.broadcast_from_pp_rank(weight, cache_key=f"{self.hf_param}_{attr}")) + if parts[0] is None: + return {} + return {self.hf_param: torch.cat([self.maybe_dequantize(part).float() for part in parts])} + + +class InitOnlyMapping(MegatronParamMapping[torch.Tensor]): + """Keep Megatron initialization for parameters that have no HF counterpart. + + Used for hyper-connection / MTP parameters that MCore creates but the source + checkpoint does not store. The mapping declares a synthetic HF key ending in + ``__init_only__`` (which never exists in any checkpoint) and returns ``None`` + on import so the random init is preserved. Export is a no-op. + + Bridges that want this behaviour should also short-circuit the HF state-dict + lookup for the synthetic key (e.g. in ``maybe_modify_loaded_hf_weight``). + """ + + def __init__(self, megatron_param: str): + super().__init__(megatron_param=megatron_param, hf_param=megatron_param + ".__init_only__") + self.allow_hf_name_mismatch = True + + def resolve(self, captures): + resolved_mg, _ = self._resolve_names(captures) + return InitOnlyMapping(resolved_mg) + + def hf_to_megatron(self, hf_weights, megatron_module): + # hf_weights is None when the (synthetic) HF key is absent -> keep init. + return None + + def megatron_to_hf(self, megatron_weights, megatron_module): + return {} diff --git a/src/megatron/bridge/models/conversion/peft_bridge.py b/src/megatron/bridge/models/conversion/peft_bridge.py index 45c4277065..dd82fb2352 100644 --- a/src/megatron/bridge/models/conversion/peft_bridge.py +++ b/src/megatron/bridge/models/conversion/peft_bridge.py @@ -371,10 +371,16 @@ def _split_qkv_linear_out_weight( """Split a fused LoRA linear_out tensor for QKV adapters.""" model = megatron_model[0] if isinstance(megatron_model, list) else megatron_model - # Pass the LoRA rank as feature_dim so split_qkv_weights doesn't - # mistake it for an FP8 compressed hidden_size. - feature_dim = linear_out_weight.shape[-1] if linear_out_weight.ndim == 2 else None - q_out, k_out, v_out = split_qkv_weights(model.config, linear_out_weight, feature_dim=feature_dim) + config = model.config + if getattr(config, "kda_two_stage_gates", False): + # Low-rank KDA gates leave a contiguous Q/K/V-only input projection. + q_dim = config.linear_num_key_heads * config.linear_key_head_dim + v_dim = config.linear_num_value_heads * config.linear_value_head_dim + q_out, k_out, v_out = linear_out_weight.split((q_dim, q_dim, v_dim), dim=0) + else: + # LoRA rank is the feature width, not an FP8-compressed hidden size. + feature_dim = linear_out_weight.shape[-1] if linear_out_weight.ndim == 2 else None + q_out, k_out, v_out = split_qkv_weights(config, linear_out_weight, feature_dim=feature_dim) return {"q_proj": q_out, "k_proj": k_out, "v_proj": v_out} def _split_gdn_in_proj_linear_out_weight( diff --git a/src/megatron/bridge/models/deepseek/deepseek_v4_bridge.py b/src/megatron/bridge/models/deepseek/deepseek_v4_bridge.py index 4bee34dd26..da52f91537 100644 --- a/src/megatron/bridge/models/deepseek/deepseek_v4_bridge.py +++ b/src/megatron/bridge/models/deepseek/deepseek_v4_bridge.py @@ -75,7 +75,7 @@ AutoMapping, ColumnParallelMapping, GatedMLPMapping, - MegatronParamMapping, + HCAlphaMapping, ReplicatedMapping, ) from megatron.bridge.models.deepseek.deepseek_v4_hybrid_provider import DeepSeekV4HybridModelProvider @@ -361,93 +361,6 @@ def _dsv4_use_mxfp4_export(hf_param: str, weight: torch.Tensor, source_scale: to # --------------------------------------------------------------------------- -class _HCAlphaMapping(MegatronParamMapping): - """Map Megatron's three scalar HC alpha parameters to/from the V4 checkpoint's - 3-element hc_*_scale tensor. - - V4 checkpoint : layers.N.hc_attn_scale shape [3] = [alpha_pre, alpha_post, alpha_res] - Megatron : three separate nn.Parameter([1]) tensors - """ - - def __init__(self, megatron_pre: str, megatron_post: str, megatron_res: str, hf_param: str): - # We register under the alpha_pre path; the others are handled inside hf_to_megatron. - super().__init__(megatron_param=megatron_pre, hf_param=hf_param) - self._megatron_post = megatron_post - self._megatron_res = megatron_res - - @staticmethod - def _resolve_single(pattern: str, captures) -> str: - result = pattern - ci = 0 - while "**" in result and ci < len(captures): - result = result.replace("**", captures[ci], 1) - ci += 1 - ci = 0 - while "*" in result and ci < len(captures): - result = result.replace("*", captures[ci], 1) - ci += 1 - return result - - def resolve(self, captures): - resolved_mg, resolved_hf = self._resolve_names(captures) - resolved_post = self._resolve_single(self._megatron_post, captures) - resolved_res = self._resolve_single(self._megatron_res, captures) - return _HCAlphaMapping( - megatron_pre=resolved_mg, - megatron_post=resolved_post, - megatron_res=resolved_res, - hf_param=resolved_hf, - ) - - def hf_to_megatron(self, hf_weights, megatron_module): - # hf_weights is hc_*_scale [3]; we write alpha_pre here (index 0). - # alpha_post and alpha_res are handled by their own mappings when registered. - target = hf_weights.to(megatron_module.alpha_pre.device) - return target[0:1] - - def megatron_to_hf(self, megatron_weights, megatron_module): - # megatron_weights is alpha_pre [1]; gather all 3 from the same module. - # With PP > 1, megatron_module may be None on non-owning ranks, - # so we broadcast alpha_post and alpha_res alongside alpha_pre. - post_tensor = megatron_module.alpha_post.detach() if megatron_module is not None else None - res_tensor = megatron_module.alpha_res.detach() if megatron_module is not None else None - megatron_weights = self.broadcast_from_pp_rank(megatron_weights, cache_key=str(self.hf_param)) - post = self.broadcast_from_pp_rank(post_tensor, cache_key=str(self.hf_param) + "_post") - res = self.broadcast_from_pp_rank(res_tensor, cache_key=str(self.hf_param) + "_res") - if megatron_weights is None: - return {} - megatron_weights = self.maybe_dequantize(megatron_weights) - return {self.hf_param: torch.cat([megatron_weights.float(), post.float(), res.float()])} - - -class _HCAlphaSecondaryMapping(MegatronParamMapping): - """Secondary mapping for alpha_post (index=1) or alpha_res (index=2). - - Import: extracts element [index] from the 3-element hc_*_scale tensor. - Export: returns {} because the primary _HCAlphaMapping (alpha_pre) already - exports all three alpha values together. This mapping just suppresses the - "No mapping found" warning for the secondary Megatron params during export. - """ - - def __init__(self, megatron_param: str, hf_scale_param: str, index: int): - super().__init__(megatron_param=megatron_param, hf_param=hf_scale_param) - self._index = index - self.allow_hf_name_mismatch = True # export is no-op; skip hf_keys check - - def hf_to_megatron(self, hf_weights, megatron_module): - attr = "alpha_post" if self._index == 1 else "alpha_res" - target = hf_weights.to(getattr(megatron_module, attr).device) - return target[self._index : self._index + 1] - - def resolve(self, captures): - resolved_mg, resolved_hf = self._resolve_names(captures) - return _HCAlphaSecondaryMapping(resolved_mg, resolved_hf, self._index) - - def megatron_to_hf(self, megatron_weights, megatron_module): - # Already handled by the primary alpha_pre _HCAlphaMapping - return {} - - class _ReplicatedOptional(ReplicatedMapping): """ReplicatedMapping for CSA-optional weights (compressor / indexer). @@ -784,14 +697,9 @@ def mapping_registry(self) -> MegatronMappingRegistry: # noqa: C901 def _hc_alpha_mappings(hc_prefix: str, hf_scale_param: str) -> list: """mHC alpha triplet: one 3-vector hc_*_scale <-> alpha_pre/post/res scalars.""" return [ - _HCAlphaMapping( - megatron_pre=f"{hc_prefix}.alpha_pre", - megatron_post=f"{hc_prefix}.alpha_post", - megatron_res=f"{hc_prefix}.alpha_res", - hf_param=hf_scale_param, - ), - _HCAlphaSecondaryMapping(f"{hc_prefix}.alpha_post", hf_scale_param, 1), - _HCAlphaSecondaryMapping(f"{hc_prefix}.alpha_res", hf_scale_param, 2), + HCAlphaMapping(f"{hc_prefix}.alpha_pre", hf_scale_param, 0), + HCAlphaMapping(f"{hc_prefix}.alpha_post", hf_scale_param, 1), + HCAlphaMapping(f"{hc_prefix}.alpha_res", hf_scale_param, 2), ] def _attention_layer_mappings(layer_prefix: str, ck: str) -> list: diff --git a/src/megatron/bridge/models/glm5next/__init__.py b/src/megatron/bridge/models/glm5next/__init__.py new file mode 100644 index 0000000000..bfe3bff21b --- /dev/null +++ b/src/megatron/bridge/models/glm5next/__init__.py @@ -0,0 +1,26 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from megatron.bridge.models.glm5next.glm53flash_bridge import ( + GLM53FlashBridge, + GLM53FlashModel, + GLM53FlashModelProvider, +) + + +__all__ = [ + "GLM53FlashBridge", + "GLM53FlashModel", + "GLM53FlashModelProvider", +] diff --git a/src/megatron/bridge/models/glm5next/glm53flash_bridge.py b/src/megatron/bridge/models/glm5next/glm53flash_bridge.py new file mode 100644 index 0000000000..fc47cd534b --- /dev/null +++ b/src/megatron/bridge/models/glm5next/glm53flash_bridge.py @@ -0,0 +1,929 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and limitations +# under the License. + +"""Unified bridge for GLM-5.3-Flash (model_type ``glm5_next``). + +GLM-5.3-Flash is a VLM: a hybrid KDA/DSA + MoE + mHC language model paired with a +replicated HF vision tower. This bridge handles both the text and vision paths in a +single registered bridge + provider + model — there is no separate "VL" split. + +Each HF decoder layer maps to two physical layers: attention at 2N and FFN at 2N+1. +Checkpoint keys use ``model.language_model``; mixed FP8 weights are dequantized on +import. Vision weights (``model.visual.**``) are loaded replicated and stay BF16 even +under FP8. The VLM wrapper degrades to text-only when ``pixel_values`` is None. +""" + +import types +from collections.abc import Mapping +from copy import deepcopy +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Optional + +import torch +from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols +from megatron.core.tensor_parallel import scatter_to_sequence_parallel_region +from megatron.core.transformer import ModuleSpec +from megatron.core.transformer.module import MegatronModule +from torch import Tensor, nn +from torch.distributed._tensor import DTensor + +from megatron.bridge.models.conversion import quantization_utils +from megatron.bridge.models.conversion.mapping_registry import MegatronMappingRegistry +from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge +from megatron.bridge.models.conversion.param_mapping import ( + AutoMapping, + ColumnParallelMapping, + GatedMLPMapping, + HCAlphaMapping, + InitOnlyMapping, + MegatronParamMapping, + ReplicatedMapping, + RowParallelMapping, + _module_uses_fsdp, +) +from megatron.bridge.models.conversion.utils import get_module_and_param_from_name +from megatron.bridge.models.hf_pretrained.causal_lm import PreTrainedCausalLM +from megatron.bridge.models.hybrid.hybrid_provider import get_default_hybrid_stack_spec +from megatron.bridge.models.hybrid_mla_provider import HybridMLAModelProvider +from megatron.bridge.training.utils.packed_seq_utils import preprocess_packed_seqs +from megatron.bridge.utils.common_utils import hook_hf_module_setattr_for_tp_grad_sync + + +try: + from transformers import Glm5NextForConditionalGeneration +except ImportError: # pragma: no cover + Glm5NextForConditionalGeneration = "Glm5NextForConditionalGeneration" # type: ignore[assignment] + + +if TYPE_CHECKING: + from megatron.core.packed_seq_params import PackedSeqParams + + +# --------------------------------------------------------------------------- +# Config helpers +# --------------------------------------------------------------------------- + + +def _glm53_text_config(hf_config): + """Return the nested text_config for the GLM-5.3-Flash VLM.""" + text_config = getattr(hf_config, "text_config", None) + return text_config if text_config is not None else hf_config + + +def _glm53_layer_types(hf_config) -> list[str]: + """Per-HF-layer attention type: 'linear_attention' (KDA) or 'deepseek_sparse_attention' (DSA).""" + text_config = _glm53_text_config(hf_config) + layer_types = list(getattr(text_config, "layer_types", [])) + num_layers = text_config.num_hidden_layers + if len(layer_types) != num_layers: + raise ValueError(f"layer_types has {len(layer_types)} entries, but num_hidden_layers is {num_layers}.") + return layer_types + + +def _glm53_mlp_layer_types(hf_config) -> list[str]: + """Per-HF-layer MLP type: 'dense' or 'sparse'.""" + text_config = _glm53_text_config(hf_config) + mlp_types = getattr(text_config, "mlp_layer_types", None) + if mlp_types is None: + first_k_dense = getattr(text_config, "first_k_dense_replace", min(3, text_config.num_hidden_layers)) + mlp_types = ["dense"] * first_k_dense + ["sparse"] * (text_config.num_hidden_layers - first_k_dense) + return list(mlp_types) + + +# --------------------------------------------------------------------------- +# Custom param mappings +# --------------------------------------------------------------------------- + + +class _ColumnParallelConcatMapping(MegatronParamMapping[dict]): + """Map separate HF Q/K/V tensors to rank-local [q_shard | k_shard | v_shard].""" + + def __init__(self, megatron_param: str, hf_params: list[str]): + super().__init__(megatron_param, {str(i): p for i, p in enumerate(hf_params)}) + self._hf_order = hf_params + self._tp_mapping = ColumnParallelMapping(megatron_param, megatron_param) + + def resolve(self, captures): + resolved_mg, resolved_hf = self._resolve_names(captures) + return _ColumnParallelConcatMapping(resolved_mg, list(resolved_hf.values())) + + def _shard_per_rank(self, merged: torch.Tensor, section_sizes: list[int]) -> list[torch.Tensor]: + if any(size % self.tp_size for size in section_sizes): + raise ValueError("Each KDA Q/K/V section must be divisible by tensor parallel size.") + sections = merged.split(section_sizes, dim=0) + return [ + torch.cat(shards, dim=0) + for shards in zip(*(section.chunk(self.tp_size, dim=0) for section in sections), strict=True) + ] + + def hf_to_megatron(self, hf_weights: dict, megatron_module: nn.Module) -> torch.Tensor: + splits = None + if self.tp_rank == 0: + sections = [hf_weights[str(i)] for i in range(len(self._hf_order))] + merged = torch.cat(sections, dim=0) + if self.tp_size == 1: + return merged + splits = self._shard_per_rank(merged, [section.shape[0] for section in sections]) + + normalized_param = self._tp_mapping._normalize_expert_param_name(self.megatron_param) + _, target = get_module_and_param_from_name(megatron_module, normalized_param) + shape = target.orig_param.shape if isinstance(target, DTensor) else target.shape + return self._tp_mapping.scatter_to_tp_ranks(splits, shape, target.dtype, target.device) + + def megatron_to_hf(self, megatron_weights: Optional[torch.Tensor], megatron_module) -> dict: + if megatron_weights is None: + return {} + weights = self.maybe_dequantize(megatron_weights) + config = self._get_config(megatron_module) if megatron_module is not None else None + sizes = self._section_sizes(config) + sharded = self.tp_size > 1 and not _module_uses_fsdp(megatron_module) + if sizes is None: + if weights.shape[0] % len(self._hf_order): + raise ValueError("Cannot infer unequal KDA Q/K/V sections without the model config.") + rows = weights.shape[0] // len(self._hf_order) * (self.tp_size if sharded else 1) + sizes = [rows] * len(self._hf_order) + + if sharded: + if any(size % self.tp_size for size in sizes): + raise ValueError("Each KDA Q/K/V section must be divisible by tensor parallel size.") + local_sizes = [size // self.tp_size for size in sizes] + shards = self._tp_mapping.gather_from_tp_ranks(weights) + parts = [ + torch.cat(section_shards, dim=0) + for section_shards in zip(*(shard.split(local_sizes, dim=0) for shard in shards), strict=True) + ] + else: + parts = weights.split(sizes, dim=0) + return dict(zip(self._hf_order, parts, strict=True)) + + def _section_sizes(self, config): + if config is None: + return None + qk = config.linear_key_head_dim * config.linear_num_key_heads + v = config.linear_value_head_dim * config.linear_num_value_heads + return [qk, qk, v] + + +class _EhProjSplitMapping(MegatronParamMapping[torch.Tensor]): + """Split HF eh_proj [hidden, 2*hidden] into column-sharded Megatron e_proj/h_proj. + + Import: split along input dim (dim 1), shard column-parallel (dim 0). + Export: gather both halves, concat along dim 1. + """ + + _export_cache: dict = {} + + def __init__(self, megatron_param: str, hf_param: str, index: int): + super().__init__(megatron_param=megatron_param, hf_param=hf_param) + self._index = index + self.allow_hf_name_mismatch = True + + def resolve(self, captures): + resolved_mg, resolved_hf = self._resolve_names(captures) + return _EhProjSplitMapping(resolved_mg, resolved_hf, self._index) + + def hf_to_megatron(self, hf_weights, megatron_module): + if hf_weights is None: + return None + half = hf_weights.shape[1] // 2 + chunk = hf_weights[:, self._index * half : (self._index + 1) * half] + return self._shard_column(chunk, megatron_module) + + def _shard_column(self, weight, megatron_module): + device = next(megatron_module.parameters()).device + weight = weight.to(device=device) + if self.tp_size == 1: + return weight + shard = weight.shape[0] // self.tp_size + start = self.tp_rank * shard + return weight[start : start + shard].contiguous() + + def megatron_to_hf(self, megatron_weights, megatron_module): + if megatron_weights is None: + return {} + megatron_weights = self.maybe_dequantize(megatron_weights) + if self.tp_size == 1: + full = megatron_weights + else: + full = torch.cat(self.gather_from_tp_ranks(megatron_weights), dim=0) + cache = _EhProjSplitMapping._export_cache + key = str(self.hf_param) + cache[f"{key}#{self._index}"] = full + other_key = f"{key}#{1 - self._index}" + if other_key not in cache: + return {} + other_w = cache.pop(other_key) + cache.pop(f"{key}#{self._index}", None) + parts = [None, None] + parts[self._index] = full + parts[1 - self._index] = other_w + return {key: torch.cat(parts, dim=1).contiguous()} + + +# --------------------------------------------------------------------------- +# Stack spec +# --------------------------------------------------------------------------- + + +def glm53_hybrid_stack_spec(config: HybridMLAModelProvider) -> ModuleSpec: + """Return the Hybrid stack with MLA q/kv norms wired into the DSA spec.""" + stack_spec = deepcopy(get_default_hybrid_stack_spec(config)) + dsa_submodules = stack_spec.submodules.dsa_layer.submodules + native_norm = dsa_submodules.input_layernorm + attention_submodules = dsa_submodules.self_attention.submodules + attention_submodules.q_layernorm = native_norm + attention_submodules.kv_layernorm = native_norm + return stack_spec + + +# --------------------------------------------------------------------------- +# Provider +# --------------------------------------------------------------------------- + + +@dataclass +class GLM53FlashModelProvider(HybridMLAModelProvider): + """Provider for the unified GLM-5.3-Flash VLM: hybrid language stack + HF vision.""" + + # Vision configuration (HF Glm5NextVisionConfig object, stored as-is). + vision_config: Any = None + spatial_merge_size: int = 2 + # Inject vision features before scattering the full sequence across TP ranks. + scatter_embedding_sequence_parallel: bool = False + # Required by HF's can_return_tuple-decorated feature extraction methods. + return_dict: bool = True + + # Multimodal token IDs (from the top-level Glm5NextConfig). + image_token_id: int = 154854 + video_token_id: int = 154855 + image_start_token_id: int = 154830 + image_end_token_id: int = 154831 + video_start_token_id: int = 154832 + video_end_token_id: int = 154833 + + # Freeze options for LoRA training. + freeze_language_model: bool = False + freeze_vision_model: bool = False + freeze_vision_projection: bool = False + + def provide( + self, pre_process: Optional[bool] = None, post_process: Optional[bool] = None, vp_stage: Optional[int] = None + ) -> "GLM53FlashModel": + """Build the VLM wrapper and apply explicitly requested freeze options.""" + model = GLM53FlashModel(self, pre_process=pre_process, post_process=post_process, vp_stage=vp_stage) + if self.freeze_language_model or self.freeze_vision_model or self.freeze_vision_projection: + model.freeze( + freeze_language_model=self.freeze_language_model, + freeze_vision_model=self.freeze_vision_model, + freeze_vision_projection=self.freeze_vision_projection, + ) + return model + + def provide_language_model( + self, pre_process: Optional[bool] = None, post_process: Optional[bool] = None, vp_stage: Optional[int] = None + ): + """Build the Megatron HybridModel language model (delegates to the parent).""" + return super().provide(pre_process=pre_process, post_process=post_process, vp_stage=vp_stage) + + +# --------------------------------------------------------------------------- +# Model (VLM wrapper: replicated HF vision tower + Megatron HybridModel) +# --------------------------------------------------------------------------- + + +class GLM53FlashModel(MegatronModule): + """Inject HF image/video features into the hybrid language model's embeddings.""" + + def __init__( + self, config, pre_process: bool = True, post_process: bool = True, vp_stage: Optional[int] = None + ) -> None: + super().__init__(config=config) + + self.pre_process = pre_process + self.post_process = post_process + self.vp_stage = vp_stage + + if pre_process: + from transformers.models.glm5_next.modeling_glm5_next import ( + Glm5NextModel, + Glm5NextVisionModel, + ) + + self.visual = Glm5NextVisionModel._from_config(config.vision_config) + hook_hf_module_setattr_for_tp_grad_sync(self.visual) + + self.language_model = self.config.provide_language_model( + pre_process=pre_process, post_process=post_process, vp_stage=vp_stage + ) + + self.share_embeddings_and_output_weights = config.share_embeddings_and_output_weights + self.shared_embedding_or_output_weight = self.language_model.shared_embedding_or_output_weight + + # Monkey-patch HF image/video feature extraction + placeholder-mask logic. + self.get_image_features = types.MethodType(Glm5NextModel.get_image_features, self) + self.get_video_features = types.MethodType(Glm5NextModel.get_video_features, self) + self.get_placeholder_mask = types.MethodType(Glm5NextModel.get_placeholder_mask, self) + + self.config.spatial_merge_size = getattr(self.config.vision_config, "spatial_merge_size", 2) + + def set_input_tensor(self, input_tensor) -> None: + """Set model chunk input tensor.""" + self.language_model.set_input_tensor(input_tensor) + + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + pixel_values: Optional[torch.Tensor] = None, + pixel_values_videos: Optional[torch.FloatTensor] = None, + image_grid_thw: Optional[torch.LongTensor] = None, + video_grid_thw: Optional[torch.LongTensor] = None, + mm_token_type_ids: Optional[torch.Tensor] = None, + labels: Optional[torch.Tensor] = None, + runtime_gather_output: Optional[bool] = None, + packed_seq_params: Optional["PackedSeqParams"] = None, + temperature: Optional[float] = None, + output_processor=None, + output_processor_context=None, + *, + loss_mask: Optional[Tensor] = None, + ) -> Tensor: + """Inject vision features before packing BSHD inputs into THD layout. + + Placeholder masks require full BSHD sequences. Packing and sequence-parallel + scattering happen afterward; NoPE attention uses only packed_seq_params. + """ + if self.pre_process: + if inputs_embeds is None: + inputs_embeds = self.language_model.embedding( + input_ids=input_ids, position_ids=None + ) # [seq_len, batch, hidden] (full seq, scatter_embedding_sequence_parallel=False) + # Transpose to HF format [batch, seq_len, hidden]. + inputs_embeds = inputs_embeds.transpose(1, 0).contiguous() + + if pixel_values is not None: + image_embeds = self.get_image_features(pixel_values, image_grid_thw).pooler_output + image_embeds = torch.cat(image_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) + image_mask, _ = self.get_placeholder_mask( + input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds + ) + inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) + + if pixel_values_videos is not None: + video_embeds = self.get_video_features(pixel_values_videos, video_grid_thw).pooler_output + video_embeds = torch.cat(video_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) + _, video_mask = self.get_placeholder_mask( + input_ids, inputs_embeds=inputs_embeds, video_features=video_embeds + ) + inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds) + + use_fp8_padding = self.config.fp8 in ("e4m3", "hybrid") + inputs_embeds_thd = preprocess_packed_seqs( + inputs_embeds, + attention_mask, + pg_collection=self.config._pg_collection, + use_fp8_padding=use_fp8_padding, + )[0] + # Megatron expects [seq, batch, hidden] with batch=1 for THD. + inputs_embeds_thd = inputs_embeds_thd.transpose(1, 0).contiguous() # [total, 1, hidden] + + # MTP re-embeds input_ids, so it also needs the packed token layout. + input_ids_thd = preprocess_packed_seqs( + input_ids, + attention_mask, + pg_collection=self.config._pg_collection, + use_fp8_padding=use_fp8_padding, + )[0] + + if self.config.sequence_parallel: + tp_group = self.config._pg_collection.tp if self.config._pg_collection is not None else None + inputs_embeds_thd = scatter_to_sequence_parallel_region(inputs_embeds_thd, group=tp_group) + + # verl's output processor needs hidden states, not the MTP auxiliary loss. + return self.language_model.forward( + input_ids=input_ids_thd, + position_ids=None, + attention_mask=None, + decoder_input=inputs_embeds_thd, + labels=labels, + loss_mask=loss_mask, + runtime_gather_output=runtime_gather_output, + packed_seq_params=packed_seq_params, + output_processor=output_processor, + output_processor_context=output_processor_context, + compute_mtp_loss=False, + ) + + def freeze( + self, + freeze_language_model: bool, + freeze_vision_model: bool, + freeze_vision_projection: bool, + ): + """Freeze model modules (set requires_grad=False).""" + modules = [] + if freeze_language_model and hasattr(self, "language_model") and self.language_model is not None: + modules.append(self.language_model) + if freeze_vision_model and hasattr(self, "visual") and self.visual is not None: + if hasattr(self.visual, "patch_embed"): + modules.append(self.visual.patch_embed) + if hasattr(self.visual, "blocks"): + modules.append(self.visual.blocks) + if freeze_vision_projection and hasattr(self, "visual") and self.visual is not None: + if hasattr(self.visual, "merger"): + modules.append(self.visual.merger) + for module in modules: + for param in module.parameters(): + param.requires_grad = False + + +# --------------------------------------------------------------------------- +# Bridge +# --------------------------------------------------------------------------- + + +@MegatronModelBridge.register_bridge( + source=Glm5NextForConditionalGeneration, + target=GLM53FlashModel, + provider=GLM53FlashModelProvider, + model_type="glm5_next", +) +class GLM53FlashBridge(MegatronModelBridge): + """Unified bridge for GLM-5.3-Flash: hybrid language model + replicated vision tower.""" + + _provider_cls = GLM53FlashModelProvider + + def provider_bridge(self, hf_pretrained: PreTrainedCausalLM): + hf_config = hf_pretrained.config + text_config = _glm53_text_config(hf_config) + + # Build base provider kwargs from the nested text_config (the VLM top-level + # config lacks hidden_size/num_layers/etc. that the text sub-model owns). + provider_kwargs = self.hf_config_to_provider_kwargs(text_config) + provider_kwargs.pop("_mla_rope_params", None) + valid_fields = self._provider_cls.__dataclass_fields__ + provider = self._provider_cls(**{k: v for k, v in provider_kwargs.items() if k in valid_fields}) + + provider.normalization = "RMSNorm" + provider.gated_linear_unit = True + provider.add_bias_linear = False + provider.share_embeddings_and_output_weights = bool(getattr(hf_config, "tie_word_embeddings", False)) + provider.qk_layernorm = True + provider.multi_latent_attention = True + provider.sequence_parallel = True + + # ---- MoE ---- + provider.moe_grouped_gemm = True + provider.moe_router_pre_softmax = True + provider.moe_token_dispatcher_type = "alltoall" + provider.moe_flex_dispatcher_backend = "hybridep" + provider.moe_flex_dispatcher_num_sms = 16 + provider.moe_router_load_balancing_type = "seq_aux_loss" + provider.moe_shared_expert_overlap = False + provider.moe_router_score_function = text_config.scoring_func # "sigmoid" + provider.moe_router_enable_expert_bias = True + provider.moe_router_dtype = "fp32" + provider.moe_permute_fusion = True + provider.moe_router_topk = text_config.num_experts_per_tok # 8 + provider.moe_router_topk_scaling_factor = getattr(text_config, "routed_scaling_factor", 1.0) # 2.5 + provider.norm_topk_prob = text_config.norm_topk_prob + provider.moe_aux_loss_coeff = getattr(text_config, "router_aux_loss_coef", 0.001) + provider.activation_func_clamp_value = getattr(text_config, "swiglu_limit", 0.0) # 10.0 + provider.moe_shared_expert_intermediate_size = text_config.moe_intermediate_size * text_config.n_shared_experts + # GLM-5.3-Flash's shared expert is always-on (no per-token gate scalar in the + # checkpoint), so disable MCore's optional shared-expert gate. + provider.moe_shared_expert_gate = False + + provider.hidden_dropout = 0.0 + provider.attention_softmax_in_fp32 = False + provider.make_vocab_size_divisible_by = 1280 + provider.actual_vocab_size = text_config.vocab_size + + # ---- Hybrid layer pattern ---- + attn_types = _glm53_layer_types(hf_config) + mlp_types = _glm53_mlp_layer_types(hf_config) + attn_sym = { + "linear_attention": Symbols.KDA, + "deepseek_sparse_attention": Symbols.DS_ATTENTION, + } + ffn_sym = {"dense": Symbols.MLP, "sparse": Symbols.MOE} + pattern_chars = [] + # Router replay indexes physical modules: each HF block becomes attention + FFN. + moe_layer_freq = [] + for at, mt in zip(attn_types, mlp_types, strict=True): + if at not in attn_sym: + raise ValueError(f"Unsupported GLM-5.3-Flash attention layer type: {at!r}") + if mt not in ffn_sym: + raise ValueError(f"Unsupported GLM-5.3-Flash MLP layer type: {mt!r}") + pattern_chars.append(attn_sym[at] + ffn_sym[mt]) + moe_layer_freq.append(0) # attention module is never MoE + moe_layer_freq.append(1 if mt == "sparse" else 0) # FFN module + main_pattern = "".join(pattern_chars) + provider.hybrid_layer_pattern = main_pattern + provider.num_layers = len(main_pattern) + provider.hybrid_stack_spec = glm53_hybrid_stack_spec + provider.moe_layer_freq = moe_layer_freq + + # ---- MLA geometry (NoPE) ---- + # GLM-5.3-Flash: mla_use_nope=True -> qk_rope_head_dim=0, qk_nope_head_dim=256, qk_head_dim=256. + provider.qk_pos_emb_head_dim = text_config.qk_head_dim - text_config.qk_nope_head_dim # 0 + provider.q_lora_rank = text_config.q_lora_rank # 1536 + provider.kv_lora_rank = text_config.kv_lora_rank # 512 + provider.v_head_dim = text_config.v_head_dim # 256 + provider.qk_head_dim = text_config.qk_head_dim # 256 + provider.num_attention_heads = text_config.num_attention_heads # 64 + provider.num_key_value_heads = text_config.num_key_value_heads # 64 + # HF head_dim=0 is a NoPE sentinel, not a valid attention channel width. + provider.kv_channels = text_config.v_head_dim # 256 + provider.rope_type = "yarn" + provider.rotary_base = 10000.0 + provider.rotary_scaling_factor = 1.0 + provider.mscale = 1.0 + provider.mscale_all_dim = 1.0 + provider.apply_rope_fusion = False + provider.rotary_interleaved = False + provider.dsa_indexer_rope_interleaved = getattr(text_config, "indexer_rope_interleave", True) + provider.cp_comm_type = "allgather" + + # ---- DSA indexer (kpool) ---- + provider.experimental_attention_variant = "dsa" + provider.dsa_indexer_head_dim = text_config.index_head_dim # 128 + provider.dsa_indexer_n_heads = text_config.index_n_heads # 32 + provider.dsa_indexer_topk = text_config.index_topk # 2048 + provider.dsa_indexer_loss_coeff = 0.0 + provider.dsa_indexer_use_sparse_loss = False + provider.dsa_indexer_rotate_activation = False + provider.dsa_indexer_kpool_fp8 = True + provider.dsa_indexer_k_norm_epsilon = 1e-6 + provider.dsa_indexer_kpool = int(getattr(text_config, "index_kpool", 1)) # 4 + provider.dsa_indexer_kpool_always_select_tail = bool( + getattr(text_config, "index_kpool_always_select_tail", False) + ) + provider.dsa_indexer_topk_freq = 1 + provider.dsa_indexer_skip_topk_offset = 0 + + # ---- KDA two-stage gates ---- + provider.kda_two_stage_gates = True + provider.kda_safe_gate = True + provider.kda_lower_bound = float(text_config.linear_attn_config.get("gate_lower_bound", -5.0)) + provider.gdn_pre_gated_delta_rule_fusion = False + provider.gdn_conv_pad_alignment = None + lin = text_config.linear_attn_config + # KDA uses equal query/key/value head geometry. + provider.linear_num_key_heads = int(lin["num_heads"]) # 64 + provider.linear_num_value_heads = int(lin["num_heads"]) # 64 (qkv equal) + provider.linear_key_head_dim = int(lin["head_dim"]) # 128 + provider.linear_value_head_dim = int(lin["head_dim"]) # 128 + provider.linear_conv_kernel_dim = int(lin["short_conv_kernel_size"]) # 4 + + # ---- mHC (multi-stream Hyper-Connections) ---- + provider.enable_mhc_connections = True + provider.mhc_num_residual_streams = int(getattr(text_config, "hc_mult", 4)) + provider.mhc_sinkhorn_iterations = int(getattr(text_config, "hc_sinkhorn_iters", 20)) + provider.mhc_learned_output_contract = False + provider.mhc_norm_eps_inside_sqrt = True + provider.mhc_keep_mappings_in_fp32 = True + + # ---- MTP (1 layer, DSA + MoE) ---- + # MTP has its own pattern and is excluded from num_layers/moe_layer_freq. + num_mtp = int(getattr(text_config, "num_nextn_predict_layers", 0) or 0) + provider.mtp_num_layers = num_mtp or None + provider.mtp_hybrid_override_pattern = "DE" + provider.mtp_use_repeated_layer = True + provider.keep_mtp_spec_in_bf16 = True + if num_mtp: + provider.mtp_loss_scaling_factor = 0.1 + + provider.persist_layer_norm = True + provider.gradient_accumulation_fusion = True + provider.bias_dropout_fusion = True + + # ---- Vision tower (replicated HF Glm5NextVisionModel) ---- + provider.vision_config = getattr(hf_config, "vision_config", None) + provider.image_token_id = getattr(hf_config, "image_token_id", 154854) + provider.video_token_id = getattr(hf_config, "video_token_id", 154855) + provider.image_start_token_id = getattr(hf_config, "image_start_token_id", 154830) + provider.image_end_token_id = getattr(hf_config, "image_end_token_id", 154831) + provider.video_start_token_id = getattr(hf_config, "video_start_token_id", 154832) + provider.video_end_token_id = getattr(hf_config, "video_end_token_id", 154833) + provider.spatial_merge_size = getattr(provider.vision_config, "spatial_merge_size", 2) + + return provider + + @classmethod + def megatron_to_hf_config(cls, provider) -> dict: + hf_config = super().megatron_to_hf_config(provider) + pattern = getattr(provider, "hybrid_layer_pattern", None) + if pattern: + main_pattern = pattern.split(Symbols.MTP_SEPARATOR)[0] + hf_config["num_hidden_layers"] = sum(1 for c in main_pattern if c in (Symbols.KDA, Symbols.DS_ATTENTION)) + return hf_config + + # ------------------------------------------------------------------ + # FP8 dequantization on import + # ------------------------------------------------------------------ + + def maybe_modify_loaded_hf_weight(self, hf_param, hf_state_dict: Mapping[str, torch.Tensor]): + # InitOnlyMapping declares a synthetic HF key ("...__init_only__") that + # never exists in the checkpoint; skip the state-dict lookup so import keeps + # the Megatron module's random initialization. + if isinstance(hf_param, str) and hf_param.endswith(".__init_only__"): + return None + hf_weights = super().maybe_modify_loaded_hf_weight(hf_param, hf_state_dict) + if isinstance(hf_weights, dict): + return { + key: self._maybe_dequantize_fp8(tensor, hf_param[key], hf_state_dict) + for key, tensor in hf_weights.items() + } + return self._maybe_dequantize_fp8(hf_weights, hf_param, hf_state_dict) + + @staticmethod + def _maybe_dequantize_fp8(weight, param_name: str, hf_state_dict: Mapping[str, torch.Tensor]): + scale_key = param_name + "_scale_inv" + return quantization_utils.maybe_dequantize_fp8_blockwise(weight, hf_state_dict.get(scale_key)) + + # ------------------------------------------------------------------ + # Weight mapping registry + # ------------------------------------------------------------------ + + def mapping_registry(self) -> MegatronMappingRegistry: + hf_config = self.hf_config + text_config = _glm53_text_config(hf_config) + attn_types = _glm53_layer_types(hf_config) + mlp_types = _glm53_mlp_layer_types(hf_config) + num_hf_layers = text_config.num_hidden_layers + + mapping_list = [ + AutoMapping( + "embedding.word_embeddings.weight", + "model.language_model.embed_tokens.weight", + ), + AutoMapping( + "decoder.final_norm.weight", + "model.language_model.norm.weight", + ), + AutoMapping( + "output_layer.weight", + "lm_head.weight", + ), + ] + + for hf_layer_idx in range(num_hf_layers): + is_kda = attn_types[hf_layer_idx] == "linear_attention" + is_moe = mlp_types[hf_layer_idx] == "sparse" + attn_layer_idx = 2 * hf_layer_idx + ffn_layer_idx = attn_layer_idx + 1 + hf_layer = f"model.language_model.layers.{hf_layer_idx}" + attn_layer = f"decoder.layers.{attn_layer_idx}" + ffn_layer = f"decoder.layers.{ffn_layer_idx}" + + # mHC hyper-connections wrap every non-MTP physical layer (attn + ffn). + mapping_list.extend(self._mhc_mappings(attn_layer, hf_layer, "attn")) + mapping_list.extend(self._mhc_mappings(ffn_layer, hf_layer, "ffn")) + + if is_kda: + mapping_list.extend(self._kda_attn_mappings(attn_layer, hf_layer)) + else: + mapping_list.extend(self._dsa_attn_mappings(attn_layer, hf_layer)) + + if is_moe: + mapping_list.extend(self._moe_ffn_mappings(ffn_layer, hf_layer)) + else: + mapping_list.extend(self._dense_ffn_mappings(ffn_layer, hf_layer)) + + # MTP uses split e/h projections and two inner layers without HF mHC weights. + num_mtp = int(getattr(text_config, "num_nextn_predict_layers", 0) or 0) + for mtp_idx in range(num_mtp): + mtp_hf = f"model.language_model.layers.{mtp_idx + num_hf_layers}" + mtp_mg = f"mtp.layers.{mtp_idx}" + eh_proj = f"{mtp_hf}.eh_proj.weight" + mapping_list.extend( + [ + AutoMapping(f"{mtp_mg}.enorm.weight", f"{mtp_hf}.enorm.weight"), + AutoMapping(f"{mtp_mg}.hnorm.weight", f"{mtp_hf}.hnorm.weight"), + # Split eh_proj [hidden, 2*hidden] along its input dimension. + _EhProjSplitMapping(f"{mtp_mg}.e_proj.weight", eh_proj, 0), + _EhProjSplitMapping(f"{mtp_mg}.h_proj.weight", eh_proj, 1), + AutoMapping( + f"{mtp_mg}.final_layernorm.weight", + "model.language_model.norm.weight", + ), + ] + ) + # Inner layer 0 is attention; layer 1 is MoE. + mtp_attn_layer = f"{mtp_mg}.mtp_model_layer.layers.0" + mtp_ffn_layer = f"{mtp_mg}.mtp_model_layer.layers.1" + mapping_list.extend(self._mhc_init_only_mappings(mtp_attn_layer)) + mapping_list.extend(self._mhc_init_only_mappings(mtp_ffn_layer)) + mapping_list.extend(self._dsa_attn_mappings(mtp_attn_layer, mtp_hf)) + mapping_list.extend(self._moe_ffn_mappings(mtp_ffn_layer, mtp_hf)) + + # The language model is a submodule of the VLM wrapper, so prefix every + # Megatron param with ``language_model.`` and append the replicated vision weights. + return self._prefix_mapping_registry( + MegatronMappingRegistry(*mapping_list), + "language_model.", + ReplicatedMapping("visual.**", "model.visual.**"), + ) + + @staticmethod + def _mhc_init_only_mappings(mg_layer: str) -> list: + """Init-only mappings for an mHC hyper_connection with no HF source. + + Used for the MTP layer's inner physical layers, whose ``hyper_connection`` + params (mapping_proj, bias, alpha_pre/post/res) are created by MCore but + have no counterpart in GLM-5.3-Flash's HF checkpoint. + """ + hc = "hyper_connection" + return [ + InitOnlyMapping(f"{mg_layer}.{hc}.mapping_proj.weight"), + InitOnlyMapping(f"{mg_layer}.{hc}.bias"), + InitOnlyMapping(f"{mg_layer}.{hc}.alpha_pre"), + InitOnlyMapping(f"{mg_layer}.{hc}.alpha_post"), + InitOnlyMapping(f"{mg_layer}.{hc}.alpha_res"), + ] + + # ------------------------------------------------------------------ + # Sub-mapping builders + # ------------------------------------------------------------------ + + @staticmethod + def _mhc_mappings(mg_layer: str, hf_layer: str, kind: str) -> list: + """mHC hyper-connection mappings for one physical layer. + + ``kind`` is 'attn' or 'ffn'. With ``enable_mhc_connections`` MCore wraps + each physical layer in a ``HyperConnectionHybridLayer`` exposing a single + ``hyper_connection`` submodule (NOT separate ``self_attention_hyper_connection`` + / ``mlp_hyper_connection``). The checkpoint stores ``hc_{kind}_{fn,base,scale}`` + as bare parameters (no ``.weight`` suffix) on the HF layer. + """ + hc = "hyper_connection" + pre = f"{mg_layer}.{hc}.alpha_pre" + post = f"{mg_layer}.{hc}.alpha_post" + res = f"{mg_layer}.{hc}.alpha_res" + scale = f"{hf_layer}.hc_{kind}_scale" + return [ + ReplicatedMapping( + f"{mg_layer}.{hc}.mapping_proj.weight", + f"{hf_layer}.hc_{kind}_fn", + ), + ReplicatedMapping( + f"{mg_layer}.{hc}.bias", + f"{hf_layer}.hc_{kind}_base", + ), + HCAlphaMapping(pre, scale, 0), + HCAlphaMapping(post, scale, 1), + HCAlphaMapping(res, scale, 2), + ] + + @staticmethod + def _kda_attn_mappings(mg_layer: str, hf_layer: str) -> list: + """KDA (linear attention, two-stage gates) weight mappings. + + MCore KDA fuses q|k|v into ``in_proj`` (order q|k|v along dim 0) and + q|k|v conv into a single ``conv1d``; HF stores them separately. + f_a/g_a are replicated; f_b/g_b, in_proj, beta_proj, conv1d, A_log, + dt_bias are TP column-sharded. + """ + sa = f"{mg_layer}.inner_layer.self_attention" + hf_attn = f"{hf_layer}.self_attn" + return [ + AutoMapping(f"{mg_layer}.inner_layer.input_layernorm.weight", f"{hf_layer}.input_layernorm.weight"), + # Fused QKV in_proj (q|k|v, column-parallel, no GQA interleave). + _ColumnParallelConcatMapping( + f"{sa}.in_proj.weight", + [ + f"{hf_attn}.q_proj.weight", + f"{hf_attn}.k_proj.weight", + f"{hf_attn}.v_proj.weight", + ], + ), + # Beta projection (separate, column-parallel). + ColumnParallelMapping(f"{sa}.beta_proj.weight", f"{hf_attn}.b_proj.weight"), + # Two-stage low-rank forget gate: f_a replicated, f_b TP-sharded. + ReplicatedMapping(f"{sa}.f_a_proj.weight", f"{hf_attn}.f_a_proj.weight"), + ColumnParallelMapping(f"{sa}.f_b_proj.weight", f"{hf_attn}.f_b_proj.weight"), + # Two-stage low-rank output gate: g_a replicated, g_b TP-sharded. + ReplicatedMapping(f"{sa}.g_a_proj.weight", f"{hf_attn}.g_a_proj.weight"), + ColumnParallelMapping(f"{sa}.g_b_proj.weight", f"{hf_attn}.g_b_proj.weight"), + # Fused q|k|v conv1d (depthwise). [conv_dim, 1, kernel] -> concat dim 0. + _ColumnParallelConcatMapping( + f"{sa}.conv1d.weight", + [ + f"{hf_attn}.q_conv1d.weight", + f"{hf_attn}.k_conv1d.weight", + f"{hf_attn}.v_conv1d.weight", + ], + ), + # Kernel parameters (TP-sharded along dim 0). + ColumnParallelMapping(f"{sa}.A_log", f"{hf_attn}.A_log"), + ColumnParallelMapping(f"{sa}.dt_bias", f"{hf_attn}.dt_bias"), + # Output norm + projection. + ReplicatedMapping(f"{sa}.out_norm.weight", f"{hf_attn}.o_norm.weight"), + RowParallelMapping(f"{sa}.out_proj.weight", f"{hf_attn}.o_proj.weight"), + ] + + @staticmethod + def _dsa_attn_mappings(mg_layer: str, hf_layer: str) -> list: + """DSA (MLA + kpool indexer) weight mappings.""" + sa = f"{mg_layer}.inner_layer.self_attention" + idx = f"{sa}.core_attention.indexer" + hf_attn = f"{hf_layer}.self_attn" + hf_idx = f"{hf_attn}.indexer" + return [ + AutoMapping(f"{mg_layer}.inner_layer.input_layernorm.weight", f"{hf_layer}.input_layernorm.weight"), + # MLA projections. The down-projections (q_a/kv_a) are TELinear with + # parallel_mode="duplicated" → replicated. The up-projections (q_b/kv_b) + # are TEColumnParallelLinear → column-parallel. + ReplicatedMapping(f"{sa}.linear_q_down_proj.weight", f"{hf_attn}.q_a_proj.weight"), + ColumnParallelMapping(f"{sa}.linear_q_up_proj.weight", f"{hf_attn}.q_b_proj.weight"), + ReplicatedMapping(f"{sa}.q_layernorm.weight", f"{hf_attn}.q_a_layernorm.weight"), + ReplicatedMapping(f"{sa}.linear_kv_down_proj.weight", f"{hf_attn}.kv_a_proj_with_mqa.weight"), + ColumnParallelMapping(f"{sa}.linear_kv_up_proj.weight", f"{hf_attn}.kv_b_proj.weight"), + ReplicatedMapping(f"{sa}.kv_layernorm.weight", f"{hf_attn}.kv_a_layernorm.weight"), + RowParallelMapping(f"{sa}.linear_proj.weight", f"{hf_attn}.o_proj.weight"), + # DSA kpool indexer. All indexer projections are parallel_mode="duplicated" + # (replicated) in MCore — see dsa.py build_module(..., parallel_mode="duplicated"). + ReplicatedMapping(f"{idx}.linear_wq_b.weight", f"{hf_idx}.wq_b.weight"), + ReplicatedMapping(f"{idx}.linear_wk.weight", f"{hf_idx}.wk.weight"), + ReplicatedMapping(f"{idx}.k_norm.weight", f"{hf_idx}.k_norm.weight"), + ReplicatedMapping(f"{idx}.k_norm.bias", f"{hf_idx}.k_norm.bias"), + ReplicatedMapping(f"{idx}.linear_weights_proj.weight", f"{hf_idx}.weights_proj.weight"), + ReplicatedMapping(f"{idx}.index_kpool_compress_ape", f"{hf_idx}.index_kpool_compress_ape"), + ReplicatedMapping(f"{idx}.index_kpool_compress_gate", f"{hf_idx}.index_kpool_compress_gate"), + ] + + @staticmethod + def _dense_ffn_mappings(mg_layer: str, hf_layer: str) -> list: + """Dense MLP FFN mappings.""" + inner = f"{mg_layer}.inner_layer" + hf_mlp = f"{hf_layer}.mlp" + return [ + AutoMapping( + f"{inner}.mlp.linear_fc1.layer_norm_weight", + f"{hf_layer}.post_attention_layernorm.weight", + ), + RowParallelMapping(f"{inner}.mlp.linear_fc2.weight", f"{hf_mlp}.down_proj.weight"), + GatedMLPMapping( + f"{inner}.mlp.linear_fc1.weight", + gate=f"{hf_mlp}.gate_proj.weight", + up=f"{hf_mlp}.up_proj.weight", + ), + ] + + @staticmethod + def _moe_ffn_mappings(mg_layer: str, hf_layer: str) -> list: + """MoE FFN (router + shared expert + routed experts) mappings.""" + inner = f"{mg_layer}.inner_layer" + mlp = f"{inner}.mlp" + hf_mlp = f"{hf_layer}.mlp" + return [ + AutoMapping( + f"{inner}.pre_mlp_layernorm.weight", + f"{hf_layer}.post_attention_layernorm.weight", + ), + AutoMapping(f"{mlp}.router.weight", f"{hf_mlp}.gate.weight"), + AutoMapping(f"{mlp}.router.expert_bias", f"{hf_mlp}.gate.e_score_correction_bias"), + # Shared expert (always-on; no per-token gate scalar in GLM-5.3-Flash). + RowParallelMapping( + f"{mlp}.shared_experts.linear_fc2.weight", + f"{hf_mlp}.shared_experts.down_proj.weight", + ), + GatedMLPMapping( + f"{mlp}.shared_experts.linear_fc1.weight", + gate=f"{hf_mlp}.shared_experts.gate_proj.weight", + up=f"{hf_mlp}.shared_experts.up_proj.weight", + ), + # Routed experts (grouped-GEMM + sequential fallback). + RowParallelMapping( + f"{mlp}.experts.linear_fc2.weight*", + f"{hf_mlp}.experts.*.down_proj.weight", + ), + RowParallelMapping( + f"{mlp}.experts.local_experts.*.linear_fc2.weight", + f"{hf_mlp}.experts.*.down_proj.weight", + ), + GatedMLPMapping( + f"{mlp}.experts.linear_fc1.weight*", + gate=f"{hf_mlp}.experts.*.gate_proj.weight", + up=f"{hf_mlp}.experts.*.up_proj.weight", + ), + GatedMLPMapping( + f"{mlp}.experts.local_experts.*.linear_fc1.weight", + gate=f"{hf_mlp}.experts.*.gate_proj.weight", + up=f"{hf_mlp}.experts.*.up_proj.weight", + ), + ] diff --git a/src/megatron/bridge/models/hybrid/hybrid_provider.py b/src/megatron/bridge/models/hybrid/hybrid_provider.py index ec629f050c..2698ee1b22 100644 --- a/src/megatron/bridge/models/hybrid/hybrid_provider.py +++ b/src/megatron/bridge/models/hybrid/hybrid_provider.py @@ -140,6 +140,7 @@ class HybridModelProvider(TransformerConfig, ModelProviderMixin[MCoreHybridModel rotary_percent: float = 1.0 rotary_base: int = 10000 seq_len_interpolation_factor: float | None = None + scatter_embedding_sequence_parallel: bool = True apply_rope_fusion: bool = True make_vocab_size_divisible_by: int = 128 gated_linear_unit: bool = False @@ -347,6 +348,7 @@ def provide(self, pre_process=None, post_process=None, vp_stage=None) -> MCoreHy post_process=post_process, pg_collection=pg_collection, vp_stage=vp_stage, + scatter_embedding_sequence_parallel=self.scatter_embedding_sequence_parallel, ) finally: self._pg_collection = pg_collection diff --git a/src/megatron/bridge/models/hybrid_mla_provider.py b/src/megatron/bridge/models/hybrid_mla_provider.py new file mode 100644 index 0000000000..e1ae8ec258 --- /dev/null +++ b/src/megatron/bridge/models/hybrid_mla_provider.py @@ -0,0 +1,148 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Hybrid model provider for Multi-Latent Attention.""" + +from dataclasses import dataclass + +from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols +from megatron.core.transformer.experimental_attention_variant.dsa import ( + is_dsa_skip_topk_layer, + source_dsa_compute_layer, +) + +from megatron.bridge.models.hybrid.hybrid_provider import HybridModelProvider +from megatron.bridge.models.transformer_config import MLATransformerConfig + + +def _validate_dsa_pipeline_segments(main_pattern: str, *, topk_freq: int, skip_topk_offset: int) -> None: + """Ensure shared DSA indexers never reference a previous pipeline stage.""" + layer_offset = 0 + for segment_idx, segment in enumerate(main_pattern.split(Symbols.PIPE)): + if segment_idx > 0 and (not segment or segment[0] != Symbols.DS_ATTENTION): + raise ValueError( + "Hybrid MLA pipeline separators must be placed before DSA layers " + "so attention and FFN physical layers stay together." + ) + for local_idx, layer_type in enumerate(segment): + layer_number = layer_offset + local_idx + 1 + if layer_type != Symbols.DS_ATTENTION or not is_dsa_skip_topk_layer( + layer_number, skip_topk_offset, topk_freq + ): + continue + source_layer = source_dsa_compute_layer(layer_number, skip_topk_offset, topk_freq) + if source_layer <= layer_offset: + raise ValueError( + "Hybrid MLA pipeline segments cannot split DSA IndexShare groups: " + f"layer {layer_number} reuses indexer output from layer {source_layer}, " + "which is on a previous pipeline stage." + ) + layer_offset += len(segment) + + +def _add_dsa_pipeline_separators( + pattern: str, + *, + pipeline_model_parallel_size: int, + topk_freq: int, + skip_topk_offset: int, +) -> str: + """Add balanced pipeline separators before full-indexer DSA layers.""" + main_pattern, separator, mtp_pattern = pattern.partition(Symbols.MTP_SEPARATOR) + if Symbols.DS_ATTENTION not in main_pattern or pipeline_model_parallel_size <= 1: + return pattern + + if Symbols.PIPE in main_pattern: + num_segments = main_pattern.count(Symbols.PIPE) + 1 + if num_segments != pipeline_model_parallel_size: + raise ValueError( + f"hybrid_layer_pattern defines {num_segments} pipeline segments, but " + f"pipeline_model_parallel_size is {pipeline_model_parallel_size}." + ) + _validate_dsa_pipeline_segments( + main_pattern, + topk_freq=topk_freq, + skip_topk_offset=skip_topk_offset, + ) + return pattern + + candidate_boundaries = [ + layer_idx + for layer_idx, layer_type in enumerate(main_pattern) + if layer_idx > 0 + and layer_type == Symbols.DS_ATTENTION + and not is_dsa_skip_topk_layer(layer_idx + 1, skip_topk_offset, topk_freq) + ] + required_boundaries = pipeline_model_parallel_size - 1 + if len(candidate_boundaries) < required_boundaries: + raise ValueError( + "Cannot split Hybrid MLA layers across pipeline stages without crossing a DSA " + f"IndexShare group: found {len(candidate_boundaries)} valid boundaries for " + f"pipeline_model_parallel_size={pipeline_model_parallel_size}." + ) + + selected_boundaries = [] + candidate_start = 0 + for stage_idx in range(1, pipeline_model_parallel_size): + remaining_boundaries = pipeline_model_parallel_size - stage_idx - 1 + choices = candidate_boundaries[candidate_start : len(candidate_boundaries) - remaining_boundaries] + target = len(main_pattern) * stage_idx / pipeline_model_parallel_size + boundary = min(choices, key=lambda candidate: (abs(candidate - target), candidate)) + selected_boundaries.append(boundary) + candidate_start = candidate_boundaries.index(boundary, candidate_start) + 1 + + boundaries = set(selected_boundaries) + segmented_pattern = "".join( + (Symbols.PIPE if layer_idx in boundaries else "") + layer_type + for layer_idx, layer_type in enumerate(main_pattern) + ) + _validate_dsa_pipeline_segments( + segmented_pattern, + topk_freq=topk_freq, + skip_topk_offset=skip_topk_offset, + ) + return segmented_pattern + (separator + mtp_pattern if separator else "") + + +@dataclass +class HybridMLAModelProvider(HybridModelProvider, MLATransformerConfig): + """Provider for Hybrid models using Multi-Latent Attention.""" + + def finalize(self) -> None: + """Finalize the provider after adding DSA-safe pipeline segments.""" + pattern_attr = "hybrid_layer_pattern" + pattern = self.hybrid_layer_pattern + if pattern is None and self.hybrid_override_pattern is not None: + pattern_attr = "hybrid_override_pattern" + pattern = self.hybrid_override_pattern + + main_pattern = pattern.split(Symbols.MTP_SEPARATOR)[0] if pattern is not None else "" + if Symbols.DS_ATTENTION in main_pattern and self.pipeline_model_parallel_size > 1: + if ( + self.num_layers_in_first_pipeline_stage is not None + or self.num_layers_in_last_pipeline_stage is not None + ): + raise ValueError( + "Hybrid MLA derives pipeline stages from hybrid_layer_pattern; " + "num_layers_in_first_pipeline_stage and num_layers_in_last_pipeline_stage are unsupported." + ) + pattern = _add_dsa_pipeline_separators( + pattern, + pipeline_model_parallel_size=self.pipeline_model_parallel_size, + topk_freq=self.dsa_indexer_topk_freq or 1, + skip_topk_offset=self.dsa_indexer_skip_topk_offset or 0, + ) + setattr(self, pattern_attr, pattern) + + super().finalize() diff --git a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py index f0d360baff..90099ca213 100644 --- a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py +++ b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py @@ -44,7 +44,6 @@ get_dist_train_vision_dp_data, get_vision_cp_data, pack_dist_train_vision_module_output, - preprocess_packed_seqs, qwen3vl_cp_split, reorganize_inputs, split_data_cp_rank, @@ -54,6 +53,7 @@ from megatron.bridge.training.utils.packed_seq_utils import ( get_packed_seq_cp_partition_indices, get_packed_seq_q_cu_seqlens, + preprocess_packed_seqs, ) from megatron.bridge.utils.vocab_utils import calculate_padded_vocab_size diff --git a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/utils.py b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/utils.py index 9f7c579322..b06476e9c8 100644 --- a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/utils.py +++ b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/utils.py @@ -17,9 +17,6 @@ from typing import Optional, Sequence, Union import torch -from megatron.core import mpu -from megatron.core.packed_seq_params import PackedSeqParams -from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.transformer.module import MegatronModule from megatron.core.transformer.spec_utils import ModuleSpec, build_module from megatron.core.utils import get_tensor_model_parallel_group_if_none @@ -698,109 +695,3 @@ def ensure_requires_grad_for_cp_collective(tensors: Sequence[torch.Tensor]) -> N for tensor in tensors: if not tensor.requires_grad: tensor.requires_grad_(True) - - -def preprocess_packed_seqs( - input_ids: torch.Tensor, - attention_mask: torch.Tensor, - pre_process: bool = True, - pg_collection: Optional[ProcessGroupCollection] = None, -) -> tuple[torch.Tensor, PackedSeqParams]: - """ - Preprocess packed sequences - CP splits sequence into CP*2 chunks, and each GPU gets 2 chunks (GPU0 gets first and last chunks, GPU1 - gets second and second last chunks, and so on), this is for load balancing with causal masking. - See https://github.com/NVIDIA/TransformerEngine/issues/1368 - """ - batch_size = input_ids.shape[0] - - # Ensure boolean dtype for correct advanced indexing (bool → mask select, - # int → fancy index which silently corrupts data when values are 0/1). - if attention_mask.dtype != torch.bool: - attention_mask = attention_mask.bool() - - seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) - if pg_collection is not None: - tp_size = pg_collection.tp.size() - cp_size = pg_collection.cp.size() - cp_rank = pg_collection.cp.rank() - else: - tp_size = mpu.get_tensor_model_parallel_world_size() - cp_size = mpu.get_context_parallel_world_size() - cp_rank = mpu.get_context_parallel_rank() - align_size = tp_size * cp_size * 2 if cp_size > 1 else tp_size - - pad_size = (align_size - seqlens_in_batch % align_size) % align_size - seqlens_in_batch_padded = seqlens_in_batch + pad_size - - cu_seqlens = torch.zeros(batch_size + 1, dtype=torch.int32, device=input_ids.device) - cu_seqlens[1:] = torch.cumsum(seqlens_in_batch, dim=0) - cu_seqlens_padded = torch.zeros(batch_size + 1, dtype=torch.int32, device=input_ids.device) - cu_seqlens_padded[1:] = torch.cumsum(seqlens_in_batch_padded, dim=0) - - # ---------------------------------------------------------------------------- - # Move the index information needed in the subsequent loop to the CPU at once, - # to avoid frequent .item() calls in the loop that cause D2H synchronization - # ---------------------------------------------------------------------------- - seqlens_in_batch_cpu: list[int] = seqlens_in_batch.tolist() # original valid lengths - seqlens_in_batch_padded_cpu: list[int] = seqlens_in_batch_padded.tolist() # lengths after padding - cu_seqlens_padded_cpu: list[int] = cu_seqlens_padded.tolist() # start positions (after padding) - - # Pure Python int calculation to avoid further synchronization - max_seqlen_in_batch = max(seqlens_in_batch_padded_cpu) - - shape = list(input_ids.shape[1:]) - shape[0] = sum(seqlens_in_batch_padded_cpu) // cp_size - if pre_process: - input_ids_rmpad = torch.zeros(shape, dtype=input_ids.dtype, device=input_ids.device) - for i in range(batch_size): - # Use Python int, so no GPU→CPU sync in the loop - if cp_size <= 1: - seqlen = seqlens_in_batch_cpu[i] - start_idx = cu_seqlens_padded_cpu[i] - input_ids_rmpad[start_idx : start_idx + seqlen] = input_ids[i, attention_mask[i]] - continue - - seqlen_padded_i = seqlens_in_batch_padded_cpu[i] - seqlen = seqlen_padded_i // cp_size - half_seqlen = seqlen // 2 - start_idx = cu_seqlens_padded_cpu[i] // cp_size - # split to 2 chunks - d = input_ids[i, attention_mask[i]] - # Chunk indices are computed in the *padded* coordinate space - # (seqlen_padded_i), while d holds only the *valid* tokens. When a row's - # valid length is shorter than the alignment padding, the first-chunk - # slice can start past the end of d and yield an empty tensor, which - # raises on assignment: - # RuntimeError: The expanded size of the tensor (2) must match the - # existing size (0) at non-singleton dimension 0 - # Clamp it the same way the remaining chunk below already is. - first_start = half_seqlen * cp_rank - first_end = min(half_seqlen * (cp_rank + 1), d.shape[0]) - first_len = max(first_end - first_start, 0) - if first_len > 0: - input_ids_rmpad[start_idx : start_idx + first_len] = d[first_start:first_end] - - remain_start = seqlen_padded_i - half_seqlen * (cp_rank + 1) - remain_end = seqlen_padded_i - half_seqlen * cp_rank - remain_end = min(remain_end, d.shape[0]) - remain_len = max(remain_end - remain_start, 0) - if remain_len > 0: - input_ids_rmpad[start_idx + half_seqlen : start_idx + half_seqlen + remain_len] = d[ - remain_start:remain_end - ] - - packed_seq_params = PackedSeqParams( - qkv_format="thd", - cu_seqlens_q=cu_seqlens_padded, - max_seqlen_q=max_seqlen_in_batch, - cu_seqlens_kv=cu_seqlens_padded, - max_seqlen_kv=max_seqlen_in_batch, - cu_seqlens_q_padded=cu_seqlens_padded, - cu_seqlens_kv_padded=cu_seqlens_padded, - ) - - if pre_process: - return input_ids_rmpad.unsqueeze(0), packed_seq_params - else: - return input_ids, packed_seq_params diff --git a/src/megatron/bridge/training/utils/packed_seq_utils.py b/src/megatron/bridge/training/utils/packed_seq_utils.py index 96bca2e283..1184769712 100644 --- a/src/megatron/bridge/training/utils/packed_seq_utils.py +++ b/src/megatron/bridge/training/utils/packed_seq_utils.py @@ -15,11 +15,13 @@ from __future__ import annotations import inspect +import math import megatron.core import torch -from megatron.core import parallel_state +from megatron.core import mpu, parallel_state from megatron.core.packed_seq_params import PackedSeqParams +from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.utils import get_batch_on_this_cp_rank from packaging.version import Version as PkgVersion @@ -366,3 +368,105 @@ def get_packed_seq_params(batch: dict[str, PackedMetadataValue]) -> PackedSeqPar else {} ), ) + + +def preprocess_packed_seqs( + input_ids: torch.Tensor, + attention_mask: torch.Tensor, + pre_process: bool = True, + pg_collection: ProcessGroupCollection | None = None, + use_fp8_padding: bool = False, +) -> tuple[torch.Tensor, PackedSeqParams]: + """Preprocess packed sequences into THD layout. + + CP splits sequence into CP*2 chunks, and each GPU gets 2 chunks (GPU0 gets + first and last chunks, GPU1 gets second and second last chunks, and so on), + for load balancing with causal masking. + See https://github.com/NVIDIA/TransformerEngine/issues/1368 + """ + batch_size = input_ids.shape[0] + + # Ensure boolean dtype for correct advanced indexing (bool → mask select, + # int → fancy index which silently corrupts data when values are 0/1). + if attention_mask.dtype != torch.bool: + attention_mask = attention_mask.bool() + + seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) + if pg_collection is not None: + tp_size = pg_collection.tp.size() + cp_size = pg_collection.cp.size() + cp_rank = pg_collection.cp.rank() + else: + tp_size = mpu.get_tensor_model_parallel_world_size() + cp_size = mpu.get_context_parallel_world_size() + cp_rank = mpu.get_context_parallel_rank() + align_size = tp_size * cp_size * 2 if cp_size > 1 else tp_size + if use_fp8_padding: + total_align = align_size * 128 + align_size = math.lcm(16, align_size) + + pad_size = (align_size - seqlens_in_batch % align_size) % align_size + seqlens_in_batch_padded = seqlens_in_batch + pad_size + + cu_seqlens = torch.zeros(batch_size + 1, dtype=torch.int32, device=input_ids.device) + cu_seqlens[1:] = torch.cumsum(seqlens_in_batch, dim=0) + cu_seqlens_padded = torch.zeros(batch_size + 1, dtype=torch.int32, device=input_ids.device) + cu_seqlens_padded[1:] = torch.cumsum(seqlens_in_batch_padded, dim=0) + + if use_fp8_padding: + pad_size_last = (total_align - cu_seqlens_padded[-1] % total_align) % total_align + cu_seqlens_padded[-1] += pad_size_last + seqlens_in_batch_padded[-1] += pad_size_last + + # Move the index information needed in the subsequent loop to the CPU at once, + # to avoid frequent .item() calls in the loop that cause D2H synchronization. + seqlens_in_batch_cpu: list[int] = seqlens_in_batch.tolist() + seqlens_in_batch_padded_cpu: list[int] = seqlens_in_batch_padded.tolist() + cu_seqlens_padded_cpu: list[int] = cu_seqlens_padded.tolist() + + max_seqlen_in_batch = max(seqlens_in_batch_padded_cpu) + + shape = list(input_ids.shape[1:]) + shape[0] = sum(seqlens_in_batch_padded_cpu) // cp_size + if pre_process: + input_ids_rmpad = torch.zeros(shape, dtype=input_ids.dtype, device=input_ids.device) + for i in range(batch_size): + if cp_size <= 1: + seqlen = seqlens_in_batch_cpu[i] + start_idx = cu_seqlens_padded_cpu[i] + input_ids_rmpad[start_idx : start_idx + seqlen] = input_ids[i, attention_mask[i]] + continue + + seqlen_padded_i = seqlens_in_batch_padded_cpu[i] + seqlen = seqlen_padded_i // cp_size + half_seqlen = seqlen // 2 + start_idx = cu_seqlens_padded_cpu[i] // cp_size + d = input_ids[i, attention_mask[i]] + first_start = half_seqlen * cp_rank + first_end = min(half_seqlen * (cp_rank + 1), d.shape[0]) + first_len = max(first_end - first_start, 0) + if first_len > 0: + input_ids_rmpad[start_idx : start_idx + first_len] = d[first_start:first_end] + + remain_start = seqlen_padded_i - half_seqlen * (cp_rank + 1) + remain_end = min(seqlen_padded_i - half_seqlen * cp_rank, d.shape[0]) + remain_len = max(remain_end - remain_start, 0) + if remain_len > 0: + input_ids_rmpad[start_idx + half_seqlen : start_idx + half_seqlen + remain_len] = d[ + remain_start:remain_end + ] + + packed_seq_params = PackedSeqParams( + qkv_format="thd", + cu_seqlens_q=cu_seqlens_padded, + max_seqlen_q=max_seqlen_in_batch, + cu_seqlens_kv=cu_seqlens_padded, + max_seqlen_kv=max_seqlen_in_batch, + cu_seqlens_q_padded=cu_seqlens_padded, + cu_seqlens_kv_padded=cu_seqlens_padded, + ) + + if pre_process: + return input_ids_rmpad.unsqueeze(0), packed_seq_params + else: + return input_ids, packed_seq_params diff --git a/tests/unit_tests/models/glm5next/test_glm53flash_mapping.py b/tests/unit_tests/models/glm5next/test_glm53flash_mapping.py new file mode 100644 index 0000000000..733785377a --- /dev/null +++ b/tests/unit_tests/models/glm5next/test_glm53flash_mapping.py @@ -0,0 +1,160 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Regression tests for GLM-5.3-Flash's section-wise KDA TP weight layout.""" + +from math import prod +from types import SimpleNamespace +from unittest.mock import PropertyMock, patch + +import pytest +import torch +from megatron.core.models.hybrid.layers.utils import create_layer_config + +from megatron.bridge.models.conversion.param_mapping import HCAlphaMapping, MegatronParamMapping +from megatron.bridge.models.glm5next.glm53flash_bridge import GLM53FlashBridge, _ColumnParallelConcatMapping + + +pytestmark = pytest.mark.unit + + +def test_vlm_provider_keeps_vision_trainable_and_copies_multimodal_config(): + config_module = pytest.importorskip("transformers.models.glm5_next.configuration_glm5_next") + config = config_module.Glm5NextConfig( + text_config={ + "scoring_func": "sigmoid", + "linear_attn_config": {"num_heads": 64, "head_dim": 128, "short_conv_kernel_size": 4}, + } + ) + provider = GLM53FlashBridge().provider_bridge(SimpleNamespace(config=config)) + assert provider.vision_config is config.vision_config + assert provider.spatial_merge_size == config.vision_config.spatial_merge_size + assert not provider.freeze_vision_model + assert not provider.freeze_vision_projection + assert not provider.scatter_embedding_sequence_parallel + for field in ( + "image_token_id", + "video_token_id", + "image_start_token_id", + "image_end_token_id", + "video_start_token_id", + "video_end_token_id", + ): + assert getattr(provider, field) == getattr(config, field) + + +def test_vlm_mapping_prefixes_text_weights_and_nested_tp_mappings(): + config_module = pytest.importorskip("transformers.models.glm5_next.configuration_glm5_next") + config = config_module.Glm5NextConfig(text_config={"scoring_func": "sigmoid"}) + bridge = GLM53FlashBridge() + bridge.hf_config = config + mappings = bridge.mapping_registry().mappings + # The unified bridge prefixes every language-model weight with ``language_model.`` + # and appends a single replicated vision mapping at the end. + assert mappings[-1].megatron_param == "visual.**" + assert mappings[-1].hf_param == "model.visual.**" + for mapping in mappings[:-1]: + assert mapping.megatron_param.startswith("language_model.") + if isinstance(mapping, _ColumnParallelConcatMapping): + assert mapping._tp_mapping.megatron_param == mapping.megatron_param + + +@pytest.mark.parametrize("gate_lower_bound", [-5.0, -1.0]) +def test_glm53_provider_precision_contract(gate_lower_bound): + config_module = pytest.importorskip("transformers.models.glm5_next.configuration_glm5_next") + config = config_module.Glm5NextConfig( + text_config={ + "scoring_func": "sigmoid", + "index_kpool": 4, + "linear_attn_config": { + "num_heads": 64, + "head_dim": 128, + "short_conv_kernel_size": 4, + "gate_lower_bound": gate_lower_bound, + }, + } + ) + provider = GLM53FlashBridge().provider_bridge(SimpleNamespace(config=config)) + assert provider.layernorm_epsilon == 1e-5 + assert provider.dsa_indexer_k_norm_epsilon == 1e-6 + assert provider.dsa_indexer_kpool_fp8 + assert create_layer_config(provider, "D").dsa_indexer_kpool_fp8 + assert provider.mhc_norm_eps_inside_sqrt + assert provider.mhc_keep_mappings_in_fp32 + assert not provider.mhc_learned_output_contract + assert provider.kda_two_stage_gates + assert provider.kda_lower_bound == gate_lower_bound + + +@pytest.mark.parametrize("index", [0, 1, 2]) +def test_hc_alpha_import_export_and_wildcard_resolution(index): + attr = ("alpha_pre", "alpha_post", "alpha_res")[index] + mapping = HCAlphaMapping(f"decoder.layers.*.hyper_connection.{attr}", "model.layers.*.hc_scale", index) + mapping = mapping.resolve(("7",)) + assert mapping.megatron_param == f"decoder.layers.7.hyper_connection.{attr}" + assert mapping.hf_param == "model.layers.7.hc_scale" + module = torch.nn.Module() + source = torch.tensor([0.1, 0.2, 0.3]) + for i, name in enumerate(("alpha_pre", "alpha_post", "alpha_res")): + module.register_parameter(name, torch.nn.Parameter(source[i : i + 1].clone())) + torch.testing.assert_close(mapping.hf_to_megatron(source, module), source[index : index + 1]) + with patch.object(mapping, "broadcast_from_pp_rank", side_effect=lambda weight, **_: weight): + exported = mapping.megatron_to_hf(getattr(module, attr), module) + if index == 0: + torch.testing.assert_close(exported[mapping.hf_param], source) + else: + assert exported == {} + + +@pytest.mark.parametrize("tp_size", [1, 2, 4, 8, 16]) +@pytest.mark.parametrize("trailing_shape", [(3,), (1, 4)]) +def test_kda_import_shards_each_projection_independently(tp_size, trailing_shape): + mapping = _ColumnParallelConcatMapping("weight", ["q.weight", "k.weight", "v.weight"]) + sections = [ + torch.arange(rows * prod(trailing_shape)).reshape(rows, *trailing_shape) + offset + for rows, offset in [(16, 0), (16, 1000), (32, 2000)] + ] + merged = torch.cat(sections) + + with patch.object(MegatronParamMapping, "tp_size", new_callable=PropertyMock, return_value=tp_size): + shards = mapping._shard_per_rank(merged, [16, 16, 32]) + + assert len(shards) == tp_size + for rank, shard in enumerate(shards): + expected = torch.cat( + [ + section[rank * (len(section) // tp_size) : (rank + 1) * (len(section) // tp_size)] + for section in sections + ] + ) + torch.testing.assert_close(shard, expected, rtol=0, atol=0) + + +@pytest.mark.parametrize("tp_size", [1, 2, 4, 8, 16]) +@pytest.mark.parametrize("trailing_shape", [(3,), (1, 4)]) +def test_kda_export_reconstructs_hf_order_from_rank_local_sections(tp_size, trailing_shape): + mapping = _ColumnParallelConcatMapping("weight", ["q.weight", "k.weight", "v.weight"]) + module = torch.nn.Module() + module.config = SimpleNamespace( + linear_key_head_dim=2, + linear_value_head_dim=4, + linear_num_key_heads=8, + linear_num_value_heads=8, + ) + sections = [torch.randn(rows, *trailing_shape) for rows in (16, 16, 32)] + # Build the distributed input independently of the import implementation. + shards = [ + torch.cat( + [section.narrow(0, rank * (len(section) // tp_size), len(section) // tp_size) for section in sections] + ) + for rank in range(tp_size) + ] + + with ( + patch.object(MegatronParamMapping, "tp_size", new_callable=PropertyMock, return_value=tp_size), + patch.object(mapping._tp_mapping, "gather_from_tp_ranks", return_value=shards), + ): + exported = mapping.megatron_to_hf(shards[0], module) + + assert list(exported) == ["q.weight", "k.weight", "v.weight"] + for name, expected in zip(exported, sections): + torch.testing.assert_close(exported[name], expected, rtol=0, atol=0) diff --git a/tests/unit_tests/models/qwen_vl/modelling_qwen3_vl/test_utils.py b/tests/unit_tests/models/qwen_vl/modelling_qwen3_vl/test_utils.py index e0afcb7521..ba45ddb1be 100644 --- a/tests/unit_tests/models/qwen_vl/modelling_qwen3_vl/test_utils.py +++ b/tests/unit_tests/models/qwen_vl/modelling_qwen3_vl/test_utils.py @@ -75,6 +75,35 @@ def test_select_sequence_handles_none_and_selects_dimension(): assert selected.is_contiguous() +@pytest.mark.parametrize("tp_size", [1, 8]) +@pytest.mark.parametrize("use_fp8_padding", [False, True]) +def test_packed_sequence_alignment_and_gradients(tp_size, use_fp8_padding): + pg_collection = SimpleNamespace( + tp=SimpleNamespace(size=lambda: tp_size), + cp=SimpleNamespace(size=lambda: 1, rank=lambda: 0), + ) + embeddings = torch.arange(2 * 19 * 4, dtype=torch.float32).view(2, 19, 4).requires_grad_() + mask = torch.arange(19)[None, :] < torch.tensor([2, 13])[:, None] + output, packed = preprocess_packed_seqs( + embeddings, mask, pg_collection=pg_collection, use_fp8_padding=use_fp8_padding + ) + alignment = 16 if use_fp8_padding else tp_size + lengths = [((length + alignment - 1) // alignment) * alignment for length in (2, 13)] + total = sum(lengths) + if use_fp8_padding: + total_alignment = tp_size * 128 + total = ((total + total_alignment - 1) // total_alignment) * total_alignment + expected = torch.zeros(1, total, 4) + expected[0, :2] = embeddings.detach()[0, :2] + expected[0, lengths[0] : lengths[0] + 13] = embeddings.detach()[1, :13] + torch.testing.assert_close(output, expected) + assert packed.cu_seqlens_q.tolist() == [0, lengths[0], total] + assert packed.cu_seqlens_q_padded.tolist() == [0, lengths[0], total] + assert packed.max_seqlen_q == max(lengths[0], total - lengths[0]) + output.sum().backward() + torch.testing.assert_close(embeddings.grad, mask[..., None].expand_as(embeddings).float()) + + def test_split_if_full_sequence_only_splits_full_length_inputs(): """Dense CP helpers should not double-slice tensors that are already local.""" full = torch.arange(16).view(1, 16) diff --git a/tests/unit_tests/models/test_model_bridge_lora.py b/tests/unit_tests/models/test_model_bridge_lora.py index 30fb2aad40..c4baba5cac 100644 --- a/tests/unit_tests/models/test_model_bridge_lora.py +++ b/tests/unit_tests/models/test_model_bridge_lora.py @@ -350,6 +350,11 @@ def test_merge_lora_adapter_weights_qkv_split(monkeypatch): hidden_size=4, attention_output_gate=False, num_moe_experts=0, + linear_num_key_heads=2, + linear_num_value_heads=2, + linear_key_head_dim=2, + linear_value_head_dim=2, + kda_two_stage_gates=False, ) megatron_model = [SimpleNamespace(config=config)] converted = { @@ -378,6 +383,21 @@ def test_merge_lora_adapter_weights_qkv_split(monkeypatch): torch.testing.assert_close(updated["v_proj.weight"], v_weight) +@pytest.mark.parametrize("trailing_shape", [(), (4,)]) +def test_split_kda_qkv_linear_out_weight(trailing_shape): + config = SimpleNamespace( + kda_two_stage_gates=True, + linear_num_key_heads=2, + linear_num_value_heads=2, + linear_key_head_dim=2, + linear_value_head_dim=2, + ) + sections = [torch.full((4, *trailing_shape), value) for value in (1.0, 2.0, 3.0)] + actual = DummyBridge()._split_qkv_linear_out_weight([SimpleNamespace(config=config)], torch.cat(sections)) + for name, expected in zip(("q_proj", "k_proj", "v_proj"), sections, strict=True): + torch.testing.assert_close(actual[name], expected) + + def test_merge_lora_adapter_weights_grouped_expert_missing_expert_idx(monkeypatch): bridge = DummyBridge() base = torch.zeros(2, 2) From f1bd9ba0bbb449b24d3b4ccddb6097023d6f9655 Mon Sep 17 00:00:00 2001 From: Hollow Man Date: Mon, 21 Sep 2026 11:01:50 +0300 Subject: [PATCH 2/8] Move config under KDALayerConfig Signed-off-by: Hollow Man --- src/megatron/bridge/models/glm5next/glm53flash_bridge.py | 3 ++- .../models/glm5next/test_glm53flash_mapping.py | 9 ++++++++- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/src/megatron/bridge/models/glm5next/glm53flash_bridge.py b/src/megatron/bridge/models/glm5next/glm53flash_bridge.py index fc47cd534b..7e700f388f 100644 --- a/src/megatron/bridge/models/glm5next/glm53flash_bridge.py +++ b/src/megatron/bridge/models/glm5next/glm53flash_bridge.py @@ -32,6 +32,7 @@ import torch from megatron.core.models.hybrid.hybrid_layer_allocation import Symbols +from megatron.core.ssm.kda_layer_config import KDALayerConfig from megatron.core.tensor_parallel import scatter_to_sequence_parallel_region from megatron.core.transformer import ModuleSpec from megatron.core.transformer.module import MegatronModule @@ -253,7 +254,7 @@ def glm53_hybrid_stack_spec(config: HybridMLAModelProvider) -> ModuleSpec: @dataclass -class GLM53FlashModelProvider(HybridMLAModelProvider): +class GLM53FlashModelProvider(HybridMLAModelProvider, KDALayerConfig): """Provider for the unified GLM-5.3-Flash VLM: hybrid language stack + HF vision.""" # Vision configuration (HF Glm5NextVisionConfig object, stored as-is). diff --git a/tests/unit_tests/models/glm5next/test_glm53flash_mapping.py b/tests/unit_tests/models/glm5next/test_glm53flash_mapping.py index 733785377a..81724c3e14 100644 --- a/tests/unit_tests/models/glm5next/test_glm53flash_mapping.py +++ b/tests/unit_tests/models/glm5next/test_glm53flash_mapping.py @@ -59,7 +59,8 @@ def test_vlm_mapping_prefixes_text_weights_and_nested_tp_mappings(): @pytest.mark.parametrize("gate_lower_bound", [-5.0, -1.0]) -def test_glm53_provider_precision_contract(gate_lower_bound): +@pytest.mark.parametrize("kda_disable_fp8", [False, True]) +def test_glm53_provider_precision_contract(gate_lower_bound, kda_disable_fp8): config_module = pytest.importorskip("transformers.models.glm5_next.configuration_glm5_next") config = config_module.Glm5NextConfig( text_config={ @@ -83,6 +84,12 @@ def test_glm53_provider_precision_contract(gate_lower_bound): assert not provider.mhc_learned_output_contract assert provider.kda_two_stage_gates assert provider.kda_lower_bound == gate_lower_bound + provider.apply_overrides_and_finalize(overrides={"kda_disable_fp8": kda_disable_fp8}) + kda_config = create_layer_config(provider, "K") + assert kda_config.kda_disable_fp8 is kda_disable_fp8 + assert kda_config.kda_two_stage_gates + assert kda_config.kda_safe_gate + assert kda_config.kda_lower_bound == gate_lower_bound @pytest.mark.parametrize("index", [0, 1, 2]) From 568b79321786936ec9a887a164fcd4eb04a7d5ad Mon Sep 17 00:00:00 2001 From: Hollow Man Date: Wed, 23 Sep 2026 01:19:31 +0300 Subject: [PATCH 3/8] Sync upstream for FP8 flags Signed-off-by: Hollow Man --- .../bridge/models/glm5next/glm53flash_bridge.py | 10 ++++++++-- .../models/glm5next/test_glm53flash_mapping.py | 10 ++++++++-- 2 files changed, 16 insertions(+), 4 deletions(-) diff --git a/src/megatron/bridge/models/glm5next/glm53flash_bridge.py b/src/megatron/bridge/models/glm5next/glm53flash_bridge.py index 7e700f388f..c0825ff6d0 100644 --- a/src/megatron/bridge/models/glm5next/glm53flash_bridge.py +++ b/src/megatron/bridge/models/glm5next/glm53flash_bridge.py @@ -547,6 +547,7 @@ def provider_bridge(self, hf_pretrained: PreTrainedCausalLM): provider.qk_head_dim = text_config.qk_head_dim # 256 provider.num_attention_heads = text_config.num_attention_heads # 64 provider.num_key_value_heads = text_config.num_key_value_heads # 64 + provider.mla_proj_disable_quantization = True # HF head_dim=0 is a NoPE sentinel, not a valid attention channel width. provider.kv_channels = text_config.v_head_dim # 256 provider.rope_type = "yarn" @@ -567,16 +568,21 @@ def provider_bridge(self, hf_pretrained: PreTrainedCausalLM): provider.dsa_indexer_loss_coeff = 0.0 provider.dsa_indexer_use_sparse_loss = False provider.dsa_indexer_rotate_activation = False - provider.dsa_indexer_kpool_fp8 = True + provider.dsa_indexer_qk_proj_disable_quantization = True + provider.dsa_indexer_kpool_use_quantization = True + # The HF FP8 config keeps the DSA weights projection in BF16. + provider.dsa_indexer_weights_proj_use_quantization = False + provider.dsa_indexer_weights_proj_output_dtype = "bf16" provider.dsa_indexer_k_norm_epsilon = 1e-6 provider.dsa_indexer_kpool = int(getattr(text_config, "index_kpool", 1)) # 4 provider.dsa_indexer_kpool_always_select_tail = bool( - getattr(text_config, "index_kpool_always_select_tail", False) + getattr(text_config, "index_kpool_always_select_tail", True) ) provider.dsa_indexer_topk_freq = 1 provider.dsa_indexer_skip_topk_offset = 0 # ---- KDA two-stage gates ---- + provider.kda_disable_fp8 = True provider.kda_two_stage_gates = True provider.kda_safe_gate = True provider.kda_lower_bound = float(text_config.linear_attn_config.get("gate_lower_bound", -5.0)) diff --git a/tests/unit_tests/models/glm5next/test_glm53flash_mapping.py b/tests/unit_tests/models/glm5next/test_glm53flash_mapping.py index 81724c3e14..5fa1ed2119 100644 --- a/tests/unit_tests/models/glm5next/test_glm53flash_mapping.py +++ b/tests/unit_tests/models/glm5next/test_glm53flash_mapping.py @@ -76,13 +76,19 @@ def test_glm53_provider_precision_contract(gate_lower_bound, kda_disable_fp8): ) provider = GLM53FlashBridge().provider_bridge(SimpleNamespace(config=config)) assert provider.layernorm_epsilon == 1e-5 + assert provider.mla_proj_disable_quantization assert provider.dsa_indexer_k_norm_epsilon == 1e-6 - assert provider.dsa_indexer_kpool_fp8 - assert create_layer_config(provider, "D").dsa_indexer_kpool_fp8 + assert provider.dsa_indexer_qk_proj_disable_quantization + assert provider.dsa_indexer_kpool_use_quantization + assert not provider.dsa_indexer_weights_proj_use_quantization + assert provider.dsa_indexer_weights_proj_output_dtype == "bf16" + assert create_layer_config(provider, "D").dsa_indexer_kpool_use_quantization + assert provider.dsa_indexer_kpool_always_select_tail assert provider.mhc_norm_eps_inside_sqrt assert provider.mhc_keep_mappings_in_fp32 assert not provider.mhc_learned_output_contract assert provider.kda_two_stage_gates + assert provider.kda_disable_fp8 assert provider.kda_lower_bound == gate_lower_bound provider.apply_overrides_and_finalize(overrides={"kda_disable_fp8": kda_disable_fp8}) kda_config = create_layer_config(provider, "K") From 983b3339d0483317214c578a516f1fd1527ad299 Mon Sep 17 00:00:00 2001 From: Hollow Man Date: Thu, 24 Sep 2026 02:27:38 +0300 Subject: [PATCH 4/8] fix MTP norm mapping Signed-off-by: Hollow Man --- .../models/glm5next/glm53flash_bridge.py | 2 +- .../glm5next/test_glm53flash_mapping.py | 24 +++++++++++++++++++ 2 files changed, 25 insertions(+), 1 deletion(-) diff --git a/src/megatron/bridge/models/glm5next/glm53flash_bridge.py b/src/megatron/bridge/models/glm5next/glm53flash_bridge.py index c0825ff6d0..52c099ff17 100644 --- a/src/megatron/bridge/models/glm5next/glm53flash_bridge.py +++ b/src/megatron/bridge/models/glm5next/glm53flash_bridge.py @@ -726,7 +726,7 @@ def mapping_registry(self) -> MegatronMappingRegistry: _EhProjSplitMapping(f"{mtp_mg}.h_proj.weight", eh_proj, 1), AutoMapping( f"{mtp_mg}.final_layernorm.weight", - "model.language_model.norm.weight", + f"{mtp_hf}.shared_head.norm.weight", ), ] ) diff --git a/tests/unit_tests/models/glm5next/test_glm53flash_mapping.py b/tests/unit_tests/models/glm5next/test_glm53flash_mapping.py index 5fa1ed2119..7833a62de9 100644 --- a/tests/unit_tests/models/glm5next/test_glm53flash_mapping.py +++ b/tests/unit_tests/models/glm5next/test_glm53flash_mapping.py @@ -58,6 +58,30 @@ def test_vlm_mapping_prefixes_text_weights_and_nested_tp_mappings(): assert mapping._tp_mapping.megatron_param == mapping.megatron_param +def test_mtp_final_layernorm_round_trips_through_the_mtp_shared_head_norm(): + config_module = pytest.importorskip("transformers.models.glm5_next.configuration_glm5_next") + config = config_module.Glm5NextConfig( + text_config={ + "scoring_func": "sigmoid", + "num_nextn_predict_layers": 1, + } + ) + bridge = GLM53FlashBridge() + bridge.hf_config = config + registry = bridge.mapping_registry() + + megatron_name = "language_model.mtp.layers.0.final_layernorm.weight" + hf_name = f"model.language_model.layers.{config.text_config.num_hidden_layers}.shared_head.norm.weight" + + imported = registry.megatron_to_hf_lookup(megatron_name) + assert imported is not None + assert imported.hf_param == hf_name + + exported = registry.hf_to_megatron_lookup(hf_name) + assert exported is not None + assert exported.megatron_param == megatron_name + + @pytest.mark.parametrize("gate_lower_bound", [-5.0, -1.0]) @pytest.mark.parametrize("kda_disable_fp8", [False, True]) def test_glm53_provider_precision_contract(gate_lower_bound, kda_disable_fp8): From d45669cb75966af377155101338047554bdd7e3d Mon Sep 17 00:00:00 2001 From: Hollow Man Date: Tue, 29 Sep 2026 16:48:41 +0300 Subject: [PATCH 5/8] Use fp32 for dsa_indexer_weights_proj_output_dtype Signed-off-by: Hollow Man --- src/megatron/bridge/models/glm5next/glm53flash_bridge.py | 5 +++-- tests/unit_tests/models/glm5next/test_glm53flash_mapping.py | 2 +- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/src/megatron/bridge/models/glm5next/glm53flash_bridge.py b/src/megatron/bridge/models/glm5next/glm53flash_bridge.py index 52c099ff17..f30806e41b 100644 --- a/src/megatron/bridge/models/glm5next/glm53flash_bridge.py +++ b/src/megatron/bridge/models/glm5next/glm53flash_bridge.py @@ -570,9 +570,10 @@ def provider_bridge(self, hf_pretrained: PreTrainedCausalLM): provider.dsa_indexer_rotate_activation = False provider.dsa_indexer_qk_proj_disable_quantization = True provider.dsa_indexer_kpool_use_quantization = True - # The HF FP8 config keeps the DSA weights projection in BF16. + # vLLM computes this projection in FP32 because BF16 rounding can change + # near-tie KPool rankings on long-context inputs. provider.dsa_indexer_weights_proj_use_quantization = False - provider.dsa_indexer_weights_proj_output_dtype = "bf16" + provider.dsa_indexer_weights_proj_output_dtype = "fp32" provider.dsa_indexer_k_norm_epsilon = 1e-6 provider.dsa_indexer_kpool = int(getattr(text_config, "index_kpool", 1)) # 4 provider.dsa_indexer_kpool_always_select_tail = bool( diff --git a/tests/unit_tests/models/glm5next/test_glm53flash_mapping.py b/tests/unit_tests/models/glm5next/test_glm53flash_mapping.py index 7833a62de9..01eee2f908 100644 --- a/tests/unit_tests/models/glm5next/test_glm53flash_mapping.py +++ b/tests/unit_tests/models/glm5next/test_glm53flash_mapping.py @@ -105,7 +105,7 @@ def test_glm53_provider_precision_contract(gate_lower_bound, kda_disable_fp8): assert provider.dsa_indexer_qk_proj_disable_quantization assert provider.dsa_indexer_kpool_use_quantization assert not provider.dsa_indexer_weights_proj_use_quantization - assert provider.dsa_indexer_weights_proj_output_dtype == "bf16" + assert provider.dsa_indexer_weights_proj_output_dtype == "fp32" assert create_layer_config(provider, "D").dsa_indexer_kpool_use_quantization assert provider.dsa_indexer_kpool_always_select_tail assert provider.mhc_norm_eps_inside_sqrt From 4a316a1bd00ef3dd07462c58161c11433fff9b3b Mon Sep 17 00:00:00 2001 From: Hollow Man Date: Wed, 30 Sep 2026 11:08:21 +0300 Subject: [PATCH 6/8] Enable dsa_indexer_k_norm_fp32 by default Signed-off-by: Hollow Man --- src/megatron/bridge/models/glm5next/glm53flash_bridge.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/megatron/bridge/models/glm5next/glm53flash_bridge.py b/src/megatron/bridge/models/glm5next/glm53flash_bridge.py index f30806e41b..cda65df367 100644 --- a/src/megatron/bridge/models/glm5next/glm53flash_bridge.py +++ b/src/megatron/bridge/models/glm5next/glm53flash_bridge.py @@ -570,6 +570,7 @@ def provider_bridge(self, hf_pretrained: PreTrainedCausalLM): provider.dsa_indexer_rotate_activation = False provider.dsa_indexer_qk_proj_disable_quantization = True provider.dsa_indexer_kpool_use_quantization = True + provider.dsa_indexer_k_norm_fp32 = True # vLLM computes this projection in FP32 because BF16 rounding can change # near-tie KPool rankings on long-context inputs. provider.dsa_indexer_weights_proj_use_quantization = False From 4180b67cb27f098dd6e0df6031b02ceeff770947 Mon Sep 17 00:00:00 2001 From: Hollow Man Date: Sat, 10 Oct 2026 11:36:50 +0300 Subject: [PATCH 7/8] Rename to dsa_indexer_qk_proj_use_quantization Signed-off-by: Hollow Man --- src/megatron/bridge/models/glm5next/glm53flash_bridge.py | 2 +- tests/unit_tests/models/glm5next/test_glm53flash_mapping.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/megatron/bridge/models/glm5next/glm53flash_bridge.py b/src/megatron/bridge/models/glm5next/glm53flash_bridge.py index cda65df367..a98270fda8 100644 --- a/src/megatron/bridge/models/glm5next/glm53flash_bridge.py +++ b/src/megatron/bridge/models/glm5next/glm53flash_bridge.py @@ -568,7 +568,7 @@ def provider_bridge(self, hf_pretrained: PreTrainedCausalLM): provider.dsa_indexer_loss_coeff = 0.0 provider.dsa_indexer_use_sparse_loss = False provider.dsa_indexer_rotate_activation = False - provider.dsa_indexer_qk_proj_disable_quantization = True + provider.dsa_indexer_qk_proj_use_quantization = False provider.dsa_indexer_kpool_use_quantization = True provider.dsa_indexer_k_norm_fp32 = True # vLLM computes this projection in FP32 because BF16 rounding can change diff --git a/tests/unit_tests/models/glm5next/test_glm53flash_mapping.py b/tests/unit_tests/models/glm5next/test_glm53flash_mapping.py index 01eee2f908..c935abfacd 100644 --- a/tests/unit_tests/models/glm5next/test_glm53flash_mapping.py +++ b/tests/unit_tests/models/glm5next/test_glm53flash_mapping.py @@ -102,7 +102,7 @@ def test_glm53_provider_precision_contract(gate_lower_bound, kda_disable_fp8): assert provider.layernorm_epsilon == 1e-5 assert provider.mla_proj_disable_quantization assert provider.dsa_indexer_k_norm_epsilon == 1e-6 - assert provider.dsa_indexer_qk_proj_disable_quantization + assert not provider.dsa_indexer_qk_proj_use_quantization assert provider.dsa_indexer_kpool_use_quantization assert not provider.dsa_indexer_weights_proj_use_quantization assert provider.dsa_indexer_weights_proj_output_dtype == "fp32" From e8202eb75f1c9b20fbe44c7edd114dae595ca4cd Mon Sep 17 00:00:00 2001 From: Hollow Man Date: Sun, 11 Oct 2026 17:30:10 +0300 Subject: [PATCH 8/8] Update to mla_proj_use_quantization Signed-off-by: Hollow Man --- src/megatron/bridge/models/glm5next/glm53flash_bridge.py | 2 +- tests/unit_tests/models/glm5next/test_glm53flash_mapping.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/megatron/bridge/models/glm5next/glm53flash_bridge.py b/src/megatron/bridge/models/glm5next/glm53flash_bridge.py index a98270fda8..a924235b89 100644 --- a/src/megatron/bridge/models/glm5next/glm53flash_bridge.py +++ b/src/megatron/bridge/models/glm5next/glm53flash_bridge.py @@ -547,7 +547,7 @@ def provider_bridge(self, hf_pretrained: PreTrainedCausalLM): provider.qk_head_dim = text_config.qk_head_dim # 256 provider.num_attention_heads = text_config.num_attention_heads # 64 provider.num_key_value_heads = text_config.num_key_value_heads # 64 - provider.mla_proj_disable_quantization = True + provider.mla_proj_use_quantization = False # HF head_dim=0 is a NoPE sentinel, not a valid attention channel width. provider.kv_channels = text_config.v_head_dim # 256 provider.rope_type = "yarn" diff --git a/tests/unit_tests/models/glm5next/test_glm53flash_mapping.py b/tests/unit_tests/models/glm5next/test_glm53flash_mapping.py index c935abfacd..18fa6acbd6 100644 --- a/tests/unit_tests/models/glm5next/test_glm53flash_mapping.py +++ b/tests/unit_tests/models/glm5next/test_glm53flash_mapping.py @@ -100,7 +100,7 @@ def test_glm53_provider_precision_contract(gate_lower_bound, kda_disable_fp8): ) provider = GLM53FlashBridge().provider_bridge(SimpleNamespace(config=config)) assert provider.layernorm_epsilon == 1e-5 - assert provider.mla_proj_disable_quantization + assert not provider.mla_proj_use_quantization assert provider.dsa_indexer_k_norm_epsilon == 1e-6 assert not provider.dsa_indexer_qk_proj_use_quantization assert provider.dsa_indexer_kpool_use_quantization