feat(model): add GLM-5.3-Flash support

Signed-off-by: Hollow Man <hollowman@opensuse.org>

Author: HollowMan6

src/megatron/bridge/models/__init__.py

@@ -85,6 +85,11 @@
     GLM45Bridge,
     GLM47FlashBridge,
 )
+from megatron.bridge.models.glm5next import (
+    GLM53FlashBridge,
+    GLM53FlashModel,
+    GLM53FlashModelProvider,
+)
 from megatron.bridge.models.glm_moe_dsa import (
     GLM5Bridge,
 )
@@ -248,6 +253,9 @@
     "GLM45Bridge",
     "GLM47FlashBridge",
     "GLM5Bridge",
+    "GLM53FlashBridge",
+    "GLM53FlashModel",
+    "GLM53FlashModelProvider",
     "GLM45VBridge",
     "GLM45VModelProvider",
     "GPTModelProvider",

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",
 ]

src/megatron/bridge/models/conversion/model_bridge.py

@@ -1010,6 +1010,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]:

src/megatron/bridge/models/conversion/param_mapping.py

@@ -4340,3 +4340,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 ``<layer>.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 {}

src/megatron/bridge/models/conversion/peft_bridge.py

@@ -370,10 +370,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(

src/megatron/bridge/models/deepseek/deepseek_v4_bridge.py

@@ -77,7 +77,7 @@
     AutoMapping,
     ColumnParallelMapping,
     GatedMLPMapping,
-    MegatronParamMapping,
+    HCAlphaMapping,
     ReplicatedMapping,
 )
 from megatron.bridge.models.hf_pretrained.causal_lm import PreTrainedCausalLM
@@ -243,93 +243,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).
 
@@ -808,41 +721,36 @@ def mapping_registry(self) -> MegatronMappingRegistry:  # noqa: C901
             ),
         ]
 
-        # HC alpha scalars need custom concatenation mapping (per-layer, both attn and ffn)
+        # HC alpha scalars: one HCAlphaMapping per alpha (index 0 exports all three;
+        # index 1/2 import only and suppress export "no mapping found" warnings).
         # These are wildcarded across all layers.
         mappings += [
-            _HCAlphaMapping(
-                megatron_pre="decoder.layers.*.self_attention_hyper_connection.alpha_pre",
-                megatron_post="decoder.layers.*.self_attention_hyper_connection.alpha_post",
-                megatron_res="decoder.layers.*.self_attention_hyper_connection.alpha_res",
-                hf_param="layers.*.hc_attn_scale",
-            ),
-            _HCAlphaMapping(
-                megatron_pre="decoder.layers.*.mlp_hyper_connection.alpha_pre",
-                megatron_post="decoder.layers.*.mlp_hyper_connection.alpha_post",
-                megatron_res="decoder.layers.*.mlp_hyper_connection.alpha_res",
-                hf_param="layers.*.hc_ffn_scale",
+            HCAlphaMapping(
+                "decoder.layers.*.self_attention_hyper_connection.alpha_pre",
+                "layers.*.hc_attn_scale",
+                0,
             ),
-        ]
-
-        # HC alpha secondary: register alpha_post and alpha_res to suppress export warnings
-        mappings += [
-            _HCAlphaSecondaryMapping(
+            HCAlphaMapping(
                 "decoder.layers.*.self_attention_hyper_connection.alpha_post",
                 "layers.*.hc_attn_scale",
                 1,
             ),
-            _HCAlphaSecondaryMapping(
+            HCAlphaMapping(
                 "decoder.layers.*.self_attention_hyper_connection.alpha_res",
                 "layers.*.hc_attn_scale",
                 2,
             ),
-            _HCAlphaSecondaryMapping(
+            HCAlphaMapping(
+                "decoder.layers.*.mlp_hyper_connection.alpha_pre",
+                "layers.*.hc_ffn_scale",
+                0,
+            ),
+            HCAlphaMapping(
                 "decoder.layers.*.mlp_hyper_connection.alpha_post",
                 "layers.*.hc_ffn_scale",
                 1,
             ),
-            _HCAlphaSecondaryMapping(
+            HCAlphaMapping(
                 "decoder.layers.*.mlp_hyper_connection.alpha_res",
                 "layers.*.hc_ffn_scale",
                 2,
@@ -936,34 +844,23 @@ def mapping_registry(self) -> MegatronMappingRegistry:  # noqa: C901
                 ),
             ]
 
-            # MTP HC alpha scalars
-            mappings += [
-                _HCAlphaMapping(
-                    megatron_pre=f"{mg_pfx}.mtp_model_layer.self_attention_hyper_connection.alpha_pre",
-                    megatron_post=f"{mg_pfx}.mtp_model_layer.self_attention_hyper_connection.alpha_post",
-                    megatron_res=f"{mg_pfx}.mtp_model_layer.self_attention_hyper_connection.alpha_res",
-                    hf_param=f"{ck_pfx}.hc_attn_scale",
-                ),
-                _HCAlphaMapping(
-                    megatron_pre=f"{mg_pfx}.mtp_model_layer.mlp_hyper_connection.alpha_pre",
-                    megatron_post=f"{mg_pfx}.mtp_model_layer.mlp_hyper_connection.alpha_post",
-                    megatron_res=f"{mg_pfx}.mtp_model_layer.mlp_hyper_connection.alpha_res",
-                    hf_param=f"{ck_pfx}.hc_ffn_scale",
-                ),
-            ]
-
-            # MTP HC alpha secondary: suppress export warnings for post/res
+            # MTP HC alpha scalars (one HCAlphaMapping per alpha; index 0 exports all three)
             for _hc_mg_sub, _hc_hf_key in [
                 ("self_attention_hyper_connection", "hc_attn_scale"),
                 ("mlp_hyper_connection", "hc_ffn_scale"),
             ]:
                 mappings += [
-                    _HCAlphaSecondaryMapping(
+                    HCAlphaMapping(
+                        f"{mg_pfx}.mtp_model_layer.{_hc_mg_sub}.alpha_pre",
+                        f"{ck_pfx}.{_hc_hf_key}",
+                        0,
+                    ),
+                    HCAlphaMapping(
                         f"{mg_pfx}.mtp_model_layer.{_hc_mg_sub}.alpha_post",
                         f"{ck_pfx}.{_hc_hf_key}",
                         1,
                     ),
-                    _HCAlphaSecondaryMapping(
+                    HCAlphaMapping(
                         f"{mg_pfx}.mtp_model_layer.{_hc_mg_sub}.alpha_res",
                         f"{ck_pfx}.{_hc_hf_key}",
                         2,