GLM-5.2 131k: dsv4-next bridge pin + vectorized FP8 dequant

.gitmodules

@@ -1,3 +1,3 @@
 [submodule "3rdparty/Megatron-LM"]
 	path = 3rdparty/Megatron-LM
-	url = https://github.com/NVIDIA/Megatron-LM.git
+	url = https://github.com/basetenlabs/Megatron-LM.git

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

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import math
+import os
 from collections.abc import Callable, Mapping
 
 import torch
@@ -80,14 +81,15 @@ def dequantize_fp8_blockwise(
     """
     M, N = weight.shape
     w = weight.float()
-    out = torch.empty_like(w)
-    sM, sN = scale_inv.shape
-    for bi in range(sM):
-        for bj in range(sN):
-            r0, r1 = bi * block_size, min((bi + 1) * block_size, M)
-            c0, c1 = bj * block_size, min((bj + 1) * block_size, N)
-            out[r0:r1, c0:c1] = w[r0:r1, c0:c1] * scale_inv[bi, bj]
-    return out.to(dtype)
+    # Vectorized block expansion: the per-(bi, bj) Python loop costs ~1k
+    # iterations per tensor (~45M across an 800B checkpoint) and made the
+    # FP8 load CPU-bound for tens of minutes. Expanding scale_inv with two
+    # repeat_interleaves and multiplying once is numerically identical
+    # (same float32 elementwise product, same scales).
+    scales = scale_inv.to(device=w.device, dtype=torch.float32)
+    scales = scales.repeat_interleave(block_size, dim=0)[:M]
+    scales = scales.repeat_interleave(block_size, dim=1)[:, :N]
+    return (w * scales).to(dtype)
 
 
 def maybe_dequantize_fp8_blockwise(
@@ -326,6 +328,14 @@ def dequantize_mxfp4_e2m1_packed(
     tensors can be passed directly; ``.to(torch.float32)`` materializes their
     power-of-two values.
     """
+    # DSV4_GPU_DEQUANT: run the unpack/dequant on GPU rather than the (much slower)
+    # CPU int64 path. Move the packed weight + scale to CUDA first; the dequantized
+    # result is copied to GPU right after load anyway, and processing one parameter
+    # at a time keeps the extra GPU memory bounded. Unset -> unchanged CPU path.
+    if os.environ.get("DSV4_GPU_DEQUANT") == "1" and torch.cuda.is_available() \
+            and weight_packed.device.type == "cpu":
+        weight_packed = weight_packed.cuda(non_blocking=True)
+        scale = scale.cuda(non_blocking=True)
     w_u8 = weight_packed.view(torch.uint8)
     lo = (w_u8 & 0xF).to(torch.int64)
     hi = (w_u8 >> 4).to(torch.int64)

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

@@ -14,6 +14,9 @@
 
 """Compatibility utilities for HuggingFace transformers 5.0+ configs."""
 
+from pathlib import Path
+
+import transformers.dynamic_module_utils as _hf_dyn
 import transformers.utils.import_utils as _hf_import_utils
 
 
@@ -24,6 +27,27 @@
     _hf_import_utils.is_torch_fx_available = lambda: True
 
 
+# transformers' get_cached_module_file() copies a remote-code module plus only
+# its *direct* (1-level) relative imports into the transformers_modules cache,
+# while get_class_in_module() resolves relative imports *recursively*. Remote
+# code with a >1-level relative-import chain (e.g. Kimi-K2's
+# modeling_kimi_k25 -> modeling_deepseek -> configuration_deepseek) therefore
+# leaves the transitive files uncopied and loading fails with a FileNotFoundError
+# for the deepest dependency. Make check_imports() return the full transitive set
+# so every relatively-imported sibling module is materialized into the cache.
+if not getattr(_hf_dyn, "_bridge_recursive_check_imports", False):
+    _orig_check_imports = _hf_dyn.check_imports
+
+    def _recursive_check_imports(filename):
+        # Preserve the original missing-third-party-package validation, then
+        # widen the copy set from direct to transitive relative imports.
+        _orig_check_imports(filename)
+        return [Path(f).stem for f in _hf_dyn.get_relative_import_files(filename)]
+
+    _hf_dyn.check_imports = _recursive_check_imports
+    _hf_dyn._bridge_recursive_check_imports = True
+
+
 def rope_theta_from_hf(config) -> float:
     """Extract rope_theta from a HuggingFace config.
 

src/megatron/bridge/models/glm_moe_dsa/glm5_bridge.py

@@ -12,8 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import json
 import logging
+import os
 
+from huggingface_hub import hf_hub_download
 from megatron.core.models.gpt.gpt_model import GPTModel
 from transformers import GlmMoeDsaForCausalLM
 
@@ -24,13 +27,41 @@
     GatedMLPMapping,
     QKVMapping,
 )
+from megatron.bridge.models.conversion.quantization_utils import maybe_dequantize_fp8_blockwise
 from megatron.bridge.models.hf_pretrained.causal_lm import PreTrainedCausalLM
 from megatron.bridge.models.mla_provider import MLAModelProvider
 
 
 logger = logging.getLogger(__name__)
 
 
+def _load_raw_hf_config(name_or_path: str) -> dict | None:
+    """Return the raw config.json for a local snapshot dir or a hub repo id.
+
+    Bypasses transformers config parsing on purpose: GlmMoeDsaConfig mangles
+    the qk head-dim split (see provider_bridge), so callers need the on-disk
+    values. For a repo id the file resolves through the hub cache —
+    transformers already fetched config.json to build the parsed config, so
+    this works offline (HF_HUB_OFFLINE) too.
+    """
+    local_path = os.path.join(name_or_path, "config.json")
+    if os.path.isfile(local_path):
+        with open(local_path) as f:
+            return json.load(f)
+    try:
+        resolved = hf_hub_download(repo_id=name_or_path, filename="config.json")
+    except Exception as exc:
+        logger.warning(
+            "Could not resolve raw config.json for %r locally or from the hub "
+            "cache: %s",
+            name_or_path,
+            exc,
+        )
+        return None
+    with open(resolved) as f:
+        return json.load(f)
+
+
 @MegatronModelBridge.register_bridge(
     source=GlmMoeDsaForCausalLM, target=GPTModel, provider=MLAModelProvider, model_type="glm_moe_dsa"
 )
@@ -74,6 +105,24 @@ def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> MLAModelProvider
         provider.qk_layernorm = True
         provider.multi_latent_attention = True
 
+        # Work around a transformers GlmMoeDsaConfig bug that collapses qk_rope_head_dim onto
+        # head_dim (e.g. it reports 192 instead of 64 for GLM-5.2), which corrupts every MLA
+        # shape derived from qk_pos_emb_head_dim (kv_a_proj, RoPE, etc.). The on-disk
+        # config.json carries the correct split dims, so read them directly — from the local
+        # directory when base_model is a snapshot path, otherwise via the hub cache, so
+        # hub-id launches (prod) get the fix too. (When qk_nope == qk_rope, as in the tiny
+        # debug model, this is a no-op.)
+        raw_config = _load_raw_hf_config(getattr(hf_config, "_name_or_path", ""))
+        if raw_config is not None:
+            provider.qk_head_dim = raw_config["qk_nope_head_dim"]
+            provider.qk_pos_emb_head_dim = raw_config["qk_rope_head_dim"]
+        else:
+            logger.warning(
+                "Skipping the GLM-5 qk head-dim workaround (raw config.json "
+                "unavailable). If qk_nope_head_dim != qk_rope_head_dim (as in "
+                "GLM-5.2), weight load will fail with a kv_a_proj shape mismatch."
+            )
+
         # Disable MTP (Multi-Token Prediction) — HF config has num_nextn_predict_layers=1
         # but Bridge does not yet have MTP weight mappings for GLM-5.
         provider.mtp_num_layers = None
@@ -111,9 +160,23 @@ def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> MLAModelProvider
         provider.dsa_indexer_head_dim = hf_config.index_head_dim
         provider.dsa_indexer_n_heads = hf_config.index_n_heads
         provider.dsa_indexer_topk = hf_config.index_topk
+        # GLM-5.2's indexer applies interleaved (GPT-J) RoPE; the trainer must
+        # match serving. Default-False-when-absent mirrors vLLM's read exactly
+        # (``is_neox_style = not getattr(config, "indexer_rope_interleave", False)``).
+        provider.dsa_indexer_rope_interleave = getattr(
+            hf_config, "indexer_rope_interleave", False
+        )
         provider.dsa_indexer_loss_coeff = 0.001
         provider.dsa_indexer_use_sparse_loss = True
 
+        # GLM-5.2 cross-layer top-k sharing (IndexShare). Architecture-level: read from
+        # the HF config. Defaults (freq=1, offset=first_k_dense_replace) reduce to the
+        # GLM-5 per-layer-indexer behaviour, so this is backward compatible.
+        provider.dsa_indexer_topk_freq = getattr(hf_config, "index_topk_freq", 1)
+        provider.dsa_indexer_skip_topk_offset = getattr(
+            hf_config, "index_skip_topk_offset", hf_config.first_k_dense_replace
+        )
+
         return provider
 
     def mapping_registry(self) -> MegatronMappingRegistry:
@@ -206,3 +269,28 @@ def mapping_registry(self) -> MegatronMappingRegistry:
         )
 
         return MegatronMappingRegistry(*mapping_list)
+
+    def maybe_modify_loaded_hf_weight(self, hf_param, hf_state_dict):
+        """Dequantize block-wise FP8 (GLM-5.2-FP8) HF weights on load.
+
+        GLM-5.2-FP8 stores linear weights as float8_e4m3fn with a companion
+        ``<param>_scale_inv`` tensor per 128x128 block (DeepSeek-style; the HF
+        config carries ``weight_block_size=[128, 128]``). Layers listed in the
+        checkpoint's ``modules_to_not_convert`` have no scale and pass through
+        unchanged. This lets the bf16 (``zai-org/GLM-5.2``) and FP8
+        (``zai-org/GLM-5.2-FP8``) checkpoints both convert through this bridge,
+        so a single Loops config can train + sample on the FP8 id.
+        """
+        hf_weights = super().maybe_modify_loaded_hf_weight(hf_param, hf_state_dict)
+        if isinstance(hf_weights, dict):
+            return {
+                key: self._maybe_dequant_fp8(tensor, hf_param[key], hf_state_dict)
+                for key, tensor in hf_weights.items()
+            }
+        return self._maybe_dequant_fp8(hf_weights, hf_param, hf_state_dict)
+
+    @staticmethod
+    def _maybe_dequant_fp8(weight, param_name, hf_state_dict):
+        """Block-wise dequant ``weight`` if FP8, using ``<param_name>_scale_inv``."""
+        scale_inv = hf_state_dict.get(param_name + "_scale_inv")
+        return maybe_dequantize_fp8_blockwise(weight, scale_inv)

src/megatron/bridge/models/qwen3_asr/hf_qwen3_asr/modeling_qwen3_asr.py

@@ -993,6 +993,8 @@ def forward(
         cache_position: Optional[torch.LongTensor] = None,
         **kwargs: Unpack[FlashAttentionKwargs],
     ) -> Union[tuple, BaseModelOutputWithPast]:
+        # TODO: Remove this local docstring workaround once this branch includes
+        # NVIDIA-NeMo/Megatron-Bridge@0a21da47 or a later upstream fix.
         r"""
         cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
             Indices depicting the position of the input sequence tokens in the sequence.
@@ -1175,6 +1177,8 @@ def forward(
         cache_position=None,
         **kwargs,
     ) -> Union[tuple, Qwen3ASRThinkerCausalLMOutputWithPast]:
+        # TODO: Remove this local docstring workaround once this branch includes
+        # NVIDIA-NeMo/Megatron-Bridge@0a21da47 or a later upstream fix.
         r"""
         cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
             Indices depicting the position of the input sequence tokens in the sequence.

src/megatron/bridge/peft/lora_layers.py

@@ -59,7 +59,7 @@ def forward(self, x: torch.Tensor, *args: Any, **kwargs: Any) -> Tuple[torch.Ten
             return linear_output, bias
         adapter_output = self.adapter_forward(self.adapter, layernorm_output.contiguous(), *args, **kwargs)
         adapter_output = adapter_output.reshape(linear_output.shape)
-        return linear_output + adapter_output, bias
+        return adapter_output.add_(linear_output), bias  # in-place: avoid a 6.25 GiB sum alloc at 131k
 
 
 class LoRATopKRouter(AdapterWrapper):

src/megatron/bridge/peft/utils.py

@@ -1138,7 +1138,15 @@ def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor:
         if self.dropout_position == "post":
             x = self.dropout(x)
 
-        x = x * (self.alpha / self.dim)
+        if torch.is_grad_enabled():
+            # Recompute/backward (or no-recompute forward): the linear_out
+            # gather-region output is a view; in-place mul_ is forbidden by
+            # autograd. Go out-of-place (costs a delta-sized alloc).
+            x = x * (self.alpha / self.dim)
+        else:
+            # no_grad checkpointed forward: in-place is safe and avoids a
+            # delta-sized (6-8 GiB) alloc at 131k.
+            x = x.mul_(self.alpha / self.dim)
 
         if pad_len > 0:
             # Remove MoE padding.

tests/unit_tests/models/glm/test_glm5_bridge_raw_config.py

@@ -0,0 +1,58 @@
+# Copyright (c) 2025, 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.
+
+"""
+Unit tests for the GLM-5 raw-config resolution used by the qk head-dim
+workaround (transformers GlmMoeDsaConfig collapses qk_rope_head_dim, so
+GLM5Bridge re-reads the on-disk config.json).
+"""
+
+import json
+
+from megatron.bridge.models.glm_moe_dsa.glm5_bridge import _load_raw_hf_config
+
+
+def test_local_snapshot_dir_reads_config_directly(tmp_path):
+    dims = {"qk_nope_head_dim": 128, "qk_rope_head_dim": 64}
+    (tmp_path / "config.json").write_text(json.dumps(dims))
+    assert _load_raw_hf_config(str(tmp_path)) == dims
+
+
+def test_hub_id_resolves_through_hub_cache(tmp_path, monkeypatch):
+    cached = tmp_path / "config.json"
+    cached.write_text(json.dumps({"qk_rope_head_dim": 64}))
+    calls = {}
+
+    def fake_download(repo_id, filename):
+        calls["repo_id"] = repo_id
+        calls["filename"] = filename
+        return str(cached)
+
+    monkeypatch.setattr(
+        "megatron.bridge.models.glm_moe_dsa.glm5_bridge.hf_hub_download",
+        fake_download,
+    )
+    assert _load_raw_hf_config("zai-org/GLM-5.2-FP8") == {"qk_rope_head_dim": 64}
+    assert calls == {"repo_id": "zai-org/GLM-5.2-FP8", "filename": "config.json"}
+
+
+def test_unresolvable_name_returns_none(monkeypatch):
+    def fake_download(repo_id, filename):
+        raise OSError("offline and not in the hub cache")
+
+    monkeypatch.setattr(
+        "megatron.bridge.models.glm_moe_dsa.glm5_bridge.hf_hub_download",
+        fake_download,
+    )
+    assert _load_raw_hf_config("zai-org/GLM-5.2-FP8") is None