From f8016ce6ab4c4902e9f598b332e9eeab462bfe49 Mon Sep 17 00:00:00 2001 From: Jack Rao Date: Wed, 8 Jul 2026 00:41:02 -0700 Subject: [PATCH 1/7] perf(ckpt): vectorize FP8 blockwise dequant (bit-exact, ~17x/tensor) The per-block Python loop made GLM-5.2-FP8's 800B load CPU-bound (8 workers pegged ~50 min). Two repeat_interleaves + one multiply, verified bit-exact against the loop. Co-Authored-By: Claude Fable 5 --- .../models/conversion/quantization_utils.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/src/megatron/bridge/models/conversion/quantization_utils.py b/src/megatron/bridge/models/conversion/quantization_utils.py index c9488467d..cd347037a 100644 --- a/src/megatron/bridge/models/conversion/quantization_utils.py +++ b/src/megatron/bridge/models/conversion/quantization_utils.py @@ -80,14 +80,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( From d48814a9a3f7750f19edb2b03bb93e8c54a0bf78 Mon Sep 17 00:00:00 2001 From: Jack Rao Date: Wed, 8 Jul 2026 19:14:01 -0700 Subject: [PATCH 2/7] fix(glm5): resolve raw config.json via the hub cache for HF-id launches MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The qk_rope_head_dim workaround (transformers GlmMoeDsaConfig collapses the qk head-dim split, corrupting kv_a_proj shapes: 704 != 576) only read /config.json as a filesystem path, so it silently no-opped when base_model is a hub repo id and the 800B weight load failed. Resolve the raw config through hf_hub_download when the local read misses — transformers has already cached config.json by the time the bridge runs, so this works offline (HF_HUB_OFFLINE) too. Warn loudly when neither path resolves. Prod launches GLM-5.2-FP8 by HF id, so this unblocks the registry row in trainers#592. Co-Authored-By: Claude Fable 5 --- .../bridge/models/glm_moe_dsa/glm5_bridge.py | 49 ++++++++++++++-- .../models/glm/test_glm5_bridge_raw_config.py | 58 +++++++++++++++++++ 2 files changed, 101 insertions(+), 6 deletions(-) create mode 100644 tests/unit_tests/models/glm/test_glm5_bridge_raw_config.py diff --git a/src/megatron/bridge/models/glm_moe_dsa/glm5_bridge.py b/src/megatron/bridge/models/glm_moe_dsa/glm5_bridge.py index 0b731cd33..7307aab41 100644 --- a/src/megatron/bridge/models/glm_moe_dsa/glm5_bridge.py +++ b/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,33 @@ 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" ) @@ -77,14 +105,23 @@ def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> MLAModelProvider provider.qk_layernorm = True provider.multi_latent_attention = True - # 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) + # 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) by default # HF config has num_nextn_predict_layers=1 diff --git a/tests/unit_tests/models/glm/test_glm5_bridge_raw_config.py b/tests/unit_tests/models/glm/test_glm5_bridge_raw_config.py new file mode 100644 index 000000000..f6e201abe --- /dev/null +++ b/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 From 8e6e3680bec0905fe5f8faee4d73e82dab2bb78f Mon Sep 17 00:00:00 2001 From: Jack Rao Date: Sat, 11 Jul 2026 17:45:41 -0700 Subject: [PATCH 3/7] build: pin Megatron-LM to the packed-CP q_causal_offsets fix (LM#14 head) --- 3rdparty/Megatron-LM | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/3rdparty/Megatron-LM b/3rdparty/Megatron-LM index 038760cd8..545ba95cb 160000 --- a/3rdparty/Megatron-LM +++ b/3rdparty/Megatron-LM @@ -1 +1 @@ -Subproject commit 038760cd89e687dafff5b94bcf33e36fb92a5775 +Subproject commit 545ba95cb63c18660a9c1effdd38c962017e9243 From 31ee2e57d91af52d3f6ded303b11b1f2aca45a6e Mon Sep 17 00:00:00 2001 From: Jack Rao Date: Mon, 13 Jul 2026 11:00:17 -0700 Subject: [PATCH 4/7] build: update Megatron-LM PR pin Keep the Bridge submodule aligned with the current LM#14 head before the dependent Bridge and Trainers PRs merge. Signed-off-by: Jack Rao --- 3rdparty/Megatron-LM | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/3rdparty/Megatron-LM b/3rdparty/Megatron-LM index 545ba95cb..ccce860e2 160000 --- a/3rdparty/Megatron-LM +++ b/3rdparty/Megatron-LM @@ -1 +1 @@ -Subproject commit 545ba95cb63c18660a9c1effdd38c962017e9243 +Subproject commit ccce860e2ff83299ca744711280e63e9740c8d7d From ce1c2855842512c4022d1bd4294014849fa46080 Mon Sep 17 00:00:00 2001 From: Jack Rao Date: Mon, 13 Jul 2026 11:00:31 -0700 Subject: [PATCH 5/7] fix(glm5): fail fast without raw config The raw config preserves GLM qk dimensions, so conversion must terminate instead of continuing with invalid MLA shapes. Remove redundant comments. Signed-off-by: Jack Rao Co-authored-by: Cursor --- .../models/conversion/quantization_utils.py | 5 --- .../bridge/models/glm_moe_dsa/glm5_bridge.py | 40 ++++--------------- .../models/glm/test_glm5_bridge_raw_config.py | 13 +++--- 3 files changed, 13 insertions(+), 45 deletions(-) diff --git a/src/megatron/bridge/models/conversion/quantization_utils.py b/src/megatron/bridge/models/conversion/quantization_utils.py index cd347037a..35cf28c4a 100644 --- a/src/megatron/bridge/models/conversion/quantization_utils.py +++ b/src/megatron/bridge/models/conversion/quantization_utils.py @@ -80,11 +80,6 @@ def dequantize_fp8_blockwise( """ M, N = weight.shape w = weight.float() - # 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] diff --git a/src/megatron/bridge/models/glm_moe_dsa/glm5_bridge.py b/src/megatron/bridge/models/glm_moe_dsa/glm5_bridge.py index 7307aab41..3b579e538 100644 --- a/src/megatron/bridge/models/glm_moe_dsa/glm5_bridge.py +++ b/src/megatron/bridge/models/glm_moe_dsa/glm5_bridge.py @@ -35,15 +35,8 @@ 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. - """ +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: @@ -51,13 +44,9 @@ def _load_raw_hf_config(name_or_path: str) -> dict | None: 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 + 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) @@ -105,23 +94,10 @@ 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.) + # transformers loses the GLM qk head-dimension split. 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." - ) + 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 diff --git a/tests/unit_tests/models/glm/test_glm5_bridge_raw_config.py b/tests/unit_tests/models/glm/test_glm5_bridge_raw_config.py index f6e201abe..c84d9348d 100644 --- a/tests/unit_tests/models/glm/test_glm5_bridge_raw_config.py +++ b/tests/unit_tests/models/glm/test_glm5_bridge_raw_config.py @@ -12,14 +12,10 @@ # 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 +import pytest + from megatron.bridge.models.glm_moe_dsa.glm5_bridge import _load_raw_hf_config @@ -47,7 +43,7 @@ def fake_download(repo_id, filename): assert calls == {"repo_id": "zai-org/GLM-5.2-FP8", "filename": "config.json"} -def test_unresolvable_name_returns_none(monkeypatch): +def test_unresolvable_name_raises(monkeypatch): def fake_download(repo_id, filename): raise OSError("offline and not in the hub cache") @@ -55,4 +51,5 @@ def fake_download(repo_id, filename): "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 + with pytest.raises(RuntimeError, match="GLM-5 requires raw config"): + _load_raw_hf_config("zai-org/GLM-5.2-FP8") From ab50d2171dd533f3a8eaa882162f0e405963de90 Mon Sep 17 00:00:00 2001 From: Jack Rao Date: Mon, 13 Jul 2026 11:03:17 -0700 Subject: [PATCH 6/7] build: advance Megatron-LM PR pin Keep the Bridge submodule aligned with the current LM#14 head before the dependent PRs merge. Signed-off-by: Jack Rao --- 3rdparty/Megatron-LM | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/3rdparty/Megatron-LM b/3rdparty/Megatron-LM index ccce860e2..a1fab1bb5 160000 --- a/3rdparty/Megatron-LM +++ b/3rdparty/Megatron-LM @@ -1 +1 @@ -Subproject commit ccce860e2ff83299ca744711280e63e9740c8d7d +Subproject commit a1fab1bb53cdc7cb21a5a77f321c25735f369ef5 From 32d432ccb1e5f366cacd9b5a35d5f1284308aa9a Mon Sep 17 00:00:00 2001 From: Jack Rao Date: Mon, 13 Jul 2026 11:14:31 -0700 Subject: [PATCH 7/7] docs(glm5): clarify raw config workaround State the Transformer configuration defect and why the raw config is required. Signed-off-by: Jack Rao Co-authored-by: Cursor --- src/megatron/bridge/models/glm_moe_dsa/glm5_bridge.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/megatron/bridge/models/glm_moe_dsa/glm5_bridge.py b/src/megatron/bridge/models/glm_moe_dsa/glm5_bridge.py index 3b579e538..3370201f0 100644 --- a/src/megatron/bridge/models/glm_moe_dsa/glm5_bridge.py +++ b/src/megatron/bridge/models/glm_moe_dsa/glm5_bridge.py @@ -94,7 +94,8 @@ def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> MLAModelProvider provider.qk_layernorm = True provider.multi_latent_attention = True - # transformers loses the GLM qk head-dimension split. + # 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 = _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"]