fix(glm5): support FP8 context-parallel launches

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

@@ -80,14 +80,10 @@ 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)
+    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(

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

@@ -16,6 +16,7 @@
 import logging
 import os
 
+from huggingface_hub import hf_hub_download
 from megatron.core.models.gpt.gpt_model import GPTModel
 from transformers import GlmMoeDsaForCausalLM
 
@@ -34,6 +35,22 @@
 logger = logging.getLogger(__name__)
 
 
+def _load_raw_hf_config(name_or_path: str) -> dict:
+    """Load raw config.json for a local snapshot or hub repo."""
+    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:
+        raise RuntimeError(
+            f"GLM-5 requires raw config.json for {name_or_path!r} to preserve qk head dimensions."
+        ) from exc
+    with open(resolved) as f:
+        return json.load(f)
+
+
 @MegatronModelBridge.register_bridge(
     source=GlmMoeDsaForCausalLM, target=GPTModel, provider=MLAModelProvider, model_type="glm_moe_dsa"
 )
@@ -79,12 +96,9 @@ def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> MLAModelProvider
 
         # Work around transformers configs that collapse qk_rope_head_dim onto
         # head_dim for GLM-5.2. The on-disk config carries the real MLA split.
-        raw_config_path = os.path.join(getattr(hf_config, "_name_or_path", ""), "config.json")
-        if os.path.isfile(raw_config_path):
-            with open(raw_config_path) as raw_config_file:
-                raw_config = json.load(raw_config_file)
-            provider.qk_head_dim = raw_config["qk_nope_head_dim"]
-            provider.qk_pos_emb_head_dim = raw_config["qk_rope_head_dim"]
+        raw_config = _load_raw_hf_config(getattr(hf_config, "_name_or_path", ""))
+        provider.qk_head_dim = raw_config["qk_nope_head_dim"]
+        provider.qk_pos_emb_head_dim = raw_config["qk_rope_head_dim"]
 
         # Disable MTP (Multi-Token Prediction) by default
         # HF config has num_nextn_predict_layers=1

tests/unit_tests/models/glm/test_glm5_bridge_raw_config.py

@@ -0,0 +1,55 @@
+# 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.
+
+import json
+
+import pytest
+
+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_raises(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,
+    )
+    with pytest.raises(RuntimeError, match="GLM-5 requires raw config"):
+        _load_raw_hf_config("zai-org/GLM-5.2-FP8")