Add GLM-5.3-Flash and GLM-5.2 support to Megatron-Bridge PTQ and HF export - #2539
kevalmorabia97 wants to merge 5 commits into
Conversation
Quantize GLM-5.3-Flash (glm5_next) with examples/megatron_bridge to the recipe behind nvidia/GLM-5.3-Flash-NVFP4 and export it as a unified HF checkpoint: - Recipe: match Megatron dense-MLP names and keep the MTP layer BF16. - Quantization: register DSAttention for FP8 KV-cache quantization. - Export: add the glm5_next mapping (mHC hyper-connections, KDA, NoPE-MLA with the DSA indexer, MTP with split e/h projections). - Exporter fixes: keep expert_bias FP32, exclude the passthrough vision tower, skip the head_dim override for MLA, and pass the generation config through unvalidated. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
|
Important Draft PR not reviewedDraft PRs are not automatically reviewed by default.
To automatically review draft PRs, update your CodeRabbit configuration: reviews:
auto_review:
drafts: trueComment |
|
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #2539 +/- ##
==========================================
- Coverage 68.89% 68.72% -0.18%
==========================================
Files 605 608 +3
Lines 67063 68908 +1845
==========================================
+ Hits 46204 47354 +1150
- Misses 20859 21554 +695
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
- export_quantized_megatron_to_hf.py: add --ep_size. Grouped-GEMM experts are already gathered across the EP group on export; SequentialMLP experts are indexed locally, so export raises for them at EP>1. - Add the GLM-5 / GLM-5.2 (glm_moe_dsa) export mapping. Megatron-Bridge does not build its MTP layer, so it is copied from the source checkpoint (dequantized, excluded from quantization) like the released NVFP4 ones. - weight_attr_names now reports TEGroupedLinear, so an experts-only quantized grouped-GEMM MoE is detected as quantized; its export used to omit hf_quant_config.json and quantization_config. - test_quantize_export: quantize with NVFP4 and export nemotron_h at EP. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
kevalmorabia97
changed the title
At EP>1 every EP rank all-gathered every expert and wrote the same shards, so host memory scaled with EP and ranks raced on the same files. A full-size GLM-5.3-Flash export at EP4 was OOM-killed (~260 GB RSS per rank). Gather experts to EP rank 0 only, and write shards from one TP/DP/EP rank per pipeline stage. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
|
/claude review |
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (gpt-6-astra) — DM the bot to share feedback.
Changes requested: Hub-ID exports silently omit GLM-5 MTP weights, and folded-layer verification mishandles depth-pruned GLM-5.3 exports.
Needs action:
- Fix Hub-ID MTP passthrough in
unified_export_megatron.py, including a mocked Hub regression test; see inline comment. - Fix the export self-check to use HF decoder depth after layer folding; test a pruned GLM-5.3 model.
- Restore FP8 integration coverage in
test_quantize_export.py, or identify equivalent retained coverage and justify its removal. - Add DSAttention regression coverage for conversion, tensor/
Nonevalues, and checkpoint round-trip intest_megatron.py. - Document the circular-import reason for the local
GroupedQuantizerimport incore_utils.py, or move it to module scope.
No action needed:
- The design extends existing DeepSeek/Qwen mappings and shared slicing rather than adding a competing export system; the PR explains the layout and memory constraints.
- Switching the Nemotron integration case to EP is justified by the new export capability; removing FP8 coverage is not explained.
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
Claude review summaryFindings: 0 CRITICAL · 1 IMPORTANT · 1 SUGGESTION
Most impactful findingRegistering What I verified and cleared
Non-blocking note
Risk assessmentModerate. The structural work is careful and well-decomposed, and the new unit tests cover the genuinely tricky mechanics: KDA fused-QKV/conv1d splitting, mHC tensor naming, decoder-resident MTP copying, and grouped-experts format detection. The documented backward-compatibility deltas — MLA |
There was a problem hiding this comment.
Claude review found 1 IMPORTANT issue (0 CRITICAL), so this is a comment review rather than an approval.
IMPORTANT (Export) — modelopt/torch/quantization/plugins/megatron.py:1055-1057: registering DSAttention makes KV-cache quantization reachable on DSA layers for the first time, but the absorbed-MLA path passes value=None, so v_bmm_quantizer is never invoked and its amax stays None. _self_attention_scaling then gates both writes behind all(s is not None for s in kv_scales) and emits neither k_scale nor v_scale, while still setting self.kv_cache_dtype from get_kv_cache_dtype(module) — producing a checkpoint that advertises kv_cache_quant_algo: FP8 with no scale tensors on those layers. The GLM recipe here is shielded only because kv_fp8_cast sets use_constant_amax: true; the plain, calibrated kv_fp8 unit is not, and the example test in this PR now uses general/ptq/nvfp4_default-kv_fp8. Suggested fixes (mirror K statistics onto V for the absorbed path, or emit scales independently and fail loudly on a missing one) are in the inline thread.
SUGGESTION — stale "only Pipeline parallelism is supported" comment and module docstring in examples/megatron_bridge/export_quantized_megatron_to_hf.py now that --ep_size exists. Non-blocking.
The rest of the change held up under review: the boolean rule flags survive with_language_model_prefix and _populate_rule_book, the folded layer-pair index arithmetic and MTP source/destination indices are correct, the new EP writer gate drops nothing (vision-tower and MTP merges both land on ep0/dp0 ranks and save_safetensors_by_layer_index tolerates an empty dict), the weight_attr_names GroupedQuantizer branch is safe for every consumer, the _QuantCoreAttention rename preserves registry keys so modelopt_state still restores, and _split_fused_projection is behavior-preserving for GatedDeltaNet. A full breakdown is in the summary comment.
…heck - DSAttention (absorbed MLA) passes value=None: calibrate V on the shared KV latent so a calibrated FP8 KV cache also exports a v_scale. - Copy GLM-5 decoder-layer MTP from a Hub-ID source by downloading only the shards that hold it. - Export self-check: compare against the HF decoder depth (folded models have 2x physical layers) and check the MTP layer at its remapped index. - test_quantize_export keeps an FP8 case; add DSAttention KV-cache, depth-pruned self-check and Hub MTP tests; document the local import; refresh the export script's parallelism notes. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
…op HAS_TE torch's state_dict()/load_state_dict() route ``_extra_state`` only through classes that override get/set_extra_state. TEDotProductAttention does, DSAttention does not, so its quantizer state (including the KV amax) was never saved and a restored model ran the KV cache with dynamic scales. Register DSAttention through a subclass that routes its extra state to the quantizer state handlers. The DSA KV test now round-trips a torch-dist checkpoint. Transformer Engine is always present in Megatron environments, so import it unconditionally like the NAS / prune plugins, and drop the HAS_TE guard. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
|
/claude review |
|
|
||
| if HAS_DSA: | ||
|
|
||
| @QuantModuleRegistry.register({DSAttention: "megatron_DSAttention"}) | ||
| class _QuantDSAttention(_QuantCoreAttention): | ||
| """DSAttention with KV-cache quantization. | ||
|
|
||
| torch's state_dict() / load_state_dict() route ``_extra_state`` only through classes that | ||
| override these; TEDotProductAttention does, DSAttention does not, so without them the | ||
| quantizer state (including amax) was dropped from checkpoints. | ||
| """ | ||
|
|
||
| def get_extra_state(self): | ||
| return quant_module_get_extra_state(self) | ||
|
|
||
| def set_extra_state(self, state): | ||
| quant_module_set_extra_state(self, state) |
There was a problem hiding this comment.
[SUGGESTION] This fixes the symptom per-class; the cause is in register_modelopt_extra_state_callbacks, and the same silent drop stays latent for every other registered Megatron module whose base class lacks these methods.
register_modelopt_extra_state_callbacks (modelopt/torch/opt/plugins/megatron.py:136-141) binds the hooks as instance attributes:
module.get_extra_state = types.MethodType(_modelopt_get_extra_state, module)but PyTorch gates extra-state entirely on a class-level check (getattr(self.__class__, "get_extra_state", Module.get_extra_state) is not Module.get_extra_state) before it ever looks the attribute up on the instance. So the instance binding only has an effect when the wrapped class already overrides get_extra_state/set_extra_state — TE's DotProductAttention does, DSAttention doesn't, which is exactly the bug this block works around.
Two consequences worth noting:
- Because the instance attribute shadows the class method (
megatron_replace_quant_module_hookbinds it on the pre-conversion module, andDynamicModuleconversion preserves__dict__),_QuantDSAttention.get_extra_stateis essentially never called — it exists only to flip torch's class check toTrue. That's non-obvious enough that a future reader may "clean up" the apparently-redundant override and silently reintroduce the amax loss. Worth saying so in the docstring if the per-class shape is kept. - A repo-wide grep finds no other class-level
get_extra_state/set_extra_stateinmodelopt/torch/, so any futureQuantModuleRegistry.registeron a Megatron module whose base lacks them (e.g. native non-TEColumnParallelLinear/RowParallelLinearin a build without TE) will hit the same silent quantizer-state drop — with no error, just missing amax after a checkpoint round-trip.
Suggested fix, in the spirit of CONTRIBUTING's "fix the bug cause, not the side effect": have the helper install the hooks on the type instead of the instance, so the class check passes for every registered module without per-class boilerplate. Since conversion already gives each converted module a fresh dynamic class, something like the following in register_modelopt_extra_state_callbacks is safe:
cls = type(module)
if cls.get_extra_state is torch.nn.Module.get_extra_state:
cls.get_extra_state = _modelopt_get_extra_state(_modelopt_get_extra_state's zero-arg super() relies on __class__ from its defining module, so if it moves onto a type keep it defined where it is and reference it, rather than re-defining it inside the class body.) If that turns out to be too invasive for this PR, an assert/warning in the helper when the class check would fail would at least make the next occurrence loud instead of silent.
| query = self.q_bmm_quantizer(query) | ||
| if value is None: | ||
| # Absorbed MLA (DSAttention) passes value=None: the key is the KV latent that both K | ||
| # and V are read from, so calibrate V on it too (output unused) to export a V scale. | ||
| self.v_bmm_quantizer(key) | ||
| else: |
There was a problem hiding this comment.
[SUGGESTION] The discarded self.v_bmm_quantizer(key) runs on every forward, not just during calibration.
Calibrating V on the KV latent is the right call for absorbed MLA, and it's what unblocks _self_attention_scaling's all(s is not None for s in kv_scales) gate so both k_scale and v_scale get written. But once amax is collected, this call still does a full fake quant-dequant of the latent tensor whose result is thrown away — per DSA layer, per forward, for the rest of the model's life (QAT steps, eval, any post-calibration forward). With fake_quant enabled that's an extra quantize+dequantize over a [seq, batch, kv_lora_rank] tensor per layer.
Two ways to avoid it:
- Gate on the stats-collection pass, e.g. only call it while the quantizer is collecting (
getattr(self.v_bmm_quantizer, "_if_calib", False), the same predicatetransformer_engine.py:54uses), or - Skip the extra forward entirely and copy K's amax into V once at the end of calibration (in
modelopt_post_restoreor a small calibration hook) — they are provably identical here, andtests/gpu_megatron/torch/quantization/plugins/test_megatron.py'storch.equal(v_bmm_quantizer.amax, k_bmm_quantizer.amax)assertion would still hold.
The second is cheaper but moves logic away from the point where the invariant is visible; the first is a one-line change and keeps the explanation next to the code. Either is fine — the current form is correct, just pays a permanent cost for a calibration-only need.
|
Claude review — re-review at Counts: 0 CRITICAL · 0 IMPORTANT · 2 SUGGESTION Full review (bare Both findings from my previous round are resolved
What I verified this round
Risk assessment Low-to-moderate, concentrated in the multi-rank export paths. The single-rank logic is well covered by the six new/extended tests, and the PR body's bit-exactness check of all 197 unquantized tensors against the released checkpoint is strong evidence for the GLM-5.3 name mappings. The residual risk is the ~114 uncovered patch lines codecov reports, which sit almost entirely in the PP/EP fan-out ( Neither SUGGESTION blocks merge. |
…y PR (#2547) ## Summary `partial-install (torch)` has been failing on **every** PR since 2026-09-24 — including PRs whose branches predate the breakage — and because it is a *collection* error rather than a test failure, it aborts the entire run: ``` ImportError while importing test module '.../tests/unit/torch/speculative/plugins/test_hf_streaming_dataset.py' E ModuleNotFoundError: No module named 'httpx' collected 2243 items / 1 error / 45 skipped !!!!!!!!!!!!!!!!!!!! Interrupted: 1 error during collection !!!!!!!!!!!!!!!!!!!! ``` `unit-pr-required-check` aggregates it, so nothing currently merges on a fresh run. ## What happened **No code changed.** `modelopt/torch/speculative/plugins/hf_streaming_dataset.py` has imported `httpx` at module scope since #1509 (2026-06-02), and `httpx` has never appeared in `pyproject.toml`. It arrived only transitively: `dev-test` → `timm` → `huggingface_hub` → `httpx`. **huggingface_hub 2.0.0**, published **2026-09-24T12:01:21Z**, replaced `httpx<1,>=0.23.0` with the separate **`httpx2<3,>=2.0.0`** distribution. Different package name, so `httpx` stopped being installed and the chain disappeared. The boundary is exact — every run *created* before that timestamp passes, every one after fails: | PR | run created | result | |---|---|---| | #2536 / #2535 | 09-23 22:02 | pass | | #2500 | 09-23 23:45 | pass — **merged 09-24 20:01 on this stale-green result** | | *hub 1.33.0 (still requires httpx)* | *09-24 09:49* | | | **hub 2.0.0 published** | **09-24 12:01** | ← | | #2539 | 09-24 16:57 | fail | | #2544 | 09-24 18:32 | fail | | #2216 | 09-25 11:58 | fail | #2500 merging afterwards is not a counterexample: GitHub does not re-run checks at merge time, so it merged on a result from ~20 hours earlier. That is also why this went unnoticed. ## The changes ### 1. Declare `httpx` in the `hf` extra `httpx` is not incidental to streaming — it is the only transport: - every fetch is HTTP: `POST /v1/completions` to the vLLM serve plus `GET /meta` and `/desc` against the connector's sidecar, all through `httpx.Client`; - there is no non-HTTP path — the base `StreamingDataset._fetch` is an abstract seam and `EagleVllmStreamingDataset._fetch` is its only implementation; - no other HTTP library appears in the module (`requests` / `urllib` / `aiohttp`: zero hits, and `requests` is not declared either); - even the retry predicate is built from it: `_TRANSIENT_FETCH_ERRORS = (httpx.HTTPError, OSError)`. It belongs in `hf` rather than in the core `dependencies`: the same module needs `transformers.trainer_pt_utils` at module scope, so one extra already gates the whole file, and a core install has no use for an HTTP client. The bound matches the 0.x API the code uses — `httpx` has no 1.0 release, and 2.x is a different distribution. This is the part that stops it recurring. `[hf]` currently gets `httpx` only because `datasets` happens to require it — the same accident with a different supplier, one release away from repeating. ### 2. Acquire `httpx` in the test through the existing skip guard The test file already intends to skip where the extra is absent — it has `pytest.importorskip("transformers")` and a comment explaining why, and `transformers` is absent in this job too. It broke only because `import httpx` sat **five lines above** that guard, where a missing module ends collection instead of skipping one file. ## Verification - With everything installed: **18 passed**, no behaviour change. - The import-order property is checked with an AST walk over the module's top-level statements: no `hf`-extra-only import precedes the first `importorskip` (which is now line 41). - A faithful local reproduction was attempted and abandoned honestly: hiding `httpx` locally also breaks `huggingface_hub` 1.28, which `modelopt.torch.opt.plugins.huggingface` imports, so the local failure is not the CI one. CI is the oracle for that half — this PR's own `partial-install (torch)` run is the check that matters. ## Scope Two files, five lines of declaration and four of test import order. Deliberately not folded into any feature PR: it blocks the whole repo, and burying a repo-wide fix inside unrelated work is how these stay invisible. 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Chores** * Optional Hugging Face installations now include `httpx`, supporting features that require HTTP communication without requiring it for all installations. * **Tests** * Hugging Face streaming dataset tests now skip when `httpx` is unavailable, allowing the remaining test suite to be collected and run without it. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This branch has not been deployed
2 participants
What does this PR do?
Type of change: new feature
Adds GLM-5.3-Flash (
glm5_next) support toexamples/megatron_bridge, covering both PTQ (quantize.py) and unified HF export (export_quantized_megatron_to_hf.py). It uses the recipe behind the publishednvidia/GLM-5.3-Flash-NVFP4, which was produced withhf_ptq.py, so the Megatron path now produces a checkpoint with the same quantization layout.GLM-5.3-Flash differs from GLM-5 /
glm_moe_dsa:Quantization
models/zai-org/GLM-5.3-Flash/ptq/nvfp4_experts_dense_mlp-kv_fp8_cast): the dense-MLP patterns only matched HF names, so on Megatron the layer 0-2 MLPs stayed BF16. Added themlp.linear_fc1/fc2patterns and a*mtp*disable (Megatron-Bridge builds the MTP layer, HF does not). Neither matches anything in the HF model, sohf_ptq.pyis unchanged.k/v_bmmquantizers toTEDotProductAttention, so on DSA layerskv_fp8_castdid nothing.DSAttentionnow registers with the same quant module, which toleratesvalue=None(the absorbed-MLA path). This also covers other DSA models on Megatron.Export (new
plugins/mcore_glm.pyplus exporter support)hc_{attn,ffn}_{fn,base,scale}.in_projandconv1dare split into HF's tensors. The split code is now shared with Qwen3.5's GatedDeltaNet.eh_projrebuilt from Megatron's splite_proj/h_proj.quantize.pyuse grouped-GEMM experts for this model.GLM-5 / GLM-5.2 export (
GlmMoeDsaForCausalLM, supported by the stock Megatron-Bridge in the 26.10 container)mtp_num_layers = None), so the exporter copies it from the source checkpoint. The copy is dequantized from FP8 and excluded asmodel.layers.<N>.*, the same waynvidia/GLM-5.2-NVFP4ships it in BF16. It also handles depth-pruned exports.Expert-parallel export:
export_quantized_megatron_to_hf.py --ep_size N. Grouped-GEMM experts are gathered across the EP group on export, so the whole MoE no longer has to fit on one GPU. The experts are gathered only to EP rank 0, and one TP/DP/EP rank per pipeline stage writes the shards. Previously every EP rank held and wrote every expert, so host memory scaled with EP, ranks raced on the same files, and a full-size GLM-5.3-Flash export at EP4 was OOM-killed. For GLM-5.3-Flash this is the only way to shard the export, because mHC rules out pipeline parallelism. SequentialMLP experts (--no_moe_grouped_gemm) are numbered by local position, so export now raises for them at EP>1 instead of writing colliding expert IDs.Exporter fixes that also affect other architectures
nvfp4_experts_only-*) wrote packed NVFP4 experts but nohf_quant_config.jsonand noquantization_config, so they would be served as unquantized weights. The cause:weight_attr_names()returned nothing forTEGroupedLinear(its weights areweight0..Nbehind oneGroupedQuantizer), which was already a documented known gap. It now reportsweight.expert_biasis exported as FP32 instead of the export dtype. Megatron holds it as an FP32 buffer, and HF checkpoints (DeepSeek, GLM) shipe_score_correction_biasas F32; it decides expert routing.exclude_modules(model.visual*), ashf_ptq.pyalready does, so deployments don't treat it as quantized.head_dimis no longer overwritten withkv_channelsfor MLA models. For MLA,kv_channelsis the V head dim, not HF'shead_dim.generation_config.jsonis written without re-validation. Transformers 5.17 refuses to re-save GLM-5.3-Flash's shipped config (top_pwithoutdo_sample).Usage
torchrun --nproc_per_node 8 examples/megatron_bridge/quantize.py \ --hf_model_name_or_path zai-org/GLM-5.3-Flash \ --recipe models/zai-org/GLM-5.3-Flash/ptq/nvfp4_experts_dense_mlp-kv_fp8_cast \ --ep_size 8 --skip_generate --export_megatron_path /tmp/GLM-5.3-Flash-NVFP4-megatron # Keep --ep_size <= GPUs per node (see Testing). mHC does not support pipeline parallelism in Megatron-Core yet, so shard the export by EP instead. torchrun --nproc_per_node 8 examples/megatron_bridge/export_quantized_megatron_to_hf.py \ --hf_model_name_or_path zai-org/GLM-5.3-Flash \ --megatron_path /tmp/GLM-5.3-Flash-NVFP4-megatron \ --export_unified_hf_path /tmp/GLM-5.3-Flash-NVFP4 --ep_size 8Testing
Environment: the NeMo
26.10.rc2container. Its Megatron-LM and Megatron-Bridge don't support GLM-5.3-Flash yet, so these unmerged PRs were cherry-picked on top. They apply cleanly to rc2's checkouts.e6f7cf28. That commit fixes the MTPfinal_layernormto load fromlayers.{N}.shared_head.norm.weightinstead of the mainnorm; the MTP layer didn't round-trip before it.transformers==5.17.0(first release withglm5_next) andtokenizers>=0.23.1,<0.24.End-to-end runs on a tiny GLM-5.3-Flash: 4 HF layers (3 KDA + 1 DSA) built from the real config, with random weights. All norms and small tensors are randomized too, so swapped tensors would show up. Each run is quantize → export, compared against the source and the published checkpoint:
nvidia/GLM-5.3-Flash-NVFP4, with layer and expert indices normalized, including the MTP layer and every scale tensor.hc_*_base/scale,A_log,dt_biasande_score_correction_biasare F32, as in the real source checkpoint.NVFP4,group_size16 andkv_cache_quant_algo: FP8match the published checkpoint.config.json: identical to the source apart fromquantization_config.GLM-5.2 on the stock
26.10.rc2container (no Megatron-LM / Megatron-Bridge patches), GB300: a tiny GLM-5.2 built from the real config (3 layers + MTP, randomized norms), quantized withgeneral/ptq/nvfp4_experts_only-kv_fp8_castand then exported:config.jsonis identical to the source, ande_score_correction_biasis F32.quantize.py's post-quantization generation check fails with a HybridEP in-place autograd error after the checkpoint is saved, so run with--skip_generate.EP export: a TP2-quantized and an EP2-quantized toy GLM-5.3-Flash checkpoint, each exported at EP2, are bit-identical to the EP1 exports: all 497 tensors,
hf_quant_config.jsonandconfig.json.Unit and GPU tests:
tests/unit/recipe/test_glm_5_3_recipe.py(new Megatron-names case).test_mcore_export_mappings.pyandtest_unified_export_megatron.py(KDA split, mHC).tests/gpu_megatron/torch/exportsuite (96 passed), including new mocked tests for the GLM-5 MTP copy (FP8 dequantize, depth-pruned remap).tests/examples/megatron_bridge/test_quantize_export.pynow quantizes with NVFP4 instead of FP8, and exportsnemotron_hat--ep_size <num_gpus>from its TP-quantized checkpoint; 3/3 passed on 2 GPUs.tests/unit/torch/export/test_get_quantization.py: new grouped-experts format and IQ detection tests, which fail without theweight_attr_namesfix.tests/unit/torch/quantization: 1133 passed;test_dbrxalso fails onmainin this container.tests/gpu_megatron/torch/quantization/plugins/test_megatron.py.Full-size GLM-5.3-Flash on GB300 (oci-jhb, 4 GPUs per node, from the real FP8-blockwise
zai-org/GLM-5.3-Flash):--ep_size 4(experts within each node, data parallel across nodes),cnn_nemotron_v2_mix, 1024 samples, sequence length 4096. Calibration took about 4 minutes and the job about 15. The 600 GB Megatron checkpoint has no NaN, Inf or zero amax. Quantized: routed experts (42 MoE layers) and the dense MLP (layers 0-2) in NVFP4, FP8 KV on the 11 DSA layers; MTP stays BF16.--ep_size 4, 39 minutes, 191 GB output. Checked againstnvidia/GLM-5.3-Flash-NVFP4from safetensors headers:NVFP4, group size 16,kv_cache_quant_algo: FP8, matching.config.json: equal to the source apart fromquantization_configand fields transformers 5.17 derives on save.--ep_size 16on 4 nodes), calibration hits NaN at an expert input quantizer on the first batch. The same EP16 layout without quantizers is NaN-free across all 16 ranks, and EP4 with quantizers is clean, so the cause is still being investigated. Keep--ep_size≤ GPUs per node.Not tested:
glm5_nextyet (0.26.0 does not), so the export isn't deployed or evaluated.Before your PR is "Ready for review"
expert_biasis now F32, VLM exports list the vision tower inexclude_modules, MLA configs keep the sourcehead_dim, and experts-only grouped-GEMM exports are now marked as quantized.CONTRIBUTING.md: N/AAdditional Information
Draft until the Megatron-LM and Megatron-Bridge PRs above land. GLM-5 / GLM-5.2 need neither.
🤖 Generated with Claude Code