From 72f9f7d9ea37aee16fe26cbe798922cc73c1b551 Mon Sep 17 00:00:00 2001 From: Shengliang Xu <106840466+shengliangxu@users.noreply.github.com> Date: Thu, 8 Oct 2026 13:37:45 -0700 Subject: [PATCH 01/35] [OmniML-3853] Backfill checkpoint aliases for Kimi-K2.7-Code and the Gemma-4 MoE releases (#2647) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? **Type of change:** New recipes (checkpoint aliases) Continues the 2026 published-checkpoint backfill started in #2643, with three more NVFP4 releases. Each entry is a thin alias importing an existing general recipe wholesale — no duplicated `quant_cfg`, and editing a base recipe flows through to every alias pointing at it. | Source model | Recipe | Published as | |---|---|---| | `moonshotai/Kimi-K2.7-Code` | `general/ptq/nvfp4_experts_only_input_scale1-kv_fp8_cast` | `nvidia/Kimi-K2.7-Code-NVFP4` | | `google/gemma-4-26B-A4B-it` | `general/ptq/nvfp4_experts_only-kv_fp8_cast` | `nvidia/Gemma-4-26B-A4B-NVFP4` | | `google/diffusiongemma-26B-A4B-it` | `general/ptq/nvfp4_experts_only-kv_fp8_cast` | `nvidia/diffusiongemma-26B-A4B-it-NVFP4` | #### How these were verified Unlike #2643, where each engineer named a recipe file outright, these three were described in prose ("NVFP4 experts only, input activation scale 1.0, calibration free"; "nvfp4 experts only, default dataset, 1024 samples"). So the mapping was checked against the published checkpoints before being recorded: - **Format and method.** All three are NVFP4 with an FP8 KV cache, and every model card documents post-training quantization with no mention of quantization-aware training or distillation. That check matters because QAD retrains the weights while leaving the quantization layout identical — a PTQ recipe could match every module and format and still not reproduce such a release. - **Scope.** Both Gemma releases list dense MLPs, router, attention, vision tower and `lm_head` in `exclude_modules`, leaving only the routed experts quantized. Kimi-K2.7-Code additionally excludes its shared experts. - **`input_scale1` vs plain experts-only.** This is what separates the Kimi recipe from the Gemma one. The `input_scale1` recipe pins expert activation amax to a constant 2688.0 (`E2M1_MAX * E4M3_MAX`), so the exported `input_scale` is exactly 1.0. Kimi-K2.7-Code's sampled expert `input_scale` tensors are exactly 1.0; the Gemma releases' are calibrated values, so they take the plain recipe. - **KV mode.** Neither Gemma release exports `k_scale` tensors, which is what cast mode produces (`use_constant_amax`, no KV calibration) — consistent with the already-merged `GLM-5.1`/`Kimi-K2.6` entries, which use `-kv_fp8_cast` recipes and likewise export none. A calibrated KV cache exports a non-1.0 `k_scale`. - **A naming quirk worth knowing:** `diffusiongemma` nests its decoder under `model.decoder.layers` rather than `model.language_model.layers`. The recipe's `*.experts.*` patterns match either way, so no model-specific body is needed. **What this cannot establish: calibration.** A checkpoint does not record whether max, MSE, or something else was used, so that rests on what each engineer said — the Gemma pair used the default, which is max. Kimi-K2.7-Code is calibration-free on the activation side by construction. #### Second commit: name the release in the Qwen3.8-27B description `models/Qwen/Qwen3.8-27B/ptq/nvfp4_w4a4_mlp_fp8_attn_local_hessian` already reproduces `nvidia/Qwen3.8-27B-NVFP4`, and `ptq.md` says so — but the recipe's own `metadata.description` did not, so on its own the file reads like an unrelated AutoQuantize artifact. Every other checkpoint entry names its release, so a coverage audit that reads the recipes can miss this one. (It missed mine.) Checked against the published checkpoint before asserting it: MLP `gate`/`up`/`down_proj` NVFP4, `self_attn` and `linear_attn` projections FP8, `lm_head` NVFP4, `mtp*` excluded, no KV quantization — and the model card names local-Hessian calibration on 2,048 samples. All matching, including the unusual NVFP4 `lm_head`. ### Usage ```bash python examples/hf_ptq/hf_ptq.py \ --pyt_ckpt_path google/gemma-4-26B-A4B-it \ --recipe models/google/gemma-4-26B-A4B-it/ptq/nvfp4_experts_only-kv_fp8_cast \ --export_path ``` ### Testing - `tests/unit/recipe/` — 460 passed. All three aliases load through `load_recipe`, inherit `algorithm: max` from their base, and resolve to 32 `quant_cfg` entries. - The aliases are swept by the shipped-recipe discovery in `tests/unit/recipe/test_loader.py`, which finds every PTQ recipe from disk, so they are covered without a bespoke test. - `pre-commit` passes on all changed files. - The 14 failures in the local run are pre-existing and environment-only — a broken `transformer_engine` `.so` in the dev venv; they reproduce identically on unmodified `main`. Not covered: **numerics**. Nothing here asserts accuracy or that running a recipe reproduces a released checkpoint's weights bit-for-bit. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ — purely additive. Three new alias files plus a docs and changelog entry; no existing recipe path, body or API changes. - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: ✅ — no new dependencies, no source changes. - Did you write any new necessary tests?: ✅ — covered by the shipped-recipe sweep. - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ - Did you get Claude approval on this PR?: ❌ — not yet run. ### Additional Information Three more releases from the same tracking list are verified for scope and KV mode but still need their calibration method confirmed by the engineer who produced them (a checkpoint cannot reveal it): `moonshotai/Kimi-K2.5`, `Qwen/Qwen3.5-397B-A17B` (the non-V2 release), and `zai-org/GLM-5.1`. They will follow separately. 🤖 Generated with [Claude Code](https://claude.com/claude-code) ## Summary by CodeRabbit * **New Features** * Added checkpoint recipes for Kimi-K2.7-Code, Gemma 4 26B A4B, and DiffusionGemma 26B A4B, with NVFP4 quantization and FP8 KV-cache options. * Added documentation identifying the published checkpoints and their available recipe configurations. * Updated the Qwen 3.8 27B recipe description to identify its published checkpoint. --------- Signed-off-by: Shengliang Xu Co-authored-by: Claude Opus 5 (1M context) Signed-off-by: Rohan Joshi --- CHANGELOG.rst | 1 + ...nvfp4_w4a4_mlp_fp8_attn_local_hessian.yaml | 2 +- .../ptq/nvfp4_experts_only-kv_fp8_cast.yaml | 30 +++++++++++++++++++ .../ptq/nvfp4_experts_only-kv_fp8_cast.yaml | 28 +++++++++++++++++ ...experts_only_input_scale1-kv_fp8_cast.yaml | 30 +++++++++++++++++++ modelopt_recipes/ptq.md | 9 ++++++ 6 files changed, 99 insertions(+), 1 deletion(-) create mode 100644 modelopt_recipes/models/google/diffusiongemma-26B-A4B-it/ptq/nvfp4_experts_only-kv_fp8_cast.yaml create mode 100644 modelopt_recipes/models/google/gemma-4-26B-A4B-it/ptq/nvfp4_experts_only-kv_fp8_cast.yaml create mode 100644 modelopt_recipes/models/moonshotai/Kimi-K2.7-Code/ptq/nvfp4_experts_only_input_scale1-kv_fp8_cast.yaml diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 98bcfd6915d..7b55df236ba 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -15,6 +15,7 @@ Changelog - Add Q8_0 weight-only quantization with 32-value GGML blocks, packed unified HF and Megatron export, and a built-in ``q8_0`` PTQ recipe. - Add Hugging Face PTQ calibration and export support for Nemotron-H MTP modules, preserving calibrated expert input scales. +- Backfill checkpoint aliases for three more published NVFP4 releases: ``moonshotai/Kimi-K2.7-Code``, ``google/gemma-4-26B-A4B-it`` and ``google/diffusiongemma-26B-A4B-it``. Each imports an existing general recipe wholesale rather than copying its body. - Backfill checkpoint aliases for nine more published NVFP4 releases, so each is reachable from its source model's hub path: ``zai-org/GLM-5.1`` and ``GLM-5.2``, ``MiniMaxAI/MiniMax-M2.5`` and ``MiniMax-M3``, ``deepseek-ai/DeepSeek-V3.1`` and ``DeepSeek-V3.2``, ``Qwen/Qwen3-235B-A22B-Instruct-2507`` and ``-Thinking-2507``, and ``Qwen/Qwen3.6-27B``. Each imports an existing general or architecture recipe wholesale rather than copying its body. - Add composed Hugging Face AutoQuantize recipes that run fixed PTQ or weight AutoQuantize before a separate KV-cache AutoQuantize stage, with independent resumable checkpoints for the weight and diff --git a/modelopt_recipes/models/Qwen/Qwen3.8-27B/ptq/nvfp4_w4a4_mlp_fp8_attn_local_hessian.yaml b/modelopt_recipes/models/Qwen/Qwen3.8-27B/ptq/nvfp4_w4a4_mlp_fp8_attn_local_hessian.yaml index 90decca0902..08894fec91d 100644 --- a/modelopt_recipes/models/Qwen/Qwen3.8-27B/ptq/nvfp4_w4a4_mlp_fp8_attn_local_hessian.yaml +++ b/modelopt_recipes/models/Qwen/Qwen3.8-27B/ptq/nvfp4_w4a4_mlp_fp8_attn_local_hessian.yaml @@ -20,7 +20,7 @@ imports: metadata: description: >- Qwen3.8-27B local-Hessian PTQ using the NVFP4 W4A4 and FP8 W8A8 assignment exported by - the 5.5-bit NVFP4-max AutoQuantize sweep. + the 5.5-bit NVFP4-max AutoQuantize sweep, as published in nvidia/Qwen3.8-27B-NVFP4. quantize: $import: base diff --git a/modelopt_recipes/models/google/diffusiongemma-26B-A4B-it/ptq/nvfp4_experts_only-kv_fp8_cast.yaml b/modelopt_recipes/models/google/diffusiongemma-26B-A4B-it/ptq/nvfp4_experts_only-kv_fp8_cast.yaml new file mode 100644 index 00000000000..d41599700b3 --- /dev/null +++ b/modelopt_recipes/models/google/diffusiongemma-26B-A4B-it/ptq/nvfp4_experts_only-kv_fp8_cast.yaml @@ -0,0 +1,30 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +# Alias recipe for nvidia/diffusiongemma-26B-A4B-it-NVFP4 (source google/diffusiongemma-26B-A4B-it). Same recipe as `general/ptq/nvfp4_experts_only-kv_fp8_cast`: NVFP4 on the routed experts +# only, max calibration, with an FP8 KV cache in cast mode -- the same scheme as the +# gemma-4-26B-A4B-it release. This model nests its decoder under `model.decoder.layers` +# rather than `model.language_model.layers`, which the recipe's `*.experts.*` patterns +# match either way. + +imports: + base: general/ptq/nvfp4_experts_only-kv_fp8_cast + +$import: base +metadata: + description: >- + google/diffusiongemma-26B-A4B-it quantized with the general expert-only NVFP4 scheme + (max calibration) and an FP8 KV cache in cast mode, as published in + nvidia/diffusiongemma-26B-A4B-it-NVFP4. diff --git a/modelopt_recipes/models/google/gemma-4-26B-A4B-it/ptq/nvfp4_experts_only-kv_fp8_cast.yaml b/modelopt_recipes/models/google/gemma-4-26B-A4B-it/ptq/nvfp4_experts_only-kv_fp8_cast.yaml new file mode 100644 index 00000000000..4f38d435a43 --- /dev/null +++ b/modelopt_recipes/models/google/gemma-4-26B-A4B-it/ptq/nvfp4_experts_only-kv_fp8_cast.yaml @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +# Alias recipe for nvidia/Gemma-4-26B-A4B-NVFP4 (source google/gemma-4-26B-A4B-it). Same recipe as `general/ptq/nvfp4_experts_only-kv_fp8_cast`: NVFP4 on the routed experts +# only, max calibration, with an FP8 KV cache in cast mode. The dense MLPs, router, +# attention, vision tower and `lm_head` stay BF16. + +imports: + base: general/ptq/nvfp4_experts_only-kv_fp8_cast + +$import: base +metadata: + description: >- + google/gemma-4-26B-A4B-it quantized with the general expert-only NVFP4 scheme (max + calibration) and an FP8 KV cache in cast mode, as published in + nvidia/Gemma-4-26B-A4B-NVFP4. diff --git a/modelopt_recipes/models/moonshotai/Kimi-K2.7-Code/ptq/nvfp4_experts_only_input_scale1-kv_fp8_cast.yaml b/modelopt_recipes/models/moonshotai/Kimi-K2.7-Code/ptq/nvfp4_experts_only_input_scale1-kv_fp8_cast.yaml new file mode 100644 index 00000000000..119cf553713 --- /dev/null +++ b/modelopt_recipes/models/moonshotai/Kimi-K2.7-Code/ptq/nvfp4_experts_only_input_scale1-kv_fp8_cast.yaml @@ -0,0 +1,30 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +# Alias recipe for nvidia/Kimi-K2.7-Code-NVFP4 (source moonshotai/Kimi-K2.7-Code). Same recipe as +# `general/ptq/nvfp4_experts_only_input_scale1-kv_fp8_cast`: NVFP4 on the routed experts +# with the expert `input_scale` pinned to 1.0 -- no activation calibration -- and an FP8 KV +# cache in cast mode. The release's exported expert `input_scale` tensors are exactly 1.0, +# and its shared experts, attention, `lm_head` and vision tower are all left unquantized. + +imports: + base: general/ptq/nvfp4_experts_only_input_scale1-kv_fp8_cast + +$import: base +metadata: + description: >- + moonshotai/Kimi-K2.7-Code quantized with the general expert-only NVFP4 scheme, expert + `input_scale` pinned to 1.0 (calibration-free), and an FP8 KV cache in cast mode, as + published in nvidia/Kimi-K2.7-Code-NVFP4. diff --git a/modelopt_recipes/ptq.md b/modelopt_recipes/ptq.md index 33c43e42359..21e5fe5744b 100644 --- a/modelopt_recipes/ptq.md +++ b/modelopt_recipes/ptq.md @@ -563,6 +563,15 @@ entry is a thin **alias** that imports that recipe wholesale and overrides only `model_type/qwen3_5_moe/ptq/nvfp4_experts_mse-fp8_rest-kv_fp8` — NVFP4 (MSE static weights) on the routed experts, ModelOpt-default FP8 elsewhere, and an FP8 KV cache — as published in `nvidia/Qwen3.5-397B-A17B-NVFP4-V2`. +- **`models/moonshotai/Kimi-K2.7-Code/ptq/nvfp4_experts_only_input_scale1-kv_fp8_cast`** + aliases `general/ptq/nvfp4_experts_only_input_scale1-kv_fp8_cast` — expert-only NVFP4 with + the expert `input_scale` pinned to 1.0 (no activation calibration) and an FP8 KV cache in + cast mode — as published in `nvidia/Kimi-K2.7-Code-NVFP4`. +- **`models/google/gemma-4-26B-A4B-it/ptq/nvfp4_experts_only-kv_fp8_cast`** and + **`models/google/diffusiongemma-26B-A4B-it/ptq/...`** alias + `general/ptq/nvfp4_experts_only-kv_fp8_cast` — expert-only NVFP4 with max calibration and an + FP8 KV cache in cast mode — as published in `nvidia/Gemma-4-26B-A4B-NVFP4` and + `nvidia/diffusiongemma-26B-A4B-it-NVFP4`. - **`models/zai-org/GLM-5.1/ptq/nvfp4_experts_only_input_scale1-kv_fp8_cast`** and **`models/zai-org/GLM-5.2/ptq/...`** alias `general/ptq/nvfp4_experts_only_input_scale1-kv_fp8_cast` — expert-only NVFP4 with the From 4c1ffaab8a861e159c11b7a5fb2956b3d4fb5bc2 Mon Sep 17 00:00:00 2001 From: Juhi Mittal <39641197+juhi10071998@users.noreply.github.com> Date: Thu, 8 Oct 2026 14:43:45 -0700 Subject: [PATCH 02/35] fix(export): Don't re-apply a reverse rename to a key that already carries its prefix during export (#2700) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Type of change: Bug fix Unified HF export misnames tensors that are **copied from the source checkpoint rather than loaded by the model**, when the checkpoint's conversion mapping renames a module to a suffix of its original name. `revert_weight_conversion_quant_aware` runs over the whole export state dict. That dict includes passthrough tensors — `unified_export_hf` logs them as `Carrying N source weight(s) the model never loaded into the export`. Those keys never went through the *forward* conversion, so they arrive at the reverse pass already in hub naming, and a reverse rule whose replacement still contains its own pattern adds its prefix a second time. C-RADIO remote code registers `WeightRenaming("radio_model.input_conditioner", "input_conditioner")`. The derived reverse is `input_conditioner` -> `radio_model.input_conditioner`, so the passthrough buffer is exported as: ``` got: vision_model.radio_model.radio_model.input_conditioner.norm_mean expected: vision_model.radio_model.input_conditioner.norm_mean ``` The tensor is present but misnamed, which is worse than missing: a strict deployment loader sees an unexpected key *and* a missing one. Only rules shaped `pattern` -> ` + pattern` can do this. `encoder.layer` -> `radio_model.model.blocks` cannot re-match its own output, so most mappings are accidentally idempotent — which is why this cost 2 names out of 43,079 and went unnoticed. **Fix:** when the replacement is exactly ` + `, compile the pattern with a negative lookbehind for that prefix, so an already-reverted key is left alone. The prefix is a literal string, so the lookbehind is fixed-width. Replacements that do not end with the pattern, including group-reference forms such as `(\.experts\.\d+\.)up_proj` -> `\g<1>up_proj`, compile unchanged. ### Usage No API change. Affected checkpoints need re-exporting to pick up the corrected names. ### Testing ```bash PYTHONPATH=. python -m pytest tests/unit/torch/export/test_quant_aware_conversion.py -q -o addopts='' ``` **57 passed** (3 new, 54 existing), transformers 5.13.0 / torch 2.12.0a0 on one node. End to end: a NemotronH omni VL checkpoint with a C-RADIO vision tower, NVFP4 PTQ via `examples/hf_ptq` on 4x184 GiB, before and after the fix. | | before | after | |---|---|---| | tensor names matching the BF16 hub checkpoint verbatim | 43,077 / 43,079 | **43,079 / 43,079** | | occurrences of `radio_model.radio_model` | 2 | **0** | Everything else in the artefact is unchanged: 0 duplicates, 42,179 NVFP4 scales, 4,120 MTP tensors, 8 shards, 457 vision tensors, 69.86 GiB. GSM8K on the re-exported checkpoint: 93.18 (1,319 questions, 5-shot, `invalid_rate` 0.0), in line with what we see for this checkpoint. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: Yes. Exported names move toward the hub layout; checkpoints already exported with the doubled prefix should be re-exported. - If you copied code from other sources or added a new PIP dependency, did you follow CONTRIBUTING.md?: N/A; no new dependency or copied code. - Did you write any new necessary tests?: Yes, 3 unit tests covering the doubling case, an already-reverted key, and a rule that cannot re-match. - Did you update Changelog?: Yes, under **Bug Fixes**, next to the related RADIO export entry. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Signed-off-by: Juhi Mittal Co-authored-by: Claude Opus 5 (1M context) Signed-off-by: Rohan Joshi --- CHANGELOG.rst | 5 ++++ .../torch/export/quant_aware_conversion.py | 18 +++++++++++- .../export/test_quant_aware_conversion.py | 29 +++++++++++++++++++ 3 files changed, 51 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 7b55df236ba..cc0431b7f3d 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -142,6 +142,11 @@ Changelog projector names instead of restoring the hub layout; deployment loaders could skip those weights. Streaming, offload, FSDP2, and layerwise export now reject reverse conversions that regroup tensors; use resident export for those models. +- Fix that same reverse rename double-prefixing tensors copied from the source checkpoint rather + than loaded by the model, when a conversion mapping renames a module to a suffix of its original + name. Those tensors already carry the hub name, so C-RADIO towers exported + ``vision_model.radio_model.radio_model.input_conditioner.norm_mean``. Re-export affected + checkpoints. - Fix shared ONNX export metadata and Diffusers attention policy: every ``NVFP4QuantExporter`` post-process now upgrades the default-domain opset to at least 23, all FP8 custom-op exports re-run ONNX shape/type inference after setting output metadata, and quantized SDPA derives FP8 MHA enablement from the live Q/K/V quantizers instead of honoring a caller-set ``_disable_fp8_mha`` attribute. - Fix ONNX FP16 conversion failing to preserve public output types when type inference changes a graph output declaration before output casts are inserted. - Fix ``examples/hf_ptq/hf_ptq.py`` discarding a completed PTQ run (no checkpoint exported) when the optional post-quantization sanity-check ``generate()`` call raised, for example because ``device_map="auto"`` placed part of the model on CPU. That failure is now caught and only skips the sanity check; export proceeds regardless. diff --git a/modelopt/torch/export/quant_aware_conversion.py b/modelopt/torch/export/quant_aware_conversion.py index 482bb95e94b..8c750702751 100644 --- a/modelopt/torch/export/quant_aware_conversion.py +++ b/modelopt/torch/export/quant_aware_conversion.py @@ -251,9 +251,25 @@ def _apply_merge_rule(state_dict: dict[str, torch.Tensor], rule: _MergeRule) -> state_dict[target_key] = merged +def _idempotent_pattern(pattern: str, repl: str) -> str: + """Stop a rename rule from re-matching its own replacement. + + Tensors copied from the source checkpoint skip the forward conversion and reach the reverse + pass already in hub naming, so a rule like ``input_conditioner`` -> + ``radio_model.input_conditioner`` would add its prefix twice. When the replacement is + `` + ``, skip keys that already carry the prefix. + """ + if len(repl) > len(pattern) and repl.endswith(pattern): + return f"(? str: diff --git a/tests/unit/torch/export/test_quant_aware_conversion.py b/tests/unit/torch/export/test_quant_aware_conversion.py index 209dafb64b3..e2efe7c7ff6 100644 --- a/tests/unit/torch/export/test_quant_aware_conversion.py +++ b/tests/unit/torch/export/test_quant_aware_conversion.py @@ -158,6 +158,35 @@ def _fp8_llama(*quantizer_names): ) +def test_reverse_rename_is_idempotent_when_repl_contains_pattern(): + """A rule whose replacement keeps the pattern must not re-apply to its own output.""" + rules = [RenameRule("input_conditioner", "radio_model.input_conditioner", ("vision_model.",))] + key = "vision_model.input_conditioner.norm_mean" + expected = "vision_model.radio_model.input_conditioner.norm_mean" + + once = apply_reverse_rules({key: torch.zeros(3)}, [], rules) + assert list(once) == [expected] + + twice = apply_reverse_rules(once, [], rules) + assert list(twice) == [expected], "reverse rename must be idempotent" + + +def test_reverse_rename_idempotent_for_already_hub_named_key(): + """An input already carrying the reverted prefix is left untouched.""" + rules = [RenameRule("summary_idxs", "radio_model.summary_idxs", ("vision_model.",))] + already = "vision_model.radio_model.summary_idxs" + out = apply_reverse_rules({already: torch.zeros(1)}, [], rules) + assert list(out) == [already] + + +def test_idempotency_guard_leaves_disjoint_rules_alone(): + """A replacement that cannot re-match keeps its original behaviour.""" + rules = [RenameRule("encoder.layer", "radio_model.model.blocks", ("vision_model.",))] + key = "vision_model.encoder.layer.0.attn.weight" + out = apply_reverse_rules({key: torch.zeros(2)}, [], rules) + assert list(out) == ["vision_model.radio_model.model.blocks.0.attn.weight"] + + def test_rename_carries_scale_siblings(): """A module rename rewrites weight + all scale siblings with identical values.""" sd = _nvfp4_linear("model.language_model.layers.10.mlp.experts.40.gate_proj", 8, 16) From 37e9123bfa8b51ab7f4f806c93398e6e36dcb4dc Mon Sep 17 00:00:00 2001 From: Frida Hou <201670829+Fridah-nv@users.noreply.github.com> Date: Thu, 8 Oct 2026 15:55:51 -0700 Subject: [PATCH 03/35] fix(gpt-oss): pin trl<1.15 and trackio>=0.41 for the example (#2718) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Type of change: Bug fix Pins two `examples/gpt-oss` dependencies that broke `tests/examples/gpt-oss/test_gpt_oss_qat.py::test_gpt_oss_sft_toy` on PR CI on 2026-10-08 (19 failed runs across 16 PRs). The example installs `requirements.txt` with open ranges, so each run picked up whatever was newest. - **`trl<1.15`**: TRL 1.15.0's `SFTTrainer.compute_loss` always calls `model(**inputs, fused_lm_head=True)`, which runs the Triton `ChunkedLogProbFunction` kernel (`trl/trainer/utils.py`). The only fallback is when Triton can't be imported, so there is no device check and no opt-out. The toy test trains on CPU (`--use_cpu True`, `CUDA_VISIBLE_DEVICES=""`), so Triton fails with `RuntimeError: 0 active drivers ([]). There should only be one.` Example: https://github.com/NVIDIA/Model-Optimizer/actions/runs/37836855877/job/113524251231 - **`trackio>=0.41`**: some runs resolved `trackio` 0.20.2, whose `commit_scheduler.py` imports `CommitOperationAdd` from `huggingface_hub.hf_api`; with `huggingface_hub` 1.33.0 that fails with `ImportError`. Example: https://github.com/NVIDIA/Model-Optimizer/actions/runs/37816170339/job/113447734730 The last passing runs had exactly `trl` 1.14.2 and `trackio` 0.41.0 (e.g. https://github.com/NVIDIA/Model-Optimizer/actions/runs/37828686284/job/113490544527). The `trl` cap should be lifted once TRL skips the fused head for CPU training. Longer term, PR CI could install example dependencies from a bot-maintained constraints file (like `bump_uv_lock.yml` does for `uv.lock`) with nightly runs on unpinned latest, so an upstream release doesn't break every open PR at once. N/A - `pip install --dry-run -r examples/gpt-oss/requirements.txt` resolves to `trl` 1.14.2 and `trackio` 0.41.0, matching the last passing CI runs. - Read the TRL 1.15.0 wheel to confirm the fused-head path has no CPU fallback or opt-out. - `pre-commit` passes on the changed file. - The `trtllm (gpt-oss)` example job on this PR is the end-to-end check. - Is this change backward compatible?: ✅ - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: N/A - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A - Did you get Claude approval on this PR?: ❌ Unblocks the `trtllm (gpt-oss)` example check on open PRs, e.g. #2440. 🤖 Generated with [Claude Code](https://claude.com/claude-code) * **Chores** * Updated setup configuration for the GPT-OSS example. This change does not alter the example’s user-facing features or behavior. --------- Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 Signed-off-by: Rohan Joshi --- examples/gpt-oss/requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/gpt-oss/requirements.txt b/examples/gpt-oss/requirements.txt index e8618d6ca5e..7462574d4d0 100644 --- a/examples/gpt-oss/requirements.txt +++ b/examples/gpt-oss/requirements.txt @@ -1,5 +1,5 @@ kernels>=0.9.0,<0.13 -trackio +trackio>=0.41 # transformers>=5.3 avoids CVE-2026-4372 (RCE via the `kernels` Hub download path, which this example installs) transformers>=5.3 # TRL 1.15 enables a GPU-only fused LM head, breaking the CPU SFT smoke test. From f8b93e79db1836735880a3d8c6eed6bb098d0723 Mon Sep 17 00:00:00 2001 From: Shengliang Xu <106840466+shengliangxu@users.noreply.github.com> Date: Thu, 8 Oct 2026 23:24:29 -0700 Subject: [PATCH 04/35] [OMNIML-3817] Move model-specific PTQ modeling into modelopt/torch/models [3/4] (#2578) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Type of change: Refactor ### Summary of the series Model-specific PTQ modeling moves out of `modelopt/torch/quantization/plugins/huggingface.py` into the per-model packages under `modelopt/torch/models//modeling_ptq.py`, next to each model's `specs.py`. By the end of the series `huggingface.py` shrinks from ~1880 to ~1500 lines and keeps only the code shared across models: sequential/fused MoE auto-detection, attention, `FP8Linear`, `CompressedLinear`, parallel linears, and the transposed-quantization helpers that gpt_oss and llama4 share. Model types that gain PTQ modeling but had no `ModelSpec` get one. - Class and function bodies move verbatim. The only rewrites: on-the-fly callbacks are added to `CUSTOM_MODEL_PLUGINS` by their own module, and `is_homogeneous_hf_model` imports `is_nemotron_h_model` lazily. - `huggingface.py` imports every `modeling_ptq` from an explicit list. The list sits after the generic wrappers these modules build on and before the homogeneous decoder discoverer is registered, so Nemotron-H's more specific discoverer still matches first. - Each package's `__init__.py` keeps importing only `specs`, so `import modelopt.torch.models` does not pull in quantization or transformers. - No public API or quantization behavior changes; only private (`_`-prefixed) names change module. Merge in order. #2576 and #2577 have merged; #2578 now targets `main`, and #2580 targets #2578's branch. 1. #2576 — Add model specs for llama4, qwen3_vl_moe, falcon, step3p5 and step3p7 (merged) 2. #2577 — Add the `modeling_ptq.py` convention; move Nemotron-H, Falcon and Llama4 (merged) 3. #2578 — Move GPT-OSS ← **this PR** 4. #2580 — Move the Step family (step3p5, shared by step3p7) The DBRX slice (#2579) was dropped: #2649 removed DBRX support, so there was nothing left to move. Likewise, the Qwen3-VL-MoE half of [3/4] was dropped: #2670 removed its pre-5.12 wrapper. ### This PR [3/4] Moves `_QuantGptOssExperts` and its registration into `gpt_oss/modeling_ptq.py`, adds `gpt_oss` to the import loop, and extends the registration test to it. This slice originally also moved the pre-5.12 Qwen3-VL-MoE wrapper; #2670 removed that wrapper, so that half was dropped. ### Usage No API change. To add PTQ support for a new model, put its wrapper in `modelopt/torch/models//modeling_ptq.py` and add `` to the import loop at the end of `quantization/plugins/huggingface.py`. ### Testing Each branch of the stack checked out and tested on `main` at 90ba9fbc65, CPU, torch 2.11. With transformers 5.18: `pytest tests/unit/torch/models tests/unit/torch/quantization tests/unit/torch/export tests/unit/recipe`, 2371 passed, 1 skipped. With transformers 5.5 (the new floor): `tests/unit/torch/models` plus `test_moe_linear.py` and `test_huggingface.py`, 206 passed, 1 skipped. GPU tests are left to CI. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ Only private (`_`-prefixed) names change module; public API and quantization behavior are unchanged. - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: N/A (verbatim move; existing tests re-pointed) - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A (internal refactor) - Did you get Claude approval on this PR?: ❌ ### Additional Information Part of a 5-PR stack; see the series list above. 🤖 Generated with [Claude Code](https://claude.com/claude-code) ## Summary by CodeRabbit * **New Features** * Added post-training quantization support for GPT-OSS models. * Added support for quantizing legacy Qwen3-VL-MoE layouts, while retaining support for newer fused layouts. * Quantization support now accommodates both legacy and newer Qwen3-VL-MoE model layouts. --------- Signed-off-by: Shengliang Xu Signed-off-by: Rohan Joshi --- modelopt/torch/models/gpt_oss/modeling_ptq.py | 121 ++++++++++++++++++ .../torch/quantization/plugins/huggingface.py | 92 +------------ .../models/test_modeling_ptq_registration.py | 48 ++++--- 3 files changed, 154 insertions(+), 107 deletions(-) create mode 100644 modelopt/torch/models/gpt_oss/modeling_ptq.py diff --git a/modelopt/torch/models/gpt_oss/modeling_ptq.py b/modelopt/torch/models/gpt_oss/modeling_ptq.py new file mode 100644 index 00000000000..e92780ef34e --- /dev/null +++ b/modelopt/torch/models/gpt_oss/modeling_ptq.py @@ -0,0 +1,121 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""GPT-OSS PTQ modeling (HF model type ``gpt_oss``).""" + +from contextlib import contextmanager +from functools import partial + +import torch +from transformers.models.gpt_oss.modeling_gpt_oss import GptOssExperts + +from modelopt.torch.quantization.nn import QuantModuleRegistry, TensorQuantizer +from modelopt.torch.quantization.plugins.custom import _QuantFunctionalMixin +from modelopt.torch.quantization.plugins.huggingface import ( + _transposed_quantize, + _TransposedExpertsCalibMixin, +) + +__all__: list[str] = [] + + +class _QuantGptOssExperts(_TransposedExpertsCalibMixin, _QuantFunctionalMixin): + """Quantized wrapper for `transformers.GptOssExperts`. + + Quantizes `gate_up_proj` and `down_proj` weights via dynamic attributes inside `quantize_weight()`. + Activations into `gate_up_proj` are quantized by `gate_up_proj_input_quantizer`. For `down_proj` + activation quantization, we intercept `torch.Tensor.__matmul__`/`torch.bmm` and quantize inputs + on every second call (since the first call computes `gate_up_proj` outputs and second call + computes `down_proj` outputs). + """ + + @staticmethod + def _get_quantized_weight(quantizer, module, weight): + # MoE weight is accessed for each expert in one forward pass. so lets cache it + if module._enable_weight_quantization: + if hasattr(quantizer, "_cached_quant_val"): + return getattr(quantizer, "_cached_quant_val") + quantizer._cached_quant_val = _transposed_quantize(weight, quantizer) + return quantizer._cached_quant_val + return weight + + def _setup_for_weight_quantization(self): + self._register_dynamic_attribute( + "gate_up_proj", partial(self._get_quantized_weight, self.gate_up_proj_weight_quantizer) + ) + self._register_dynamic_attribute( + "down_proj", partial(self._get_quantized_weight, self.down_proj_weight_quantizer) + ) + + def _setup(self): + assert not hasattr(self, "kernel_layer_name"), ( + "ModelOpt quantization does not support patched forward for kernel_hub" + ) + self.gate_up_proj_input_quantizer = TensorQuantizer() + self.gate_up_proj_weight_quantizer = TensorQuantizer() + self.down_proj_input_quantizer = TensorQuantizer() + self.down_proj_weight_quantizer = TensorQuantizer() + + self._register_temp_attribute("_enable_weight_quantization", False) + self._register_temp_attribute("_down_proj_mul", False) + self._setup_for_weight_quantization() + + @property + def functionals_to_replace(self): + # Use torch.ops.aten to bypass Python dispatch and avoid RecursionError + # (torch.matmul / __matmul__ can dispatch to each other) + _aten_bmm = torch.ops.aten.bmm + _aten_matmul = torch.ops.aten.matmul + + def _quantized_bmm(batch1, batch2, *, out=None): + batch1 = self.down_proj_input_quantizer(batch1) if self._down_proj_mul else batch1 + self._down_proj_mul = not self._down_proj_mul # toggle the flag + if out is not None: + return torch.ops.aten.bmm.out(batch1, batch2, out=out) + return _aten_bmm(batch1, batch2) + + def _tensor_matmul(self_t, other): + self_t = self.down_proj_input_quantizer(self_t) if self._down_proj_mul else self_t + self._down_proj_mul = not self._down_proj_mul + return _aten_matmul(self_t, other) + + return [ + (torch, "bmm", _quantized_bmm), + (torch.Tensor, "__matmul__", _tensor_matmul), + ] + + @contextmanager + def quantize_weight(self): + """Context in which MoE weight is quantized.""" + self._enable_weight_quantization = True + try: + yield + finally: + for module in self.modules(): + if isinstance(module, TensorQuantizer) and hasattr(module, "_cached_quant_val"): + delattr(module, "_cached_quant_val") + self._enable_weight_quantization = False + + def forward( + self, hidden_states: torch.Tensor, router_indices=None, routing_weights=None + ) -> torch.Tensor: + """Forward method to add quantization.""" + hidden_states = self.gate_up_proj_input_quantizer(hidden_states) + with self.quantize_weight(): + return super().forward(hidden_states, router_indices, routing_weights) + + +if GptOssExperts not in QuantModuleRegistry: + QuantModuleRegistry.register({GptOssExperts: "hf.GptOssExperts"})(_QuantGptOssExperts) diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index 5eae0ae937a..05e589835ac 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -30,7 +30,6 @@ from torch import Tensor from torch.nn.functional import linear from transformers.integrations.finegrained_fp8 import FP8Linear -from transformers.models.gpt_oss.modeling_gpt_oss import GptOssExperts from transformers.models.t5.modeling_t5 import T5Attention from modelopt.torch.kernels.common.attention import IS_AVAILABLE as TRITON_FA_AVAILABLE @@ -1206,96 +1205,6 @@ def unpack_weight(self): QuantModuleRegistry.register({FP8Linear: "hf.FP8Linear"})(_QuantFP8Linear) -class _QuantGptOssExperts(_TransposedExpertsCalibMixin, _QuantFunctionalMixin): - """Quantized wrapper for `transformers.GptOssExperts`. - - Quantizes `gate_up_proj` and `down_proj` weights via dynamic attributes inside `quantize_weight()`. - Activations into `gate_up_proj` are quantized by `gate_up_proj_input_quantizer`. For `down_proj` - activation quantization, we intercept `torch.Tensor.__matmul__`/`torch.bmm` and quantize inputs - on every second call (since the first call computes `gate_up_proj` outputs and second call - computes `down_proj` outputs). - """ - - @staticmethod - def _get_quantized_weight(quantizer, module, weight): - # MoE weight is accessed for each expert in one forward pass. so lets cache it - if module._enable_weight_quantization: - if hasattr(quantizer, "_cached_quant_val"): - return getattr(quantizer, "_cached_quant_val") - quantizer._cached_quant_val = _transposed_quantize(weight, quantizer) - return quantizer._cached_quant_val - return weight - - def _setup_for_weight_quantization(self): - self._register_dynamic_attribute( - "gate_up_proj", partial(self._get_quantized_weight, self.gate_up_proj_weight_quantizer) - ) - self._register_dynamic_attribute( - "down_proj", partial(self._get_quantized_weight, self.down_proj_weight_quantizer) - ) - - def _setup(self): - assert not hasattr(self, "kernel_layer_name"), ( - "ModelOpt quantization does not support patched forward for kernel_hub" - ) - self.gate_up_proj_input_quantizer = TensorQuantizer() - self.gate_up_proj_weight_quantizer = TensorQuantizer() - self.down_proj_input_quantizer = TensorQuantizer() - self.down_proj_weight_quantizer = TensorQuantizer() - - self._register_temp_attribute("_enable_weight_quantization", False) - self._register_temp_attribute("_down_proj_mul", False) - self._setup_for_weight_quantization() - - @property - def functionals_to_replace(self): - # Use torch.ops.aten to bypass Python dispatch and avoid RecursionError - # (torch.matmul / __matmul__ can dispatch to each other) - _aten_bmm = torch.ops.aten.bmm - _aten_matmul = torch.ops.aten.matmul - - def _quantized_bmm(batch1, batch2, *, out=None): - batch1 = self.down_proj_input_quantizer(batch1) if self._down_proj_mul else batch1 - self._down_proj_mul = not self._down_proj_mul # toggle the flag - if out is not None: - return torch.ops.aten.bmm.out(batch1, batch2, out=out) - return _aten_bmm(batch1, batch2) - - def _tensor_matmul(self_t, other): - self_t = self.down_proj_input_quantizer(self_t) if self._down_proj_mul else self_t - self._down_proj_mul = not self._down_proj_mul - return _aten_matmul(self_t, other) - - return [ - (torch, "bmm", _quantized_bmm), - (torch.Tensor, "__matmul__", _tensor_matmul), - ] - - @contextmanager - def quantize_weight(self): - """Context in which MoE weight is quantized.""" - self._enable_weight_quantization = True - try: - yield - finally: - for module in self.modules(): - if isinstance(module, TensorQuantizer) and hasattr(module, "_cached_quant_val"): - delattr(module, "_cached_quant_val") - self._enable_weight_quantization = False - - def forward( - self, hidden_states: torch.Tensor, router_indices=None, routing_weights=None - ) -> torch.Tensor: - """Forward method to add quantization.""" - hidden_states = self.gate_up_proj_input_quantizer(hidden_states) - with self.quantize_weight(): - return super().forward(hidden_states, router_indices, routing_weights) - - -if GptOssExperts not in QuantModuleRegistry: - QuantModuleRegistry.register({GptOssExperts: "hf.GptOssExperts"})(_QuantGptOssExperts) - - def _has_num_experts(obj): # n_routed_experts: NemotronH-style MoE return hasattr(obj, "num_experts") or hasattr(obj, "n_routed_experts") @@ -1572,6 +1481,7 @@ def _is_param_grad_enabled_for_auto_quantize(pname, model): # Nemotron-H's is the more specific one. for _model_type in ( "falcon", + "gpt_oss", "llama4", "nemotron_h", ): diff --git a/tests/unit/torch/models/test_modeling_ptq_registration.py b/tests/unit/torch/models/test_modeling_ptq_registration.py index b97a2294a67..fbaa761d38d 100644 --- a/tests/unit/torch/models/test_modeling_ptq_registration.py +++ b/tests/unit/torch/models/test_modeling_ptq_registration.py @@ -21,8 +21,10 @@ so an in-process test would only ever see the import order of whichever test ran first. """ +import os import subprocess import sys +from concurrent.futures import ThreadPoolExecutor import pytest @@ -30,6 +32,7 @@ FIRST_IMPORTS = [ "modelopt.torch.models.falcon.modeling_ptq", + "modelopt.torch.models.gpt_oss.modeling_ptq", "modelopt.torch.models.llama4.modeling_ptq", "modelopt.torch.models.nemotron_h.modeling_ptq", "modelopt.torch.quantization", @@ -63,6 +66,13 @@ else: assert Llama4TextExperts in QuantModuleRegistry, "Llama4TextExperts not registered" +try: + from transformers.models.gpt_oss.modeling_gpt_oss import GptOssExperts +except ImportError: + pass +else: + assert GptOssExperts in QuantModuleRegistry, "GptOssExperts not registered" + assert register_falcon_linears_on_the_fly in CUSTOM_MODEL_PLUGINS, "Falcon callback missing" # The first matching decoder discoverer wins, so Nemotron-H's must precede the generic one. @@ -73,23 +83,29 @@ """ +def _run_check(first_import): + # The child only asserts; tracing its heavy imports for coverage (inherited through + # COVERAGE_PROCESS_START) would just slow it down. + env = {k: v for k, v in os.environ.items() if k != "COVERAGE_PROCESS_START"} + result = subprocess.run( + [sys.executable, "-c", CHECK, first_import], + capture_output=True, + text=True, + env=env, + check=False, + ) + return first_import, result + + +# Several fresh interpreters each import torch and transformers, which outlasts the default +# 60 s per-test cap on a small CI runner. +@pytest.mark.timeout(300) def test_registration_is_independent_of_import_order(): - # Each first import needs its own interpreter, but starting them one after another would - # add over a minute to the unit-test job, so launch them all and then collect the results. - procs = { - first_import: subprocess.Popen( - [sys.executable, "-c", CHECK, first_import], - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - text=True, - ) - for first_import in FIRST_IMPORTS - } - failures = {} - for first_import, proc in procs.items(): - _, stderr = proc.communicate(timeout=600) - if proc.returncode != 0: - failures[first_import] = stderr + # Each first import needs its own interpreter; run them in parallel, at most one per CPU, + # rather than one after another. + with ThreadPoolExecutor(max_workers=os.cpu_count() or 1) as pool: + results = list(pool.map(_run_check, FIRST_IMPORTS)) + failures = {name: result.stderr for name, result in results if result.returncode != 0} assert not failures, "\n\n".join( f"importing {name} first:\n{err}" for name, err in failures.items() ) From 752b51bd0223edd33c80dcb78e31585a99e6187e Mon Sep 17 00:00:00 2001 From: Shengliang Xu <106840466+shengliangxu@users.noreply.github.com> Date: Fri, 9 Oct 2026 12:33:22 -0700 Subject: [PATCH 05/35] [OMNIML-3817] Move model-specific PTQ modeling into modelopt/torch/models [4/4] (#2580) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Type of change: Refactor ### Summary of the series Model-specific PTQ modeling moves out of `modelopt/torch/quantization/plugins/huggingface.py` into the per-model packages under `modelopt/torch/models//modeling_ptq.py`, next to each model's `specs.py`. By the end of the series `huggingface.py` shrinks from ~1880 to ~1500 lines and keeps only the code shared across models: sequential/fused MoE auto-detection, attention, `FP8Linear`, `CompressedLinear`, parallel linears, and the transposed-quantization helpers that gpt_oss and llama4 share. Model types that gain PTQ modeling but had no `ModelSpec` get one. - Class and function bodies move verbatim. The only rewrites: on-the-fly callbacks are added to `CUSTOM_MODEL_PLUGINS` by their own module, and `is_homogeneous_hf_model` imports `is_nemotron_h_model` lazily. - `huggingface.py` imports every `modeling_ptq` from an explicit list. The list sits after the generic wrappers these modules build on and before the homogeneous decoder discoverer is registered, so Nemotron-H's more specific discoverer still matches first. - Each package's `__init__.py` keeps importing only `specs`, so `import modelopt.torch.models` does not pull in quantization or transformers. - No public API or quantization behavior changes; only private (`_`-prefixed) names change module. Merge in order. #2576, #2577 and #2578 have merged; this last slice, #2580, now targets `main`. 1. #2576 — Add model specs for llama4, qwen3_vl_moe, falcon, step3p5 and step3p7 (merged) 2. #2577 — Add the `modeling_ptq.py` convention; move Nemotron-H, Falcon and Llama4 (merged) 3. #2578 — Move GPT-OSS (merged) 4. #2580 — Move the Step family (step3p5, shared by step3p7) ← **this PR** The DBRX slice (#2579) was dropped: #2649 removed DBRX support, so there was nothing left to move. Likewise, the Qwen3-VL-MoE half of [3/4] was dropped: #2670 removed its pre-5.12 wrapper. ### This PR [4/4] Moves `_QuantMoELinear`, `_is_expert_indexed_moe_linear`, `_is_step_family_model`, `register_moe_linear_on_the_fly` and `_reconstruct_fused_moe_linear` into `step3p5/modeling_ptq.py`, which also serves `step3p7`. This carries over #2569's switch of `_is_step_family_model` to `hf_model_type`; `huggingface.py` drops its now-unused `hf_model_type` and `re` imports. Export (`hf_export_handlers.py`, `layerwise_export.py`, `unified_export_hf.py`, `unified_export_hf_streaming.py`) imports these names from the new module. The imports stay deferred, because `modelopt.torch.quantization` imports `modelopt.torch.export` and importing at module scope would risk an import cycle; the comments now give that reason. The PTQ skill references (`unsupported-models.md`, `checkpoint-validation.md`) now point agents at `modelopt/torch/models//modeling_ptq.py` for model-specific patches. Rebasing onto #2649 also replaces the README's DBRX example of a model-specific wrapper with Llama4's fused BMM experts. `step3p7` also gets its own `modeling_ptq.py`, which imports the Step-3.5 module: Step-3.7 remote-code checkpoints use the same `MoELinear`, and today they are covered only because the HF plugin imports every `modeling_ptq`. With its own module, a future loader that imports PTQ modeling by model type finds it too. ### Usage No API change. To add PTQ support for a new model, put its wrapper in `modelopt/torch/models//modeling_ptq.py` and add `` to the import loop at the end of `quantization/plugins/huggingface.py`. ### Testing Each branch of the stack checked out and tested on `main` at 90ba9fbc65, CPU, torch 2.11. With transformers 5.18: `pytest tests/unit/torch/models tests/unit/torch/quantization tests/unit/torch/export tests/unit/recipe`, 2372 passed, 1 skipped. With transformers 5.5 (the new floor): `tests/unit/torch/models` plus `test_moe_linear.py` and `test_huggingface.py`, 207 passed, 1 skipped. GPU tests are left to CI. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ Only private (`_`-prefixed) names change module; public API and quantization behavior are unchanged. - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: N/A (verbatim move; existing tests re-pointed) - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A (internal refactor) - Did you get Claude approval on this PR?: ❌ ### Additional Information Part of a 5-PR stack; see the series list above. 🤖 Generated with [Claude Code](https://claude.com/claude-code) ## Summary by CodeRabbit * **New Features** * Added post-training quantization support for Step-family models with expert-indexed MoE layers. * Checkpoint exports preserve the original layout of quantized expert weights and scales. * **Documentation** * Updated MoE architecture examples and guidance for configuring model-specific quantization support, including Step model revisions. --------- Signed-off-by: Shengliang Xu Signed-off-by: Rohan Joshi --- modelopt/torch/export/hf_export_handlers.py | 5 +- modelopt/torch/export/layerwise_export.py | 4 +- modelopt/torch/export/unified_export_hf.py | 7 +- .../export/unified_export_hf_streaming.py | 12 +- modelopt/torch/models/README.md | 9 +- modelopt/torch/models/step3p5/modeling_ptq.py | 264 ++++++++++++++++++ modelopt/torch/models/step3p7/modeling_ptq.py | 26 ++ modelopt/torch/models/step3p7/specs.py | 4 +- .../torch/quantization/plugins/huggingface.py | 231 +-------------- .../ptq/references/checkpoint-validation.md | 2 +- .../ptq/references/unsupported-models.md | 10 +- .../models/test_modeling_ptq_registration.py | 4 + .../test_step3p5_modeling_ptq.py} | 4 +- 13 files changed, 326 insertions(+), 256 deletions(-) create mode 100644 modelopt/torch/models/step3p5/modeling_ptq.py create mode 100644 modelopt/torch/models/step3p7/modeling_ptq.py rename tests/unit/torch/{quantization/plugins/test_moe_linear.py => models/test_step3p5_modeling_ptq.py} (99%) diff --git a/modelopt/torch/export/hf_export_handlers.py b/modelopt/torch/export/hf_export_handlers.py index 7c526d8255f..200ec6e5dae 100644 --- a/modelopt/torch/export/hf_export_handlers.py +++ b/modelopt/torch/export/hf_export_handlers.py @@ -108,12 +108,13 @@ def _is_quant_moe_linear(module: nn.Module) -> bool: name gets uniquified — would bypass this handler and export without the input-amax fallback. The name check is kept as a fallback so stand-in modules match too. - The wrapper lives in the optional transformers plugin, hence the lazy import. + The wrapper lives in the Step PTQ modeling module, imported lazily like other export + modules do to avoid an import cycle with ``modelopt.torch.quantization``. """ if not hasattr(module, "experts"): return False try: - from modelopt.torch.quantization.plugins.huggingface import _QuantMoELinear + from modelopt.torch.models.step3p5.modeling_ptq import _QuantMoELinear if isinstance(module, _QuantMoELinear): return True diff --git a/modelopt/torch/export/layerwise_export.py b/modelopt/torch/export/layerwise_export.py index 8bd9e6cc3f8..92f31011ea3 100644 --- a/modelopt/torch/export/layerwise_export.py +++ b/modelopt/torch/export/layerwise_export.py @@ -286,8 +286,8 @@ def export_layer( a fusing format can rediscover which modules share an input; omit them only when nothing fuses. """ - # Local, as in every other export module: the plugin imports transformers. - from modelopt.torch.quantization.plugins.huggingface import _reconstruct_fused_moe_linear + # Local, as in every other export module: see pack_quantized_weights. + from modelopt.torch.models.step3p5.modeling_ptq import _reconstruct_fused_moe_linear assert self._bound, "export_layer() before bind()" assert not self._finalized, "export_layer() called after finalize()" diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 4c39674f585..f6116f4746a 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -1027,10 +1027,9 @@ def _prepare_model_for_export(model, dtype, is_modelopt_qlora): def pack_quantized_weights(model, dtype, is_modelopt_qlora: bool = False) -> None: """Quantize every module's weight in place, then rebuild the fused MoE linears.""" - # Deferred: modelopt.torch.quantization.plugins.huggingface imports transformers at module - # scope, and transformers is an optional extra -- importing it here keeps - # ``import modelopt.torch.export`` working without it. - from modelopt.torch.quantization.plugins.huggingface import _reconstruct_fused_moe_linear + # Deferred: modelopt.torch.quantization imports modelopt.torch.export, so importing the + # per-model PTQ module at module scope here would risk an import cycle. + from modelopt.torch.models.step3p5.modeling_ptq import _reconstruct_fused_moe_linear _process_quantized_modules(model, dtype, is_modelopt_qlora) _reconstruct_fused_moe_linear(model) diff --git a/modelopt/torch/export/unified_export_hf_streaming.py b/modelopt/torch/export/unified_export_hf_streaming.py index 7f692194fa6..9188bb5780e 100644 --- a/modelopt/torch/export/unified_export_hf_streaming.py +++ b/modelopt/torch/export/unified_export_hf_streaming.py @@ -346,9 +346,9 @@ def _export_transformers_checkpoint_streaming( NotImplementedError: if the model's conversion mapping regroups tensors. RuntimeError: if decoder layers cannot be discovered for layer-wise materialization. """ - # Deferred: the huggingface plugin imports transformers at module scope, and transformers - # is an optional extra -- keep ``import modelopt.torch.export`` working without it. - from modelopt.torch.quantization.plugins.huggingface import _reconstruct_fused_moe_linear + # Deferred: modelopt.torch.quantization imports modelopt.torch.export, so importing the + # per-model PTQ module at module scope here would risk an import cycle. + from modelopt.torch.models.step3p5.modeling_ptq import _reconstruct_fused_moe_linear export_dir = Path(export_dir) # Materialization dispatch walks the module tree from the root; without these maps each @@ -545,9 +545,9 @@ def collect_export_tensors( All the gathers finish before this returns, so the caller can write or postprocess without stalling anyone. Returning a list rather than a generator is what guarantees that. """ - # Deferred: the huggingface plugin imports transformers at module scope, and transformers - # is an optional extra -- keep ``import modelopt.torch.export`` working without it. - from modelopt.torch.quantization.plugins.huggingface import _reconstruct_fused_moe_linear + # Deferred: modelopt.torch.quantization imports modelopt.torch.export, so importing the + # per-model PTQ module at module scope here would risk an import cycle. + from modelopt.torch.models.step3p5.modeling_ptq import _reconstruct_fused_moe_linear _assert_fsdp2_owns_every_mesh_dim(model) my_rank, world = _dist.rank(), _dist.size() diff --git a/modelopt/torch/models/README.md b/modelopt/torch/models/README.md index 287fd7f7da1..c0caf7f6e01 100644 --- a/modelopt/torch/models/README.md +++ b/modelopt/torch/models/README.md @@ -53,18 +53,19 @@ an intra-family detail, not a package API. ## PTQ modeling Quantized-module wrappers and registrations that only one model needs (e.g. Llama4's fused -BMM experts) live in `/modeling_ptq.py`; generic +BMM experts, Step's expert-indexed `MoELinear`) live in `/modeling_ptq.py`; generic ones (fused/sequential MoE auto-detection, attention, `FP8Linear`) stay in `modelopt/torch/quantization/plugins/huggingface.py`. The HF plugin imports every `modeling_ptq` from an explicit list at its end, so add the new model type there. Leave the model's `__init__.py` importing only `specs`: `modeling_ptq` pulls in quantization and transformers, which `import modelopt.torch.models` must not. -A `modeling_ptq.py` registers on import — static classes via `QuantModuleRegistry.register` -behind a `try`/`except ImportError`, remote-code classes via a callback added to +A `modeling_ptq.py` registers on import — static classes via `QuantModuleRegistry.register`, +imported from transformers at module scope; remote-code classes via a callback added to `CUSTOM_MODEL_PLUGINS` — and exports nothing (`__all__ = []`). Every package still has a `specs.py`, even when its spec only records `model_type` and where the modeling code comes -from. +from. A family revision that shares another type's PTQ code (`step3p7` uses `step3p5`'s) still +gets a `modeling_ptq.py` that imports the shared one, so loading it by model type works. ## Sections diff --git a/modelopt/torch/models/step3p5/modeling_ptq.py b/modelopt/torch/models/step3p5/modeling_ptq.py new file mode 100644 index 00000000000..9e3d40d49b8 --- /dev/null +++ b/modelopt/torch/models/step3p5/modeling_ptq.py @@ -0,0 +1,264 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""Step-family PTQ modeling (HF model types ``step3p5``, ``step3p7``). + +The package is named for the first Step release; later revisions share its expert-indexed +``MoELinear`` layout and are matched by :func:`_is_step_family_model`. +""" + +import inspect +import re + +import torch +import torch.nn as nn +from torch import Tensor + +from modelopt.torch.models import hf_model_type +from modelopt.torch.quantization.nn import QuantModule, QuantModuleRegistry +from modelopt.torch.quantization.plugins.custom import CUSTOM_MODEL_PLUGINS + +__all__: list[str] = [] + + +class _QuantMoELinear(QuantModule): + """Quantization wrapper for expert-indexed MoELinear modules (fused expert weights). + + MoELinear has weight shape [num_experts, out_features, in_features] with + forward(x, expert_id). We expand it into per-expert nn.Linear modules so + each expert gets its own weight_quantizer and input_quantizer, calibrated + only on tokens actually routed to that expert. + + On export, _reconstruct_fused_moe_linear() stacks the per-expert quantized + weights and scales back into the original 3D format. + + Note: we use expansion-then-reconstruction rather than the add_module() approach + because vLLM requires stacked 3D scaling factors; per-expert expanded keys are + not accepted by the downstream serving engine. + """ + + def _setup(self): + # Deferred: accelerate is a separate optional dependency, and the HF plugin (which + # imports this module) must stay importable without it. + from accelerate import init_empty_weights + + # Accelerate's CPU/disk offload (`device_map="auto"`, `--offload_folder`) leaves + # `weight` as a meta tensor and keeps the real value in the module's offload hook, + # keyed on the original `weight` name. Expanding that would copy meta storage into + # every expert and then delete the key the hook restores into, silently producing a + # checkpoint of zeros. Refuse instead of corrupting. + if self.weight.is_meta or getattr(getattr(self, "_hf_hook", None), "offload", False): + raise NotImplementedError( + f"{type(self).__name__}: expert-indexed MoELinear weights cannot be quantized " + "while offloaded by Accelerate (the weight is a meta tensor whose value lives " + "in the offload hook). Load the model without CPU/disk offload — more GPUs, or " + "a device_map that keeps the MoE layers resident — and re-run." + ) + + dtype, device = self.weight.dtype, self.weight.device + + with init_empty_weights(): + experts = nn.ModuleList( + [ + nn.Linear(self.in_features, self.out_features, bias=False) + for _ in range(self.num_experts) + ] + ) + + for i in range(self.num_experts): + experts[i].to_empty(device=device) + with torch.no_grad(): + experts[i].weight.data = self.weight[i].detach().to(dtype=dtype, device=device) + + delattr(self, "weight") + self.experts = experts + + def forward(self, x, expert_id): + # experts[expert_id] is a _QuantLinear after quantization wrapping, providing + # per-expert input_quantizer and weight_quantizer. + # + # MoELinear.forward always promotes to fp32 for the matmul regardless of storage + # dtype (`F.linear(x.float(), self.weight[expert_id].float())`), so leaving the + # expert's weight at its native storage dtype (e.g. bf16) and downcasting x to + # match before the matmul would compute in bf16 and change the model's output even + # with every quantizer disabled -- the bf16 rounding this class exists to quantize + # *past*, not to reintroduce as a side effect of conversion. + # + # A prior version of this fix instead expanded every expert's weight in fp32 + # permanently in `_setup`. That reproduces Step's fp32 compute but turns a per-call + # transient promotion into persistent model state: on Step-3.7's full routed-expert + # set (42 layers x 3 projections x 288 experts x 4096 x 1280), doubling from bf16 to + # fp32 adds roughly 354 GiB held throughout calibration, on top of device placement + # already sized for bf16 -- a model that loaded successfully can then OOM. It also + # left disabled/unquantized experts reconstructed at fp32 in the exported checkpoint. + # + # Instead, only the one expert actually being called is promoted, transiently, for + # the duration of this one call -- matching Step's own per-call `.float()` memory + # profile instead of Step-3.7's full expert set. `expert.weight` is read here + # outside any `quantize_weight()` context, so `_get_quantized_weight` passes it + # through unchanged and this is the real underlying nn.Parameter (the same pattern + # `_setup` above uses), not a value computed by the quantizer -- so reassigning its + # `.data` genuinely mutates the persisted storage, not a transient wrapper. + # + # This must keep calling `expert(x)` (`__call__`, not `.forward()`) rather than + # reimplementing the input/weight-quantize/output-quantize sequence inline: some + # calibration algorithms (e.g. `local_hessian_calibrate`) register a + # `forward_pre_hook` directly on the quantized Linear module, which only fires + # through standard `nn.Module.__call__` dispatch. + expert = self.experts[expert_id] + original_weight = expert.weight.data + with torch.no_grad(): + expert.weight.data = original_weight.float() + try: + out = expert(x.float()) + finally: + with torch.no_grad(): + expert.weight.data = original_weight + return out.float() + + +def _is_expert_indexed_moe_linear(module: nn.Module) -> bool: + """Whether ``module`` packs one projection's experts into an expert-indexed 3-D weight. + + The Step family (``stepfun-ai/Step-3.5-Flash``, ``stepfun-ai/Step-3.7-Flash``) ships a + custom ``MoELinear`` via ``trust_remote_code``: a plain ``nn.Module`` holding a single + ``weight`` of shape ``[num_experts, out_features, in_features]``, whose + ``forward(x, expert_id)`` runs ``F.linear`` against the selected expert's slice. The + weights therefore live on the projection submodule rather than on the expert container, + which is what :func:`_fused_experts_wrapper_class` looks for, and the module is not an + ``nn.Linear``, so neither the fused-experts path nor the plain linear path claims it. + + Detection is structural rather than keyed on class names so new Step revisions are picked + up without another hardcoded name, but the shape alone is not a sufficient contract: the + replacement forward indexes ``self.experts[expert_id]``, so it only works for callers that + pass a **scalar expert index**. Grouped-GEMM MoE layers share the exact same 3-D weight and + attribute set while passing a per-expert token-count *tensor* instead (e.g. Moondream3's + ``MoeFusedLinear.forward(input, m_sizes)``, which would raise ``TypeError: only integer + tensors of a single element can be converted to an index`` on the first calibration + forward). The second parameter must therefore be named ``expert_id``, which is the + scalar-index contract both Step revisions declare. + """ + weight = getattr(module, "weight", None) + if not isinstance(weight, (nn.Parameter, Tensor)) or weight.dim() != 3: + return False + if not all(hasattr(module, attr) for attr in ("num_experts", "in_features", "out_features")): + return False + # The wrapper rebuilds the weight as `num_experts` Linears of (out_features, in_features), + # so a 3-D weight laid out any other way would silently copy the wrong slices. + if tuple(weight.shape) != (module.num_experts, module.out_features, module.in_features): + return False + try: + params = list(inspect.signature(type(module).forward).parameters.values())[1:] + except (TypeError, ValueError): + return False + # The replacement forward is exactly `(x, expert_id)`, so anything the caller could pass + # beyond those two — a keyword-only `router_state`, *args, **kwargs — would raise once + # converted. Require the signature to match what the wrapper can honour. + return ( + len(params) == 2 + and all(p.kind is p.POSITIONAL_OR_KEYWORD and p.default is p.empty for p in params) + and params[1].name == "expert_id" + ) + + +_STEP_FAMILY_RE = re.compile(r"(?i)^step\d") + + +def _is_step_family_model(model: nn.Module) -> bool: + """Whether ``model`` is a Step-family root model (Step-3.5, Step-3.7, or a future revision). + + Matched against the ``step`` convention shared by ``model_type`` (``"step3p5"``, + ``"step3p7"``) and the remote-code class name (``Step3p5ForCausalLM``, + ``Step3p7ForConditionalGeneration``), not an exact revision, so a new Step release is + still picked up without another hardcoded name. This is deliberately narrower than the + shape/signature check in :func:`_is_expert_indexed_moe_linear` alone: that check accepts + any module with a matching 3-D weight and an ``(x, expert_id)`` forward, which is a + coincidence risk on its own -- an unrelated architecture happening to reuse the parameter + name ``expert_id`` with different semantics (a per-expert bias or post-scale, say) would + be claimed and have that behavior silently dropped by the replacement wrapper. Gating on + the model family keeps the shape check doing what it is actually good at: telling + Step revisions apart without a class-name allowlist, rather than distinguishing Step + from arbitrary third-party MoE code. + """ + model_type = hf_model_type(model) or "" + return bool(_STEP_FAMILY_RE.match(model_type) or _STEP_FAMILY_RE.match(type(model).__name__)) + + +def register_moe_linear_on_the_fly(model): + """Register expert-indexed ``MoELinear`` modules (Step-3.5 / Step-3.7) for quantization. + + Without this the routed experts carry no quantizer at all: an experts-only recipe matches + nothing and the export writes a checkpoint with ``quant_algo: null``. + """ + if not _is_step_family_model(model): + return + visited_types = set() + for name, module in model.named_modules(): + mod_type = type(module) + if mod_type in visited_types or QuantModuleRegistry.get(mod_type) is not None: + continue + visited_types.add(mod_type) + + if _is_expert_indexed_moe_linear(module): + print( + f"\033[1mDetected expert-indexed MoE linear '{name}' of type " + f"{mod_type.__name__}, registering with _QuantMoELinear.\033[0m" + ) + QuantModuleRegistry.register({mod_type: f"hf.{mod_type.__name__}"})(_QuantMoELinear) + + +def _reconstruct_fused_moe_linear(model: nn.Module) -> None: + """Reconstruct :class:`_QuantMoELinear` per-expert weights back to the 3-D MoELinear format. + + After _process_quantized_modules, each expert's nn.Linear inside the wrapper has: + - weight: fp4-quantized tensor [out_features, in_features] + - weight_scale, weight_scale_2: per-block / global scales + - input_scale: activation scale (if calibrated) + + This stacks them back into the original MoELinear layout so the exported state_dict + uses the original key names (e.g. moe.up_proj.weight with shape [N, out, in]). + + Matched by wrapper type rather than by the dynamically generated class name (``Quant`` + + the model's own class name): a model whose class is not spelled ``MoELinear`` would + otherwise quantize normally but export unusable per-expert keys. + """ + for _name, module in model.named_modules(): + if not isinstance(module, _QuantMoELinear): + continue + + n = module.num_experts + experts = module.experts + + # Reconstruct 3D weight: [num_experts, out_features, in_features] + module.weight = nn.Parameter( + torch.stack([experts[i].weight.data for i in range(n)]), + requires_grad=False, + ) + + # Stack per-expert scales back under the original attribute names. + # Check all experts: some may lack input_scale if they were never routed + # during calibration, so only stack when every expert has the attribute. + for attr in ("weight_scale", "weight_scale_2", "input_scale"): + if all(hasattr(experts[i], attr) for i in range(n)): + module.register_buffer( + attr, + torch.stack([getattr(experts[i], attr) for i in range(n)]), + ) + + # Remove expanded experts — the reconstructed 3D tensors replace them + del module.experts + + +CUSTOM_MODEL_PLUGINS.add(register_moe_linear_on_the_fly) diff --git a/modelopt/torch/models/step3p7/modeling_ptq.py b/modelopt/torch/models/step3p7/modeling_ptq.py new file mode 100644 index 00000000000..e43c5237278 --- /dev/null +++ b/modelopt/torch/models/step3p7/modeling_ptq.py @@ -0,0 +1,26 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""Step-3.7 PTQ modeling (HF model type ``step3p7``). + +Step-3.7 shares Step-3.5's expert-indexed ``MoELinear`` layout, so its PTQ support is the +Step-family code in ``step3p5/modeling_ptq.py``, which matches every Step revision. This +module imports it, so that loading the PTQ modeling for model type ``step3p7`` by name +registers that support too. +""" + +from ..step3p5 import modeling_ptq as _step3p5_modeling_ptq # noqa: F401 + +__all__: list[str] = [] diff --git a/modelopt/torch/models/step3p7/specs.py b/modelopt/torch/models/step3p7/specs.py index f38c0339989..74db68aed4b 100644 --- a/modelopt/torch/models/step3p7/specs.py +++ b/modelopt/torch/models/step3p7/specs.py @@ -15,8 +15,8 @@ """Step-3.7 specs (HF model type ``step3p7``). -Shares Step-3.5's architecture; its PTQ modeling (in the HF quantization plugin) -matches every Step revision. +Shares Step-3.5's architecture; its PTQ modeling lives in ``step3p5/modeling_ptq.py``, +which matches every Step revision. """ from ..specs import ModelSpec, register diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index 05e589835ac..4110a95e5bd 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -18,7 +18,6 @@ import importlib import inspect import logging -import re import warnings from contextlib import contextmanager from functools import partial @@ -38,7 +37,6 @@ validate_triton_attention_envelope, ) from modelopt.torch.kernels.quantization.gemm import IS_AVAILABLE as IS_TRITON_AVAILABLE -from modelopt.torch.models import hf_model_type from modelopt.torch.opt.dynamic import DynamicModule from modelopt.torch.utils.distributed import ParallelState @@ -1484,6 +1482,8 @@ def _is_param_grad_enabled_for_auto_quantize(pname, model): "gpt_oss", "llama4", "nemotron_h", + "step3p5", + "step3p7", ): importlib.import_module(f"modelopt.torch.models.{_model_type}.modeling_ptq") @@ -1492,235 +1492,8 @@ def _is_param_grad_enabled_for_auto_quantize(pname, model): ) -class _QuantMoELinear(QuantModule): - """Quantization wrapper for expert-indexed MoELinear modules (fused expert weights). - - MoELinear has weight shape [num_experts, out_features, in_features] with - forward(x, expert_id). We expand it into per-expert nn.Linear modules so - each expert gets its own weight_quantizer and input_quantizer, calibrated - only on tokens actually routed to that expert. - - On export, _reconstruct_fused_moe_linear() stacks the per-expert quantized - weights and scales back into the original 3D format. - - Note: we use expansion-then-reconstruction rather than the add_module() approach - because vLLM requires stacked 3D scaling factors; per-expert expanded keys are - not accepted by the downstream serving engine. - """ - - def _setup(self): - from accelerate import init_empty_weights - - # Accelerate's CPU/disk offload (`device_map="auto"`, `--offload_folder`) leaves - # `weight` as a meta tensor and keeps the real value in the module's offload hook, - # keyed on the original `weight` name. Expanding that would copy meta storage into - # every expert and then delete the key the hook restores into, silently producing a - # checkpoint of zeros. Refuse instead of corrupting. - if self.weight.is_meta or getattr(getattr(self, "_hf_hook", None), "offload", False): - raise NotImplementedError( - f"{type(self).__name__}: expert-indexed MoELinear weights cannot be quantized " - "while offloaded by Accelerate (the weight is a meta tensor whose value lives " - "in the offload hook). Load the model without CPU/disk offload — more GPUs, or " - "a device_map that keeps the MoE layers resident — and re-run." - ) - - dtype, device = self.weight.dtype, self.weight.device - - with init_empty_weights(): - experts = nn.ModuleList( - [ - nn.Linear(self.in_features, self.out_features, bias=False) - for _ in range(self.num_experts) - ] - ) - - for i in range(self.num_experts): - experts[i].to_empty(device=device) - with torch.no_grad(): - experts[i].weight.data = self.weight[i].detach().to(dtype=dtype, device=device) - - delattr(self, "weight") - self.experts = experts - - def forward(self, x, expert_id): - # experts[expert_id] is a _QuantLinear after quantization wrapping, providing - # per-expert input_quantizer and weight_quantizer. - # - # MoELinear.forward always promotes to fp32 for the matmul regardless of storage - # dtype (`F.linear(x.float(), self.weight[expert_id].float())`), so leaving the - # expert's weight at its native storage dtype (e.g. bf16) and downcasting x to - # match before the matmul would compute in bf16 and change the model's output even - # with every quantizer disabled -- the bf16 rounding this class exists to quantize - # *past*, not to reintroduce as a side effect of conversion. - # - # A prior version of this fix instead expanded every expert's weight in fp32 - # permanently in `_setup`. That reproduces Step's fp32 compute but turns a per-call - # transient promotion into persistent model state: on Step-3.7's full routed-expert - # set (42 layers x 3 projections x 288 experts x 4096 x 1280), doubling from bf16 to - # fp32 adds roughly 354 GiB held throughout calibration, on top of device placement - # already sized for bf16 -- a model that loaded successfully can then OOM. It also - # left disabled/unquantized experts reconstructed at fp32 in the exported checkpoint. - # - # Instead, only the one expert actually being called is promoted, transiently, for - # the duration of this one call -- matching Step's own per-call `.float()` memory - # profile instead of Step-3.7's full expert set. `expert.weight` is read here - # outside any `quantize_weight()` context, so `_get_quantized_weight` passes it - # through unchanged and this is the real underlying nn.Parameter (the same pattern - # `_setup` above uses), not a value computed by the quantizer -- so reassigning its - # `.data` genuinely mutates the persisted storage, not a transient wrapper. - # - # This must keep calling `expert(x)` (`__call__`, not `.forward()`) rather than - # reimplementing the input/weight-quantize/output-quantize sequence inline: some - # calibration algorithms (e.g. `local_hessian_calibrate`) register a - # `forward_pre_hook` directly on the quantized Linear module, which only fires - # through standard `nn.Module.__call__` dispatch. - expert = self.experts[expert_id] - original_weight = expert.weight.data - with torch.no_grad(): - expert.weight.data = original_weight.float() - try: - out = expert(x.float()) - finally: - with torch.no_grad(): - expert.weight.data = original_weight - return out.float() - - -def _is_expert_indexed_moe_linear(module: nn.Module) -> bool: - """Whether ``module`` packs one projection's experts into an expert-indexed 3-D weight. - - The Step family (``stepfun-ai/Step-3.5-Flash``, ``stepfun-ai/Step-3.7-Flash``) ships a - custom ``MoELinear`` via ``trust_remote_code``: a plain ``nn.Module`` holding a single - ``weight`` of shape ``[num_experts, out_features, in_features]``, whose - ``forward(x, expert_id)`` runs ``F.linear`` against the selected expert's slice. The - weights therefore live on the projection submodule rather than on the expert container, - which is what :func:`_fused_experts_wrapper_class` looks for, and the module is not an - ``nn.Linear``, so neither the fused-experts path nor the plain linear path claims it. - - Detection is structural rather than keyed on class names so new Step revisions are picked - up without another hardcoded name, but the shape alone is not a sufficient contract: the - replacement forward indexes ``self.experts[expert_id]``, so it only works for callers that - pass a **scalar expert index**. Grouped-GEMM MoE layers share the exact same 3-D weight and - attribute set while passing a per-expert token-count *tensor* instead (e.g. Moondream3's - ``MoeFusedLinear.forward(input, m_sizes)``, which would raise ``TypeError: only integer - tensors of a single element can be converted to an index`` on the first calibration - forward). The second parameter must therefore be named ``expert_id``, which is the - scalar-index contract both Step revisions declare. - """ - weight = getattr(module, "weight", None) - if not isinstance(weight, (nn.Parameter, Tensor)) or weight.dim() != 3: - return False - if not all(hasattr(module, attr) for attr in ("num_experts", "in_features", "out_features")): - return False - # The wrapper rebuilds the weight as `num_experts` Linears of (out_features, in_features), - # so a 3-D weight laid out any other way would silently copy the wrong slices. - if tuple(weight.shape) != (module.num_experts, module.out_features, module.in_features): - return False - try: - params = list(inspect.signature(type(module).forward).parameters.values())[1:] - except (TypeError, ValueError): - return False - # The replacement forward is exactly `(x, expert_id)`, so anything the caller could pass - # beyond those two — a keyword-only `router_state`, *args, **kwargs — would raise once - # converted. Require the signature to match what the wrapper can honour. - return ( - len(params) == 2 - and all(p.kind is p.POSITIONAL_OR_KEYWORD and p.default is p.empty for p in params) - and params[1].name == "expert_id" - ) - - -_STEP_FAMILY_RE = re.compile(r"(?i)^step\d") - - -def _is_step_family_model(model: nn.Module) -> bool: - """Whether ``model`` is a Step-family root model (Step-3.5, Step-3.7, or a future revision). - - Matched against the ``step`` convention shared by ``model_type`` (``"step3p5"``, - ``"step3p7"``) and the remote-code class name (``Step3p5ForCausalLM``, - ``Step3p7ForConditionalGeneration``), not an exact revision, so a new Step release is - still picked up without another hardcoded name. This is deliberately narrower than the - shape/signature check in :func:`_is_expert_indexed_moe_linear` alone: that check accepts - any module with a matching 3-D weight and an ``(x, expert_id)`` forward, which is a - coincidence risk on its own -- an unrelated architecture happening to reuse the parameter - name ``expert_id`` with different semantics (a per-expert bias or post-scale, say) would - be claimed and have that behavior silently dropped by the replacement wrapper. Gating on - the model family keeps the shape check doing what it is actually good at: telling - Step revisions apart without a class-name allowlist, rather than distinguishing Step - from arbitrary third-party MoE code. - """ - model_type = hf_model_type(model) or "" - return bool(_STEP_FAMILY_RE.match(model_type) or _STEP_FAMILY_RE.match(type(model).__name__)) - - -def register_moe_linear_on_the_fly(model): - """Register expert-indexed ``MoELinear`` modules (Step-3.5 / Step-3.7) for quantization. - - Without this the routed experts carry no quantizer at all: an experts-only recipe matches - nothing and the export writes a checkpoint with ``quant_algo: null``. - """ - if not _is_step_family_model(model): - return - visited_types = set() - for name, module in model.named_modules(): - mod_type = type(module) - if mod_type in visited_types or QuantModuleRegistry.get(mod_type) is not None: - continue - visited_types.add(mod_type) - - if _is_expert_indexed_moe_linear(module): - print( - f"\033[1mDetected expert-indexed MoE linear '{name}' of type " - f"{mod_type.__name__}, registering with _QuantMoELinear.\033[0m" - ) - QuantModuleRegistry.register({mod_type: f"hf.{mod_type.__name__}"})(_QuantMoELinear) - - -def _reconstruct_fused_moe_linear(model: nn.Module) -> None: - """Reconstruct :class:`_QuantMoELinear` per-expert weights back to the 3-D MoELinear format. - - After _process_quantized_modules, each expert's nn.Linear inside the wrapper has: - - weight: fp4-quantized tensor [out_features, in_features] - - weight_scale, weight_scale_2: per-block / global scales - - input_scale: activation scale (if calibrated) - - This stacks them back into the original MoELinear layout so the exported state_dict - uses the original key names (e.g. moe.up_proj.weight with shape [N, out, in]). - - Matched by wrapper type rather than by the dynamically generated class name (``Quant`` + - the model's own class name): a model whose class is not spelled ``MoELinear`` would - otherwise quantize normally but export unusable per-expert keys. - """ - for _name, module in model.named_modules(): - if not isinstance(module, _QuantMoELinear): - continue - - n = module.num_experts - experts = module.experts - - # Reconstruct 3D weight: [num_experts, out_features, in_features] - module.weight = nn.Parameter( - torch.stack([experts[i].weight.data for i in range(n)]), - requires_grad=False, - ) - - # Stack per-expert scales back under the original attribute names. - # Check all experts: some may lack input_scale if they were never routed - # during calibration, so only stack when every expert has the attribute. - for attr in ("weight_scale", "weight_scale_2", "input_scale"): - if all(hasattr(experts[i], attr) for i in range(n)): - module.register_buffer( - attr, - torch.stack([getattr(experts[i], attr) for i in range(n)]), - ) - - # Remove expanded experts — the reconstructed 3D tensors replace them - del module.experts - - CUSTOM_MODEL_PLUGINS.update( [ - register_moe_linear_on_the_fly, register_fused_experts_on_the_fly, force_eager_experts_impl_on_the_fly, register_sparse_moe_on_the_fly, diff --git a/plugins/modelopt/skills/ptq/references/checkpoint-validation.md b/plugins/modelopt/skills/ptq/references/checkpoint-validation.md index 3d755933134..c836681147f 100644 --- a/plugins/modelopt/skills/ptq/references/checkpoint-validation.md +++ b/plugins/modelopt/skills/ptq/references/checkpoint-validation.md @@ -201,6 +201,6 @@ Layers silently skipped because the quantization config patterns don't match the ## What to do when warnings appear -- **Layers should have been quantized** (e.g., MoE experts with `nvfp4_mlp_only`): the quantization config patterns missed them. Fix by adding the missing pattern to the config and re-running PTQ. Check if ModelOpt already has a plugin for the model in `modelopt/torch/quantization/plugins/huggingface.py`. +- **Layers should have been quantized** (e.g., MoE experts with `nvfp4_mlp_only`): the quantization config patterns missed them. Fix by adding the missing pattern to the config and re-running PTQ. Check if ModelOpt already has a plugin for the model in `modelopt/torch/models//modeling_ptq.py`. - **Layers are intentionally unquantized** (e.g., attention layers with `nvfp4_mlp_only`): they should be in the `exclude_modules` list but the export didn't add them. Add them manually to both `hf_quant_config.json` and `config.json` `quantization_config.ignore` in the checkpoint to prevent deployment failures. diff --git a/plugins/modelopt/skills/ptq/references/unsupported-models.md b/plugins/modelopt/skills/ptq/references/unsupported-models.md index 4ed7959c71b..09514ebf4be 100644 --- a/plugins/modelopt/skills/ptq/references/unsupported-models.md +++ b/plugins/modelopt/skills/ptq/references/unsupported-models.md @@ -66,7 +66,7 @@ Manual dequantization is only needed for **non-standard parameter names** (e.g., Read the model source to identify how weights are stored. **If all linear layers are plain `nn.Linear`, no custom code is needed** — ModelOpt quantizes them automatically. -**For HuggingFace models**, check `modelopt/torch/quantization/plugins/huggingface.py` first — it already registers patches for common non-standard modules (`Llama4TextExperts`, `FP8Linear`, `FalconLinear`, `Conv1D`, `Qwen3_5MoeExperts`, etc.). If your model's non-standard class is already registered there, no extra code is needed. +**For HuggingFace models**, check `modelopt/torch/models//modeling_ptq.py` (model-specific patches, e.g. `Llama4TextExperts`, `FalconLinear`) and `modelopt/torch/quantization/plugins/huggingface.py` (generic ones: `FP8Linear`, `Conv1D`, auto-detected MoE such as `Qwen3_5MoeExperts`) first. If your model's non-standard class is already registered, no extra code is needed. Custom patches are required when: @@ -94,7 +94,7 @@ Compare against the `enable`/`disable` patterns in the config. Add custom overri After Steps A-D: - **No patches needed** (all standard modules) → run `hf_ptq.py` with a smoke test (`--calib_size 4`). If it succeeds, proceed with full calibration. If it fails, read the error and revisit Steps C/D. -- **Patches needed** → patch ModelOpt directly using the patterns below (add `QuantModule` in `modelopt/torch/quantization/plugins/huggingface.py`, update `modelopt/torch/export/` if needed), then run `hf_ptq.py` with a smoke test. This is preferred over writing a standalone script because it reuses all existing `hf_ptq.py` logic. Debug failures iteratively — quantization errors often reveal additional modules that need patching. +- **Patches needed** → patch ModelOpt directly using the patterns below (add `QuantModule` in `modelopt/torch/models//modeling_ptq.py`, update `modelopt/torch/export/` if needed), then run `hf_ptq.py` with a smoke test. This is preferred over writing a standalone script because it reuses all existing `hf_ptq.py` logic. Debug failures iteratively — quantization errors often reveal additional modules that need patching. --- @@ -138,7 +138,7 @@ class QuantCustomModule(OriginalModule): - **Fused experts** (`gate_up_proj` + `down_proj` 3D tensors, all native transformers MoEs) → auto-detected by `register_fused_experts_on_the_fly`, handled by `_QuantFusedExperts`. Covers Mixtral, Qwen, DeepSeek, Jamba, OlMoE, etc. - **Sequential per-expert `nn.Linear`** with `gate` + `experts` (e.g. remote-code models) → auto-detected by `register_sparse_moe_on_the_fly`. -**Custom MoE** (non-standard layout not matching auto-detection) requires patching. Find the closest pattern in the plugin (`modelopt/torch/quantization/plugins/huggingface.py`): +**Custom MoE** (non-standard layout not matching auto-detection) requires patching. Find the closest pattern in `modelopt/torch/models/*/modeling_ptq.py`: | MoE design | Strategy | Plugin example | | --- | --- | --- | @@ -163,7 +163,7 @@ for name, module in model.named_modules(): ## Pattern 3: Registering with ModelOpt -**When patching the plugin directly** (preferred): Use `QuantModuleRegistry.register` in `modelopt/torch/quantization/plugins/huggingface.py`, following existing examples: +**When patching ModelOpt directly** (preferred): Use `QuantModuleRegistry.register` in `modelopt/torch/models//modeling_ptq.py` and add `` to the `modeling_ptq` import loop at the end of `modelopt/torch/quantization/plugins/huggingface.py`, following existing examples: ```python from modelopt.torch.quantization.nn import QuantModuleRegistry @@ -179,6 +179,8 @@ def register_my_model_on_the_fly(model): if QuantModuleRegistry.get(mod_type) is None: QuantModuleRegistry.register({mod_type: f"hf.{mod_type.__name__}"})(QuantCustomModule) break + +CUSTOM_MODEL_PLUGINS.add(register_my_model_on_the_fly) # from ...quantization.plugins.custom ``` **When writing a standalone script** (fallback): Use `mtq.register()`: diff --git a/tests/unit/torch/models/test_modeling_ptq_registration.py b/tests/unit/torch/models/test_modeling_ptq_registration.py index fbaa761d38d..6782b64c344 100644 --- a/tests/unit/torch/models/test_modeling_ptq_registration.py +++ b/tests/unit/torch/models/test_modeling_ptq_registration.py @@ -35,6 +35,8 @@ "modelopt.torch.models.gpt_oss.modeling_ptq", "modelopt.torch.models.llama4.modeling_ptq", "modelopt.torch.models.nemotron_h.modeling_ptq", + "modelopt.torch.models.step3p5.modeling_ptq", + "modelopt.torch.models.step3p7.modeling_ptq", "modelopt.torch.quantization", "modelopt.torch.export", ] @@ -48,6 +50,7 @@ from modelopt.torch.models.falcon.modeling_ptq import register_falcon_linears_on_the_fly from modelopt.torch.models.nemotron_h.modeling_ptq import is_nemotron_h_model +from modelopt.torch.models.step3p5.modeling_ptq import register_moe_linear_on_the_fly from modelopt.torch.quantization.nn import QuantModuleRegistry from modelopt.torch.quantization.plugins.custom import CUSTOM_MODEL_PLUGINS from modelopt.torch.quantization.plugins.huggingface import is_homogeneous_hf_model @@ -74,6 +77,7 @@ assert GptOssExperts in QuantModuleRegistry, "GptOssExperts not registered" assert register_falcon_linears_on_the_fly in CUSTOM_MODEL_PLUGINS, "Falcon callback missing" +assert register_moe_linear_on_the_fly in CUSTOM_MODEL_PLUGINS, "Step callback missing" # The first matching decoder discoverer wins, so Nemotron-H's must precede the generic one. predicates = [is_supported for is_supported, _ in LayerActivationCollector._decoder_layer_support] diff --git a/tests/unit/torch/quantization/plugins/test_moe_linear.py b/tests/unit/torch/models/test_step3p5_modeling_ptq.py similarity index 99% rename from tests/unit/torch/quantization/plugins/test_moe_linear.py rename to tests/unit/torch/models/test_step3p5_modeling_ptq.py index cceaacd44d5..66b52325c24 100644 --- a/tests/unit/torch/quantization/plugins/test_moe_linear.py +++ b/tests/unit/torch/models/test_step3p5_modeling_ptq.py @@ -27,13 +27,13 @@ import modelopt.torch.quantization as mtq from modelopt.torch.export.hf_export_handlers import _export_moe_linear from modelopt.torch.export.registry import ExportContext, ExportModuleRegistry -from modelopt.torch.quantization.nn import QuantModuleRegistry -from modelopt.torch.quantization.plugins.huggingface import ( +from modelopt.torch.models.step3p5.modeling_ptq import ( _is_expert_indexed_moe_linear, _QuantMoELinear, _reconstruct_fused_moe_linear, register_moe_linear_on_the_fly, ) +from modelopt.torch.quantization.nn import QuantModuleRegistry NUM_EXPERTS = 4 HIDDEN_SIZE = 32 From 836a2239370cc0515bf66b5fba74c66c5c231273 Mon Sep 17 00:00:00 2001 From: Shengliang Xu <106840466+shengliangxu@users.noreply.github.com> Date: Fri, 9 Oct 2026 15:38:33 -0700 Subject: [PATCH 06/35] Fix the GLM-5.1 alias: it is not the input_scale1 recipe (#2711) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? **Type of change:** Bug fix (checkpoint alias pointed at the wrong recipe) `models/zai-org/GLM-5.1/` aliases `general/ptq/nvfp4_experts_only_input_scale1-kv_fp8_cast`. It should alias the plain `general/ptq/nvfp4_experts_only-kv_fp8_cast`. The `input_scale1` recipe pins the expert activation amax to a constant 2688.0 (`E2M1_MAX * E4M3_MAX`), so every exported expert `input_scale` is exactly 1.0 and no activation statistics are collected. That is a visible property of the published checkpoint, and GLM-5.1 does not have it: | Release | Expert `input_scale`, 8 randomly sampled | | |---|---|---| | `nvidia/GLM-5.2-NVFP4` | 8 of 8 exactly `1.0` | matches `input_scale1` | | `nvidia/GLM-5.1-NVFP4` | 0 of 8 equal to 1.0, spanning `3.4e-4` … `1.3e-2` | ordinary calibration | So GLM-5.1 used ordinary activation calibration and takes the plain experts-only recipe. **GLM-5.2 is correct and is not touched.** Scope and KV mode were re-checked against the release and are unchanged — routed experts only, FP8 KV cache in cast mode. Only the activation handling was wrong. Verified after the change that the alias now resolves identically to its base, that the pinned `2688.0` amax is gone from the resolved config, and that GLM-5.2 still carries it. #### How this happened The recipes were collected by asking the engineer behind each release. GLM-5.1 and GLM-5.2 were answered together in one line naming `input_scale1`. That was recorded for both without checking the one property the recipe is named for. #2643 merged it. #### One thing left open **Max versus MSE is not recoverable from a checkpoint** — calibration leaves no trace in the exported tensors, which is why the sampling above can settle `input_scale1` but not this. The `max` variant here reflects the general default, not a confirmed fact. If GLM-5.1 used MSE, the follow-up is a one-line change of the import to `general/ptq/nvfp4_experts_only_mse-kv_fp8_cast`; scope, KV mode and activation handling all stay as they are in this PR. This PR is still strictly closer to the release than what is on `main` today, which misstates activation handling outright. ### Usage ```bash python examples/hf_ptq/hf_ptq.py \ --pyt_ckpt_path zai-org/GLM-5.1 \ --recipe models/zai-org/GLM-5.1/ptq/nvfp4_experts_only-kv_fp8_cast \ --export_path ``` ### Testing - `tests/unit/recipe/` — 461 passed. - Verified the new path resolves to a config identical to `general/ptq/nvfp4_experts_only-kv_fp8_cast`, that `constant_amax: 2688.0` no longer appears in it, that the old path now raises `ValueError`, and that GLM-5.2 still resolves with the pinned amax. - `pre-commit` passes on both changed files. - The 14 failures in the local run are pre-existing and environment-only — a broken `transformer_engine` `.so` in the dev venv; they reproduce identically on unmodified `main`. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ⚠️ — `--recipe models/zai-org/GLM-5.1/ptq/nvfp4_experts_only_input_scale1-kv_fp8_cast` stops resolving and becomes `.../nvfp4_experts_only-kv_fp8_cast`. That is the fix: the old path named a recipe that does not produce this release. - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: ✅ — no new dependencies, no source changes. - Did you write any new necessary tests?: ✅ — covered by the shipped-recipe sweep in `tests/unit/recipe/test_loader.py`. - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ❌ — corrects an alias added in the same unreleased cycle (#2643), which `CONTRIBUTING.md` says to skip. Happy to add one if you would rather it be recorded. - Did you get Claude approval on this PR?: ❌ — not yet run. 🤖 Generated with [Claude Code](https://claude.com/claude-code) ## Summary by CodeRabbit * **Documentation** * Clarified that GLM-5.1 uses calibrated expert activation scales, while GLM-5.2 pins them to 1.0. * Updated the GLM-5.1 recipe description to note max calibration and an FP8 KV cache in cast mode. Signed-off-by: Shengliang Xu Co-authored-by: Claude Opus 5 (1M context) Signed-off-by: Rohan Joshi --- ...ml => nvfp4_experts_only-kv_fp8_cast.yaml} | 19 ++++++++++++------- modelopt_recipes/ptq.md | 10 +++++++--- 2 files changed, 19 insertions(+), 10 deletions(-) rename modelopt_recipes/models/zai-org/GLM-5.1/ptq/{nvfp4_experts_only_input_scale1-kv_fp8_cast.yaml => nvfp4_experts_only-kv_fp8_cast.yaml} (54%) diff --git a/modelopt_recipes/models/zai-org/GLM-5.1/ptq/nvfp4_experts_only_input_scale1-kv_fp8_cast.yaml b/modelopt_recipes/models/zai-org/GLM-5.1/ptq/nvfp4_experts_only-kv_fp8_cast.yaml similarity index 54% rename from modelopt_recipes/models/zai-org/GLM-5.1/ptq/nvfp4_experts_only_input_scale1-kv_fp8_cast.yaml rename to modelopt_recipes/models/zai-org/GLM-5.1/ptq/nvfp4_experts_only-kv_fp8_cast.yaml index 876a47584e0..dbb98cac5f5 100644 --- a/modelopt_recipes/models/zai-org/GLM-5.1/ptq/nvfp4_experts_only_input_scale1-kv_fp8_cast.yaml +++ b/modelopt_recipes/models/zai-org/GLM-5.1/ptq/nvfp4_experts_only-kv_fp8_cast.yaml @@ -13,16 +13,21 @@ # See the License for the specific language governing permissions and # limitations under the License. -# Alias recipe for nvidia/GLM-5.1-NVFP4 (source zai-org/GLM-5.1). Same recipe as `general/ptq/nvfp4_experts_only_input_scale1-kv_fp8_cast`: NVFP4 on the -# routed experts with the expert `input_scale` pinned to 1.0, and an FP8 KV cache in cast -# mode. +# Alias recipe for nvidia/GLM-5.1-NVFP4 (source zai-org/GLM-5.1). Same recipe as +# `general/ptq/nvfp4_experts_only-kv_fp8_cast`: NVFP4 on the routed experts with calibrated +# activation scales, and an FP8 KV cache in cast mode. +# +# Not the `input_scale1` variant, despite the sibling GLM-5.2 entry using it. That recipe +# pins the expert activation amax to a constant 2688.0 (E2M1_MAX * E4M3_MAX) so every +# exported expert `input_scale` is exactly 1.0. GLM-5.2's are; GLM-5.1's are calibrated +# values (sampled 0 of 8 equal to 1.0, spanning 3.4e-4 to 1.3e-2), so this release used +# ordinary activation calibration. imports: - base: general/ptq/nvfp4_experts_only_input_scale1-kv_fp8_cast + base: general/ptq/nvfp4_experts_only-kv_fp8_cast $import: base metadata: description: >- - zai-org/GLM-5.1 quantized with the general expert-only NVFP4 scheme, expert - `input_scale` pinned to 1.0, and an FP8 KV cache in cast mode, as published in - nvidia/GLM-5.1-NVFP4. + zai-org/GLM-5.1 quantized with the general expert-only NVFP4 scheme (max calibration) + and an FP8 KV cache in cast mode, as published in nvidia/GLM-5.1-NVFP4. diff --git a/modelopt_recipes/ptq.md b/modelopt_recipes/ptq.md index 21e5fe5744b..d0db470af53 100644 --- a/modelopt_recipes/ptq.md +++ b/modelopt_recipes/ptq.md @@ -572,11 +572,15 @@ entry is a thin **alias** that imports that recipe wholesale and overrides only `general/ptq/nvfp4_experts_only-kv_fp8_cast` — expert-only NVFP4 with max calibration and an FP8 KV cache in cast mode — as published in `nvidia/Gemma-4-26B-A4B-NVFP4` and `nvidia/diffusiongemma-26B-A4B-it-NVFP4`. -- **`models/zai-org/GLM-5.1/ptq/nvfp4_experts_only_input_scale1-kv_fp8_cast`** and - **`models/zai-org/GLM-5.2/ptq/...`** alias +- **`models/zai-org/GLM-5.2/ptq/nvfp4_experts_only_input_scale1-kv_fp8_cast`** aliases `general/ptq/nvfp4_experts_only_input_scale1-kv_fp8_cast` — expert-only NVFP4 with the expert `input_scale` pinned to 1.0 and an FP8 KV cache in cast mode — as published in - `nvidia/GLM-5.1-NVFP4` and `nvidia/GLM-5.2-NVFP4`. + `nvidia/GLM-5.2-NVFP4`. +- **`models/zai-org/GLM-5.1/ptq/nvfp4_experts_only-kv_fp8_cast`** aliases + `general/ptq/nvfp4_experts_only-kv_fp8_cast` — expert-only NVFP4 with calibrated + activation scales and an FP8 KV cache in cast mode — as published in + `nvidia/GLM-5.1-NVFP4`. Unlike GLM-5.2, this release's exported expert `input_scale` + tensors are calibrated rather than pinned to 1.0. - **`models/MiniMaxAI/MiniMax-M2.5/ptq/nvfp4_experts_only-kv_fp8_cast`** aliases `general/ptq/nvfp4_experts_only-kv_fp8_cast` — expert-only NVFP4 with an FP8 KV cache in cast mode — as published in `nvidia/MiniMax-M2.5-NVFP4`. From e70cbfa3956c195e4185cfe12b6dd8790c863b4e Mon Sep 17 00:00:00 2001 From: Chenjie Luo <108829653+cjluo-nv@users.noreply.github.com> Date: Mon, 12 Oct 2026 16:55:42 -0700 Subject: [PATCH 07/35] Fix the FSDP2 export deadlock at world > 1 (#2731) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Type of change: bug fix **HF export with FSDP2 deadlocks at 2 or more ranks.** It hangs with rank 0 in `collect_export_tensors` and every other rank at the barrier after it. This hits both the streaming export (`owner="share"`) and the gather-to-rank-0 export (`owner="rank0"`). In practice that's `hf_ptq --use_fsdp2` with more than one rank, and the nightly 2-GPU `gpu` job, which has hit its 60-minute limit every night for at least a week: `test_fsdp2_streaming_export_matches_reference[nvfp4]` and `[fp8]` each time out at 120s. PR CI runs GPU tests on a 1-GPU runner, where these tests skip, so it never saw this. **Root cause:** 1. A forward pass leaves the root FSDP module's parameters unsharded, because FSDP2 skips the root's post-forward reshard. Calibration and the export's own dummy forward both do this. 2. FSDP2 registers a `state_dict()` pre-hook on every module that holds FSDP parameters, and the hook reshards that module's whole parameter group. 3. In `collect_export_tensors`, only a unit's owner calls `state_dict()`. So the owner of the non-layer unit (`embed_tokens`, `norm`, `lm_head`) reshards the root on its rank alone. 4. At the next module, the gather window all-gathers a parameter on that rank which the other ranks still hold as a plain tensor. The others never join the collective. A per-rank trace of every collective shows the two ranks agreeing up to that one extra all-gather. **Fix:** before the loop, every rank reshards every FSDP module. That's a local operation, so the owner's pre-hook becomes a no-op. ### Usage No API change. ### Testing - **New unit test:** `test_streaming_export_two_cpu_ranks` in `tests/unit/torch/export/test_fsdp2_parallel_export.py`. It shards a tiny Llama over two CPU ranks with Gloo, calibrates it (which leaves the root unsharded), runs the streaming export, and compares the merged checkpoint with a single-process export. - It runs in PR CI, which so far had no multi-rank FSDP2 export coverage. - Without the fix it deadlocks. The worker dumps its stack and exits after 45s, so the test fails in about 50s instead of hanging the job. - With the fix it passes in about 10s on torch 2.12 and torch 2.8 (the oldest in the CI matrix). - **GPU tests:** all 38 of `tests/gpu/torch/export/test_fsdp2_export.py` pass in CI's image (`nvcr.io/nvidia/pytorch:26.05-py3`) on 2× RTX PRO 6000, the nightly runner's GPU. That includes `test_fsdp2_streaming_export_matches_reference[nvfp4]` (4.5s) and `[fp8]` (3.0s), which time out without the fix. - **Rest of the unit file:** the other 5 tests in `test_fsdp2_parallel_export.py` (world=1) pass. - **Pre-commit:** passes. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: ✅ - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A (the round-robin FSDP2 export and this deadlock both arrived in the unreleased 0.48 cycle) - Did you get Claude approval on this PR?: ❌ (pending) ### Additional Information Found while adding FSDP2 coverage to an earlier revision of #2710. 🤖 Generated with [Claude Code](https://claude.com/claude-code) ## Summary by CodeRabbit * **Bug Fixes** * Improved reliability of distributed model exports with FSDP2 by resharding modules across all ranks before collecting export tensors. * Multi-rank exports are checked against an unsharded reference to help ensure the exported tensors are consistent. --------- Signed-off-by: Chenjie Luo Co-authored-by: Claude Opus 5.5 (1M context) Signed-off-by: Rohan Joshi --- .../export/unified_export_hf_streaming.py | 9 +++++ .../export/test_fsdp2_parallel_export.py | 40 ++++++++++++++++++- 2 files changed, 47 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/export/unified_export_hf_streaming.py b/modelopt/torch/export/unified_export_hf_streaming.py index 9188bb5780e..26896da92f5 100644 --- a/modelopt/torch/export/unified_export_hf_streaming.py +++ b/modelopt/torch/export/unified_export_hf_streaming.py @@ -550,6 +550,15 @@ def collect_export_tensors( from modelopt.torch.models.step3p5.modeling_ptq import _reconstruct_fused_moe_linear _assert_fsdp2_owns_every_mesh_dim(model) + # A forward leaves the root's parameters unsharded (FSDP2 skips the root's post-forward + # reshard), and FSDP2's state_dict() pre-hook reshards a module's whole parameter group. Only + # the owner calls state_dict() below, so it would reshard the root on its rank alone, and the + # next gather window would then gather parameters on that rank that the others see as plain + # tensors -- a collective they never join. Resharding locally on every rank first makes the + # pre-hook a no-op. + for module in model.modules(): + if isinstance(module, FSDPModule): + module.reshard() my_rank, world = _dist.rank(), _dist.size() names = module_name_maps(model) ctx = ExportContext( diff --git a/tests/unit/torch/export/test_fsdp2_parallel_export.py b/tests/unit/torch/export/test_fsdp2_parallel_export.py index 18846c9844d..997a23020e5 100644 --- a/tests/unit/torch/export/test_fsdp2_parallel_export.py +++ b/tests/unit/torch/export/test_fsdp2_parallel_export.py @@ -18,16 +18,22 @@ The pipeline is world-agnostic: with ``torch.distributed`` uninitialized, ``rank/size`` are ``0/1``, the unshard window is a no-op, and rank 0 owns every unit. So these cover unit enumeration, per-tensor export, the shard writer, index naming and tied-weight dropping without a -GPU. The multi-rank unshard collective needs real FSDP2 and is covered in ``tests/gpu``. +GPU. ``test_streaming_export_two_cpu_ranks`` adds one two-rank gloo run; the multi-rank paths +on GPU are covered in ``tests/gpu``. """ import copy +import faulthandler import json +from functools import partial from pathlib import Path import pytest import torch +from _test_utils.torch.distributed.utils import spawn_multiprocess_job from safetensors.torch import load_file +from torch.distributed.device_mesh import init_device_mesh +from torch.distributed.fsdp import fully_shard import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq @@ -38,7 +44,7 @@ from transformers import AutoModelForCausalLM, LlamaConfig -def _tiny_quantized_llama(quant_cfg=None, tie=True): +def _tiny_quantized_llama(quant_cfg=None, tie=True, mesh=None): mto.enable_huggingface_checkpointing() torch.manual_seed(0) cfg = LlamaConfig( @@ -52,6 +58,10 @@ def _tiny_quantized_llama(quant_cfg=None, tie=True): architectures=["LlamaForCausalLM"], ) model = AutoModelForCausalLM.from_config(cfg).eval() + if mesh is not None: + for layer in model.model.layers: + fully_shard(layer, mesh=mesh) + fully_shard(model, mesh=mesh) calib = [torch.randint(0, 128, (1, 8)) for _ in range(4)] return mtq.quantize(model, quant_cfg or mtq.FP8_DEFAULT_CFG, lambda m: [m(x) for x in calib]) @@ -143,3 +153,29 @@ def test_streaming_export_extra_state_dict_mtp(tmp_path): loaded = _load_all(d) assert "model.layers.99.mtp.weight" in loaded assert torch.equal(loaded["model.layers.99.mtp.weight"], extra["model.layers.99.mtp.weight"]) + + +def _two_rank_streaming_export(rank, size, export_dir): + # A deadlocked rank exits with its stack instead of hanging the run on the join. + faulthandler.dump_traceback_later(45, exit=True) + # The unsharded reference is built and exported on every rank, so any collective its + # calibration or export issues has a partner. + expected, _ = _export_transformers_checkpoint(_tiny_quantized_llama(tie=False), torch.bfloat16) + # Calibrating after sharding leaves the root's parameters unsharded: FSDP2 skips the root's + # post-forward reshard. That is the state in which the owner's state_dict() used to reshard + # them on its rank alone, so the next gather window ran a collective the other rank skipped. + model = _tiny_quantized_llama(tie=False, mesh=init_device_mesh("cpu", (size,))) + _export_fsdp2_checkpoint_streaming(model, torch.bfloat16, export_dir=export_dir) + torch.distributed.barrier() + faulthandler.cancel_dump_traceback_later() + if rank != 0: + return + exported = _load_all(Path(export_dir)) + assert set(exported) == set(expected) + for key, tensor in expected.items(): + assert torch.equal(exported[key].float(), tensor.float()), key + + +def test_streaming_export_two_cpu_ranks(tmp_path, skip_on_windows): + """Two gloo ranks export what a single process does, without a rank waiting on the other.""" + spawn_multiprocess_job(2, partial(_two_rank_streaming_export, export_dir=str(tmp_path))) From 59c86687f2d580699009a9f98da7c15ce70f046f Mon Sep 17 00:00:00 2001 From: Shengliang Xu <106840466+shengliangxu@users.noreply.github.com> Date: Mon, 12 Oct 2026 17:30:13 -0700 Subject: [PATCH 08/35] Fix megatron_preprocess_data multi-worker segfault: use forkserver instead of fork (#2726) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Type of change: Bug fix `megatron_preprocess_data(..., workers>1)` built its tokenization pool with the default `multiprocessing.Pool`, i.e. the **fork** start method on Linux. Forking a process that has already initialized CUDA is unsafe. In the failing `gpu_megatron` CI jobs the forked worker died with `Fatal Python error: Segmentation fault` in garbage collection while the pool `initializer` loaded the HF fast tokenizer ([example job](https://github.com/NVIDIA/Model-Optimizer/actions/runs/37437843465/job/112188656139)). The pool kept respawning the dying workers and the main process hung until `pytest-timeout` fired at 120 s. The crash is state-dependent: roughly 4% of lane runs in 10-02..10-08 and none before, because the tests run late in the lane, after in-process CUDA tests. **Fix:** - Create the pool from a **`forkserver`** context and preload this module into the server (`_worker_context()`). Each worker is forked from a separate server process that never held the parent's CUDA objects, and the server imports torch/megatron **once** instead of once per worker. Windows has no `forkserver`, so it falls back to `spawn`, which is what the default context already used there. - The server's first start cold-imports torch/megatron, and `Pool()` blocks until that is done, so the wait would land inside the caller. `_prestart_forkserver()` starts the server without waiting: at the top of `megatron_preprocess_data` when `workers > 1` (overlapping the tokenizer and dataset loading), and in a session fixture in `tests/gpu_megatron/conftest.py` (overlapping the earlier tests; the repo sets `timeout_func_only`, and `_prebuild_quant_cuda_extensions` already uses this pattern). - The `workers` docstring notes the `if __name__ == "__main__":` requirement (see Usage). Why not `fork` or `spawn`: - `fork`: the segfault above. - `spawn`: safe, but every worker cold-imports the full torch/megatron stack. In CI that took about 1.6 min per pool ([job](https://github.com/NVIDIA/Model-Optimizer/actions/runs/37981490141/job/113998426449)), and the two `*_stops_at_max_tokens` tests, which build two pools each, timed out at the 120 s cap. ### Usage No API change. Behavior notes: - With `workers > 1`, call from under `if __name__ == "__main__":`. Workers re-import the main script, so a script without the guard now hangs (the pool keeps respawning workers that die on the re-import), where it ran under fork. In-repo callers are guarded. - The first multi-worker call in a process starts a forkserver (about 1 GB resident) that stays until the process exits. ### Testing CPU only (Python 3.13, CUDA hidden, tiny local tokenizer plus the cached `gpt2`): - Pre-commit hooks on the changed files: ruff, ruff-format, mypy, bandit, license. - `workers=2` output is byte-identical to `workers=1`: full run over three files, `max_tokens` early stop, two calls in one process, and the `python -m` CLI. - `test_megatron_preprocess_data.py` under pytest on CPU: `with_jsonl_path`, `jsonl_stops_at_max_tokens` (`workers=2`), `with_gzip_input`, and the two new `_worker_context` tests pass. - Early start: with 25 s of earlier tests, the `workers=2` test's call took 1.8 s with the fixture and 15.0 s without. With a simulated 20 s tokenizer/dataset load, `megatron_preprocess_data` took 20.3 s with the early start and 41.5 s without. (The import takes about 14 s on a dev box; the lane is slower.) - Spawn fallback (`get_all_start_methods` reporting only `spawn`): identical output, no forkserver started. Not run yet: the `gpu_megatron` lane on GPUs (Python 3.12, nemo:26.08). Kept in **draft** until that validates the change. Given the roughly 4% crash rate, one green run proves little, so it needs repeated runs. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ⚠️ API unchanged, but scripts that call `workers > 1` without an `if __name__ == "__main__":` guard now hang (documented in the docstring) - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: ✅ - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ - Did you get Claude approval on this PR?: ❌ ### Additional Information Addresses the intermittent `gpu_megatron` failure in `test_megatron_preprocess_data_*_stops_at_max_tokens` (fork after CUDA initialization). Whether it fails depends on what ran earlier in the lane's process, which is why it passes on some runs and segfaults on others. --------- Signed-off-by: Shengliang Xu Signed-off-by: Rohan Joshi --- CHANGELOG.rst | 1 + .../utils/plugins/megatron_preprocess_data.py | 47 ++++++++++++++++--- tests/gpu_megatron/conftest.py | 12 +++++ .../plugins/test_megatron_preprocess_data.py | 20 +++++++- 4 files changed, 72 insertions(+), 8 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index cc0431b7f3d..51deb4e6117 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -138,6 +138,7 @@ Changelog - Fix ``hf_ptq`` overwriting a model's existing ``pad_token`` with ``eos_token`` when the model already has a valid padding token. The exported tokenizer now preserves the source model's padding configuration. - Fix Hugging Face exports dropping off-index safetensors such as GLM-4.7's ``mtp.safetensors``. - Fix Megatron-Core checkpoint saving for quantized grouped MoE experts when tensor and expert parallelism are both enabled. +- Fix ``megatron_preprocess_data`` with ``workers > 1`` segfaulting when run in a process that has already initialized CUDA or other native threads (for example in-process after GPU work): the tokenization worker pool now uses the ``forkserver`` start method instead of ``fork``. - Fix unified HuggingFace export of RADIO-based VLMs retaining post-conversion vision and projector names instead of restoring the hub layout; deployment loaders could skip those weights. Streaming, offload, FSDP2, and layerwise export now reject reverse conversions that regroup diff --git a/modelopt/torch/utils/plugins/megatron_preprocess_data.py b/modelopt/torch/utils/plugins/megatron_preprocess_data.py index 50768abb0d8..a3601b02e67 100644 --- a/modelopt/torch/utils/plugins/megatron_preprocess_data.py +++ b/modelopt/torch/utils/plugins/megatron_preprocess_data.py @@ -90,6 +90,7 @@ import multiprocessing import time import warnings +from multiprocessing.context import BaseContext from pathlib import Path from datasets import get_dataset_config_names, get_dataset_split_names, load_dataset @@ -117,6 +118,35 @@ def _is_main_or_first_worker() -> bool: return not identity or identity[0] == 1 +_PRELOAD_MODULE = "modelopt.torch.utils.plugins.megatron_preprocess_data" + + +def _worker_context() -> BaseContext: + """Return the multiprocessing context for the tokenization worker pool. + + ``fork`` can segfault the children once CUDA is initialized in this process, so workers are + forked from a separate ``forkserver`` process instead. It preloads this module, so the heavy + ``torch`` / ``megatron`` import is paid once rather than per worker as with ``spawn``. + Windows has no ``forkserver`` and uses ``spawn``. + """ + if "forkserver" not in multiprocessing.get_all_start_methods(): + return multiprocessing.get_context("spawn") + ctx = multiprocessing.get_context("forkserver") + ctx.set_forkserver_preload([_PRELOAD_MODULE]) + return ctx + + +def _prestart_forkserver() -> None: + """Start the forkserver without waiting for it, so its preload overlaps with other work. + + The later ``Pool()`` then blocks only for what remains of it. No-op where ``spawn`` is used. + """ + if _worker_context().get_start_method() == "forkserver": + from multiprocessing import forkserver # POSIX-only module + + forkserver.ensure_running() + + class _Encoder: tokenizer: AutoTokenizer = None _chat_template_logged: set[str] = set() @@ -280,12 +310,10 @@ def _encode_in_batches(pool, encoder: "_Encoder", lines, batch_size: int): yield from encoded_batch def _encode_docs(self, encoder: "_Encoder", lines, may_stop_early: bool = False): - """Tokenize ``lines``, forking worker processes only when ``workers > 1``. + """Tokenize ``lines``, using worker processes only when ``workers > 1``. - ``multiprocessing.Pool`` always ``fork()``s, even for a single worker. Forking a - process that has already initialized a CUDA context / many native threads (e.g. when - this is called in-process after GPU work) is unsafe and can segfault the children. - The single-worker path avoids the fork entirely by tokenizing inline in this process. + Workers start through ``_worker_context``; ``_Encoder`` holds only picklable fields, as each + worker loads its tokenizer in ``initializer``. The single-worker path tokenizes inline. When ``may_stop_early`` is true, wait for finite batches so no worker results are pending if the caller stops consuming documents after reaching its token limit. @@ -295,7 +323,7 @@ def _encode_docs(self, encoder: "_Encoder", lines, may_stop_early: bool = False) if self.workers == 1: encoder.initializer() return None, map(encoder.encode, lines) - pool = multiprocessing.Pool(self.workers, initializer=encoder.initializer) + pool = _worker_context().Pool(self.workers, initializer=encoder.initializer) if may_stop_early: batch_size = self.workers * 4 # Balance throughput against unused final-batch work. encoded_docs = self._encode_in_batches(pool, encoder, lines, batch_size) @@ -556,7 +584,8 @@ def megatron_preprocess_data( json_keys: Key or list of keys to extract from json. Defaults to ["text"]. append_eod: Append an token to the end of a document. Defaults to False. max_sequence_length: Maximum tokenized sequence length. Defaults to None. - workers: Number of worker processes to launch. Defaults to 1. + workers: Number of worker processes to launch. Defaults to 1. With more than one, call from + under an ``if __name__ == "__main__":`` guard, since workers re-import the main script. log_interval: Interval between progress updates. Defaults to 100000. reasoning_content: How to handle the ``reasoning_content`` field present in many Nemotron Post-Training v3 datasets. One of: @@ -591,6 +620,10 @@ def megatron_preprocess_data( stacklevel=2, ) + if workers > 1: + # Start the worker server now so its import overlaps the tokenizer and dataset loading below. + _prestart_forkserver() + Path(output_dir).mkdir(parents=True, exist_ok=True) vocab_size = AutoTokenizer.from_pretrained(tokenizer_name_or_path).vocab_size diff --git a/tests/gpu_megatron/conftest.py b/tests/gpu_megatron/conftest.py index df061553001..882b553d0da 100644 --- a/tests/gpu_megatron/conftest.py +++ b/tests/gpu_megatron/conftest.py @@ -23,6 +23,7 @@ import modelopt.torch.quantization.extensions as ext import modelopt.torch.utils.distributed as dist +from modelopt.torch.utils.plugins.megatron_preprocess_data import _prestart_forkserver @pytest.fixture(scope="session") @@ -54,6 +55,17 @@ def _prebuild_quant_cuda_extensions(): ext.precompile() +@pytest.fixture(scope="session", autouse=True) +def _prestart_preprocess_forkserver(): + """Start the forkserver that ``megatron_preprocess_data`` workers fork from, at session start. + + Its first start cold-imports torch/megatron, which would otherwise land inside the first + multi-worker preprocess test and its per-test timeout. Starting it here, without waiting, + overlaps that import with the earlier tests. The server's environment is fixed at start. + """ + _prestart_forkserver() + + def megatron_worker_teardown(rank, world_size): """Clean up model-parallel state between tests in persistent workers.""" if dist.is_initialized(): diff --git a/tests/gpu_megatron/torch/utils/plugins/test_megatron_preprocess_data.py b/tests/gpu_megatron/torch/utils/plugins/test_megatron_preprocess_data.py index 009c638f8bd..5dd0815080c 100644 --- a/tests/gpu_megatron/torch/utils/plugins/test_megatron_preprocess_data.py +++ b/tests/gpu_megatron/torch/utils/plugins/test_megatron_preprocess_data.py @@ -15,11 +15,29 @@ import gzip import json +import multiprocessing +from multiprocessing import forkserver from pathlib import Path import pytest -from modelopt.torch.utils.plugins.megatron_preprocess_data import megatron_preprocess_data +from modelopt.torch.utils.plugins.megatron_preprocess_data import ( + _prestart_forkserver, + _worker_context, + megatron_preprocess_data, +) + + +def test_worker_context_uses_forkserver(): + assert _worker_context().get_start_method() == "forkserver" + + +def test_worker_context_falls_back_to_spawn_without_forkserver(monkeypatch): + monkeypatch.setattr(multiprocessing, "get_all_start_methods", lambda: ["spawn"]) + monkeypatch.setattr(forkserver, "ensure_running", lambda: pytest.fail("started a forkserver")) + + assert _worker_context().get_start_method() == "spawn" + _prestart_forkserver() # Depends on the external HF datasets-server, which intermittently returns 5xx From 13589a1bfb0c874ec447b03dc5a9d12fa52076e5 Mon Sep 17 00:00:00 2001 From: yingguo-trt <244492186+yingguo-trt@users.noreply.github.com> Date: Tue, 13 Oct 2026 13:05:02 +0800 Subject: [PATCH 09/35] [https://nvbugspro.nvidia.com/bug/6886263] Fix FlashInfer sparse attention packed KV cache support (#2697) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Type of change: Bug fix Fixes FlashInfer skip-softmax calibration and serving with packed KV caches, including the layout used by vLLM 0.30.0 and 0.31.0. The adapter assumed a five-dimensional `[blocks, 2, page, heads, dim]` cache, but these versions provide `[blocks, heads, page, 2 * dim]`. As a result, calibration fails during engine warmup with: ```text ValueError: FlashInfer KV cache must have logical shape [blocks, 2, page, heads, dim] ``` A shared helper now extracts logical `[blocks, page, heads, dim]` K/V views for both calibration and serving. Packed caches use `transpose(1, 2).split(head_size, dim=-1)` without copying the cache; legacy five-dimensional caches retain their existing indexing. ### Usage Use the existing FlashInfer calibration command. `` should be a private writable checkpoint copy/view because `--update_checkpoint_config` updates its `config.json`; `` is the prepared calibration data directory. ```bash python examples/vllm_serve/calibrate_sparse_attn.py "" \ --calib_data_dir "" \ --calib_samples 64 --calib_max_seqlen 16384 \ --target_sparse_ratio 0.5 --decode_tokens 32 \ --max_model_len 32768 --dtype bfloat16 \ --attention_backend FLASHINFER --tensor_parallel_size 4 \ --update_checkpoint_config ``` ### Testing Environment: BF16 Llama-3.1-8B-Instruct, 4 H200 GPUs, TP4, `FLASHINFER`, with 32 decode steps. **Before the fix:** [Jenkins #4044](http://dlswqa-nas:18880/job/modelopt_release_testing/4044/) reproduced the cache-shape error during engine warmup with vLLM 0.30.0. **With the fix:** the FlashInfer calibration-and-serving E2E passed on both versions: | vLLM | CI run | Result | | --- | --- | --- | | 0.30.0 | [#4045](http://dlswqa-nas:18880/job/modelopt_release_testing/4045/) | 1 passed, 0 failed, 0 skipped | | 0.31.0 | [#4047](http://dlswqa-nas:18880/job/modelopt_release_testing/4047/) | 1 passed, 0 failed, 0 skipped | Each run completed prefill/decode calibration, exported and merged the sparse configuration, then loaded the checkpoint in a fresh HTTP server and generated a response. The 0.31.0 artifacts also confirm `ModelOptSparseFlashInferImpl` during calibration and serving on all 32 attention layers. Both passing runs used commit `698154ff0d0266270235f98d57246f616e20eeff`. Current PR head `92c48f240781dd70af4581f1a059089d4d288c2d` contains the identical fix patch and an identical `plugins/vllm.py` file after rebase; the rebased head has not been rerun. Pre-commit checks passed on the tested, pre-rebase fix. The E2E validates the calibration/serving flow; it does not measure actual serving tile skips, accuracy, or performance. ### Before your PR is "*Ready for review*" Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors). - Is this change backward compatible?: ✅ — preserves the legacy five-dimensional cache layout. - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A — no copied code or new dependency. - Did you write any new necessary tests?: ❌ — no tests added in this PR; external QA E2E validation is listed above. - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ❌ — no changelog entry included. - Did you get Claude approval on this PR?: ❌ — not yet obtained. ### Additional Information - Related feature: #1992. - Tracking: [NVBug 6886263](https://nvbugspro.nvidia.com/bug/6886263). - The legacy cache writer remains appropriate for older vLLM versions. In [vLLM 0.30.0](https://github.com/vllm-project/vllm/blob/v0.30.0/vllm/v1/attention/backends/flashinfer.py#L551) and [0.31.0](https://github.com/vllm-project/vllm/blob/v0.31.0/vllm/v1/attention/backends/flashinfer.py#L555), `forward_includes_kv_cache_update=False`, so native vLLM cache updates run before the adapter and the legacy writer is bypassed. ## Summary by CodeRabbit * **Bug Fixes** * Improved compatibility with FlashInfer KV caches using supported layouts, helping calibration and regular operation handle more cache configurations. Signed-off-by: yingguo-trt <244492186+yingguo-trt@users.noreply.github.com> Signed-off-by: Rohan Joshi --- .../attention_sparsity/plugins/vllm.py | 30 ++++++++------- .../test_sparse_attn_worker.py | 33 ++++++++++++++--- .../test_vllm_calibration.py | 37 +++++++++++++++---- 3 files changed, 74 insertions(+), 26 deletions(-) diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py index 3c327d48bd0..c4f4fd79afa 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/vllm.py @@ -844,6 +844,19 @@ def build(*args, **kwargs): return True +def _flashinfer_kv_cache_views(kv_cache: torch.Tensor, head_size: int): + """Return logical [blocks, page, heads, dim] K/V views for FlashInfer.""" + if kv_cache.ndim == 5 and kv_cache.shape[1] == 2: + return kv_cache[:, 0], kv_cache[:, 1] + if kv_cache.ndim == 4 and kv_cache.shape[-1] == 2 * head_size: + return kv_cache.transpose(1, 2).split(head_size, dim=-1) + raise ValueError( + "FlashInfer KV cache must have logical shape [blocks, 2, page, heads, dim] " + "or [blocks, heads, page, 2 * dim], " + f"got {tuple(kv_cache.shape)}" + ) + + def _flashinfer_cache_write(layer, key, value, kv_cache, attn_metadata, impl) -> None: """Issue FlashInfer's native paged K/V cache write.""" torch.ops._C_cache_ops.reshape_and_cache_flash( @@ -924,18 +937,15 @@ def prepare_modelopt(): "FlashInfer metadata is missing the ModelOpt calibration " f"fields: {', '.join(missing)}" ) - if kv_cache.ndim != 5 or kv_cache.shape[1] != 2: - raise ValueError( - "FlashInfer KV cache must have logical shape [blocks, 2, page, heads, dim]" - ) + key_cache, value_cache = _flashinfer_kv_cache_views(kv_cache, impl.head_size) # Order matters: releases that update the KV cache inside forward must # write the current K/V before the calibrate kernel reads the cache. prepare_modelopt() return _forward_calibrate( impl, query=query, - key_cache=kv_cache[:, 0], - value_cache=kv_cache[:, 1], + key_cache=key_cache, + value_cache=value_cache, block_table=attn_metadata._modelopt_block_table, seq_lens=attn_metadata._modelopt_seq_lens, cu_seqlens_q=attn_metadata._modelopt_query_start_loc, @@ -954,13 +964,7 @@ def prepare_modelopt(): if resolved is None: return dense_fallback() - if kv_cache.ndim != 5 or kv_cache.shape[1] != 2: - raise ValueError( - "FlashInfer KV cache must have logical shape [blocks, 2, page, heads, dim]" - ) - - key_cache = kv_cache[:, 0] - value_cache = kv_cache[:, 1] + key_cache, value_cache = _flashinfer_kv_cache_views(kv_cache, impl.head_size) max_query_len = attn_metadata._modelopt_max_query_len is_decode_only = max_query_len <= 1 common_kw = { diff --git a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py index e45c4dcf54f..10af727381f 100644 --- a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_sparse_attn_worker.py @@ -395,12 +395,32 @@ def fake_decode(query, key_cache, value_cache, block_table, seq_lens, **kwargs): assert finalize_kw["page_size"] == page_size -def test_flashinfer_sparse_prefill_uses_shared_triton_kernel(monkeypatch): +@pytest.mark.parametrize("layout", ["HND", "packed"]) +def test_flashinfer_sparse_prefill_uses_shared_triton_kernel(monkeypatch, layout): """FlashInfer must route 2:4 prefill through the current ModelOpt kernel.""" impl = _make_flashinfer_impl(sparse=True) page_size, num_heads, head_dim = 16, 2, 64 - physical = torch.zeros(4, 2, num_heads, page_size, head_dim, dtype=torch.float16) - kv_cache = physical.permute(0, 1, 3, 2, 4) + expected_key = ( + torch.arange(4 * page_size * num_heads * head_dim) + .remainder(127) + .reshape(4, page_size, num_heads, head_dim) + .to(torch.float16) + ) + expected_value = -expected_key - 1 + if layout == "packed": + kv_cache = torch.cat((expected_key, expected_value), dim=-1).transpose(1, 2).contiguous() + expected_stride = ( + 2 * num_heads * page_size * head_dim, + 2 * head_dim, + 2 * page_size * head_dim, + 1, + ) + else: + physical = ( + torch.stack((expected_key, expected_value), dim=1).permute(0, 1, 3, 2, 4).contiguous() + ) + kv_cache = physical.permute(0, 1, 3, 2, 4) + expected_stride = kv_cache[:, 0].stride() metadata = _flashinfer_metadata(query_lens=(2, 2), max_query_len=4) query = torch.zeros(4, num_heads, head_dim, dtype=torch.float16) captured = {} @@ -419,8 +439,11 @@ def fake_attention(query, **kwargs): assert captured["sparsity_n"] == 2 assert captured["sparsity_m"] == 4 assert captured["page_size"] == page_size - assert captured["k_cache"].stride() == kv_cache[:, 0].stride() - assert captured["v_cache"].stride() == kv_cache[:, 1].stride() + for name, expected in (("k_cache", expected_key), ("v_cache", expected_value)): + cache = captured[name] + torch.testing.assert_close(cache, expected) + assert cache.stride() == expected_stride + assert cache.untyped_storage().data_ptr() == kv_cache.untyped_storage().data_ptr() assert captured["block_table"] is metadata._modelopt_block_table assert captured["b_seq_len_k"] is metadata._modelopt_seq_lens diff --git a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_calibration.py b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_calibration.py index 4ed2637ef5d..5661ffa43fd 100644 --- a/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_calibration.py +++ b/tests/gpu_vllm/torch/sparsity/attention_sparsity/test_vllm_calibration.py @@ -487,7 +487,7 @@ def test_rejects_non_16bit_cache(self, monkeypatch): # FlashInfer adapter: cache write must precede the calibrate-kernel read # --------------------------------------------------------------------------- class TestFlashInferCalibrationOrdering: - @pytest.mark.parametrize("layout", ["NHD", "HND"]) + @pytest.mark.parametrize("layout", ["NHD", "HND", "packed"]) def test_cache_write_happens_before_calibrate_read(self, monkeypatch, layout): calls = [] monkeypatch.setattr( @@ -499,6 +499,11 @@ def test_cache_write_happens_before_calibrate_read(self, monkeypatch, layout): def fake_calibrate(q, *args, **kwargs): calls.append("calibrate") calls.append(kwargs["k_cache"].stride()) + for name, expected in (("k_cache", expected_key), ("v_cache", expected_value)): + cache = kwargs[name] + torch.testing.assert_close(cache, expected) + assert cache.stride() == expected_stride + assert cache.untyped_storage().data_ptr() == kv_cache.untyped_storage().data_ptr() counters = torch.zeros(len(TRIALS), 2, dtype=torch.int64) return torch.zeros_like(q), counters @@ -513,12 +518,28 @@ def fake_calibrate(q, *args, **kwargs): _calib_threshold_trials=list(TRIALS), _calib_records=[], ) - shape = (3, 2, page, num_kv_heads, head_dim) - kv_cache = torch.zeros(shape, dtype=torch.bfloat16) - if layout == "HND": - kv_cache = torch.zeros( - shape[0], shape[1], shape[3], shape[2], shape[4], dtype=torch.bfloat16 - ).permute(0, 1, 3, 2, 4) + expected_key = ( + torch.arange(3 * page * num_kv_heads * head_dim) + .remainder(127) + .reshape(3, page, num_kv_heads, head_dim) + .to(torch.bfloat16) + ) + expected_value = -expected_key - 1 + if layout == "packed": + kv_cache = ( + torch.cat((expected_key, expected_value), dim=-1).transpose(1, 2).contiguous() + ) + expected_stride = ( + 2 * num_kv_heads * page * head_dim, + 2 * head_dim, + 2 * page * head_dim, + 1, + ) + else: + kv_cache = torch.stack((expected_key, expected_value), dim=1) + if layout == "HND": + kv_cache = kv_cache.permute(0, 1, 3, 2, 4).contiguous().permute(0, 1, 3, 2, 4) + expected_stride = kv_cache[:, 0].stride() attn_metadata = SimpleNamespace( _modelopt_block_table=torch.zeros(1, 1, dtype=torch.int32), _modelopt_seq_lens=torch.tensor([8], dtype=torch.int32), @@ -543,7 +564,7 @@ def fake_calibrate(q, *args, **kwargs): output=torch.empty_like(q), ) - assert calls == ["cache_write", "calibrate", kv_cache[:, 0].stride()] + assert calls == ["cache_write", "calibrate", expected_stride] assert torch.isfinite(out).all() assert len(impl._calib_records) == 1 assert impl._calib_records[0]["phase"] == "decode" From 903a23548647e263007365d6f515a957a03e72cc Mon Sep 17 00:00:00 2001 From: Wei-Ming Chen <17592131+meenchen@users.noreply.github.com> Date: Mon, 12 Oct 2026 22:14:27 -0700 Subject: [PATCH 10/35] [NVBugs 6914520, 6914493] Right-pad labeled calibration for finite attention gradients (#2769) Type of change: Bug fix Right-pad labeled calibration samples per tokenizer call to avoid fully masked causal attention rows that can produce non-finite cuDNN gradients during AutoQuantize. The tokenizer's generation padding setting and unlabeled calibration behavior remain unchanged. Generation previews remove either left or right padding before decoding. cuDNN stays enabled, and the existing non-finite-gradient fail-fast guard remains intact. A separate documentation change explains that KV-only recipes correctly leave the weight quantization algorithm null and should be validated through KV metadata and scaling tensors. No caller changes are required. Labeled calibration now selects right padding: ```python from modelopt.torch.utils.dataset_utils import get_dataset_dataloader calibration = get_dataset_dataloader(tokenizer=tokenizer, include_labels=True) ``` - Split branch: 99 dataset and HF PTQ model-type/preview tests passed. - Applicable pre-commit hooks and `git diff --check` passed. - GB200 and GB300 operator probes reproduce non-finite query gradients with left padding and cuDNN, while right padding keeps gradients finite with cuDNN enabled. - Prior full-model validation of the original combined patch used four GB200 GPUs, Transformers 5.5.4, 32 calibration samples, batch 1, and sequence length 512. Unmodified main reproduces the non-finite-gradient diagnostic. The corrected composed recipe completes gradient scoring, weight optimization at 5.40 bits, KV search, and export. - KV-only export succeeds with null weight algorithm, an empty weight-layer map, and mixed-precision KV metadata for 16 attention layers. - The composed-export validation also included the companion export-accounting fix. Both fixes are needed for the fully corrected release checkpoints. GPU validation was not repeated solely for the split; executable changes are unchanged. Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors) (e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(..., weights_only=False)`, `pickle`, etc.). - Is this change backward compatible?: Yes - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: Yes - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: Yes, under 0.48.0 bug fixes - Did you get Claude approval on this PR?: Not run The KV-only documentation describes existing correct behavior, not a new weight-quantization feature. A validator that requires a non-null weight algorithm for KV-only recipes needs to check `kv_cache_quant_algo`, `kv_cache_quantized_layers`, and exported K/V scales instead. * **Bug Fixes** * AutoQuantize calibration with labeled samples now uses right padding to avoid fully masked causal rows and non-finite attention gradients. * Generation previews now remove padding before decoding, including for left-padded inputs. * **Documentation** * Clarified how to verify KV-cache quantization in exported configuration when model weights are not quantized. **Companion PR:** [#2768](https://github.com/NVIDIA/Model-Optimizer/pull/2768). --------- Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com> Signed-off-by: Rohan Joshi --- CHANGELOG.rst | 1 + examples/hf_ptq/README.md | 4 +++ examples/hf_ptq/hf_ptq.py | 3 ++ modelopt/torch/utils/dataset_utils.py | 2 ++ .../examples/hf_ptq/test_hf_ptq_model_type.py | 14 ++++++++++ tests/unit/torch/utils/test_dataset_utils.py | 28 +++++++++++++++++++ 6 files changed, 52 insertions(+) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 51deb4e6117..9657532eda7 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -126,6 +126,7 @@ Changelog **Bug Fixes** - Fix Hugging Face exports classifying registered Qwen VLM parameters as unplaced source weights, which overwrote quantized tensors with original weights and incorrectly excluded their modules from quantization. Re-export affected checkpoints. +- Right-pad labeled calibration samples to avoid fully masked causal rows that can produce non-finite cuDNN attention gradients during AutoQuantize. Generation previews remove padding before decoding. - Fix calibration silently dropping samples when an out-of-memory error forces the batch to be split and a smaller working batch size is already known: the slices stopped tiling the batch, so the rows between the old and the new width never reached the model and the collected diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index 1d4690bf00a..7e3c7d8fc44 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -516,6 +516,10 @@ formats in `kv_cache_quantized_layers`. `mtq.auto_quantize` returns the sensitiv selected recipe in its search state; `--kv_auto_quantize_checkpoint` stores that resumable state, including the candidate quantizer tensors needed for replay. +KV-only recipes leave weights in their original precision, so `hf_quant_config.json` contains +`quantization.quant_algo: null` and `quantized_layers: {}`. Validate their quantization through +`kv_cache_quant_algo`, `kv_cache_quantized_layers`, and the exported K/V scaling tensors. + KV sensitivity scoring runs one reference forward plus one forward per eligible-layer candidate for every scoring step. Choose the recipe's `auto_quantize.score_size` with the number of eligible layers and candidates in mind. Peak scoring memory also grows with the number of selected tokens diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index db29d5fbf61..e75c7a4d70c 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -757,6 +757,9 @@ def pre_quantize( preview_input_ids = batch[input_key][0:1] # Pass attention_mask to generate(): HF cannot infer it when pad_token == eos_token. preview_attention_mask = batch["attention_mask"][0:1] if "attention_mask" in batch else None + if input_key == "input_ids" and preview_attention_mask is not None: + preview_input_ids = preview_input_ids[:, preview_attention_mask[0].bool()] + preview_attention_mask = torch.ones_like(preview_input_ids) # Generate preview before quantization if args.skip_generate: diff --git a/modelopt/torch/utils/dataset_utils.py b/modelopt/torch/utils/dataset_utils.py index 6202fd96076..f882a6e7364 100644 --- a/modelopt/torch/utils/dataset_utils.py +++ b/modelopt/torch/utils/dataset_utils.py @@ -888,6 +888,8 @@ def get_dataset_dataloader( all_samples, return_tensors="pt", padding=True, + # Left padding creates fully masked causal rows, which can yield NaNs in cuDNN backward. + padding_side="right" if include_labels else None, truncation=True, max_length=max_sample_length, ) diff --git a/tests/examples/hf_ptq/test_hf_ptq_model_type.py b/tests/examples/hf_ptq/test_hf_ptq_model_type.py index 0ab585b19d0..328dcd5e564 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_model_type.py +++ b/tests/examples/hf_ptq/test_hf_ptq_model_type.py @@ -53,6 +53,20 @@ def test_pre_quantize_generation_skip(hf_ptq, model_type: str, skips_generation: assert len(generate_calls) == (0 if skips_generation else 1) +@pytest.mark.parametrize( + ("ids", "mask"), [([0, 0, 1, 2], [0, 0, 1, 1]), ([1, 2, 0, 0], [1, 1, 0, 0])] +) +def test_pre_quantize_preview_removes_padding(hf_ptq, ids, mask): + calls = [] + model = SimpleNamespace(generate=lambda ids, **kwargs: calls.append((ids, kwargs)) or ids) + args = SimpleNamespace(specdec_offline_dataset=None, skip_generate=False) + batch = {"input_ids": torch.tensor([ids]), "attention_mask": torch.tensor([mask])} + prompt, attention_mask, _ = hf_ptq.pre_quantize(args, model, "llama", None, [batch], False) + assert torch.equal(prompt, torch.tensor([[1, 2]])) + assert torch.equal(attention_mask, torch.ones_like(prompt)) + assert torch.equal(calls[0][0], prompt) + + @pytest.mark.parametrize( ("model_type", "expected"), [ diff --git a/tests/unit/torch/utils/test_dataset_utils.py b/tests/unit/torch/utils/test_dataset_utils.py index 35543f8e9b9..1666308e1d4 100644 --- a/tests/unit/torch/utils/test_dataset_utils.py +++ b/tests/unit/torch/utils/test_dataset_utils.py @@ -18,6 +18,7 @@ import pytest import torch +from _test_utils.torch.transformers_models import get_tiny_tokenizer from torch.utils.data import DataLoader from modelopt.torch.utils import dataset_utils @@ -59,6 +60,33 @@ def test_prepare_messages_for_chat_template(): ) +@pytest.mark.parametrize("include_labels", [False, True]) +def test_calibration_padding_avoids_masked_gradient_rows(monkeypatch, include_labels): + tokenizer = get_tiny_tokenizer() + monkeypatch.setattr( + dataset_utils, "get_dataset_samples", lambda *args, **kwargs: ["a", "a b c"] + ) + batch = next( + iter( + get_dataset_dataloader( + tokenizer=tokenizer, + batch_size=2, + num_samples=2, + include_labels=include_labels, + device="cpu", + ) + ) + ) + mask = batch["attention_mask"].bool() + assert tokenizer.padding_side == "left" + if include_labels: + causal_mask = torch.ones(mask.shape[-1], mask.shape[-1], dtype=torch.bool).tril() + assert (causal_mask & mask[:, None, :]).any(dim=-1).all() + assert torch.equal(batch["labels"][~mask], torch.full_like(batch["labels"][~mask], -100)) + else: + assert not mask[0, 0] and mask[0, -1] + + def setup_test_data(): # Create sample batch data batch_data = { From 2ceb7334a11dbbb4359fa67d7a38fe4a2aefcb98 Mon Sep 17 00:00:00 2001 From: yueshen2016 <39203804+yueshen2016@users.noreply.github.com> Date: Tue, 13 Oct 2026 00:04:16 -0600 Subject: [PATCH 11/35] Add freeze_mtp to Megatron distillation to keep the MTP head fixed (#2770) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Type of change: New feature Adds `freeze_mtp` to the Megatron `DistillationConfig`, exposed as `--freeze_mtp` in `examples/megatron_bridge/distill.py`. With it, the student's MTP head keeps its input weights during distillation: its loss is skipped and its parameters are frozen. Today the plugin skips the MTP loss only when the MTP head is left out of quantization. A recipe that quantizes it, such as `model_type/nemotron_h/ptq/nvfp4-aggressive-mse` (MTP block-1 experts in NVFP4), keeps the MTP cross-entropy even at `kd_loss_alpha=1.0`. That loss trains the MTP head and, because the head reads the backbone's hidden states and shares the embedding and output layer, it trains the rest of the student too: QAD is then no longer distillation-only. `freeze_mtp` keeps the MTP head out of training: the student trains on the distillation loss alone, and the MTP head is exported with its PTQ weights. The default behavior is unchanged. ```bash python examples/megatron_bridge/distill.py ... --freeze_mtp ``` ```python DistillationConfig(kd_loss_alpha=1.0, freeze_mtp=True) ``` `tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py` passes on GB300 (23 passed): - `test_skip_lm_loss_with_mtp` gains `freeze_mtp` cases, with and without the LM loss. In both, the MTP head's parameters are frozen and its loss is zero. - New `test_freeze_mtp_with_quantized_mtp` covers the motivating case, with the MTP head quantized as well as the body. It runs a forward pass, KD backward and an AdamW step. The total loss follows `kd_loss_alpha`, and the backbone gets gradients. With `freeze_mtp`, the MTP head gets no gradient and is bit-identical after the step. Without it, the MTP head trains. In progress: QAD of a NemotronH omni VL checkpoint quantized with the recipe above (8 nodes), checking that the log reports the frozen MTP head and that the MTP tensors of the QAD checkpoint match the PTQ checkpoint. - Is this change backward compatible?: ✅ - Did you write any new necessary tests?: ✅ - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ 🤖 Generated with [Claude Code](https://claude.com/claude-code) - **New Features** - Added an option for Megatron distillation to freeze the student’s MTP head and skip its loss, allowing distillation to update the rest of the model. --------- Signed-off-by: Yue Co-authored-by: Claude Opus 5.5 Signed-off-by: Rohan Joshi --- CHANGELOG.rst | 1 + examples/megatron_bridge/distill.py | 6 + modelopt/torch/distill/plugins/megatron.py | 22 +++- .../distill/plugins/test_distill_megatron.py | 120 +++++++++++++++++- 4 files changed, 138 insertions(+), 11 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 9657532eda7..4e6b06bd10f 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -48,6 +48,7 @@ Changelog *Megatron Framework (M-LM / M-Bridge)* +- Add ``freeze_mtp`` to the Megatron ``DistillationConfig`` (``--freeze_mtp`` in ``examples/megatron_bridge/distill.py``): the student's MTP head keeps its input weights and its loss is skipped, so distillation trains the rest of the model on the distillation loss alone. - Add optional Top-P (nucleus) truncation to the Megatron top-k logits KD loss (now ``TopLogitsKLLoss``) via ``logit_kl_top_p`` and ``logit_kl_top_p_min_k`` in ``DistillationConfig``: after the global Top-K selection, only the smallest prefix whose cumulative teacher probability reaches ``top_p`` (with a floor of ``min_k`` entries) contributes to the KL, mirroring ``--logits-save-top-p`` / ``--logits-save-top-p-min-k`` in Megatron-LM's logits saver. - Add an end-to-end W4A4 NVFP4 PTQ and QAD tutorial for Qwen3.6-35B-A3B also covering evaluation and vLLM throughput benchmarking. See `examples/megatron_bridge/tutorials/Qwen3.6-35B-A3B/README.md `_ for details. - Add ``--mlflow `` to the ``examples/megatron_bridge`` scripts that write a checkpoint -- ``prune_minitron.py``, ``quantize.py``, ``distill.py``, ``export_quantized_megatron_to_hf.py`` and ``export_distilled_megatron_to_hf.py`` (MLflow's own ``MLFLOW_TRACKING_URI`` is honoured too). Each run records the invocation, its arguments as searchable params and its log -- a distillation records its training metrics instead of the rank-0 log -- and writes ``.experiment.json`` into the checkpoint it produced, so a pruning, a quantization, the distillation that refines its checkpoint and the export that deploys it can be traced to one another; uploading the checkpoints themselves stays off unless ``--mlflow_log_checkpoints`` is passed. diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py index f71114cf1f3..aec3425295a 100644 --- a/examples/megatron_bridge/distill.py +++ b/examples/megatron_bridge/distill.py @@ -213,6 +213,11 @@ def get_args(): default=1.0, help="KD loss weight alpha in (1 - alpha) * lm_loss + alpha * kd_loss. 1.0 skips the LM loss entirely.", ) + parser.add_argument( + "--freeze_mtp", + action="store_true", + help="Keep the student's MTP head at its input weights: skip its loss and freeze it.", + ) parser.add_argument( "--no_async_save", action="store_true", @@ -493,6 +498,7 @@ def _build_model_provider(hf_path, load_weights=True, moe_grouped_gemm=True): logit_kl_topk=args.logit_kl_top_k, logit_kl_top_p=args.logit_kl_top_p, logit_kl_top_p_min_k=args.logit_kl_top_p_min_k, + freeze_mtp=args.freeze_mtp, ) # HF VLM configs expose ``vision_config``; Megatron-Bridge nests the text model under diff --git a/modelopt/torch/distill/plugins/megatron.py b/modelopt/torch/distill/plugins/megatron.py index 3aac2666650..26b3c9cc5c8 100644 --- a/modelopt/torch/distill/plugins/megatron.py +++ b/modelopt/torch/distill/plugins/megatron.py @@ -71,6 +71,8 @@ class DistillationConfig: Only the smallest prefix of the (sorted) Top-K whose cumulative teacher probability reaches this value contributes to the loss. Requires ``logit_kl_topk``. Must be in (0, 1]. logit_kl_top_p_min_k: Minimum number of Top-K entries kept per token when top-P is active. + freeze_mtp: Keep the student's MTP head fixed: skip its loss, which otherwise also trains the + shared layers it reads from, and freeze its parameters. Default: ``False``. """ intermediate_layer_pairs: list[tuple[str, ...]] = field(default_factory=list) @@ -80,6 +82,7 @@ class DistillationConfig: logit_kl_topk: int | None = None logit_kl_top_p: float | None = None logit_kl_top_p_min_k: int = 1 + freeze_mtp: bool = False criterion: Criterion | None = None loss_balancer: mtd.DistillationLossBalancer | None = None @@ -701,21 +704,32 @@ def _sharded_state_dict(self, *args, **kwargs) -> "ShardedStateDict": # skip_lm_loss=True — only the final main-head call is zeroed. # An MTP head left out of quantization is exempt from that: there is no quantization # error to recover there, and its CE materialises an fp32 [seq, vocab] tensor. - skip_mtp_loss = _mtp_excluded_from_quantization(model) + # ``freeze_mtp`` skips the MTP loss unconditionally. skip_lm_loss = distill_cfg.skip_lm_loss - if skip_lm_loss and skip_mtp_loss: + skip_mtp_loss = distill_cfg.freeze_mtp or ( + skip_lm_loss and _mtp_excluded_from_quantization(model) + ) + if skip_mtp_loss: # Freeze the untrained MTP head: DDP's overlapped grad reduce asserts on params with no grad. - warn_rank_0("MTP head is outside quantization and its loss is skipped: freezing it.") + num_frozen = 0 with model.hide_teacher_model(): for name, param in model.named_parameters(): if "mtp" in name.split("."): param.requires_grad_(False) + num_frozen += 1 + if num_frozen: + warn_rank_0( + f"MTP head is frozen ({num_frozen} params) and its loss is skipped" + + (" (freeze_mtp)." if distill_cfg.freeze_mtp else ": it is outside quantization.") + ) + elif parallel_state.is_pipeline_last_stage(): # earlier stages never hold the MTP head + warn_rank_0("freeze_mtp is set but the model has no MTP head: nothing to freeze.") def _compute_student_lm_loss(self, labels, logits) -> Tensor: self._lm_loss_call_count += 1 mtp_num_layers = self.config.mtp_num_layers or 0 is_mtp_call = self._lm_loss_call_count <= mtp_num_layers - if skip_lm_loss and self.training and (not is_mtp_call or skip_mtp_loss): + if self.training and (skip_mtp_loss if is_mtp_call else skip_lm_loss): return torch.zeros_like(labels, dtype=logits.dtype) return type(self).compute_language_model_loss(self, labels, logits) diff --git a/tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py b/tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py index 9985f9778b0..2bbcc108d69 100644 --- a/tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py +++ b/tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py @@ -14,6 +14,7 @@ # limitations under the License. import dataclasses +import warnings from functools import partial from types import SimpleNamespace @@ -233,8 +234,11 @@ def _test_topk_logits_kl_loss(kd_kwargs, rank, size): loss["kd_loss"].backward() -def _test_skip_lm_loss_with_mtp(quantized, skip_lm_loss, rank, size): - """Test that skip_lm_loss only zeroes the main LM head, unless MTP is left out of quantization.""" +def _test_skip_lm_loss_with_mtp(quantized, skip_lm_loss, freeze_mtp, rank, size): + """Test that skip_lm_loss only zeroes the main LM head, unless MTP is left out of quantization. + + ``freeze_mtp`` zeroes the MTP loss and freezes the MTP head regardless. + """ set_seed(SEED) num_layers = 2 @@ -282,7 +286,9 @@ def _test_skip_lm_loss_with_mtp(quantized, skip_lm_loss, rank, size): distill_cfg = setup_distillation_config( # kd_loss_alpha == 1.0 skips the LM loss; any alpha < 1 keeps it. - config_or_path=DistillationConfig(kd_loss_alpha=1.0 if skip_lm_loss else 0.5), + config_or_path=DistillationConfig( + kd_loss_alpha=1.0 if skip_lm_loss else 0.5, freeze_mtp=freeze_mtp + ), student_cfg=student_model.config, teacher_cfg=teacher_model.config, ) @@ -295,7 +301,7 @@ def _test_skip_lm_loss_with_mtp(quantized, skip_lm_loss, rank, size): adjust_distillation_model_for_mcore(distillation_model, distill_cfg) # An untrained MTP head is frozen so DDP doesn't wait for its grads. - mtp_skipped = quantized and skip_lm_loss + mtp_skipped = freeze_mtp or (quantized and skip_lm_loss) with distillation_model.hide_teacher_model(): mtp_params = [p for n, p in distillation_model.named_parameters() if "mtp" in n.split(".")] assert mtp_params @@ -337,6 +343,93 @@ def _recording_loss(labels, logits): assert recorded_losses[-1].any() != skip_lm_loss, "Main LM head loss is zero iff skip_lm_loss" +def _test_freeze_mtp_with_quantized_mtp(freeze_mtp, kd_loss_alpha, rank, size): + """With the MTP head quantized too, only ``freeze_mtp`` keeps it out of a training step.""" + set_seed(SEED) + + vocab_size = 32 + batch_size = 2 + max_sequence_length = 8 + model_kwargs = { + "tensor_model_parallel_size": size, + "pipeline_model_parallel_size": 1, + "num_layers": 2, + "hidden_size": 8, + "num_attention_heads": 4, + "num_query_groups": 2, + "ffn_hidden_size": 8, + "max_sequence_length": max_sequence_length, + "vocab_size": vocab_size, + "activation_func": "squared_relu", + "mtp_num_layers": 1, + } + teacher_model = get_mcore_gpt_model(initialize_megatron=True, **model_kwargs).cuda() + student_model = get_mcore_gpt_model(initialize_megatron=False, **model_kwargs).cuda() + # A recipe that quantizes the MTP head as well as the body. + student_model.decoder.weight_quantizer = TensorQuantizer() + student_model.mtp.weight_quantizer = TensorQuantizer() + assert not _mtp_excluded_from_quantization(student_model) + + distill_cfg = setup_distillation_config( + config_or_path=DistillationConfig(kd_loss_alpha=kd_loss_alpha, freeze_mtp=freeze_mtp), + student_cfg=student_model.config, + teacher_cfg=teacher_model.config, + ) + kd_config = { + "teacher_model": teacher_model, + "criterion": distill_cfg.criterion, + "loss_balancer": distill_cfg.loss_balancer, + } + distillation_model = mtd.convert(student_model, mode=[("kd_loss", kd_config)]) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + adjust_distillation_model_for_mcore(distillation_model, distill_cfg) + if rank == 0: + frozen_msg = [w for w in caught if "MTP head is frozen" in str(w.message)] + assert bool(frozen_msg) == freeze_mtp + + with distillation_model.hide_teacher_model(): + named_params = dict(distillation_model.named_parameters()) + mtp_params = {n: p for n, p in named_params.items() if "mtp" in n.split(".")} + body_params = {n: p for n, p in named_params.items() if n.startswith("decoder.")} + assert mtp_params and body_params + mtp_before = {n: p.detach().clone() for n, p in mtp_params.items()} + optimizer = torch.optim.AdamW( + [p for p in named_params.values() if p.requires_grad], lr=1e-2, weight_decay=0.1 + ) + + distillation_model.train() + prompt_tokens = torch.randint(0, vocab_size, (batch_size, max_sequence_length)).cuda() + labels = torch.randint(0, vocab_size, (batch_size, max_sequence_length)).cuda() + position_ids = ( + torch.arange(max_sequence_length, dtype=torch.long) + .unsqueeze(0) + .repeat(batch_size, 1) + .cuda() + ) + attention_mask = torch.tril( + torch.ones((batch_size, 1, max_sequence_length, max_sequence_length), dtype=torch.bool) + ).cuda() + student_loss = distillation_model(prompt_tokens, position_ids, attention_mask, labels=labels) + loss = distillation_model.compute_kd_loss( + student_loss=student_loss.mean(), loss_reduction_fn=lambda x: x[0].mean() + ) + loss["kd_loss"].backward() + optimizer.step() + + # The main LM loss follows kd_loss_alpha either way. + assert student_loss.any() == (kd_loss_alpha < 1) + expected = (1 - kd_loss_alpha) * student_loss.mean() + kd_loss_alpha * loss["logits_loss"] + assert torch.allclose(loss["kd_loss"], expected) + # The rest of the student keeps training. + assert any(p.grad is not None and p.grad.any() for p in body_params.values()) + # The MTP head trains unless frozen; frozen, it gets no gradient and the step leaves it as is. + mtp_trained = any(p.grad is not None and p.grad.any() for p in mtp_params.values()) + assert mtp_trained != freeze_mtp + if freeze_mtp: + assert all(torch.equal(p, mtp_before[n]) for n, p in mtp_params.items()) + + def test_logits_kl_loss(dist_workers): """Test LogitsKLLoss with TP parallelism.""" dist_workers.run(_test_logits_kl_loss) @@ -460,11 +553,24 @@ def test_distillation_config_top_p_validation(): @pytest.mark.parametrize( - ("quantized", "skip_lm_loss"), [(False, True), (True, True), (True, False)] + ("quantized", "skip_lm_loss", "freeze_mtp"), + [ + (False, True, False), + (True, True, False), + (True, False, False), + (False, True, True), + (False, False, True), + ], ) -def test_skip_lm_loss_with_mtp(dist_workers, quantized, skip_lm_loss): +def test_skip_lm_loss_with_mtp(dist_workers, quantized, skip_lm_loss, freeze_mtp): """Test that skip_lm_loss only zeroes the main LM head, unless MTP is left out of quantization.""" - dist_workers.run(partial(_test_skip_lm_loss_with_mtp, quantized, skip_lm_loss)) + dist_workers.run(partial(_test_skip_lm_loss_with_mtp, quantized, skip_lm_loss, freeze_mtp)) + + +@pytest.mark.parametrize(("freeze_mtp", "kd_loss_alpha"), [(True, 1.0), (True, 0.5), (False, 1.0)]) +def test_freeze_mtp_with_quantized_mtp(dist_workers, freeze_mtp, kd_loss_alpha): + """With the MTP head quantized too, only ``freeze_mtp`` keeps it out of a training step.""" + dist_workers.run(partial(_test_freeze_mtp_with_quantized_mtp, freeze_mtp, kd_loss_alpha)) def test_mtp_excluded_from_quantization(): From 5ca3b5b824ef81f51ec5caba1013eae2d059b550 Mon Sep 17 00:00:00 2001 From: joe <222375587+joe-nv@users.noreply.github.com> Date: Wed, 14 Oct 2026 23:26:42 +0800 Subject: [PATCH 12/35] fix(quantization): isolate optional Megatron imports (#2719) ### What does this PR do? Type of change: Bug fix Prevent a broken optional Megatron dependency from aborting ModelOpt initialization for unrelated workflows. Target: `release/0.48.0`. `import modelopt.torch.opt` also initializes quantization, whose compression module imports the Megatron plugin. In the reported environment, the dependency chain reaches Transformer Engine, FlashAttention 4 and Quack, where accessing the unavailable `cutlass.cute.core.ThrMma` raises `AttributeError`. The plugin initializers already contain this failure through `import_plugin`, but `compress.py` retries the import with an `ImportError`-only handler. This second import allows the exception to escape. Reuse `import_plugin("megatron")` at that optional import boundary. Initialize `mcore_available` to `False`, setting it to `True` only after all required classes load. This follows the existing warning and exception policy without introducing a duplicate handler or a package-version-specific workaround. Successful Megatron imports retain their existing behavior. A broken backend remains unavailable and emits a warning; this change does not fix the underlying Quack/CUTLASS incompatibility. No dependency versions, public APIs or quantization algorithms change. ### Usage No new API or flag. The existing `import modelopt.torch.opt` entry point is unchanged. ### Testing - During development, local CPU import checks covered successful import, missing dependency and dependency AttributeError. Before the fix, the AttributeError case failed; after the fix, all three passed, including on the release branch. - These checks were development-only validation. The test file has been removed from the final PR; only the production import fix remains. - Targeted pre-commit hooks passed for the unchanged production code, including Ruff, mypy, license checks and Bandit. - After removing the test file, `git diff --check` and local SSH commit-signature verification passed. - The failing TRT-LLM container and real Megatron GPU workflows have not been validated with this patch. ## Summary by CodeRabbit * **Bug Fixes** * Improved compatibility when loading optional Megatron support, including in environments where Megatron is unavailable. This helps quantization workflows handle Megatron availability more consistently. --------- Signed-off-by: joe-nv <222375587+joe-nv@users.noreply.github.com> Signed-off-by: Rohan Joshi --- modelopt/torch/quantization/compress.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/modelopt/torch/quantization/compress.py b/modelopt/torch/quantization/compress.py index 2a5cbbee9fc..ee7e642a0e8 100644 --- a/modelopt/torch/quantization/compress.py +++ b/modelopt/torch/quantization/compress.py @@ -27,6 +27,7 @@ from modelopt.torch.opt.conversion import ModelLikeModule, ModeloptStateManager from modelopt.torch.opt.dynamic import _DMRegistryCls from modelopt.torch.opt.mode import ConvertReturnType, MetadataDict +from modelopt.torch.utils.import_utils import import_plugin from .backends.gemm_registry import disable_real_quant_gemm, enable_real_quant_gemm from .config import CompressCfgType, CompressConfig @@ -35,7 +36,8 @@ from .qtensor import QTensorWrapper, pack_real_quantize_weight from .utils import is_quantized_linear -try: +mcore_available = False +with import_plugin("megatron"): from .plugins.megatron import ( _MegatronColumnParallelLinear, _MegatronRowParallelLinear, @@ -44,8 +46,6 @@ ) mcore_available = True -except ImportError: - mcore_available = False RealQuantModuleRegistry = _DMRegistryCls("RealQuant") From 51cbc97bb8f7af821a0d3a6922013c6b36455dee Mon Sep 17 00:00:00 2001 From: Ajinkya Rasane <131806219+ajrasane@users.noreply.github.com> Date: Wed, 14 Oct 2026 15:44:18 -0400 Subject: [PATCH 13/35] [6915406] Fix TensorRT 11 plugin registry enumeration (#2796) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Type of change: Bug fix Loading a TensorRT plugin library during ONNX preprocessing raises `AttributeError` on TensorRT 11 because `get_custom_layers()` evaluates a debug f-string using the removed `IPluginRegistry.plugin_creator_list` property, even when debug logging is disabled. Use `IPluginRegistry.all_creators`, which is available in TensorRT 10 and 11. Add a real-runtime GPU regression test that loads the shipped TensorRT plugin library and verifies ONNX parsing and returned tensor metadata at WARNING and DEBUG log levels. Add a changelog entry for the compatibility fix. Existing ONNX preprocessing calls with `trt_plugins` continue without API changes: ```python from modelopt.onnx.trt_utils import get_custom_layers custom_layers, tensor_info = get_custom_layers( "model.onnx", trt_plugins=["/path/to/plugin.so"] ) ``` - TensorRT 11.1.0.106: `python -m pytest -o addopts= -o cache_dir=/tmp/pytest-cache tests/gpu/onnx/test_trt_utils.py -q` — both cases failed at the removed property before the fix and passed after the fix. - TensorRT 10.16.1.11: the same test with `--noconftest` — 2 passed. Suite conftests were disabled because that environment lacks unrelated Hugging Face dependencies. - Reran the original minimal reproduction with the fixed source on TensorRT 11; plugin loading and ONNX parsing completed without the reported exception. - All applicable pre-commit hooks passed for the three changed files, including Ruff, mypy, and Bandit. - Full custom-op PTQ and engine validation were not run. Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors) (e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(..., weights_only=False)`, `pickle`, etc.). - Is this change backward compatible?: ✅ - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: ✅ - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ - Did you get Claude approval on this PR?: N/A — small, focused compatibility fix. * **Bug Fixes** * Fixed ONNX preprocessing with custom TensorRT plugin libraries on TensorRT 11. * Improved reporting of registered TensorRT plugin creators. Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> Signed-off-by: Rohan Joshi --- CHANGELOG.rst | 1 + modelopt/onnx/trt_utils.py | 2 +- tests/gpu/onnx/test_trt_utils.py | 49 ++++++++++++++++++++++++++++++++ 3 files changed, 51 insertions(+), 1 deletion(-) create mode 100644 tests/gpu/onnx/test_trt_utils.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 4e6b06bd10f..2fb53f9f4e5 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -167,6 +167,7 @@ Changelog - Fix two issues in the vLLM offline hidden-state dump (``examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py``) that only surface on large runs. **Resume:** the filter that skips conversations whose ``.pt`` already exists now runs with ``load_from_cache_file=False``. It depends on on-disk state, which is not part of the fingerprint ``datasets`` computes from the function and the dataset, so with a persistent HF cache reused across a resumed or requeued run the cached "keep everything" result from an earlier run was replayed and the dump re-generated and overwrote conversations it had already finished (observed: tens of thousands of ``.pt`` rewritten while the output count stayed flat). **Staging:** generation is now chunked (``--save-chunk-size``, default 256), so each chunk is saved and its staged hidden states freed before the next chunk is generated. Previously the whole dataset was generated before anything was saved, which kept every conversation staged in the connector's ``shared_storage_path`` (``/dev/shm``, i.e. RAM, by default) at once and exhausted it partway through large dumps. Chunking also makes the dump incrementally durable, so an interrupted run keeps its finished conversations and resumes from them. The save path now also frees each conversation's staged hidden states in a ``finally``, so a conversation skipped mid-loop (e.g. a short ``loss_mask``) can no longer leak its staging file, and conversation ids are validated as plain filenames before being used to build output paths. - Fix YAML config I/O decoding with the locale codepage instead of UTF-8, which made a config containing any non-ASCII byte fail to load on a machine whose locale is not UTF-8 (notably Windows, where the default is cp1252). ``modelopt/recipe/loader.py``, the two ONNX autotune state files, the two transformers config readers, the distill config and the puzzletron profile now pass ``encoding="utf-8"`` explicitly. Only the YAML config paths are covered: these are the files most likely to carry non-ASCII text in comments, model names or paths, and the only ones read inside a user's process. - Hybrid (e.g. Nemotron-H) checkpoints saved by the ``examples/megatron_bridge`` scripts now record their layer spec in ``run_config.yaml`` in a form that reloads, so they can be converted to HuggingFace; a checkpoint saved by an earlier release still needs its ``model.hybrid_stack_spec`` block replaced by hand. +- Fix ONNX preprocessing with custom TensorRT plugin libraries on TensorRT 11. 0.47.0 (2026-09-23) ^^^^^^^^^^^^^^^^^^^ diff --git a/modelopt/onnx/trt_utils.py b/modelopt/onnx/trt_utils.py index b407fdc5411..88e199b1788 100644 --- a/modelopt/onnx/trt_utils.py +++ b/modelopt/onnx/trt_utils.py @@ -139,7 +139,7 @@ def get_custom_layers( for plugin in trt_plugins: _load_trt_plugin(plugin, registry) - logger.debug(f"Total registered plugin creators: {len(registry.plugin_creator_list)}") + logger.debug(f"Total registered plugin creators: {len(registry.all_creators)}") builder = trt.Builder(trt_logger) network = ( builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)) diff --git a/tests/gpu/onnx/test_trt_utils.py b/tests/gpu/onnx/test_trt_utils.py new file mode 100644 index 00000000000..a1dae0466ff --- /dev/null +++ b/tests/gpu/onnx/test_trt_utils.py @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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 ctypes.util +import logging + +import onnx +import pytest + +from modelopt.onnx import trt_utils + +pytest.importorskip("tensorrt") + + +@pytest.mark.parametrize("log_level", [logging.WARNING, logging.DEBUG]) +def test_get_custom_layers_with_plugin_library(caplog, log_level): + """Loading a plugin library must preserve ONNX parsing at either log level.""" + plugin_library = ctypes.util.find_library("nvinfer_plugin") + assert plugin_library is not None + model = onnx.helper.make_model( + onnx.helper.make_graph( + [onnx.helper.make_node("Identity", ["x"], ["y"])], + "identity", + [onnx.helper.make_tensor_value_info("x", onnx.TensorProto.FLOAT, [1])], + [onnx.helper.make_tensor_value_info("y", onnx.TensorProto.FLOAT, [1])], + ), + opset_imports=[onnx.helper.make_opsetid("", 13)], + ir_version=10, + ) + + with caplog.at_level(log_level, logger=trt_utils.logger.name): + custom_layers, tensor_info = trt_utils.get_custom_layers(model, [plugin_library]) + + assert custom_layers == [] + assert tensor_info == { + name: {"shape": [1], "dtype": trt_utils.trt.float32} for name in ["x", "y"] + } From 62f040979ea05649ce38616f51a37953e8159312 Mon Sep 17 00:00:00 2001 From: Ajinkya Rasane <131806219+ajrasane@users.noreply.github.com> Date: Wed, 14 Oct 2026 15:57:25 -0400 Subject: [PATCH 14/35] [6919103] Fix custom-op plugin compilation with TensorRT 11 (#2801) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Type of change: Bug fix The ONNX PTQ custom-op plugin fails to compile with TensorRT 11 because `isOutputBroadcastAcrossBatch` and `canBroadcastInputAcrossBatch` no longer override methods in the base interface. Guard their declarations and definitions with `NV_TENSORRT_MAJOR < 11`, preserving the legacy callbacks for older TensorRT versions. Add a regression in the existing TensorRT example test lane that runs the real CMake build and checks that it produces `libidentity_conv_plugin.so`. ### Usage ```bash cmake -S examples/onnx_ptq/custom_op_plugin/plugin -B /tmp/identity-plugin-build cmake --build /tmp/identity-plugin-build --parallel 2 ``` ### Testing - Confirmed the new build regression fails before the fix with both callback override errors on TensorRT 11.1. - Ran `python -m pytest --noconftest -o addopts= tests/examples/torch_trt/test_custom_op_plugin.py -q`: passed with TensorRT 11.1 and TensorRT 10.13.3. This isolates the C++ build test from unrelated Python suite fixtures and options. - Reran the documented CMake configure/build commands on the `0.48.0rc0` source with only this fix applied: passed. - All applicable pre-commit checks passed for the three changed files. - No GPU inference, engine building, or INT8 calibration was run. Validation covers plugin compilation and linking. ### Before your PR is "*Ready for review*" Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors) (e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(..., weights_only=False)`, `pickle`, etc.). - Is this change backward compatible?: ✅ - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: ✅ - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A — minor example build compatibility fix; no critical stable-release regression or public API change. - Did you get Claude approval on this PR?: N/A — small compatibility fix. ## Summary by CodeRabbit * **Bug Fixes** * Improved compatibility of the custom operation plugin with TensorRT 11 and later. * **Tests** * Added an automated build check for the custom operation plugin when the required build tools are available. Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> Signed-off-by: Rohan Joshi --- .../plugin/IdentityConvPlugin.cpp | 2 ++ .../plugin/IdentityConvPlugin.h | 2 ++ .../torch_trt/test_custom_op_plugin.py | 32 +++++++++++++++++++ 3 files changed, 36 insertions(+) create mode 100644 tests/examples/torch_trt/test_custom_op_plugin.py diff --git a/examples/onnx_ptq/custom_op_plugin/plugin/IdentityConvPlugin.cpp b/examples/onnx_ptq/custom_op_plugin/plugin/IdentityConvPlugin.cpp index 7b8a41d7b90..ecd7809b7c7 100644 --- a/examples/onnx_ptq/custom_op_plugin/plugin/IdentityConvPlugin.cpp +++ b/examples/onnx_ptq/custom_op_plugin/plugin/IdentityConvPlugin.cpp @@ -179,12 +179,14 @@ nvinfer1::DataType IdentityConv::getOutputDataType(int32_t index, return inputTypes[0]; } +#if NV_TENSORRT_MAJOR < 11 bool IdentityConv::isOutputBroadcastAcrossBatch(int32_t outputIndex, bool const *inputIsBroadcasted, int32_t nbInputs) const noexcept { return false; } bool IdentityConv::canBroadcastInputAcrossBatch(int32_t inputIndex) const noexcept { return false; } +#endif int32_t IdentityConv::enqueue(int32_t batchSize, void const *const *inputs, void *const *outputs, void *workspace, cudaStream_t stream) noexcept { diff --git a/examples/onnx_ptq/custom_op_plugin/plugin/IdentityConvPlugin.h b/examples/onnx_ptq/custom_op_plugin/plugin/IdentityConvPlugin.h index da6128424e0..1f111fb5a05 100644 --- a/examples/onnx_ptq/custom_op_plugin/plugin/IdentityConvPlugin.h +++ b/examples/onnx_ptq/custom_op_plugin/plugin/IdentityConvPlugin.h @@ -88,10 +88,12 @@ class IdentityConv : public nvinfer1::IPluginV2IOExt { char const *getPluginNamespace() const noexcept override; +#if NV_TENSORRT_MAJOR < 11 bool isOutputBroadcastAcrossBatch(int32_t outputIndex, bool const *inputIsBroadcasted, int32_t nbInputs) const noexcept override; bool canBroadcastInputAcrossBatch(int32_t inputIndex) const noexcept override; +#endif private: void deserialize(uint8_t const *data, size_t length); diff --git a/tests/examples/torch_trt/test_custom_op_plugin.py b/tests/examples/torch_trt/test_custom_op_plugin.py new file mode 100644 index 00000000000..7af15527b57 --- /dev/null +++ b/tests/examples/torch_trt/test_custom_op_plugin.py @@ -0,0 +1,32 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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 shutil +import subprocess +from pathlib import Path + +import pytest + + +@pytest.mark.skipif( + shutil.which("cmake") is None or shutil.which("nvcc") is None, + reason="Requires CMake, CUDA Toolkit, and TensorRT development files", +) +def test_identity_conv_plugin_build(tmp_path): + """Build the custom-op plugin against the installed TensorRT C++ interfaces.""" + source = Path(__file__).resolve().parents[3] / "examples/onnx_ptq/custom_op_plugin/plugin" + subprocess.run(["cmake", "-S", str(source), "-B", str(tmp_path)], check=True, timeout=60) + subprocess.run(["cmake", "--build", str(tmp_path), "--parallel", "2"], check=True, timeout=120) + assert (tmp_path / "libidentity_conv_plugin.so").is_file() From 7fbdeb4cd00c973c94d925203b41a2544fc334e2 Mon Sep 17 00:00:00 2001 From: Wei-Ming Chen <17592131+meenchen@users.noreply.github.com> Date: Wed, 14 Oct 2026 14:18:33 -0700 Subject: [PATCH 15/35] [NVBug 6914119] Fix native NemotronH export config mappings (#2794) ### What does this PR do? Type of change: Bug fix Fix native NemotronH exports retaining legacy remote-code config/model mappings. PTQ selects the built-in Transformers implementation, which serializes `layers_block_type`, but the retained `auto_map` directs trusted config reloads to a legacy class exposing that field as a read-only property. This reproduces `AttributeError: property 'layers_block_type' of 'NemotronHConfig' object has no setter`. Remove only the stale `AutoConfig`, `AutoModel`, and `AutoModelForCausalLM` mappings when the loaded config is native Transformers NemotronH. Preserve the native layer schema, genuine remote-code exports, tokenizer mappings, copied sidecars, and the in-memory config. Add offline export/reload regression coverage. ### Usage No new API or flags. Re-export affected checkpoints with the existing PTQ recipe (for example, `general/ptq/fp8_default-kv_fp8`) before deployment. ### Testing - Reproduced the exact error on main `2758c046fbe0` with NVIDIA's published Nemotron-3-Nano-30B-A3B-BF16 config files. Config-only reproduction needs no GPU. - Tests-only on unmodified main: two native export cases fail with the reported error; two remote-code cases pass. All four pass with the fix, including tokenizer-map and source-config preservation. - 150 focused checkpoint/export/load tests passed on each of Transformers 5.5.4 and 5.18.0. The four final regression cases were rerun on 5.5.4. - Published 52-layer config export/reload passed on 5.5.4, preserving all layer types and FP8 KV metadata. - Tiny native NemotronH Mamba/MoE/attention CPU smoke on 5.18.0: BF16 export/model reload preserves all 31 source state tensors; the reported FP8/KV recipe calibrates, exports weight shards, and reloads its native config successfully. - Pre-commit hooks and `git diff --check` passed. ```sh python -m pytest -q --tb=short \ tests/unit/torch/utils/test_hf_checkpoint_utils.py \ tests/unit/torch/export/test_unified_export_hf.py \ tests/unit/torch/export/test_convert_hf_config.py \ tests/unit/torch/utils/test_model_load_utils.py ``` Full 30B PTQ, GPU inference, and TRT-LLM serving were not run. The FP8 smoke validates calibration/export/config reload, not FP8 inference. The 5.5.4 tiny-model smoke cannot initialize the HF MoE adapter on macOS without Triton; the config reproducer and offline regression tests are unaffected. ### Before your PR is "*Ready for review*" Contributor and security guidelines reviewed. Remote-code execution in the new tests is confined to generated local fixtures with `local_files_only=True`. - Is this change backward compatible?: Yes; genuine remote-code exports and tokenizer mappings are preserved. Previously affected native exports need re-exporting. - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A; no copied third-party code or new dependency. - Did you write any new necessary tests?: Yes; four offline export/reload regression cases. - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: Yes; 0.48 bug-fix entry with re-export guidance. - Did you get Claude approval on this PR?: N/A; not requested or run. ### Additional Information NVBug: 6914119. ## Summary by CodeRabbit * **Bug Fixes** * Fixed Nemotron-H exports for use with the built-in Transformers implementation by removing incompatible legacy model mappings from deployment configuration. Affected checkpoints must be re-exported before deployment. Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com> Signed-off-by: Rohan Joshi --- CHANGELOG.rst | 1 + .../utils/plugins/hf_checkpoint_utils.py | 12 ++++ .../torch/utils/test_hf_checkpoint_utils.py | 72 +++++++++++++++++++ 3 files changed, 85 insertions(+) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 2fb53f9f4e5..80b5a643d22 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -126,6 +126,7 @@ Changelog **Bug Fixes** +- Fix exports using the built-in Transformers NemotronH implementation retaining incompatible legacy remote-code config and model mappings. Re-export affected checkpoints before deployment. - Fix Hugging Face exports classifying registered Qwen VLM parameters as unplaced source weights, which overwrote quantized tensors with original weights and incorrectly excluded their modules from quantization. Re-export affected checkpoints. - Right-pad labeled calibration samples to avoid fully masked causal rows that can produce non-finite cuDNN attention gradients during AutoQuantize. Generation previews remove padding before decoding. - Fix calibration silently dropping samples when an out-of-memory error forces the batch to be diff --git a/modelopt/torch/utils/plugins/hf_checkpoint_utils.py b/modelopt/torch/utils/plugins/hf_checkpoint_utils.py index de8532e5297..dc8138d4286 100644 --- a/modelopt/torch/utils/plugins/hf_checkpoint_utils.py +++ b/modelopt/torch/utils/plugins/hf_checkpoint_utils.py @@ -187,10 +187,22 @@ def sanitize_hf_config_for_deployment(config_data: dict[str, Any], model: Any) - Fix conservative deployment-only config incompatibilities: + * remove legacy config/model auto mappings from native NemotronH exports; * add missing llama3 ``rope_theta`` metadata when available; * trim trailing MTP/next-token-prediction ``layer_types`` entries only when the mismatch is exactly explained by next-token-prediction metadata. """ + model_config = getattr(model, "config", None) + auto_map = config_data.get("auto_map") + if isinstance(auto_map, dict) and type(model_config).__module__.startswith( + "transformers.models.nemotron_h." + ): + # The native schema cannot reload through the source's legacy read-only config. + for auto_class in ("AutoConfig", "AutoModel", "AutoModelForCausalLM"): + auto_map.pop(auto_class, None) + if not auto_map: + config_data.pop("auto_map") + _sanitize_llama3_rope_config(config_data, model) num_hidden_layers = _as_nonnegative_int(config_data.get("num_hidden_layers")) diff --git a/tests/unit/torch/utils/test_hf_checkpoint_utils.py b/tests/unit/torch/utils/test_hf_checkpoint_utils.py index 490ca79d7ed..b436a9cb838 100644 --- a/tests/unit/torch/utils/test_hf_checkpoint_utils.py +++ b/tests/unit/torch/utils/test_hf_checkpoint_utils.py @@ -227,6 +227,78 @@ def test_sanitize_hf_config_for_deployment_trims_nextn_layer_types(): assert config_data["layer_types"] == hidden_layer_types +@pytest.mark.parametrize("native", [False, True]) +@pytest.mark.parametrize("keep_tokenizer", [False, True]) +def test_export_config_reloads_nemotron_h_with_source_auto_map(tmp_path, native, keep_tokenizer): + # Native NemotronH is optional on the minimum-Transformers test lane. + transformers = pytest.importorskip("transformers", minversion="5.5") + from modelopt.torch.export.unified_export_hf import ( + _write_hf_export_config, + save_non_weight_artifacts, + ) + + source, export = tmp_path / "source", tmp_path / "export" + source.mkdir() + (source / "configuration_legacy.py").write_text( + "from transformers import PretrainedConfig\n" + "class LegacyConfig(PretrainedConfig):\n" + " model_type = 'nemotron_h'\n" + " def __init__(self, hybrid_override_pattern='M*', **kwargs):\n" + " self.hybrid_override_pattern = hybrid_override_pattern\n" + " super().__init__(**kwargs)\n" + " @property\n" + " def layers_block_type(self):\n" + " return ['mamba', 'attention']\n" + ) + auto_map = { + "AutoConfig": "configuration_legacy.LegacyConfig", + "AutoModel": "modeling_legacy.LegacyModel", + "AutoModelForCausalLM": "modeling_legacy.LegacyModel", + } + if keep_tokenizer: + auto_map["AutoTokenizer"] = [None, "tokenization_legacy.LegacyTokenizer"] + original_auto_map = dict(auto_map) + (source / "config.json").write_text( + json.dumps({"model_type": "nemotron_h", "auto_map": auto_map}) + ) + config = ( + transformers.NemotronHConfig(layers_block_type=["mamba", "attention"], auto_map=auto_map) + if native + else transformers.AutoConfig.from_pretrained( + source, trust_remote_code=True, local_files_only=True + ) + ) + model = torch.nn.Module() + model.config = config + model._modelopt_source_checkpoint = str(source) + quant_config = {"quantization": {"quant_algo": "FP8", "kv_cache_quant_algo": "FP8"}} + + save_non_weight_artifacts(model, export) + _write_hf_export_config(model, quant_config, export) + reloaded = transformers.AutoConfig.from_pretrained( + export, trust_remote_code=True, local_files_only=True + ) + + assert isinstance(reloaded, transformers.NemotronHConfig) == native + assert reloaded.layers_block_type == config.layers_block_type + assert reloaded.quantization_config["quant_algo"] == "FP8" + assert reloaded.quantization_config["kv_cache_scheme"] == { + "dynamic": False, + "num_bits": 8, + "type": "float", + } + assert config.auto_map == original_auto_map + expected_map = { + key: value + for key, value in original_auto_map.items() + if not native or key == "AutoTokenizer" + } + assert getattr(reloaded, "auto_map", {}) == expected_map + assert (export / "configuration_legacy.py").read_text() == ( + source / "configuration_legacy.py" + ).read_text() + + def test_sanitize_hf_config_for_deployment_adds_rope_theta_to_llama3_rope_parameters(): """Transformers 5.x requires rope_theta inside llama3 rope_parameters.""" config_data = { From c5870499659fb31213d6ba6b638582ab6e68f555 Mon Sep 17 00:00:00 2001 From: Chenjie Luo <108829653+cjluo-nv@users.noreply.github.com> Date: Wed, 14 Oct 2026 15:50:29 -0700 Subject: [PATCH 16/35] [NVBugs 6918092] Fix FP8 fake quant zeroing a small nonzero amax (static NVFP4 blocks) (#2787) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Type of change: Bug fix Fixes [NVBug 6918092](https://nvbugspro.nvidia.com/bug/6918092): static NVFP4 fake quantization (`NVFP4StaticQuantizer`) silently zeroes every block of a tensor whose global amax is small. **Root cause.** The FP8 fake quant `scaled_e4m3_impl` (CUDA extension `fake_e4m3fy` / `fake_e4m3fy_with_axis` and the eager fallback) treats any `amax <= 2**-24` as zero and normalizes with a unit amax. Static NVFP4 FP8-quantizes its per-block scales through it with `amax = global_amax / 6` (`compute_fp4_scales`). So when `global_amax < 6 * 2**-24` (~3.6e-7), every block scale rounds to 0 and the whole tensor fake-quantizes to zeros, without an error. The dynamic NVFP4 Triton kernel (after #2549), the CUDA mx fallback and the NVFP4 export path only treat a zero global amax as degenerate. The same floor also hits the other callers that pass a small but nonzero amax: - **LSQ NVFP4 block scales** (`StaticBlockScaleQuantizer._block_scale_from_amax`): the same 2^-22 block is zeroed. - **Plain FP8** tensors or channels with `0 < amax <= 2**-24`: fake quant zeroes them, while HF/TRT-LLM export only clamps amax at fp32's smallest normal (`export_amax`) and keeps them. **Fix.** Lower the zero-amax threshold in `scaled_e4m3_impl` (both CUDA functions and `_fp8_eager`) from `2**-24` to `448 / FLT_MAX`, the smallest amax whose FP8 scale `448 / amax` is still finite. An amax of exactly zero still uses the unit-amax fallback, and NaN/inf amax behave as before. Nothing else in the arithmetic changes, so results are bit-identical wherever the old floor did not apply. The INT fake-quant path keeps its own `2**-24` floor, which `test_cuda_ext_tiny_amax` pins; it is out of scope here. ### Usage ```python # Reproducer from the bug: one constant 2**-22 block, exact amax / global_amax. x = torch.full((1, 16), 2.0**-22, device="cuda") q = NVFP4StaticQuantizer(quant_attribute_cfg=QuantizerAttributeConfig( num_bits=(2, 1), block_sizes={-1: 16, "type": "static", "scale_bits": (4, 3)})).cuda() q.amax = x.abs().amax(-1) q.global_amax = x.abs().amax() assert torch.equal(q(x), x) # was all zeros before this PR ``` ### Testing The B200 nodes on computelab were unavailable, so this was run on a B300 (sm_103) in `nvcr.io/nvidia/tensorrt-llm/release:1.4.0rc0` (torch 2.14, triton 3.8): - **Repro:** the reporter's unchanged `static_probe.py` against `main` (2758c046f) reproduces the bug exactly. The 2^-22 block comes out as 16 zeros with block scale 0 in FP32 and BF16, while the dynamic kernel returns the input. With this PR, all four cases return the input, with the expected block scale of 3.97e-8. The LSQ path (tied and untied amax) also returns the input now. - **Bit-exactness against `main`:** 192 outputs were dumped with both builds, covering FP8 `scaled_e4m3` per-tensor, per-channel, multi-axis eager and CPU, plus `compute_fp4_scales` and static NVFP4 fake quant with `fp8_max_for_normalization` 448 and 256, for magnitudes 2^-30 to 2^39. 174 are bit-identical. The other 18 are exactly the cases where `main` returned all zeros (FP8 amax <= 2^-24, NVFP4 global amax < 6 * 2^-24). - **New tests**, all failing on `main` and passing here: - `TestScaledE4M3::test_small_amax_scales_output`: per-tensor, per-channel and multi-axis, on CUDA and CPU. - `Testfp4::test_static_fp4_kernel_scales_with_small_inputs`: FP32 and BF16; it also checks that an all-zero tensor still gives zeros. - `test_lsq_fp4_small_global_amax`: tied and untied amax. - **GPU suites:** `tests/gpu/torch/quantization/` `test_tensor_quant_cuda`, `test_nvfp4_static_quantizer_cuda`, `test_nvfp4_fp8_sweep_kernel`, `test_lsq_cuda`, `test_gptq`, `test_qtensor_cuda`, `test_gpt_oss_mxfp4_nvfp4_cast_cuda`, `test_quantize_cuda`, `test_real_quantize_cuda`, `test_tensor_quantizer_cuda`, `test_quantize_mxformats_cuda` and `test_torch_export`, plus `tests/gpu/torch/export/test_export_weight_gpu` and `test_export_fp8_linear`: 793 passed, 34 skipped. In the first run, 56 `TestCudaExt` tests (the unchanged INT extension) failed while that extension was building for the first time; rerunning the file with it built gives 147 passed, 2 skipped. - **CPU:** `tests/unit/torch/quantization` (excluding plugins): 1254 passed, 1 skipped. - **Performance:** the kernels and arithmetic are unchanged. `compute_fp4_scales` on 4M blocks takes 60.3 us before and 62.4 us after (run-to-run noise), with the same 32 MiB peak. - **Not run:** the ONNX export tests (the container lacks `onnxruntime` and `onnx_graphsurgeon`), multi-GPU and FSDP2 paths (no distributed code changed), and B200 (no allocation available). The CUDA extension source changed, so `modelopt_cuda_ext_fp8` rebuilds once on first use, as with any change to it. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ (bit-identical except for a positive amax <= 2^-24, which goes from all zeros to correct values) - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: ✅ - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ (0.48.0 Bug Fixes; the bug is committed for 0.48.0, labeled `cherry-pick-0.48.0`) - Did you get Claude approval on this PR?: ❌ ### Additional Information - NVBug: https://nvbugspro.nvidia.com/bug/6918092 - Related: #2549, which removed an equivalent floor from the dynamic NVFP4 Triton kernel. 🤖 Generated with [Claude Code](https://claude.com/claude-code) ## Summary by CodeRabbit * **Bug Fixes** * Corrected FP8 fake quantization for very small positive values, improving consistency with exported checkpoints across tensor and channel quantization. * Improved handling of small-scale static FP4 weights, including cases where all inputs and scales are zero. --------- Signed-off-by: Chenjie Luo Co-authored-by: Claude Opus 5.5 (1M context) Signed-off-by: Rohan Joshi --- CHANGELOG.rst | 1 + .../quantization/gemm/tensor_quant_gpu_fp8.cu | 11 +++-- modelopt/torch/quantization/tensor_quant.py | 4 +- tests/gpu/torch/quantization/test_lsq_cuda.py | 18 ++++++++ .../quantization/test_tensor_quant_cuda.py | 42 +++++++++++++++++++ 5 files changed, 70 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 80b5a643d22..e45c509337d 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -141,6 +141,7 @@ Changelog - Fix ``hf_ptq`` overwriting a model's existing ``pad_token`` with ``eos_token`` when the model already has a valid padding token. The exported tokenizer now preserves the source model's padding configuration. - Fix Hugging Face exports dropping off-index safetensors such as GLM-4.7's ``mtp.safetensors``. - Fix Megatron-Core checkpoint saving for quantized grouped MoE experts when tensor and expert parallelism are both enabled. +- Fix FP8 fake quantization treating a positive amax at or below 2**-24 as zero, which silently zeroed such tensors and channels and every block of a static NVFP4 weight (``NVFP4StaticQuantizer``, including LSQ) with a global amax below ~3.6e-7. Fake quantization now matches the exported checkpoint. - Fix ``megatron_preprocess_data`` with ``workers > 1`` segfaulting when run in a process that has already initialized CUDA or other native threads (for example in-process after GPU work): the tokenization worker pool now uses the ``forkserver`` start method instead of ``fork``. - Fix unified HuggingFace export of RADIO-based VLMs retaining post-conversion vision and projector names instead of restoring the hub layout; deployment loaders could skip those weights. diff --git a/modelopt/torch/kernels/quantization/gemm/tensor_quant_gpu_fp8.cu b/modelopt/torch/kernels/quantization/gemm/tensor_quant_gpu_fp8.cu index 9db141f0cb4..b4493ac8939 100644 --- a/modelopt/torch/kernels/quantization/gemm/tensor_quant_gpu_fp8.cu +++ b/modelopt/torch/kernels/quantization/gemm/tensor_quant_gpu_fp8.cu @@ -17,12 +17,17 @@ #include #include +#include #include #include #include #define BLOCK_SIZE 128 +// Only an amax whose FP8 scale (448 / amax) overflows is treated as zero. Any other positive amax +// is valid, e.g. the tiny global_amax / 6 used to quantize NVFP4 block scales. +#define ZERO_AMAX_THRESHOLD (448.f / FLT_MAX) + #define AT_DISPATCH_CASE_FLOATING_TYPES(...) \ AT_DISPATCH_CASE(at::ScalarType::Double, __VA_ARGS__) \ AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \ @@ -71,8 +76,7 @@ at::Tensor fake_e4m3fy_with_axis(at::Tensor inputs, at::Tensor amax, int axis) { int axis_size = inputs.size(axis); int outer_size = inputs.stride(axis); - const float epsilon = 1.0f / (1 << 24); - auto zero_amax_mask = amax <= epsilon; + auto zero_amax_mask = amax <= ZERO_AMAX_THRESHOLD; auto safe_amax = at::where(zero_amax_mask, torch::ones_like(amax), amax); auto scale = 448.f / safe_amax; auto inv_scale = 1.f / scale; @@ -91,8 +95,7 @@ at::Tensor fake_e4m3fy(at::Tensor inputs, at::Tensor amax) { inputs = inputs.contiguous(); amax = amax.view(-1).to(at::kFloat); size_t numel = inputs.numel(); - const float epsilon = 1.0f / (1 << 24); - auto zero_amax_mask = amax <= epsilon; + auto zero_amax_mask = amax <= ZERO_AMAX_THRESHOLD; auto safe_amax = at::where(zero_amax_mask, torch::ones_like(amax), amax); at::Tensor scale = 448.f / safe_amax; auto inv_scale = 1.f / scale; diff --git a/modelopt/torch/quantization/tensor_quant.py b/modelopt/torch/quantization/tensor_quant.py index 20e083491aa..e9c942aac27 100644 --- a/modelopt/torch/quantization/tensor_quant.py +++ b/modelopt/torch/quantization/tensor_quant.py @@ -47,8 +47,8 @@ def _fp8_eager(x, amax=None): dtype = x.dtype if amax is not None: amax = amax.to(torch.float32) - epsilon = 1.0 / (1 << 24) - zero_amax_mask = amax <= epsilon + # Like the CUDA extension, only an amax whose FP8 scale overflows is treated as zero. + zero_amax_mask = amax <= 448.0 / torch.finfo(torch.float32).max safe_amax = torch.where(zero_amax_mask, torch.ones_like(amax), amax) scale = 448.0 / safe_amax scale_inv = 1 / scale diff --git a/tests/gpu/torch/quantization/test_lsq_cuda.py b/tests/gpu/torch/quantization/test_lsq_cuda.py index ffb6fce95d4..4e9890737ec 100644 --- a/tests/gpu/torch/quantization/test_lsq_cuda.py +++ b/tests/gpu/torch/quantization/test_lsq_cuda.py @@ -182,6 +182,24 @@ def test_lsq_fp4_fake_quantize_differentiable(): assert sbsq._amax_post.grad is not None +@pytest.mark.parametrize("tied_amax", [False, True]) +def test_lsq_fp4_small_global_amax(tied_amax): + """A block representable in NVFP4 survives LSQ fake quant even when global_amax < 6 * 2**-24.""" + device = torch.device("cuda") + x = torch.full((1, 16), 2.0**-22, device=device) + tq = TensorQuantizer() + tq._num_bits = (2, 1) + tq._unsigned = False + tq._narrow_range = True + tq._disabled = False + tq._block_sizes = {-1: 16, "type": "static", "scale_bits": (4, 3)} + tq.register_buffer("_amax", x.abs().amax(-1)) + sbsq = StaticBlockScaleQuantizer.from_tensor_quantizer(tq, global_amax=x.abs().amax()) + sbsq.enable_lsq(quantize_scales=True, tied_amax=tied_amax) + + assert torch.equal(sbsq._fake_quantize(x), x) + + def test_lsq_fp4_cast_ste(): """Test fp4_cast_ste on GPU.""" device = torch.device("cuda") diff --git a/tests/gpu/torch/quantization/test_tensor_quant_cuda.py b/tests/gpu/torch/quantization/test_tensor_quant_cuda.py index cc3262e1917..c1cbd6a09fa 100644 --- a/tests/gpu/torch/quantization/test_tensor_quant_cuda.py +++ b/tests/gpu/torch/quantization/test_tensor_quant_cuda.py @@ -178,6 +178,23 @@ def test_zero_amax_per_channel_is_finite(self): xq = tensor_quant.scaled_e4m3(x, amax, None, 4, 3) assert torch.isfinite(xq).all() + @pytest.mark.parametrize("device", ["cuda", "cpu"]) + @pytest.mark.parametrize( + "amax_dims", [None, (1, 2), (2,)], ids=["per_tensor", "per_channel", "multi_axis"] + ) + def test_small_amax_scales_output(self, device, amax_dims): + """Scaling the input and amax by a power of two scales the output by the same factor. + + A positive amax below 2**-24 is not zero and must not flush the tensor to zero. + """ + torch.manual_seed(0) + x = torch.randn(4, 3, 16, device=device) * 10 + amax = x.abs().amax() if amax_dims is None else x.abs().amax(amax_dims, keepdim=True) + reference = tensor_quant.scaled_e4m3(x, amax, None, 4, 3) + for factor in (2.0**-30, 2.0**-60): + scaled = tensor_quant.scaled_e4m3(x * factor, amax * factor, None, 4, 3) + assert torch.equal(scaled, reference * factor) + class Testfp4: @pytest.mark.skipif(not NATIVE_E4M3_AVAILABLE, reason="Native E4M3 requires compute >= 8.9") @@ -373,6 +390,31 @@ def test_fp4_kernel_scales_with_small_inputs(self, dtype): scaled = triton_kernel.fp4_fake_quant_block(x * factor, x.float().abs().amax() * factor) assert torch.equal(scaled, reference * factor) + @pytest.mark.skipif(not triton_kernel.IS_AVAILABLE, reason="triton kernel is not available") + @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) + def test_static_fp4_kernel_scales_with_small_inputs(self, dtype): + """Static counterpart of ``test_fp4_kernel_scales_with_small_inputs``. + + A global amax below 6 * 2**-24 must not zero the FP8-quantized block scales, while an + all-zero tensor still quantizes to zeros. + """ + torch.manual_seed(0) + x = (torch.randn(8, 64, device="cuda") * 10).to(dtype) + amax = x.float().abs().view(-1, 16).amax(-1) + reference = triton_kernel.static_blockwise_fp4_fake_quant(x, amax, amax.max()) + assert reference.abs().amax() > 0 + for factor in (2.0**-10, 2.0**-20, 2.0**-30): + scaled = triton_kernel.static_blockwise_fp4_fake_quant( + x * factor, amax * factor, amax.max() * factor + ) + assert torch.equal(scaled, reference * factor) + + zeros = torch.zeros_like(x) + output = triton_kernel.static_blockwise_fp4_fake_quant( + zeros, torch.zeros_like(amax), torch.zeros_like(amax.max()) + ) + assert torch.equal(output, zeros) + @pytest.mark.skipif( not hasattr(triton_kernel, "fp4_fake_quant_block"), reason="fp4_fake_quant_block requires compute >= 8.9", From c353ad530b7212b9c4e8f9ad575f70702fb97d95 Mon Sep 17 00:00:00 2001 From: Ajinkya Rasane <131806219+ajrasane@users.noreply.github.com> Date: Wed, 14 Oct 2026 19:49:38 -0400 Subject: [PATCH 17/35] [1/4] [6466805] Make quantization carriers safe for CPU export capture (#2682) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Type of change: new feature Part 1 of the four-PR Dynamo ONNX export stack replacing #2418. Adds export metadata as optional trailing custom-op schema arguments, keeps amax detached during compiler capture, and routes CPU integer carriers through the existing tensor implementation. Legacy positional calls and eager numerical behavior are preserved. This is capture plumbing only; the supported quantized ONNX helper workflow is enabled in part 4. Automatic CPU INT8/INT4 carrier routing is deferred until the ONNX helper supplies its translation table. CPU integer quantizers retain the existing tensor-math path; the registered carriers can be captured explicitly. ### Usage No new public interface in this prerequisite. Existing quantization calls remain unchanged. ### Testing With CUDA hidden on PyTorch 2.14 development: **91 passed** across CPU quantization and ONNX helper tests. ```bash CUDA_VISIBLE_DEVICES="" python -m pytest -o addopts= -p no:cacheprovider \ tests/unit/torch/quantization/test_tensor_quant_cpu.py \ tests/unit/torch/deploy/utils/test_torch_onnx_utils.py -q ``` On PyTorch 2.9 development, the focused CPU quantization suite also passed **40 tests**, and both INT8/INT4 Dynamo helper calls retained their working legacy fallback. Coverage includes legacy positional schemas, CPU numerical parity, real strict `torch.export` capture of the registered INT8/INT4 carriers, and FP8 `TensorQuantizer` capture. Changed-file pre-commit checks and `git diff --check` passed. GPU and full multi-version CI were not run for this revision. Compatibility probes compared the original PR with its pinned `main` base on PyTorch 2.9, 2.13, and 2.14 development builds. For fixed-amax CPU INT8/INT4 quantizers, the original PR captured a carrier but failed ONNX translation; the base failed earlier on data-dependent capture. The PyTorch 2.9 helper succeeded for both through legacy fallback. These probes do not establish that nobody uses CPU Dynamo export, so this revision preserves the existing CPU routing instead of relying on that assumption. ### Before your PR is "*Ready for review*" Contributor guidelines: https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md. Commits are signed and signed off with `git commit -s -S`. Security Best Practices: https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors. - Is this change backward compatible?: ✅ - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A — neither added. - Did you write any new necessary tests?: ✅ - Did you update Changelog?: N/A — private prerequisite for the same unreleased feature; the user-facing feature entry lands with the helper in part 4. - Did you get Claude approval on this PR?: Review feedback addressed; awaiting re-review. ### Additional Information Merge order: capture → lowering → normalization → integration. This PR targets `main`; part 2 targets this branch. Parts 1–3 intentionally do not expose a new export workflow. The complete feature remains helper-owned and supports FP8, INT8, INT4 AWQ, MXFP8, NVFP4, and mixed FP8 + NVFP4. > 🤖 _Generated by Codex (AI agent)._ ### Replacement stack Merge in order; each PR targets the preceding branch. Retarget the next PR when its dependency merges. 1. [#2682: capture and CPU routing](https://github.com/NVIDIA/Model-Optimizer/pull/2682) 2. [#2683: private ONNXScript lowering](https://github.com/NVIDIA/Model-Optimizer/pull/2683) 3. [#2684: canonical weight normalization](https://github.com/NVIDIA/Model-Optimizer/pull/2684) 4. [#2685: helper and deployment integration](https://github.com/NVIDIA/Model-Optimizer/pull/2685) ## Summary by CodeRabbit * **New Features** * Dynamo ONNX export now supports quantization and dequantization for FP8, INT8/UINT8, INT4, NVFP4, and MXFP8 formats, with validation for supported configurations. * Quantization operators accept optional TensorRT dtype, block-size, axis, and ONNX quantizer metadata. * **Bug Fixes** * CPU quantization now matches eager fake-quantization behavior for supported inputs. * Quantization scale values are detached during export. --------- Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> Signed-off-by: Ajinkya Rasane <131806219+ajrasane@users.noreply.github.com> Signed-off-by: Rohan Joshi --- modelopt/torch/quantization/_dynamo_onnx.py | 266 ++++++++++++++++++ .../nn/modules/tensor_quantizer.py | 5 +- modelopt/torch/quantization/tensor_quant.py | 49 ++-- .../quantization/test_dynamo_onnx_export.py | 226 +++++++++++++++ .../quantization/test_tensor_quant_cpu.py | 66 ++++- 5 files changed, 591 insertions(+), 21 deletions(-) create mode 100644 modelopt/torch/quantization/_dynamo_onnx.py create mode 100644 tests/unit/torch/quantization/test_dynamo_onnx_export.py diff --git a/modelopt/torch/quantization/_dynamo_onnx.py b/modelopt/torch/quantization/_dynamo_onnx.py new file mode 100644 index 00000000000..55439488184 --- /dev/null +++ b/modelopt/torch/quantization/_dynamo_onnx.py @@ -0,0 +1,266 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""Private ONNXScript translations for quantized Dynamo export.""" + +import onnx +import onnxscript +import torch +from onnxscript.function_libs.torch_lib.tensor_typing import TFloat +from onnxscript.onnx_types import FLOAT4E2M1, FLOAT8E4M3FN + +from .export_onnx import onnx_dtype_map + +__all__ = [] + +_OPSET = onnxscript.opset23 +_TRT_OPSET = onnxscript.values.Opset(domain="trt", version=1) + + +@onnxscript.script(_TRT_OPSET) +def _fp8_qdq(inputs: TFloat, scale: TFloat) -> TFloat: + quantized = _TRT_OPSET.TRT_FP8QuantizeLinear(inputs, scale) + return _TRT_OPSET.TRT_FP8DequantizeLinear(quantized, scale) + + +@onnxscript.script(_TRT_OPSET) +def _int4_dq(inputs: TFloat, scale: TFloat, axis: int, block_size: int) -> TFloat: + return _TRT_OPSET.DequantizeLinear(inputs, scale, axis=axis, block_size=block_size) + + +@onnxscript.script(_TRT_OPSET) +def _fp4_qdq(inputs: TFloat, block_size: int) -> TFloat: + return _TRT_OPSET.TRT_FP4QDQ(inputs, block_size=block_size) + + +@onnxscript.script(_TRT_OPSET) +def _fp4_dynamic_quantize( + inputs: TFloat, scale: TFloat, block_size: int +) -> tuple[FLOAT4E2M1[...], FLOAT8E4M3FN[...]]: + quantized, dynamic_scale = _TRT_OPSET.TRT_FP4DynamicQuantize( + inputs, scale, axis=-1, block_size=block_size, scale_type=17 + ) + return quantized, dynamic_scale + + +@onnxscript.script(_TRT_OPSET) +def _mxfp8_dynamic_qdq(inputs: TFloat, block_size: int, output_dtype: int) -> TFloat: + quantized, scale = _TRT_OPSET.TRT_MXFP8DynamicQuantize( + inputs, axis=-1, block_size=block_size, output_dtype=17 + ) + return _TRT_OPSET.TRT_MXFP8DequantizeLinear( + quantized, + scale, + axis=-1, + block_size=block_size, + output_dtype=output_dtype, + ) + + +@onnxscript.script(_TRT_OPSET) +def _mxfp8_static_dq(inputs: TFloat, scale: TFloat, block_size: int, output_dtype: int) -> TFloat: + return _TRT_OPSET.TRT_MXFP8DequantizeLinear( + inputs, + scale, + axis=-1, + block_size=block_size, + output_dtype=output_dtype, + ) + + +def _cast(inputs, dtype: int): + return inputs if int(inputs.dtype) == dtype else _OPSET.Cast(inputs, to=dtype) + + +def _resolve_dtype(inputs, high_precision_dtype: str | None) -> int: + return ( + int(inputs.dtype) if high_precision_dtype is None else onnx_dtype_map[high_precision_dtype] + ) + + +def _static_shape(value) -> list[int]: + try: + return [int(dim) for dim in value.shape] + except (AttributeError, TypeError, ValueError): + raise NotImplementedError("Dynamo ONNX export does not support dynamic shapes.") from None + + +def _block_activation_shape( + inputs, block_size: int, quantizer_type: str | None +) -> list[int] | None: + if quantizer_type != "dynamic": + return None + shape = _static_shape(inputs) + if len(shape) not in (2, 3): + raise NotImplementedError("Dynamo ONNX block activation export supports rank 2 or 3 only.") + if shape[-1] % block_size: + raise NotImplementedError( + f"Dynamo ONNX block size {block_size} must divide the last dimension {shape[-1]}." + ) + return shape + + +def _translate_quantize_op( + inputs, + amax, + num_bits: int, + exponent_bits: int, + unsigned: bool, + narrow_range: bool, + high_precision_dtype: str | None = None, + block_size: int | None = None, + axis: int | None = None, +): + if num_bits == 8 and exponent_bits == 4: + if amax is None: + scale = _OPSET.CastLike(1.0, inputs) + elif any(dim != 1 for dim in _static_shape(amax)): + raise AssertionError( + "E4M3 supports ONNX export only for per-tensor quantization with scalar amax." + ) + else: + scale = _OPSET.Div(_OPSET.CastLike(amax, inputs), _OPSET.CastLike(448.0, inputs)) + return _fp8_qdq(inputs, scale) + + output_dtype = _resolve_dtype(inputs, high_precision_dtype) + if num_bits == 8 and exponent_bits == 0: + source_dtype = int(inputs.dtype) + input_shape = _static_shape(inputs) + if not unsigned: + assert not narrow_range, "ONNX does not support signed narrow-range INT8." + assert output_dtype in ( + source_dtype, + onnx.TensorProto.FLOAT, + onnx.TensorProto.BFLOAT16, + ), "TensorRT strongly typed mode requires Q/DQ in the input dtype, FP32, or BF16." + inputs = _cast(inputs, output_dtype) + input_rank = len(input_shape) + amax_shape = _static_shape(amax) + quantized_axes = [index for index, dim in enumerate(amax_shape) if dim != 1] + if len(quantized_axes) > 1: + raise AssertionError("ONNX does not support multi-axis quantization.") + if axis is not None: + if not -input_rank <= axis < input_rank: + raise ValueError(f"INT8 axis {axis} is out of bounds for input rank {input_rank}.") + axis %= input_rank + if not quantized_axes: + axis = None + elif axis is None: + if len(amax_shape) != input_rank: + raise ValueError("INT8 per-channel amax requires an explicit input axis.") + axis = quantized_axes[0] + elif len(amax_shape) == input_rank and quantized_axes[0] != axis: + inferred_axis = quantized_axes[0] + raise ValueError(f"INT8 axis {axis} does not match amax axis {inferred_axis}.") + amax = _OPSET.Cast(amax, to=output_dtype) + amax = _OPSET.Squeeze(amax) if axis is None else _OPSET.Reshape(amax, [-1]) + scale = _OPSET.Div(amax, float((1 << (7 + int(unsigned))) - 1)) + scale = _OPSET.Where(_OPSET.Equal(scale, 0.0), _OPSET.CastLike(1.0, scale), scale) + zero_point_dtype = onnx.TensorProto.UINT8 if unsigned else onnx.TensorProto.INT8 + zero_point = _OPSET.Cast(_OPSET.Mul(amax, 0.0), to=zero_point_dtype) + if axis is None: + quantized = _OPSET.QuantizeLinear(inputs, scale, zero_point) + output = _OPSET.DequantizeLinear(quantized, scale, zero_point) + else: + quantized = _OPSET.QuantizeLinear(inputs, scale, zero_point, axis=axis) + output = _OPSET.DequantizeLinear(quantized, scale, zero_point, axis=axis) + return output if output_dtype == source_dtype else _OPSET.Cast(output, to=source_dtype) + + if num_bits == 4 and exponent_bits == 0: + if unsigned: + raise NotImplementedError("Dynamo ONNX export supports signed INT4 only.") + if block_size is None or axis is None: + raise ValueError("INT4 ONNX export requires block_size and axis.") + if amax is None: + raise ValueError("INT4 ONNX export requires amax.") + source_dtype = int(inputs.dtype) + amax = _OPSET.Cast(amax, to=output_dtype) + scale = _OPSET.Div(amax, _OPSET.CastLike(7.0, amax)) + output = _int4_dq(inputs, scale, axis, block_size) + # The TRT carrier has no schema to infer metadata when input and scale dtypes differ. + output.dtype = onnxscript.ir.DataType(output_dtype) + output.shape = inputs.shape + return output if output_dtype == source_dtype else _OPSET.Cast(output, to=source_dtype) + + raise NotImplementedError( + f"Unsupported num_bits={num_bits}, exponent_bits={exponent_bits} for ONNX export." + ) + + +def _translate_dynamic_block_quantize_op( + inputs, + block_size: int, + amax, + num_bits: int, + exponent_bits: int, + scale_num_bits: int, + scale_exponent_bits: int, + high_precision_dtype: str | None = None, + quantizer_type: str | None = None, +): + activation_shape = _block_activation_shape(inputs, block_size, quantizer_type) + format_bits = (num_bits, exponent_bits, scale_num_bits, scale_exponent_bits) + + if format_bits == (4, 2, 8, 4): + if quantizer_type != "dynamic": + # Keep a standard-domain op so older ONNXScript retains the default opset import. + return _OPSET.Identity(_fp4_qdq(inputs, block_size)) + assert activation_shape is not None + source_dtype = int(inputs.dtype) + output_dtype = _resolve_dtype(inputs, high_precision_dtype) + inputs = _cast(inputs, output_dtype) + if amax is None: + scale = _OPSET.Constant(value_float=1.0) + else: + scale = _OPSET.Div(_OPSET.Cast(amax, to=onnx.TensorProto.FLOAT), 2688.0) + scale = _OPSET.Where(_OPSET.Equal(scale, 0.0), _OPSET.CastLike(1.0, scale), scale) + quantized, dynamic_scale = _fp4_dynamic_quantize(inputs, scale, block_size) + # ONNXScript before 0.6 may drop custom multi-output metadata during inlining. + quantized.dtype = onnxscript.ir.DataType.FLOAT4E2M1 + dynamic_scale.dtype = onnxscript.ir.DataType.FLOAT8E4M3FN + quantized.shape = onnxscript.ir.Shape(activation_shape) + scale_shape = [*activation_shape[:-1], activation_shape[-1] // block_size] + dynamic_scale.shape = onnxscript.ir.Shape(scale_shape) + dequantized_scale = _OPSET.DequantizeLinear(dynamic_scale, scale) + output = _OPSET.DequantizeLinear( + quantized, dequantized_scale, axis=-1, block_size=block_size + ) + return _OPSET.Cast(output, to=source_dtype) + + if format_bits == (8, 4, 9, 8): + output_dtype = int(inputs.dtype) + if quantizer_type == "dynamic": + output = _mxfp8_dynamic_qdq(inputs, block_size, output_dtype) + else: + output = _mxfp8_static_dq( + inputs, _OPSET.CastLike(1.0, inputs), block_size, output_dtype + ) + # Keep a standard-domain op so older ONNXScript retains the default opset import. + return _OPSET.Identity(output) + + raise NotImplementedError( + "Unsupported block format " + f"({exponent_bits}, {num_bits - exponent_bits - 1}) with scale format " + f"({scale_exponent_bits}, {scale_num_bits - scale_exponent_bits - 1})." + ) + + +def _get_dynamo_onnx_translation_table() -> dict: + dynamic_op = torch.ops.tensorrt.dynamic_block_quantize_op + return { + torch.ops.tensorrt.quantize_op.default: _translate_quantize_op, + dynamic_op.default: _translate_dynamic_block_quantize_op, + dynamic_op.overload: _translate_dynamic_block_quantize_op, + } diff --git a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py index a698514d53a..1893ed3e800 100644 --- a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py +++ b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py @@ -749,8 +749,9 @@ def _get_amax(self, inputs): reduce_axis = quant_utils.convert_quantization_axis_to_reduce_axis(inputs, self._axis) amax = quant_utils.reduce_amax(inputs, axis=reduce_axis, keepdims=True).detach() - amax = amax.detach() if is_torch_export_mode() else amax.data - return amax + if is_torch_export_mode() or torch.compiler.is_exporting(): + return amax.detach() + return amax.data def validate_attr( self, attr_value=None, attr_name="amax", raise_error=False, warn_error=False, name="" diff --git a/modelopt/torch/quantization/tensor_quant.py b/modelopt/torch/quantization/tensor_quant.py index e9c942aac27..7ba13264576 100644 --- a/modelopt/torch/quantization/tensor_quant.py +++ b/modelopt/torch/quantization/tensor_quant.py @@ -119,11 +119,16 @@ def _quantize_impl( exponent_bits: int = 0, unsigned: bool = False, narrow_range: bool = True, + trt_high_precision_dtype: str | None = None, + block_size: int | None = None, + axis: int | None = None, ): if num_bits == 8 and exponent_bits == 4: return scaled_e4m3_impl(inputs=inputs, amax=amax) elif isinstance(num_bits, int): - return fake_quant_impl( + # CPU export executes this carrier directly; use the same math as the eager CPU fallback. + quantize_impl = fake_quant_impl if inputs.is_cuda else _tensor_quant + return quantize_impl( inputs=inputs, amax=amax, num_bits=num_bits, @@ -143,14 +148,15 @@ def _quantize_impl_abstract( exponent_bits: int = 0, unsigned: bool = False, narrow_range: bool = True, + trt_high_precision_dtype: str | None = None, + block_size: int | None = None, + axis: int | None = None, ) -> torch.Tensor: """Register an abstract implementation for quantizing tensor. This abstract function returns an empty tensor with the same shape and dtype. """ - output = torch.empty_like(input) - - return output + return torch.empty_like(input) # Argument types: Tensor, int, NoneType, int, int, int, int, @@ -162,6 +168,8 @@ def _dynamic_block_quantize_impl( exponent_bits: int, scale_num_bits: int, scale_exponent_bits: int, + trt_high_precision_dtype: str | None = None, + onnx_quantizer_type: str | None = None, ): scale_bits = (scale_exponent_bits, scale_num_bits - scale_exponent_bits - 1) if exponent_bits != 0: @@ -203,14 +211,14 @@ def _dynamic_block_quantize_impl_abstract( exponent_bits: int, scale_num_bits: int, scale_exponent_bits: int, + trt_high_precision_dtype: str | None = None, + onnx_quantizer_type: str | None = None, ): """Register an abstract implementation for dynamic block quantization. This abstract function returns an empty tensor with the same shape and dtype. """ - output = torch.empty_like(inputs) - - return output + return torch.empty_like(inputs) quantize_op = _quantize_impl @@ -223,17 +231,20 @@ def _dynamic_block_quantize_impl_abstract( torch.library.define( "tensorrt::quantize_op", "(Tensor input, Tensor amax, int num_bits, int exponent_bits, " - "bool unsigned, bool narrow_range) -> Tensor", + "bool unsigned, bool narrow_range, str? trt_high_precision_dtype=None, " + "int? block_size=None, int? axis=None) -> Tensor", ) torch.library.define( "tensorrt::dynamic_block_quantize_op", "(Tensor input, int block_size, Tensor amax, int num_bits, int exponent_bits, " - "int scale_num_bits, int scale_exponent_bits) -> Tensor", + "int scale_num_bits, int scale_exponent_bits, str? trt_high_precision_dtype=None, " + "str? onnx_quantizer_type=None) -> Tensor", ) torch.library.define( "tensorrt::dynamic_block_quantize_op.overload", "(Tensor input, int block_size, None amax, int num_bits, int exponent_bits, " - "int scale_num_bits, int scale_exponent_bits) -> Tensor", + "int scale_num_bits, int scale_exponent_bits, str? trt_high_precision_dtype=None, " + "str? onnx_quantizer_type=None) -> Tensor", ) # Implement the None amax case @@ -245,6 +256,8 @@ def _dynamic_block_quantize_impl_none_amax( exponent_bits: int, scale_num_bits: int, scale_exponent_bits: int, + trt_high_precision_dtype: str | None = None, + onnx_quantizer_type: str | None = None, ): return torch.empty_like(inputs) @@ -368,8 +381,7 @@ def forward( def legacy_quant_func(): # The LegacyFakeTensorQuantFunction support cpu and amax with any shape that can be broadcasted to inputs. - outputs = _tensor_quant(inputs, amax, num_bits, unsigned, narrow_range) - return outputs + return _tensor_quant(inputs, amax, num_bits, unsigned, narrow_range) if not inputs.is_cuda: outputs = legacy_quant_func() @@ -382,6 +394,9 @@ def legacy_quant_func(): exponent_bits=0, unsigned=unsigned, narrow_range=narrow_range, + trt_high_precision_dtype=trt_high_precision_dtype, + block_size=block_size, + axis=axis, ) except (AttributeError, ValueError): # AttributeError: cuda_ext is not imported, possibly due to CPU only installation @@ -447,6 +462,7 @@ def forward( exponent_bits=4, unsigned=False, narrow_range=False, + trt_high_precision_dtype=trt_high_precision_dtype, ) if bias is not None: @@ -461,7 +477,6 @@ def backward(ctx, grad_outputs): def _dynamic_block_quantize_forward( - ctx, inputs, block_size, amax, @@ -469,7 +484,6 @@ def _dynamic_block_quantize_forward( scale_bits, trt_high_precision_dtype=None, onnx_quantizer_type="dynamic", - pass_through_bwd=True, ): """Forward method.""" if isinstance(num_bits, int): @@ -482,7 +496,7 @@ def _dynamic_block_quantize_forward( assert isinstance(scale_bits, tuple) and len(scale_bits) == 2 scale_exponent_bits = scale_bits[0] scale_num_bits = scale_bits[0] + scale_bits[1] + 1 - outputs = dynamic_block_quantize_op( + return dynamic_block_quantize_op( inputs, block_size, amax, @@ -490,8 +504,9 @@ def _dynamic_block_quantize_forward( exponent_bits, scale_num_bits, scale_exponent_bits, + trt_high_precision_dtype, + onnx_quantizer_type, ) - return outputs class DynamicBlockQuantizationFunction(Function): @@ -551,7 +566,6 @@ def forward( """Forward method.""" _save_for_backward_if_needed(ctx, pass_through_bwd, inputs, amax) return _dynamic_block_quantize_forward( - ctx, inputs, block_size, amax, @@ -559,7 +573,6 @@ def forward( scale_bits, trt_high_precision_dtype, onnx_quantizer_type, - pass_through_bwd, ) @staticmethod diff --git a/tests/unit/torch/quantization/test_dynamo_onnx_export.py b/tests/unit/torch/quantization/test_dynamo_onnx_export.py new file mode 100644 index 00000000000..1c0b64b7210 --- /dev/null +++ b/tests/unit/torch/quantization/test_dynamo_onnx_export.py @@ -0,0 +1,226 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""CPU contract tests for the supported Dynamo ONNX export path.""" + +import inspect + +import pytest +import torch +from torch import nn + +onnx = pytest.importorskip("onnx") +pytest.importorskip("onnxscript") + +from modelopt.torch.quantization._dynamo_onnx import _get_dynamo_onnx_translation_table + +_NVFP4_ARGS = (4, 2, 8, 4, "Float") +_MXFP8_ARGS = (8, 4, 9, 8, "Float") + + +class _StrictQuantOp(nn.Module): + def __init__(self, quant_format, axis=0): + super().__init__() + self.quant_format = quant_format + self.axis = axis + + def forward(self, x, amax): + if self.quant_format.startswith("fp8"): + return torch.ops.tensorrt.quantize_op.default( + x, + None if self.quant_format == "fp8_none" else amax, + 8, + 4, + False, + False, + "Float", + ) + if self.quant_format in {"int8", "uint8"}: + return torch.ops.tensorrt.quantize_op.default( + x, amax, 8, 0, self.quant_format == "uint8", False, "Float", None, self.axis + ) + if self.quant_format == "int4": + return torch.ops.tensorrt.quantize_op.default( + x, amax, 4, 0, False, True, "Float", 16, -1 + ) + if self.quant_format.startswith("nvfp4"): + return torch.ops.tensorrt.dynamic_block_quantize_op.default( + x, 16, amax, *_NVFP4_ARGS, self.quant_format.removeprefix("nvfp4_") + ) + return torch.ops.tensorrt.dynamic_block_quantize_op.overload( + x, 32, None, *_MXFP8_ARGS, self.quant_format.removeprefix("mxfp8_") + ) + + +def _assert_valid_without_custom_ops(model): + onnx.checker.check_model(model, full_check=True) + assert not model.functions + assert not any( + node.domain == "tensorrt" or node.op_type in {"quantize_op", "dynamic_block_quantize_op"} + for node in model.graph.node + ) + return model + + +def _raw_export(model, inputs, path, opset=23): + compatibility_options = ( + {"fallback": False} if "fallback" in inspect.signature(torch.onnx.export).parameters else {} + ) + torch.onnx.export( + model, + inputs, + path, + dynamo=True, + opset_version=opset, + custom_translation_table=_get_dynamo_onnx_translation_table(), + **compatibility_options, + ) + exported = onnx.load(path) + return _assert_valid_without_custom_ops(exported) + + +def _attribute(node, name): + return onnx.helper.get_attribute_value( + next(attr for attr in node.attribute if attr.name == name) + ) + + +def _constant(model, name): + initializer = next((item for item in model.graph.initializer if item.name == name), None) + if initializer is not None: + return onnx.numpy_helper.to_array(initializer) + node = next( + item for item in model.graph.node if item.op_type == "Constant" and name in item.output + ) + return onnx.numpy_helper.to_array(_attribute(node, "value")) + + +_STRICT_CASES = [ + ("fp8_none", ("trt", "TRT_FP8QuantizeLinear")), + ("fp8", ("trt", "TRT_FP8QuantizeLinear")), + ("int8", ("", "QuantizeLinear")), + ("uint8", ("", "QuantizeLinear")), + ("int4", ("trt", "DequantizeLinear")), + ("nvfp4_static", ("trt", "TRT_FP4QDQ")), + ("mxfp8_static", ("trt", "TRT_MXFP8DequantizeLinear")), + ("mxfp8_dynamic", ("trt", "TRT_MXFP8DynamicQuantize")), +] + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16]) +@pytest.mark.parametrize(("quant_format", "expected_node"), _STRICT_CASES) +def test_private_table_strict_formats(tmp_path, quant_format, expected_node, dtype): + sample_input = torch.randn(4, 32, dtype=dtype) + if quant_format in {"int8", "uint8"}: + amax = torch.ones(4, 1) + elif quant_format == "int4": + amax = torch.ones(4, 2) + else: + amax = torch.tensor(1.0) + exported_program = torch.export.export( + _StrictQuantOp(quant_format), (sample_input, amax), strict=True + ) + exported = _raw_export(exported_program, (), tmp_path / f"{quant_format}.onnx") + assert expected_node in {(node.domain, node.op_type) for node in exported.graph.node} + expected_dtype = onnx.TensorProto.FLOAT if dtype == torch.float32 else onnx.TensorProto.FLOAT16 + assert exported.graph.output[0].type.tensor_type.elem_type == expected_dtype + tensor_types = { + value.name: value.type.tensor_type.elem_type + for value in (*exported.graph.input, *exported.graph.value_info, *exported.graph.output) + } + if quant_format in {"int8", "uint8"}: + quantizer = next(node for node in exported.graph.node if node.op_type == "QuantizeLinear") + assert _attribute(quantizer, "axis") == 0 + unsigned = quant_format == "uint8" + assert tensor_types[quantizer.input[2]] == ( + onnx.TensorProto.UINT8 if unsigned else onnx.TensorProto.INT8 + ) + divisor = next(node for node in exported.graph.node if node.op_type == "Div").input[1] + assert _constant(exported, divisor).item() == (255.0 if unsigned else 127.0) + elif quant_format == "int4": + dequantizer = next( + node for node in exported.graph.node if node.op_type == "DequantizeLinear" + ) + assert _attribute(dequantizer, "axis") == -1 + assert _attribute(dequantizer, "block_size") == 16 + assert tensor_types[dequantizer.output[0]] == onnx.TensorProto.FLOAT + divisor = next(node for node in exported.graph.node if node.op_type == "Div").input[1] + assert _constant(exported, divisor).item() == 7.0 + elif quant_format == "fp8": + divisor = next(node for node in exported.graph.node if node.op_type == "Div").input[1] + assert _constant(exported, divisor).item() == 448.0 + + +@pytest.mark.parametrize("amax_shape", [(), (1,), (1, 1)]) +@pytest.mark.parametrize("axis", [0, -1]) +def test_private_table_int8_singleton_amax_is_per_tensor(tmp_path, amax_shape, axis): + inputs = torch.randn(4, 32) + exported_program = torch.export.export( + _StrictQuantOp("int8", axis), (inputs, torch.ones(amax_shape)), strict=True + ) + exported = _raw_export(exported_program, (), tmp_path / "int8_per_tensor.onnx") + tensor_shapes = { + value.name: value.type.tensor_type.shape + for value in (*exported.graph.input, *exported.graph.value_info, *exported.graph.output) + } + for node in exported.graph.node: + if node.op_type in {"QuantizeLinear", "DequantizeLinear"}: + assert len(tensor_shapes[node.input[1]].dim) == 0 + assert len(tensor_shapes[node.input[2]].dim) == 0 + + +def test_private_table_rejects_int4_without_amax(tmp_path): + exported_program = torch.export.export( + _StrictQuantOp("int4"), (torch.randn(4, 32), None), strict=True + ) + with pytest.raises(torch.onnx.OnnxExporterError, match="INT4 ONNX export requires amax"): + _raw_export(exported_program, (), tmp_path / "int4_without_amax.onnx") + + +def test_private_table_preserves_nvfp4_carrier_dtype(tmp_path): + inputs = torch.randn(4, 32, dtype=torch.float16) + exported_program = torch.export.export( + _StrictQuantOp("nvfp4_dynamic"), (inputs, torch.tensor(1.0)), strict=True + ) + + exported = _raw_export(exported_program, (), tmp_path / "nvfp4_dtype.onnx", 23) + + assert exported.graph.output[0].type.tensor_type.elem_type == onnx.TensorProto.FLOAT16 + assert {node.domain for node in exported.graph.node if node.op_type == "DequantizeLinear"} == { + "" + } + + +def test_private_table_rejects_int8_axis_disagreement(tmp_path): + inputs = torch.randn(2, 3, 4) + amax = torch.ones(1, 3, 1) + exported_program = torch.export.export(_StrictQuantOp("int8"), (inputs, amax), strict=True) + + with pytest.raises(torch.onnx.OnnxExporterError, match="axis 0 does not match amax axis 1"): + _raw_export(exported_program, (), tmp_path / "int8_axis.onnx", 23) + + +@pytest.mark.parametrize( + ("fmt", "inputs", "amax", "message"), + [ + ("fp8", torch.randn(2, 4), torch.ones(2, 1), "scalar amax"), + ("int8", torch.randn(2, 3, 4), torch.ones(2, 1, 4), "multi-axis"), + ("nvfp4_dynamic", torch.randn(1, 2, 3, 32), torch.tensor(1.0), "rank 2 or 3"), + ], +) +def test_private_table_rejects_shapes(tmp_path, fmt, inputs, amax, message): + exported_program = torch.export.export(_StrictQuantOp(fmt), (inputs, amax), strict=True) + with pytest.raises(torch.onnx.OnnxExporterError, match=message): + _raw_export(exported_program, (), tmp_path / f"{fmt}.onnx") diff --git a/tests/unit/torch/quantization/test_tensor_quant_cpu.py b/tests/unit/torch/quantization/test_tensor_quant_cpu.py index d72bdca6e1c..26f048d0d1d 100644 --- a/tests/unit/torch/quantization/test_tensor_quant_cpu.py +++ b/tests/unit/torch/quantization/test_tensor_quant_cpu.py @@ -23,7 +23,7 @@ import modelopt.torch.quantization as mtq import modelopt.torch.quantization.nn.modules.tensor_quantizer as tensor_quantizer_module -from modelopt.torch.quantization import QuantModuleRegistry +from modelopt.torch.quantization import QuantModuleRegistry, tensor_quant from modelopt.torch.quantization.config import QuantizerAttributeConfig, RotateConfig from modelopt.torch.quantization.nn import ( SequentialQuantizer, @@ -37,6 +37,70 @@ class TestFakeTensorQuantCPU(FakeTensorQuantTester): device = "cpu" +def test_custom_op_schemas_keep_legacy_positional_calls(): + inputs = torch.randn(4, 32) + amax = torch.tensor(1.0) + fp8 = torch.ops.tensorrt.quantize_op(inputs, amax, 8, 4, False, False) + mxfp8 = torch.ops.tensorrt.dynamic_block_quantize_op.overload(inputs, 32, None, 8, 4, 9, 8) + + assert (fp8.shape, fp8.dtype) == (inputs.shape, inputs.dtype) + assert (mxfp8.shape, mxfp8.dtype) == (inputs.shape, inputs.dtype) + + +@pytest.mark.parametrize("num_bits", [8, 4]) +def test_cpu_integer_carrier_strict_capture_preserves_custom_op(num_bits): + class QuantizeCarrier(torch.nn.Module): + def forward(self, inputs, amax): + return torch.ops.tensorrt.quantize_op(inputs, amax, num_bits, 0, False, False, "Float") + + inputs = torch.randn(4, 32) + amax = torch.tensor(1.0) + exported = torch.export.export(QuantizeCarrier(), (inputs, amax), strict=True) + + assert any( + node.target == torch.ops.tensorrt.quantize_op.default for node in exported.graph.nodes + ) + expected = tensor_quant.fake_tensor_quant(inputs, amax, None, num_bits, False, False) + torch.testing.assert_close(exported.module()(inputs, amax), expected) + + +def test_cpu_fp8_quantizer_strict_capture_preserves_custom_op(): + quantizer = TensorQuantizer(QuantizerAttributeConfig(num_bits=(4, 3), narrow_range=False)) + quantizer.amax = torch.tensor(1.0) + inputs = (torch.randn(4, 32),) + + exported = torch.export.export(quantizer, inputs, strict=True) + + assert any( + node.target == torch.ops.tensorrt.quantize_op.default for node in exported.graph.nodes + ) + torch.testing.assert_close(exported.module()(*inputs), quantizer(*inputs)) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16]) +@pytest.mark.parametrize( + ("amax_shape", "num_bits", "block_size", "axis"), + [([], 8, None, None), ([2, 1], 8, None, 0), ([4, 1], 4, 2, 1)], + ids=["scalar", "per_axis", "block"], +) +def test_registered_cpu_quantize_matches_eager(amax_shape, num_bits, block_size, axis, dtype): + inputs = torch.linspace(-3, 3, 8, dtype=dtype).reshape(2, 4) + if block_size is not None: + inputs = inputs.reshape(-1, block_size) + amax = torch.full(amax_shape, 2.0, dtype=dtype) + if amax.numel() > 1: + amax[0] = 0 + expected = tensor_quant.fake_tensor_quant(inputs, amax, None, num_bits, False, True) + + legacy = torch.ops.tensorrt.quantize_op(inputs, amax, num_bits, 0, False, True) + with_metadata = torch.ops.tensorrt.quantize_op( + inputs, amax, num_bits, 0, False, True, "Float", block_size, axis + ) + + torch.testing.assert_close(legacy, expected, rtol=0, atol=0) + torch.testing.assert_close(with_metadata, expected, rtol=0, atol=0) + + class TestQuantizerAttributeConfig: def test_scaled_mode(self): num_bits = np.random.randint(1, 16) From 925ce341334a030dbc398d6f31bfcea9ab48efd8 Mon Sep 17 00:00:00 2001 From: Chenjie Luo <108829653+cjluo-nv@users.noreply.github.com> Date: Wed, 14 Oct 2026 22:00:29 -0700 Subject: [PATCH 18/35] [NVBugs 6918319] Keep INT8 real-quant scales in FP32 so small FP16 channels survive compress (#2793) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Type of change: Bug fix Fixes [NVBug 6918319](https://nvbugspro.nvidia.com/bug/6918319): INT8 real quantization (`mtq.compress`) silently zeroes small FP16 weight channels. **Root cause.** `INT8QTensor.quantize` computes `scales = amax / 127.0` in the weight dtype, and `dequantize` casts the scale to the output dtype before multiplying. For FP16 weights: - a channel with amax < 2^-18 gets a zero scale; it is divided by zero, saturates, and dequantizes to zeros; - a channel with amax < 127 * 2^-14 (~7.75e-3) gets a subnormal FP16 scale with few significant bits. At amax = 1e-5 the scale rounds down by ~25%, so the largest values are clipped by ~24% (fake quant: 0.6%). Even an FP32 scale underflows at dequantize. Separately, with a 0-dim (per-tensor) scale `input / scale` stays in the input dtype, so BF16 quotients are rounded before `round()` and some codes are off by one (4 / (7/127) = 72.57 -> 72.5 -> 72). Fake quantization computes in FP32 and keeps all of these, so the compressed model diverged from the calibrated one. **Fix** (`INT8QTensor`): - compute and store the scales in FP32 (`amax.float() / 127`), as the INT8 HF export already does (`amax.float() / maxbound`); - quantize an explicit FP32 copy of the input (in-place div/round/clamp, so the peak stays at one FP32 copy); - dequantize by multiplying in FP32 and writing straight into a `dtype` output (`torch.mul(..., out=...)`), so each product is cast as it is stored and no FP32 copy of the weight is materialized; - divide all-zero channels or blocks by 1 instead of 0 (previously NaN -> int8, which is undefined). Checkpoints compressed before this change still load: their FP16 `_scale` dequantizes bit-identically, because an int8 x FP16 product is exact in FP32 and is rounded once, like the old FP16 multiply. INT8 real quantization now equals fake quantization on the BF16 per-channel and 2D-block cases in `test_qtensor_cuda.py`; their hard-coded expected outputs carried the old BF16 rounding errors and are updated. ### Usage The reproducer from the bug: ```python import copy import torch import modelopt.torch.quantization as mtq model = torch.nn.Sequential(torch.nn.Linear(4, 2, bias=False)).half().cuda() with torch.no_grad(): model[0].weight.copy_(torch.tensor([[2**-20, -(2**-20), 0, 2**-21], [1, -1, 0, 0.5]])) x = torch.tensor([[1.0, 0, 0, 0]], dtype=torch.float16, device="cuda") mtq.quantize(model, copy.deepcopy(mtq.INT8_DEFAULT_CFG), lambda m: m(x)) print(model(x)) # [2^-20, 1] mtq.compress(model) print(model(x)) # before: [0, 1]; after: [2^-20, 1] ``` ### Testing - New `tests/unit/torch/quantization/test_int8_tensor.py`: FP16/BF16 real quantization equals the FP32 result cast to that dtype (per channel, per block and per tensor, with tiny and all-zero channels), and `mtq.compress` keeps a small FP16 channel. 4 of its 7 cases fail on `main`. - NVIDIA B300 SXM6 AC, TRT-LLM 1.4.0rc0 container: the reproducer gives `[0, 1]` after compress on `main` and `[2^-20, 1]` with this PR; FP32 and BF16 are unchanged. `test_int8_tensor.py`, `tests/gpu/torch/quantization/test_qtensor_cuda.py` and `tests/gpu/torch/quantization/test_real_quantize_cuda.py`: 250 passed, 3 skipped. - CPU: `test_tensor_quantizer_cpu.py`, `test_tensor_quant_cpu.py`, `test_quantize_cpu.py` and `test_q8_0.py` pass; pre-commit is clean. - Cost on B300 for one 8192x28672 FP16 weight (448 MiB): quantize peak extra memory 1120 MiB (`main`: 896), 2.6 ms (1.7), once per layer at compress; dequantize 448 MiB, i.e. only the output (`main`: 896), 0.75 ms (1.06). The `out=` dequantize is bit-identical to `(int8 * fp32_scale).to(dtype)` on B300 for FP16/BF16/FP32 with per-channel, per-tensor and full-size scales; the B300 test run above was repeated on the second commit with the same result. ### Before your PR is "*Ready for review*" Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors) (e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(..., weights_only=False)`, `pickle`, etc.). - Is this change backward compatible?: ✅ (previously compressed checkpoints dequantize bit-identically; newly compressed INT8 weights store an FP32 `_scale`) - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: ✅ - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ - Did you get Claude approval on this PR?: ❌ (not run yet) ### Additional Information - Bug is committed for 0.48.0, hence `cherry-pick-0.48.0`. - Related, not fixed here: `FP8QTensor.dequantize` casts its scale to FP16 the same way. With a calibrated FP32 amax a small FP16 channel dequantizes to zeros; without one, the FP16 scale underflows and the channel becomes NaN. #2694 is open on the same file, so this will be a separate follow-up. 🤖 Generated with [Claude Code](https://claude.com/claude-code) ## Summary by CodeRabbit * **Bug Fixes** * Fixed INT8 quantization of FP16 and BF16 weights, which could underflow scales and zero or clip channels. Low-precision weights now retain values more accurately through quantization and dequantization, including zero values and channel outputs. * Improved consistency between low-precision quantization results and FP32 reference results across per-channel, block, and per-tensor layouts. --------- Signed-off-by: Chenjie Luo Co-authored-by: Claude Opus 5.5 (1M context) Signed-off-by: Rohan Joshi --- CHANGELOG.rst | 1 + .../torch/quantization/qtensor/int8_tensor.py | 20 +++-- .../torch/quantization/test_qtensor_cuda.py | 4 +- .../torch/quantization/test_int8_tensor.py | 76 +++++++++++++++++++ 4 files changed, 94 insertions(+), 7 deletions(-) create mode 100644 tests/unit/torch/quantization/test_int8_tensor.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index e45c509337d..db909867703 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -127,6 +127,7 @@ Changelog **Bug Fixes** - Fix exports using the built-in Transformers NemotronH implementation retaining incompatible legacy remote-code config and model mappings. Re-export affected checkpoints before deployment. +- Fix INT8 real quantization (``mtq.compress``) of FP16 weights zeroing or clipping channels with a small amax: their scales underflowed in FP16. Scales of newly compressed INT8 weights are now stored in FP32. - Fix Hugging Face exports classifying registered Qwen VLM parameters as unplaced source weights, which overwrote quantized tensors with original weights and incorrectly excluded their modules from quantization. Re-export affected checkpoints. - Right-pad labeled calibration samples to avoid fully masked causal rows that can produce non-finite cuDNN attention gradients during AutoQuantize. Generation previews remove padding before decoding. - Fix calibration silently dropping samples when an out-of-memory error forces the batch to be diff --git a/modelopt/torch/quantization/qtensor/int8_tensor.py b/modelopt/torch/quantization/qtensor/int8_tensor.py index 890b88e9ae9..fe3126279d8 100644 --- a/modelopt/torch/quantization/qtensor/int8_tensor.py +++ b/modelopt/torch/quantization/qtensor/int8_tensor.py @@ -60,11 +60,15 @@ def quantize( else: reduce_axis = convert_quantization_axis_to_reduce_axis(input, axis) amax = reduce_amax(input, axis=reduce_axis) - scales = amax / 127.0 + # Keep the scales in FP32: amax / 127 of a small FP16 channel lands in the FP16 + # subnormal range, where it loses precision or rounds to zero. + scales = amax.float() / 127.0 # Calculate the scale shape and make sure it aligns with input and block_sizes expected_shape = list(input.shape) - expanded_scales = scales.clone() + # A zero scale means an all-zero channel or block; divide it by one so it stays zero + # instead of becoming NaN. + expanded_scales = torch.where(scales == 0, 1.0, scales) if block_sizes: for dim, block_size in block_sizes.items(): dim = dim if dim >= 0 else len(input.shape) + dim # Convert negative index @@ -84,8 +88,10 @@ def quantize( for dim, block_size in block_sizes.items(): expanded_scales = expanded_scales.repeat_interleave(block_size, dim=dim) - # Quantization - quantized_data = (input / expanded_scales).round().clamp(-128, 127).to(torch.int8) + # Quantize in FP32. `input / expanded_scales` alone would run in the input dtype when the + # scale is 0-dim (per-tensor), rounding the quotient before round(). + quantized_data = input.to(torch.float32, copy=True).div_(expanded_scales) + quantized_data = quantized_data.round_().clamp_(-128, 127).to(torch.int8) return cls(original_input.shape, original_input.dtype, quantized_data), scales @@ -121,4 +127,8 @@ def dequantize(self, dtype: torch.dtype = None, **kwarg): slices = tuple(slice(0, dim) for dim in self.metadata["shape"]) scales = scales.to(self._quantized_data.device) - return (self._quantized_data.view(torch.int8).to(dtype) * scales.to(dtype))[slices] + # Multiply in FP32 and cast once: casting a small scale to FP16 first underflows it. + # Writing into a `dtype` output casts each product as it is stored, with no FP32 copy. + dequantized = torch.empty(shape, dtype=dtype, device=self._quantized_data.device) + torch.mul(self._quantized_data.view(torch.int8), scales.float(), out=dequantized) + return dequantized[slices] diff --git a/tests/gpu/torch/quantization/test_qtensor_cuda.py b/tests/gpu/torch/quantization/test_qtensor_cuda.py index cfdf38864fb..c5d89b17d44 100644 --- a/tests/gpu/torch/quantization/test_qtensor_cuda.py +++ b/tests/gpu/torch/quantization/test_qtensor_cuda.py @@ -175,7 +175,7 @@ def test_amax_from_tensor_quantizer( 0, torch.tensor([[0, 1, 2, 3, 4, 5, 6, 7]], dtype=torch.bfloat16), torch.tensor( - [[0.0000, 0.9922, 1.9844, 2.9844, 3.9688, 4.9688, 5.9688, 7.0000]], + [[0.0000, 0.9922, 1.9844, 2.9688, 4.0312, 5.0312, 6.0000, 7.0000]], dtype=torch.bfloat16, ), ), @@ -186,7 +186,7 @@ def test_amax_from_tensor_quantizer( None, torch.tensor([[0, 1, 2, 3], [4, 5, 6, 7]], dtype=torch.bfloat16), torch.tensor( - [[0.0000, 1.0234, 1.9844, 2.9844], [4.0000, 5.0000, 5.9688, 7.0000]], + [[0.0000, 0.9844, 1.9844, 2.9688], [4.0312, 5.0000, 6.0000, 7.0000]], dtype=torch.bfloat16, ), ), diff --git a/tests/unit/torch/quantization/test_int8_tensor.py b/tests/unit/torch/quantization/test_int8_tensor.py new file mode 100644 index 00000000000..47e190fe0bd --- /dev/null +++ b/tests/unit/torch/quantization/test_int8_tensor.py @@ -0,0 +1,76 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""Tests for INT8QTensor real quantization of small FP16/BF16 channels on CPU.""" + +import copy + +import pytest +import torch + +import modelopt.torch.quantization as mtq +from modelopt.torch.quantization.qtensor import INT8QTensor + +# Rows: amax 2^-20 (its FP16 scale amax / 127 rounds to zero), amax 1e-5 (FP16 subnormal +# scale, rounded down by ~25%), all zeros, and a normal channel. +WEIGHT = [ + [2**-20, -(2**-20), 0.0, 2**-21], + [1e-5, -5e-6, 2e-6, 0.0], + [0.0, 0.0, 0.0, 0.0], + [1.0, -1.0, 0.0, 0.5], +] + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize( + ("layout", "weight"), + [ + ({"axis": 0}, WEIGHT), + ({"block_sizes": {-1: 2}}, WEIGHT), + # A per-tensor scale is 0-dim and does not promote `weight / scale` out of the dtype. + ({}, [list(range(8))]), + ], +) +def test_int8_quantize_matches_fp32(dtype, layout, weight): + """Real-quantizing low-precision weights must give the FP32 result, cast to that dtype.""" + weight = torch.tensor(weight, dtype=dtype) + block_sizes = layout.get("block_sizes") + + def round_trip(w): + qtensor, scale = INT8QTensor.quantize(w, **layout) + return qtensor.dequantize(dtype=dtype, scale=scale, block_sizes=block_sizes) + + dequantized = round_trip(weight) + + assert torch.equal(dequantized, round_trip(weight.float())) + # No nonzero value here is small enough to round to zero against its own amax. + assert torch.equal(dequantized == 0, weight == 0) + + +def test_int8_compress_keeps_small_fp16_channel(): + """mtq.compress must not zero a small FP16 channel that fake quantization keeps.""" + model = torch.nn.Sequential(torch.nn.Linear(4, 2, bias=False)).half() + with torch.no_grad(): + model[0].weight.copy_(torch.tensor([WEIGHT[0], WEIGHT[3]])) + x = torch.tensor([[1.0, 0.0, 0.0, 0.0]], dtype=torch.float16) + + mtq.quantize(model, copy.deepcopy(mtq.INT8_DEFAULT_CFG), lambda m: m(x)) + with torch.no_grad(): + fake_quant_output = model(x) + mtq.compress(model) + real_quant_output = model(x) + + assert fake_quant_output[0, 0] == 2**-20 + assert torch.equal(real_quant_output, fake_quant_output) From 4afeac19006327c6f28e1626194f4a45faad88a1 Mon Sep 17 00:00:00 2001 From: yueshen2016 <39203804+yueshen2016@users.noreply.github.com> Date: Thu, 15 Oct 2026 01:41:05 -0400 Subject: [PATCH 19/35] [1/2] Add a chat JSONL output mode to prepare_megatron_data_blend (#2791) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Type of change: New feature `prepare_megatron_data_blend` writes tokenized `.bin`/`.idx` data, which has no loss mask, so `--data_paths` training computes the loss on every token. Recipes that distill only on assistant responses use Megatron-Bridge's chat path instead (`distill.py --sft --sft_hf_dataset ... --sft_loss_mode assistant`), which reads chat JSONL and applies the model's chat template and assistant mask itself. Until now each such recipe needed its own blend script. `output_format: chat_jsonl` writes the same YAML-defined weighted sources as `train.jsonl` / `validation.jsonl` (one `{"messages": [...]}` per line) for that path: - `target_samples` sets the row budget; each source contributes `round(target_samples * weight / sum(weights))` rows. It is required, since the selected rows are mixed in memory. - `content_field` must be `messages` or `text`. Rows are kept only if Megatron-Bridge's `is_text_only_chat_example` accepts them, judged on the converted conversation plus the source row's top-level media keys. A `text` row without markers becomes one assistant turn, so all of it is trained. A `text` row with the `User` / `Assistant` / `System` markers is split into turns. It is dropped if it has text before the first marker, a marker has any other role (a Tool turn, a SteerLM label line), a role name is not followed by a newline, or a turn other than `System` is empty. Any row with no assistant turn is dropped, since assistant-only loss trains nothing in it. - `messages` go through `prepare_messages_for_chat_template`, like the tokenizing mode, so `tool_calls` arguments are normalized, and a row's `tools` list and top-level `system` prompt are kept for the chat path. `reasoning_content` (`native` / `inline` / `strip`) defaults to `native`, unlike the tokenizing mode's `inline`: Bridge applies the model's chat template at training time, and templates such as Nemotron 3's render `reasoning_content` themselves. Use `inline` for a template that ignores the field. - `shuffle: false` takes the first rows of each split in order (default: a buffer shuffle); `seed` seeds that shuffle, the cross-source mix and the `min(1000, n // 20)` validation hold-out. `files` may be `.jsonl` or `.jsonl.gz`; malformed records are skipped with a warning, as in the tokenizing reader. - A source that fails to load, has no usable rows or runs short stops the run; a source whose weight rounds to zero rows is skipped. Each source is cached atomically under `/parts`, keyed by the settings that select its rows, so an interrupted run resumes and a changed configuration selects afresh. `parts` can be deleted once `train.jsonl` is written. The tokenizing mode is unchanged, except that it now rejects the chat-only keys instead of ignoring them. Each entry point rejects a config whose `output_format` belongs to the other. This is [1/2]; #2720 [2/2] uses it for the NemotronH omni VL QAD tutorial (replacing that tutorial's `build_qad_blend.py`) and merges after this one. ```yaml output_dir: /path/to/chat_blend output_format: chat_jsonl target_samples: 36_800 shuffle: false seed: 1234 sources: - hf_dataset: nvidia/Nemotron-Pretraining-SFT-v1 config: Nemotron-SFT-General split: train content_field: text weight: 20 - hf_dataset: nvidia/Nemotron-SFT-Math-v3 config: default split: train content_field: messages weight: 17 ``` ```bash python -m modelopt.torch.utils.plugins.prepare_megatron_data_blend --config blend.yaml python examples/megatron_bridge/distill.py ... --sft --sft_hf_dataset /path/to/chat_blend/train.jsonl \ --sft_hf_validation /path/to/chat_blend/validation.jsonl --sft_loss_mode assistant ``` - `tests/gpu_megatron/torch/utils/plugins/test_prepare_megatron_data_blend.py`: 28 passed in the NeMo 26.08 container with Megatron-Bridge `main` (2 existing, 26 new, consolidated per review). They cover the row shapes (leading text, unknown-role and label markers, roles without a newline, empty turns, rows without an assistant turn), `tools` and `reasoning_content` handling, the text-only filter and give-up rule, sources without usable rows, `.jsonl.gz` and malformed records in `files`, the weighted split with a zero-weight source, config validation (budgets, weights, a missing `target_samples`, `content_field`, `output_format`, the other mode's keys), cache reuse and rebuild, no reuse after a changed selection, atomic writes, and failing or short sources. Reverting each review fix of the last three rounds (14 in all) one at a time makes at least one test fail. - End to end: the 11-source blend of #2720, built from scratch with this mode, is byte-identical (sha256 of `train.jsonl` and `validation.jsonl`) to the data that tutorial's published QAD run trained on; re-checked on the latest commit. - Is this change backward compatible?: ✅ The default output and the tokenizing path are unchanged; only configs that set the new chat-only keys without `output_format: chat_jsonl` are now rejected. - Did you write any new necessary tests?: ✅ - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ 🤖 Generated with [Claude Code](https://claude.com/claude-code) * **New Features** * Added chat JSONL output for weighted dataset blends, creating `train.jsonl` and `validation.jsonl` files for chat training workflows with assistant-only loss configuration. * Supports conversations from message or text fields, configurable sample targets, shuffling, seeds, and reasoning-content handling. * Reuses cached source selections and filters unusable records. Reports errors when sources or output splits lack usable data. * **Documentation** * Added guidance and configuration examples for preparing chat JSONL blends. --------- Signed-off-by: Yue Co-authored-by: Claude Opus 5.5 Signed-off-by: Rohan Joshi --- CHANGELOG.rst | 1 + examples/dataset/MEGATRON_DATA_PREP.md | 49 +++ .../plugins/prepare_megatron_data_blend.py | 303 +++++++++++++++++- .../test_prepare_megatron_data_blend.py | 276 ++++++++++++++++ 4 files changed, 617 insertions(+), 12 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index db909867703..87930a8aaf4 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -49,6 +49,7 @@ Changelog *Megatron Framework (M-LM / M-Bridge)* - Add ``freeze_mtp`` to the Megatron ``DistillationConfig`` (``--freeze_mtp`` in ``examples/megatron_bridge/distill.py``): the student's MTP head keeps its input weights and its loss is skipped, so distillation trains the rest of the model on the distillation loss alone. +- Add ``output_format: chat_jsonl`` to ``prepare_megatron_data_blend``: it writes a weighted blend as ``train.jsonl`` / ``validation.jsonl`` chat data for ``distill.py --sft --sft_hf_dataset``, so the distillation loss can be restricted to assistant tokens. See the `Megatron data preparation guide `_. - Add optional Top-P (nucleus) truncation to the Megatron top-k logits KD loss (now ``TopLogitsKLLoss``) via ``logit_kl_top_p`` and ``logit_kl_top_p_min_k`` in ``DistillationConfig``: after the global Top-K selection, only the smallest prefix whose cumulative teacher probability reaches ``top_p`` (with a floor of ``min_k`` entries) contributes to the KL, mirroring ``--logits-save-top-p`` / ``--logits-save-top-p-min-k`` in Megatron-LM's logits saver. - Add an end-to-end W4A4 NVFP4 PTQ and QAD tutorial for Qwen3.6-35B-A3B also covering evaluation and vLLM throughput benchmarking. See `examples/megatron_bridge/tutorials/Qwen3.6-35B-A3B/README.md `_ for details. - Add ``--mlflow `` to the ``examples/megatron_bridge`` scripts that write a checkpoint -- ``prune_minitron.py``, ``quantize.py``, ``distill.py``, ``export_quantized_megatron_to_hf.py`` and ``export_distilled_megatron_to_hf.py`` (MLflow's own ``MLFLOW_TRACKING_URI`` is honoured too). Each run records the invocation, its arguments as searchable params and its log -- a distillation records its training metrics instead of the rank-0 log -- and writes ``.experiment.json`` into the checkpoint it produced, so a pruning, a quantization, the distillation that refines its checkpoint and the export that deploys it can be traced to one another; uploading the checkpoints themselves stays off unless ``--mlflow_log_checkpoints`` is passed. diff --git a/examples/dataset/MEGATRON_DATA_PREP.md b/examples/dataset/MEGATRON_DATA_PREP.md index 91b89375907..c920ed02567 100644 --- a/examples/dataset/MEGATRON_DATA_PREP.md +++ b/examples/dataset/MEGATRON_DATA_PREP.md @@ -5,6 +5,7 @@ | From JSONL files | Tokenize local JSONL files | \[[Link](#from-jsonl-files)\] | | From Hugging Face Hub | Stream or download HF datasets and tokenize | \[[Link](#from-hugging-face-hub)\] | | Token-budgeted data blends | Prepare weighted subsets for fast experiments | \[[Link](#prepare-token-budgeted-data-blends)\] | +| Chat JSONL blends | Weighted chat data for assistant-only loss | \[[Link](#chat-jsonl-blends-for-assistant-only-loss)\] | | `reasoning_content` for Post-Training v3 | Control how chain-of-thought traces are handled | \[[Link](#reasoning_content-for-post-training-v3-datasets)\] | | Nemotron Pre/Post-Training Datasets | Ready-to-run commands for all Nemotron datasets | \[[Link](#ready-to-run-tokenization-commands)\] | @@ -172,6 +173,54 @@ The output contains tokenized Megatron `.bin`/`.idx` files, `data_blend.txt` wit and `config.yaml` recording how the blend was generated. The final token count can slightly exceed the target because the final document from each source is kept whole. +### Chat JSONL blends for assistant-only loss + +Tokenized `.bin`/`.idx` data carries no loss mask, so `--data_paths` training computes the loss on every token. +To compute it on assistant tokens only, set `output_format: chat_jsonl`. The same sources are then written, +untokenized, as `train.jsonl` and `validation.jsonl` with one `{"messages": [...]}` conversation per line, for +Megatron-Bridge's chat path (`examples/megatron_bridge/distill.py --sft --sft_hf_dataset /train.jsonl +--sft_hf_validation /validation.jsonl --sft_loss_mode assistant`), which applies the model's chat template +when it loads the data. Run it in an environment with Megatron-Bridge, whose text-only chat filter selects the rows. + +```yaml +output_dir: /path/to/chat_blend +output_format: chat_jsonl +target_samples: 36_800 # Total conversations; required, since the selected rows are mixed in memory. +shuffle: false # Optional; take the first rows of each split in order (default: buffer shuffle). +seed: 1234 # Optional; seeds the shuffle, the cross-source mix and the validation hold-out. +reasoning_content: native # Optional; native (default), inline or strip. +sources: + - hf_dataset: nvidia/Nemotron-Pretraining-SFT-v1 + config: Nemotron-SFT-General + split: train + content_field: text + weight: 20 + - hf_dataset: nvidia/Nemotron-SFT-Math-v3 + config: default + split: train + content_field: messages + weight: 17 +``` + +Each source contributes `round(target_samples * weight / sum(weights))` conversations that contain only text +and have an assistant turn; a source that rounds to zero is skipped. `content_field` must be `messages` or `text`: + +- A `text` row becomes one assistant turn, or is split into turns if it uses the `User` / + `Assistant` / `System` chat markers. A row with a marker for any other role (such as a SteerLM + label line) or with an empty turn other than `System` is dropped. +- `messages` and a row's `tools` are prepared as in the tokenizing mode (JSON-string `tool_calls` arguments + become dicts), and a top-level `system` prompt is kept for Megatron-Bridge to prepend. The one difference is + that `reasoning_content` stays on the message by default (`native`), because the chat template is applied + later, at training time. Use `inline` if your model's chat template ignores `reasoning_content`, so the + reasoning is kept as a `` block in `content` (see [below](#reasoning_content-for-post-training-v3-datasets)), + or `strip` to drop it. + +`files` may be `.jsonl` or `.jsonl.gz`; malformed records in them are skipped with a warning. The conversations are +mixed and `min(1000, n // 20)` of them are held out for validation. A source that fails to load, has no usable +rows, or runs out of them stops the run. Each source's selection is cached under `/parts`, keyed by +the settings that choose its rows, so an interrupted run resumes and a changed configuration selects afresh. +Nothing prunes `parts`; delete it once `train.jsonl` is written to reclaim the space. + ## `reasoning_content` for Post-Training v3 Datasets v3 datasets include a `reasoning_content` field in assistant messages (chain-of-thought separate from diff --git a/modelopt/torch/utils/plugins/prepare_megatron_data_blend.py b/modelopt/torch/utils/plugins/prepare_megatron_data_blend.py index fa5875a5e2c..2bc7aacb8c9 100644 --- a/modelopt/torch/utils/plugins/prepare_megatron_data_blend.py +++ b/modelopt/torch/utils/plugins/prepare_megatron_data_blend.py @@ -16,17 +16,42 @@ """Prepare a weighted Megatron data blend from a YAML configuration.""" import argparse +import gzip +import hashlib +import json +import math import os +import random +import re import shutil +from collections.abc import Iterator from pathlib import Path from typing import Any, cast import huggingface_hub import yaml +from datasets import load_dataset + +from modelopt.torch.utils import ( + REASONING_CONTENT_MODES, + ReasoningContentMode, + prepare_messages_for_chat_template, +) from .megatron_preprocess_data import megatron_preprocess_data -__all__ = ["prepare_megatron_data_blend"] +__all__ = ["prepare_chat_data_blend", "prepare_megatron_data_blend"] + +# A turn of a ``text`` row in the NeMo chat format: ``User\n...``. Any marker ends the +# turn, so a trailing delimiter is not kept as content. The role must be followed by a newline, a +# marker or the end, so ``Userguide`` is no role at all. +_CHAT_TURN = re.compile( + r"(User|Assistant|System)(?:\n|(?=|\Z))(.*?)(?=|\Z)", + re.DOTALL, +) +_CHAT_ROLES = {"User": "user", "Assistant": "assistant", "System": "system"} +_CHAT_MARKER = re.compile(r"") +_CHAT_TRAILING_MARKER = re.compile(r"\s*\Z") def load_config(path: Path) -> dict[str, Any]: @@ -53,6 +78,10 @@ def load_config(path: Path) -> dict[str, Any]: optional ``target_tokens``. Each source has ``hf_dataset``, ``content_field``, and ``weight``; it uses ``split`` with optional ``config`` and ``max_samples``, or selects repository ``files``. + + With ``output_format: chat_jsonl``, the blend is written as chat JSONL instead (see + :func:`prepare_chat_data_blend`): ``tokenizer`` and ``target_tokens`` are not used, and the + optional ``target_samples``, ``shuffle``, ``seed`` and ``reasoning_content`` keys apply. """ with path.open(encoding="utf-8") as stream: return cast("dict[str, Any]", yaml.safe_load(stream)) @@ -76,6 +105,18 @@ def _copy_config(source: Path, destination: Path) -> None: shutil.copyfile(source, destination) +def _download_files(source: dict[str, Any], output_dir: Path) -> list[str]: + """Download the repository ``files`` of a source and return their local paths.""" + dataset = source["hf_dataset"] + raw_dir = output_dir.parent / "raw" / dataset.replace("/", "--") + return [ + huggingface_hub.hf_hub_download( + repo_id=dataset, filename=file, repo_type="dataset", local_dir=raw_dir + ) + for file in source["files"] + ] + + def _prepare_sources( sources: list[dict[str, Any]], output_dir: Path, @@ -104,17 +145,7 @@ def _prepare_sources( content_field = source["content_field"] input_args: dict[str, Any] if "files" in source: - raw_dir = output_dir.parent / "raw" / dataset.replace("/", "--") - paths = [ - huggingface_hub.hf_hub_download( - repo_id=dataset, - filename=file, - repo_type="dataset", - local_dir=raw_dir, - ) - for file in source["files"] - ] - input_args = {"jsonl_paths": paths} + input_args = {"jsonl_paths": _download_files(source, output_dir)} else: input_args = { "hf_dataset": dataset, @@ -146,9 +177,252 @@ def _prepare_sources( return blend +def _to_chat( + row: dict[str, Any], content_field: str, reasoning_content: ReasoningContentMode = "native" +) -> dict[str, Any] | None: + """Convert a source row to ``{"messages": [...]}``, or ``None`` if it has nothing to train on. + + ``messages`` go through :func:`prepare_messages_for_chat_template` with ``reasoning_content``, + and non-empty ``tools`` and ``system`` fields are kept for the chat path. A ``text`` row in the NeMo chat + format (``User`` / ``Assistant`` markers) is split into turns; any other ``text`` + row becomes one assistant turn, so all of it is trained. + """ + if content_field == "messages": + messages = row.get("messages") + if not messages: + return None + messages = prepare_messages_for_chat_template(messages, reasoning_content=reasoning_content) + chat: dict[str, Any] = {"messages": messages} + # As in the tokenizing mode, the chat template renders the tool definitions into the prompt; + # Megatron-Bridge turns a top-level system prompt into the leading system message. + for key in ("tools", "system"): + if row.get(key): + chat[key] = row[key] + else: + text = row.get("text") or "" + if not text: + return None + if " delimiters as assistant content. Text before the first + # marker, a marker without a known role (a Tool turn, or a SteerLM label line) or a known turn + # left empty would silently cut it out of the conversation, so drop the row. Only a bare + # delimiter at the very end needs no role, and only the System block, which real rows leave + # empty, may be empty. + roleless = len(_CHAT_MARKER.findall(text)) - len(matches) + if ( + not _CHAT_MARKER.match(text.lstrip()) + or roleless > (1 if _CHAT_TRAILING_MARKER.search(text) else 0) + or any(not content.strip() and role != "System" for role, content in matches) + ): + return None + turns = [ + {"role": _CHAT_ROLES[role], "content": content.strip()} + for role, content in matches + if content.strip() + ] + chat = {"messages": turns} + # Nothing in a row without an assistant turn is trained under assistant-only loss. + if not any(turn.get("role") == "assistant" for turn in chat["messages"]): + return None + return chat + + +def _iter_rows( + source: dict[str, Any], output_dir: Path, shuffle: bool, seed: int +) -> Iterator[dict[str, Any]]: + """Yield the rows of a source: its ``files`` in order, or its Hugging Face split, streamed.""" + if "files" in source: + for path in _download_files(source, output_dir): + opener = gzip.open if path.endswith(".gz") else open + with opener(path, "rt", encoding="utf-8") as f: + for line in f: + if not line.strip(): + continue + try: + yield json.loads(line) + except ValueError as e: + # As in the tokenizing reader: some Nemotron files carry a malformed record. + print(f"[WARN] Skipping malformed JSONL record in {path}: {e}", flush=True) + return + dataset = load_dataset( + source["hf_dataset"], source.get("config"), split=source["split"], streaming=True + ) + yield from dataset.shuffle(seed=seed) if shuffle else dataset + + +def _take_chat_rows( + source: dict[str, Any], + output_dir: Path, + want: int, + shuffle: bool, + seed: int, + reasoning_content: ReasoningContentMode, +) -> list[dict[str, Any]]: + """Return the first ``want`` usable chat rows of a source. + + Fails if the source cannot be read or runs out of usable rows, so the blend cannot change + silently; gives up after ``10 * want + 100`` rejected rows, a floor that keeps a small budget from + misjudging a source by its first few rows. + """ + # Megatron-Bridge's chat SFT path consumes this output; keep exactly the rows it accepts. Imported + # here, not at module scope, because the tokenizing mode does not need Megatron-Bridge installed. + from megatron.bridge.data.conversation_processing import is_text_only_chat_example + + name = f"{source['hf_dataset']} {source.get('split') or source.get('files')}" + rows: list[dict[str, Any]] = [] + rejected = 0 + max_rejected = 10 * want + 100 + try: + for row in _iter_rows(source, output_dir, shuffle, seed): + chat = _to_chat(row, source["content_field"], reasoning_content) + # Judge the source row's top-level media keys too, which the conversion does not keep. + if chat is None or not is_text_only_chat_example({**row, **chat}): + rejected += 1 + if rejected > max_rejected: + break + continue + rows.append(chat) + if len(rows) >= want: + break + except Exception as e: # e.g. a gated split, or a schema the streaming reader cannot cast + raise RuntimeError(f"{name}: failed to read or convert") from e + if len(rows) < want: + gave_up = rejected > max_rejected + reason = f"gave up after {rejected} rejected rows" if gave_up else "source exhausted" + raise RuntimeError(f"{name}: only {len(rows)}/{want} usable rows ({reason})") + print(f"{name}: {len(rows)} rows ({rejected} rejected)", flush=True) + return rows + + +def _write_jsonl(path: Path, rows: list[dict[str, Any]]) -> None: + """Write ``rows`` atomically, so an interrupted write never leaves a partial file.""" + tmp = path.with_name(path.name + ".tmp") + with tmp.open("w", encoding="utf-8") as f: + f.writelines(json.dumps(row) + "\n" for row in rows) + os.replace(tmp, path) + + +def _read_jsonl(path: Path) -> list[dict[str, Any]] | None: + """Read a cached part, or return ``None`` if it is missing or unreadable.""" + try: + with path.open(encoding="utf-8") as f: + return [json.loads(line) for line in f] + except (OSError, ValueError): + return None + + +def _part_path(parts_dir: Path, source: dict[str, Any], settings: dict[str, Any]) -> Path: + """Return the cache file of a source's selected rows, keyed by every input that selects them. + + A part holds the first usable rows of the source, so it serves any budget it covers. + """ + selection = {key: source.get(key) for key in ("hf_dataset", "config", "split", "files")} + selection.update(settings, content_field=source["content_field"]) + digest = hashlib.sha256(json.dumps(selection, sort_keys=True).encode()).hexdigest()[:16] + label = source.get("split") or Path(source["files"][0]).name.split(".")[0] + return parts_dir / f"{source['hf_dataset'].replace('/', '--')}_{label}_{digest}.jsonl" + + +def _chat_budgets(config: dict[str, Any]) -> list[int]: + """Return each source's row budget; validate the inputs.""" + weights = [float(source["weight"]) for source in config["sources"]] + if not all(math.isfinite(w) and w >= 0 for w in weights) or sum(weights) <= 0: + raise ValueError(f"Source weights must be finite, nonnegative and not all zero: {weights}.") + # Required: the selected rows are held in memory to be mixed, which a whole source may not fit. + target = config.get("target_samples") + if isinstance(target, bool) or not isinstance(target, int) or target <= 0: + raise ValueError( + f"target_samples is required and must be a positive integer, got {target!r}." + ) + return [round(target * w / sum(weights)) for w in weights] + + +def prepare_chat_data_blend(config_path: Path) -> dict[str, int]: + """Write a weighted chat blend as ``train.jsonl`` / ``validation.jsonl``; return their sizes. + + For Megatron-Bridge's chat SFT path (``distill.py --sft --sft_hf_dataset``), which renders each + conversation with the model's chat template and can restrict the loss to assistant tokens. + Each source contributes ``round(target_samples * weight / sum(weights))`` rows (a source rounded + to zero rows is skipped) that the chat path accepts as text-only chat. ``content_field`` must be + ``messages`` or ``text``; ``messages`` keep ``reasoning_content`` for the chat template by + default (``reasoning_content: native``; set ``inline`` for a template that ignores the field, or + ``strip``). Hugging Face splits are buffer-shuffled with ``seed`` unless ``shuffle: false``; + ``files`` (``.jsonl`` or ``.jsonl.gz``) are read from the beginning, skipping malformed records. + The rows are mixed with ``seed`` and ``min(1000, n // 20)`` of them are held out for validation. + Each source is cached under ``/parts``, keyed by the settings that select its rows, + so an interrupted run resumes. + """ + config = load_config(config_path) + if config.get("output_format") != "chat_jsonl": + raise ValueError( + "prepare_chat_data_blend writes output_format: chat_jsonl, " + f"got {config.get('output_format')!r}." + ) + ignored = {"tokenizer", "target_tokens"} & config.keys() + if ignored: + raise ValueError(f"output_format: chat_jsonl does not use {sorted(ignored)}.") + if any("max_samples" in source for source in config["sources"]): + raise ValueError("output_format: chat_jsonl does not use max_samples; set target_samples.") + fields = {source["content_field"] for source in config["sources"]} - {"messages", "text"} + if fields: + raise ValueError( + f"output_format: chat_jsonl reads content_field messages or text: {fields}." + ) + reasoning_content = config.get("reasoning_content", "native") + if reasoning_content not in REASONING_CONTENT_MODES: + raise ValueError( + f"reasoning_content must be one of {sorted(REASONING_CONTENT_MODES)}, " + f"got {reasoning_content!r}." + ) + budgets = _chat_budgets(config) + output_dir = Path(config["output_dir"]) + parts_dir = output_dir / "parts" + parts_dir.mkdir(parents=True, exist_ok=True) + shuffle = bool(config.get("shuffle", True)) + seed = int(config.get("seed", 1234)) + # Everything besides the source itself that decides which rows a part holds. + settings = {"shuffle": shuffle, "reasoning_content": reasoning_content} + if shuffle: + settings["seed"] = seed + + rows: list[dict[str, Any]] = [] + for source, want in zip(config["sources"], budgets): + if want == 0: + print(f"{source['hf_dataset']}: weight rounds to 0 rows, skipped", flush=True) + continue + part = _part_path(parts_dir, source, settings) + got = _read_jsonl(part) + if got is None or len(got) < want: + got = _take_chat_rows(source, output_dir, want, shuffle, seed, reasoning_content) + _write_jsonl(part, got) + rows.extend(got[:want]) + + random.Random(seed).shuffle(rows) + num_validation = min(1000, len(rows) // 20) + splits = {"validation": rows[:num_validation], "train": rows[num_validation:]} + if not splits["train"] or not splits["validation"]: + raise RuntimeError(f"Too few rows ({len(rows)}) for a train/validation split.") + for name, split_rows in splits.items(): + _write_jsonl(output_dir / f"{name}.jsonl", split_rows) + _copy_config(config_path, output_dir / "config.yaml") + return {name: len(split_rows) for name, split_rows in splits.items()} + + def prepare_megatron_data_blend(config_path: Path) -> list[tuple[float, str]]: """Download and tokenize the configured weighted data sources.""" config = load_config(config_path) + output_format = config.get("output_format", "megatron") + if output_format != "megatron": + raise ValueError( + f"Unknown output_format {output_format!r} for a tokenized blend; use megatron, or " + "chat_jsonl with prepare_chat_data_blend." + ) + chat_keys = {"target_samples", "shuffle", "seed", "reasoning_content"} & config.keys() + if chat_keys: + raise ValueError(f"Tokenized blends do not use {sorted(chat_keys)}.") output_dir = Path(config["output_dir"]) output_dir.mkdir(parents=True, exist_ok=True) target_tokens = config.get("target_tokens") @@ -165,6 +439,11 @@ def main() -> None: """Prepare a data blend from the supplied configuration.""" parser = _build_parser() args = parser.parse_args() + # Each entry point rejects a config written for the other. + if load_config(args.config).get("output_format") == "chat_jsonl": + sizes = prepare_chat_data_blend(args.config) + print(f"Prepared chat JSONL: {sizes}. See train.jsonl, validation.jsonl and config.yaml.") + return blend = prepare_megatron_data_blend(args.config) print(f"Prepared {len(blend)} data paths. See data_blend.txt and config.yaml in the output.") diff --git a/tests/gpu_megatron/torch/utils/plugins/test_prepare_megatron_data_blend.py b/tests/gpu_megatron/torch/utils/plugins/test_prepare_megatron_data_blend.py index f531be27ad4..6306ecef292 100644 --- a/tests/gpu_megatron/torch/utils/plugins/test_prepare_megatron_data_blend.py +++ b/tests/gpu_megatron/torch/utils/plugins/test_prepare_megatron_data_blend.py @@ -13,8 +13,11 @@ # See the License for the specific language governing permissions and # limitations under the License. +import gzip +import importlib import json import sys +from collections import Counter from pathlib import Path from unittest.mock import Mock @@ -24,6 +27,9 @@ from modelopt.torch.utils.plugins.prepare_megatron_data_blend import main +# The plugins package re-exports the function under the module's name, so import the module itself. +blend_module = importlib.import_module("modelopt.torch.utils.plugins.prepare_megatron_data_blend") + def _setup_test( tiny_tokenizer_path: str, tmp_path: Path, monkeypatch, target_tokens: int | None @@ -103,3 +109,273 @@ def test_prepare_megatron_data_blend_with_split_and_files_sources( f"competitive_programming_python_00_messages{token_suffixes[1]}", ] assert (output_dir / "config.yaml").read_bytes() == config_path.read_bytes() + + +def _chat(i: int) -> dict: + return { + "messages": [ + {"role": "user", "content": f"q{i}"}, + {"role": "assistant", "content": f"a{i}"}, + ] + } + + +def _image_chat(i: int) -> dict: + return { + "messages": [ + {"role": "user", "content": [{"type": "image", "image": "x.png"}]}, + {"role": "assistant", "content": f"a{i}"}, + ] + } + + +def _source(name: str, **fields) -> dict: + source = { + "hf_dataset": f"org/{name}", + "split": "train", + "content_field": "messages", + "weight": 1, + } + source.update(fields) + return {k: v for k, v in source.items() if v is not None} + + +def _read_jsonl(path: Path) -> list[dict]: + return [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines()] + + +@pytest.fixture +def chat_blend(tmp_path: Path, monkeypatch): + """Sources of 100 ``messages`` rows each; returns (config writer, stream counts per source).""" + calls = Counter() + + def iter_rows(source, output_dir, shuffle, seed): + calls[source["hf_dataset"]] += 1 + yield from (_chat(i) for i in range(100)) + + monkeypatch.setattr(blend_module, "_iter_rows", iter_rows) + + def write_config(**overrides) -> Path: + config = { + "output_dir": str(tmp_path / "blend"), + "output_format": "chat_jsonl", + "target_samples": 100, + "shuffle": False, + "seed": 0, + "sources": [_source("a"), _source("b")], + } + config.update(overrides) + path = tmp_path / "config.yaml" + path.write_text(yaml.safe_dump(config), encoding="utf-8") + return path + + return write_config, calls + + +def test_to_chat_handles_all_row_shapes(): + msgs = [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "hello"}] + assert blend_module._to_chat({"messages": msgs}, "messages") == {"messages": msgs} + assert blend_module._to_chat({"messages": []}, "messages") is None + # Nothing would be trained under assistant-only loss. + assert blend_module._to_chat({"messages": msgs[:1]}, "messages") is None + # Tool definitions and a top-level system prompt reach the chat path. + tools = [{"type": "function", "function": {"name": "add"}}] + row = {"messages": msgs, "tools": tools, "system": "Be brief.", "id": 7} + chat = blend_module._to_chat(row, "messages") + assert chat == {"messages": msgs, "tools": tools, "system": "Be brief."} + # Unlike the tokenizing mode, reasoning_content stays on the message by default (native). + reasoning = [{"role": "assistant", "content": "4", "reasoning_content": "2+2"}] + assert blend_module._to_chat({"messages": reasoning}, "messages")["messages"] == reasoning + inline = blend_module._to_chat({"messages": reasoning}, "messages", "inline") + assert inline["messages"] == [{"role": "assistant", "content": "\n2+2\n\n4"}] + + expected = { + "messages": [ + {"role": "user", "content": "hello"}, + {"role": "assistant", "content": "world"}, + ] + } + text = "System\nUser\nhello\nAssistant\nworld\n" + assert blend_module._to_chat({"text": text}, "text") == expected + # Any marker ends a turn, so a trailing marker is not trained as content. + assert blend_module._to_chat({"text": text + ""}, "text") == expected + assert blend_module._to_chat({"text": "def f(): pass"}, "text") == { + "messages": [{"role": "assistant", "content": "def f(): pass"}] + } + assert blend_module._to_chat({"text": ""}, "text") is None + # Markers that parse to no turns: dropped, never trained as assistant text. + assert blend_module._to_chat({"text": "\nraw body\n"}, "text") is None + # A role without a mapping would cut its turn out of the middle: the whole row is dropped. + text = "User\nq\nTool\nresult\nAssistant\na\n" + assert blend_module._to_chat({"text": text}, "text") is None + for text in ( + # Text before the first marker would be cut from the conversation. + "preamble\nUser\nq\nAssistant\na\n", + # A SteerLM label line is no role either. + "User\nq\nAssistant\nquality:4\na\n", + # A role must end at a newline: this is no User turn. + "Userguide\nq\nAssistant\na\n", + # An empty known turn would be cut out of the middle of the conversation. + "User\nq1\nAssistant\nUser\nq2\nAssistant\na2\n", + # No assistant turn to train on. + "User\nq\nAssistant\n", + ): + assert blend_module._to_chat({"text": text}, "text") is None, text + + +def test_take_chat_rows_skips_non_text_rows_and_fails_without_enough(tmp_path: Path, monkeypatch): + source = _source("a") + # Media is rejected inside message content and as a top-level key the conversion drops. + rows = [_chat(0), _image_chat(1), {**_chat(9), "images": ["a.png"]}, _chat(2), _chat(3)] + monkeypatch.setattr(blend_module, "_iter_rows", lambda *a: iter(rows)) + got = blend_module._take_chat_rows(source, tmp_path, 3, False, 0, "native") + assert [c["messages"][0]["content"] for c in got] == ["q0", "q2", "q3"] + + images = [_image_chat(i) for i in range(200)] + monkeypatch.setattr(blend_module, "_iter_rows", lambda *a: iter(images)) + with pytest.raises(RuntimeError, match=r"only 0/2 usable rows \(gave up after 121 rejected"): + blend_module._take_chat_rows(source, tmp_path, 2, False, 0, "native") + + +def test_chat_blend_splits_weighted_rows(tmp_path: Path, chat_blend, monkeypatch): + write_config, calls = chat_blend + sources = [_source("a", weight=3), _source("b"), _source("c", weight=0)] + config_path = write_config(sources=sources) + monkeypatch.setattr( + sys, "argv", ["prepare_megatron_data_blend.py", "--config", str(config_path)] + ) + main() + out = tmp_path / "blend" + assert len(_read_jsonl(out / "validation.jsonl")) == 5 + rows = _read_jsonl(out / "train.jsonl") + _read_jsonl(out / "validation.jsonl") + # 3:1 weights: the first 75 rows of org/a and 25 of org/b; org/c rounds to zero and is skipped. + expected = Counter({f"q{i}": 1 for i in range(75)}) + Counter(f"q{i}" for i in range(25)) + assert Counter(r["messages"][0]["content"] for r in rows) == expected + assert calls["org/c"] == 0 + assert (out / "config.yaml").read_bytes() == config_path.read_bytes() + + +def test_chat_blend_reuses_and_rebuilds_cached_parts(tmp_path: Path, chat_blend): + write_config, calls = chat_blend + config_path = write_config() + blend_module.prepare_chat_data_blend(config_path) + first = (tmp_path / "blend" / "train.jsonl").read_text(encoding="utf-8") + blend_module.prepare_chat_data_blend(config_path) + assert calls == {"org/a": 1, "org/b": 1} # the second run streams nothing + + (part,) = (tmp_path / "blend" / "parts").glob("org--a_*.jsonl") + part.write_text(part.read_text(encoding="utf-8")[:-20], encoding="utf-8") # cut mid-record + blend_module.prepare_chat_data_blend(config_path) + assert calls == {"org/a": 2, "org/b": 1} + assert (tmp_path / "blend" / "train.jsonl").read_text(encoding="utf-8") == first + + +@pytest.mark.parametrize( + ("first", "second"), + [ + ({}, {"sources": [_source("a", split="test")]}), + ( + {"sources": [_source("a", split=None, files=["a.jsonl"])]}, + {"sources": [_source("a", split=None, files=["b.jsonl"])]}, + ), + ({}, {"shuffle": True}), + ({"shuffle": True, "seed": 0}, {"shuffle": True, "seed": 1}), + ({}, {"reasoning_content": "inline"}), + ], + ids=["split", "files", "shuffle", "seed", "reasoning_content"], +) +def test_chat_blend_reselects_after_a_selection_change(chat_blend, first, second): + write_config, calls = chat_blend + blend_module.prepare_chat_data_blend(write_config(**first)) + blend_module.prepare_chat_data_blend(write_config(**second)) + assert calls["org/a"] == 2 + + +@pytest.mark.parametrize( + ("overrides", "match"), + [ + ({"target_samples": 0}, "positive integer"), + ({"target_samples": 1.5}, "positive integer"), + # A whole source would be held in memory. + ({"target_samples": None}, "target_samples is required"), + ({"output_format": "megatron"}, "writes output_format: chat_jsonl"), + ({"sources": [_source("a", weight=-1)]}, "nonnegative"), + ({"sources": [_source("a", weight=0)]}, "not all zero"), + ({"sources": [_source("a", weight=float("inf"))]}, "finite"), + ({"tokenizer": "x"}, "does not use"), + ({"target_tokens": 1000}, "does not use"), + ({"reasoning_content": "drop"}, "reasoning_content must be one of"), + # A row with both columns must not silently be read from ``text``. + ({"sources": [_source("a", content_field="conversations")]}, "messages or text"), + ], +) +def test_chat_blend_rejects_invalid_configs(chat_blend, overrides, match): + write_config, calls = chat_blend + with pytest.raises(ValueError, match=match): + blend_module.prepare_chat_data_blend(write_config(**overrides)) + assert not calls + + +@pytest.mark.parametrize( + ("overrides", "match"), + [ + ({}, "Unknown output_format 'chat_jsonl'"), + ({"output_format": "chat-jsonl"}, "Unknown output_format 'chat-jsonl'"), + ({"output_format": "megatron"}, r"Tokenized blends do not use \['seed', 'shuffle'"), + ], +) +def test_tokenized_blend_rejects_chat_options(chat_blend, overrides, match): + write_config, _ = chat_blend + with pytest.raises(ValueError, match=match): + blend_module.prepare_megatron_data_blend(write_config(**overrides)) + + +def test_chat_blend_fails_on_a_bad_short_or_tiny_source(tmp_path: Path, chat_blend, monkeypatch): + write_config, _ = chat_blend + with pytest.raises(RuntimeError, match="org/a train: only 100/250 usable rows"): + blend_module.prepare_chat_data_blend(write_config(target_samples=500)) + with pytest.raises(RuntimeError, match="Too few rows"): + blend_module.prepare_chat_data_blend(write_config(target_samples=10)) + + def iter_rows(source, output_dir, shuffle, seed): + yield from (_chat(i) for i in range(10)) # yields some rows, then the reader fails + raise TypeError("cannot cast") + + monkeypatch.setattr(blend_module, "_iter_rows", iter_rows) + with pytest.raises(RuntimeError, match="org/a train: failed to read or convert"): + blend_module.prepare_chat_data_blend(write_config()) + assert not (tmp_path / "blend" / "train.jsonl").exists() + + +def test_write_jsonl_is_atomic(tmp_path: Path): + path = tmp_path / "part.jsonl" + with pytest.raises(TypeError): + blend_module._write_jsonl(path, [_chat(0), {"bad": object()}]) + assert not path.exists() + + +def test_chat_blend_reads_files_sources(tmp_path: Path, monkeypatch): + plain, packed = tmp_path / "a.jsonl", tmp_path / "b.jsonl.gz" + lines = [json.dumps(_chat(i)) for i in range(20)] + # A malformed record is skipped with a warning, as in the tokenizing reader. + plain.write_text("\n".join([*lines[:10], '{"messages": [', *lines[10:]]) + "\n", "utf-8") + with gzip.open(packed, "wt", encoding="utf-8") as f: + f.writelines(json.dumps(_chat(i)) + "\n" for i in range(20, 40)) + paths = {"data/a.jsonl": str(plain), "data/b.jsonl.gz": str(packed)} + monkeypatch.setattr( + huggingface_hub, "hf_hub_download", lambda repo_id, filename, **kwargs: paths[filename] + ) + config_path = tmp_path / "config.yaml" + config = { + "output_dir": str(tmp_path / "blend"), + "output_format": "chat_jsonl", + "target_samples": 40, + "sources": [_source("a", split=None, files=list(paths))], + } + config_path.write_text(yaml.safe_dump(config), encoding="utf-8") + assert blend_module.prepare_chat_data_blend(config_path) == {"validation": 2, "train": 38} + out = tmp_path / "blend" + rows = _read_jsonl(out / "train.jsonl") + _read_jsonl(out / "validation.jsonl") + assert sorted(r["messages"][0]["content"] for r in rows) == sorted(f"q{i}" for i in range(40)) + # The part is named after its first file. + assert len(list((out / "parts").glob("org--a_a_*.jsonl"))) == 1 From 8689869d5bce0f658954e4ae77f836116035aba4 Mon Sep 17 00:00:00 2001 From: Colin McDonough Date: Thu, 15 Oct 2026 03:58:23 -0400 Subject: [PATCH 20/35] docs(dms): fix broken link to the small-budget training script (#2792) ## What does this PR do? **Type of change:** Documentation `experimental/dms/README.md` links to `scripts/train_small_debug.sh`, which does not exist in `experimental/dms/scripts/` (only `train.sh`, `train_small.sh`, `evaluate.sh`). `train_small.sh` is the single-GPU, limited-memory debug launcher the note describes (its header: "Example debug launch script for single-GPU training with limited memory"), so the link now points there. ## Testing Docs-only, one line. Confirmed the target file exists at `experimental/dms/scripts/train_small.sh`. Commit is signed and signed-off (DCO). ## Summary by CodeRabbit * **Documentation** * Updated the training guide to direct users with smaller compute budgets to `scripts/train_small.sh`. Signed-off-by: ColinM-sys Signed-off-by: Rohan Joshi --- experimental/dms/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/experimental/dms/README.md b/experimental/dms/README.md index 2107c80079a..edaa34bd3d6 100644 --- a/experimental/dms/README.md +++ b/experimental/dms/README.md @@ -48,7 +48,7 @@ This single install provides everything needed for training and evaluation (incl ### Train DMS adapters -**Note:** The number of GPUs determines the effective batch size. The configuration below was tested on a DGX H100 with 8× H100 80GB GPUs. For debugging with a smaller compute budget (e.g., a single RTX 5090), see [`scripts/train_small_debug.sh`](scripts/train_small_debug.sh). +**Note:** The number of GPUs determines the effective batch size. The configuration below was tested on a DGX H100 with 8× H100 80GB GPUs. For debugging with a smaller compute budget (e.g., a single RTX 5090), see [`scripts/train_small.sh`](scripts/train_small.sh). ```bash bash scripts/train.sh configs/qwen3_8b.yaml From 6c89743df1de8ce146bd901e0d425dac50faf26b Mon Sep 17 00:00:00 2001 From: Juhi Mittal <39641197+juhi10071998@users.noreply.github.com> Date: Thu, 15 Oct 2026 09:58:38 -0700 Subject: [PATCH 21/35] [NVBug -6914213] fix(export): drop reverse renames that invent names the source checkpoint lacks (#2766) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## What changed and why A transformers conversion mapping is reversed on save so a quantized export keeps the original hub checkpoint's names. That reverse must be a **no-op on names already in source form** — the invariant `_idempotent_pattern` already enforces for prefix rules. A model breaks it by declaring a backwards-compatibility *load* rename as a reversible `WeightRenaming`. NemotronH does exactly this: ```python register_checkpoint_conversion_mapping("nemotron_h", [ WeightRenaming("backbone.", "model."), WeightRenaming("embedding.weight", "embeddings.weight"), # BC-only, but reversible ], overwrite=True) ``` The second rule mirrors a transformers-4 `load_hook`, which was load-only. Under transformers 5 it is bidirectional, so reversing it rewrites the canonical `backbone.embeddings.weight` to a legacy `backbone.embedding.weight` that the published checkpoint never used. vLLM's NemotronH loader maps the substring `embeddings` to `embed_tokens`, so the singular name reaches it unmapped and serving fails at load: ``` ValueError: There is no module or parameter named 'model.embedding' in NemotronHForCausalLM ``` The published checkpoint is plural, so the reverse was not restoring a name — it was inventing one. This surfaced on a 0.48.0rc0 export of a 120B NemotronH checkpoint (NVBug 6914213); 0.47 was unaffected only because it did not reverse the mapping at all. ### The fix `_source_checkpoint_keys()` reads the tensor names of the checkpoint the model was loaded from. `_SourceAwareRenamer` then resolves **each key** against them: ``` apply the rename chain result is a name the checkpoint carries? -> use it (unchanged behavior) result absent, but the chain without the name-inventing rules gives one it carries? -> use that (the fix) neither names a source tensor? -> use the chain (unchanged behavior) ``` Reversing a rename is genuinely ambiguous: given a canonical in-memory name, both the canonical and a legacy spelling are valid pre-images of the forward rule, and nothing inside the model breaks the tie. The source checkpoint does, which is why the evidence has to come from it. The decision is **per key, never per rule**. No rename is ever discarded, so a rule that is required for one namespace keeps applying there even if it would invent a name elsewhere — see the mixed-namespace case under Testing. Tensors the export adds, such as quantizer scales, are resolved by the weight of the module they belong to, so a module is never split across two namespaces. Rules are flagged as name-inventing only on positive evidence: a rule must rewrite at least one source key, *and* every source key it rewrites must land on a name absent from the checkpoint. When the source checkpoint cannot be read (a config-built model, a hub id never materialized locally, an unrecognized layout) behavior is unchanged. The same resolution drives `build_reverse_name_mapper`, so `exclude_modules` / `quantized_layers` move with the tensors and stay aligned. ## Testing **Unit** — `tests/unit/torch/export/test_quant_aware_conversion.py`: 74 pass, of which 17 items cover this change — the suppression, the two apply-anyway cases, the inert path, the mixed-namespace case (weights, scale tensors and the config mapping), the sharded-index and single-file reads, the malformed-index fallback, the cached-hub-id resolution and its revision pinning, the empty-conversion-mapping early return, and the resolution end to end through `_build_reverse_rules` (asserting the exported tensor name, `quantized_layers` and `exclude_modules` all stay canonical while an unrelated prefix rename still reverses). **Mixed-namespace checkpoint (review follow-up)** — a published checkpoint can store one container name for some layers and the in-memory name for others, so a single rule is both needed and wrong within the same model. Resolving per rule would discard it wholesale. Checked against a real published 23,416-key index whose sparse layers are stored under `block_sparse_moe` while its dense and vision layers are already stored under `mlp`, driving every published key through the reverse chain: ``` round-trip of all 23,416 published keys: ALL CORRECT layers.N.mlp.experts.0.w1{.weight,.weight_scale,.input_scale} -> layers.N.block_sparse_moe.experts.0.w1{...} (renamed) layers.M.mlp.down_proj{.weight,.weight_scale,.input_scale} -> unchanged (already canonical) ``` Every key lands on a name the checkpoint actually carries, and a module's scale tensors follow its weight. **Real checkpoint, before/after** — tiny NemotronH built from the published `nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16` config and its own remote code, with `name_or_path` pointing at that checkpoint so the guard reads its real 42,683 names. Same tree, same venv (transformers 5.17.0), one file apart: ``` BEFORE (b) modelopt reverse : ['backbone.embedding.weight'] exclude_modules: backbone.embedding AFTER (b) modelopt reverse : ['backbone.embeddings.weight'] exclude_modules: backbone.embeddings source keys read: 42683 canonical in source: True legacy in source: False ``` Renaming that key in the real 119 GB failing artefact clears the reported error: weight loading completes (`Model loading took 28.72 GiB memory and 33.76 seconds`). **Hub-id `name_or_path`** — `from_pretrained("nvidia/...")` leaves `name_or_path` as the repo id, which is not a directory, so the first revision of this guard read no keys and silently did nothing. Resolved through the local HF cache and re-checked on the same checkpoint, two arms one file apart: ``` without the resolver : hub resolver unavailable -> no keys with the resolver : local snapshot dir : keys=42683 canonical=True legacy=False hub id via HF cache : keys=42683 canonical=True legacy=False ``` **Full export, end to end** — the 120B re-quantized from the published BF16 checkpoint with this commit on `PYTHONPATH`, same recipe as the original report (`general/ptq/fp8_default-kv_fp8_cast`), 4×GPU, 119 GB out: ``` backbone.embeddings.weight (canonical) in export : True backbone.embedding.weight (legacy) in export : False source names absent from export : 0 export-only non-scale names : 0 quantization_config exclude_modules embedding entry : backbone.embeddings ``` 125,147 export tensors against the source's 42,683, reconciling exactly: every published name present, none invented. vLLM's `nemotron_h` mapper turns `backbone.embeddings` into `model.embed_tokens`, which exists, so the reported `ValueError` cannot arise from this artefact. Not booted in vLLM — the environment used for the export has no nvcc, so a serve would stop in FlashInfer JIT for unrelated reasons. **No-regression on an unaffected model** — a separate model whose conversion mapping carries several reverse renames, all of which do restore real source names, was re-exported end to end from an identical tree plus this commit: ``` tensor-name delta vs the pre-guard export : 0 source-checkpoint names absent from export: 0 shard md5 : all shards byte-identical ``` No rule is dropped when every rename is sound, so the export is unchanged. ## Note for maintainers This guards ModelOpt's export path. A plain `save_pretrained` round-trip of that checkpoint under transformers 5 still produces the legacy name, because HF's own `revert_weight_conversion` is unguarded — the model card's `WeightRenaming` is worth fixing separately at the source. 🤖 Generated with [Claude Code](https://claude.com/claude-code) ## Summary by CodeRabbit * **Bug Fixes** * Fixed unified Hugging Face exports so tensor names and quantization configuration references are not incorrectly renamed when reversing a rule would produce a name absent from the source checkpoint. * Reverse renaming rules are checked against known checkpoint contents, helping preserve valid names for local and cached checkpoints. If checkpoint contents are unavailable or cannot be read, existing renaming behavior remains unchanged. --------- Signed-off-by: Juhi Mittal Co-authored-by: Claude Opus 5 (1M context) Signed-off-by: Rohan Joshi --- CHANGELOG.rst | 1 + .../torch/export/quant_aware_conversion.py | 206 +++++++++++- .../export/unified_export_hf_streaming.py | 2 +- .../export/test_quant_aware_conversion.py | 297 +++++++++++++++++- 4 files changed, 490 insertions(+), 16 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 87930a8aaf4..120b7015037 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -127,6 +127,7 @@ Changelog **Bug Fixes** +- Fix unified HF export writing a weight name the source checkpoint does not use when a model declares a backwards-compatibility *load* rename as a reversible ``WeightRenaming``. Reversing such a rule rewrote the canonical name to a legacy spelling, and a deployment runtime keyed on the canonical name then failed to load the quantized checkpoint. Each exported name is now resolved against the names the source checkpoint actually carries, for both the exported tensors and the quantization config's module references. - Fix exports using the built-in Transformers NemotronH implementation retaining incompatible legacy remote-code config and model mappings. Re-export affected checkpoints before deployment. - Fix INT8 real quantization (``mtq.compress``) of FP16 weights zeroing or clipping channels with a small amax: their scales underflowed in FP16. Scales of newly compressed INT8 weights are now stored in FP32. - Fix Hugging Face exports classifying registered Qwen VLM parameters as unplaced source weights, which overwrote quantized tensors with original weights and incorrectly excluded their modules from quantization. Re-export affected checkpoints. diff --git a/modelopt/torch/export/quant_aware_conversion.py b/modelopt/torch/export/quant_aware_conversion.py index 8c750702751..ad186df9300 100644 --- a/modelopt/torch/export/quant_aware_conversion.py +++ b/modelopt/torch/export/quant_aware_conversion.py @@ -62,13 +62,21 @@ (in-memory-name) behavior rather than emit a silently-wrong checkpoint. """ +import logging import re from collections.abc import Iterable from dataclasses import dataclass from fnmatch import fnmatchcase +from pathlib import Path from typing import Any import torch +from huggingface_hub import try_to_load_from_cache +from huggingface_hub.errors import HFValidationError + +from modelopt.torch.utils.plugins.hf_checkpoint_utils import indexed_weight_map + +logger = logging.getLogger(__name__) __all__ = [ "QuantConversionUnsupportedError", @@ -321,6 +329,7 @@ def apply_reverse_rules( split_rules: list[SplitRule], rename_rules: list[RenameRule], merge_rules: list[_MergeRule] | None = None, + source_keys: set[str] | None = None, ) -> dict[str, torch.Tensor]: """Apply quant-aware reverse conversion: tensor transforms first, then renames. @@ -329,6 +338,10 @@ def apply_reverse_rules( ``merge_rules`` is an internal hook for model-derived rules: unlike the public rename/split descriptors, each rule retains a live transformers converter for scoped matching. Call :func:`revert_weight_conversion_quant_aware` to derive them. + + ``source_keys`` are the tensor names of the checkpoint the model was loaded from. + They let a rename that would otherwise invent a name resolve to the spelling the + checkpoint actually carries; omitting them applies the rules unchanged. """ out = dict(state_dict) for rule in split_rules: @@ -336,10 +349,10 @@ def apply_reverse_rules( for merge_rule in merge_rules or []: _apply_merge_rule(out, merge_rule) - compiled = _compile_rename_rules(rename_rules) + rename = _SourceAwareRenamer(rename_rules, source_keys) renamed: dict[str, torch.Tensor] = {} for key, value in out.items(): - new_key = _apply_rename_rules(key, compiled) + new_key = rename(key) if new_key in renamed: raise QuantConversionUnsupportedError(f"rename collision on '{new_key}'") renamed[new_key] = value @@ -354,11 +367,13 @@ def revert_weight_conversion_quant_aware(model, state_dict: dict[str, torch.Tens when the mapping uses an op that cannot be reversed quant-aware yet, so the caller can fall back to the legacy behavior. """ - split_rules, merge_rules, rename_rules, expert_fused_leaves = _build_reverse_rules(model) + split_rules, merge_rules, rename_rules, expert_fused_leaves, source_keys = _build_reverse_rules( + model + ) if not split_rules and not merge_rules and not rename_rules: return state_dict _assert_experts_pre_expanded(state_dict, expert_fused_leaves) - return apply_reverse_rules(state_dict, split_rules, rename_rules, merge_rules) + return apply_reverse_rules(state_dict, split_rules, rename_rules, merge_rules, source_keys) def build_reverse_name_mapper(model, *, tensor_keys: bool = False): @@ -382,15 +397,15 @@ def build_reverse_name_mapper(model, *, tensor_keys: bool = False): :class:`QuantConversionUnsupportedError` when the mapping can't be reversed, so the caller can keep the in-memory names for BOTH weights and config (mutually consistent). """ - _, merge_rules, rename_rules, _ = _build_reverse_rules(model) + _, merge_rules, rename_rules, _, source_keys = _build_reverse_rules(model) if not merge_rules and not rename_rules: return None - compiled = _compile_rename_rules(rename_rules) + rename = _SourceAwareRenamer(rename_rules, source_keys) def _apply(text: str) -> str: for rule in merge_rules: text, _ = rule.converter.rename_source_key(text) - return _apply_rename_rules(text, compiled) + return rename(text) def _map(name: str) -> str: if tensor_keys: @@ -504,6 +519,160 @@ def _scope_prefixes(rev) -> tuple[str, ...]: return tuple(seen) +_SAFETENSORS_INDEX = "model.safetensors.index.json" +_SAFETENSORS_SINGLE = "model.safetensors" + + +def _checkpoint_dir(path: str, revision: str | None = None) -> Path | None: + """Local directory holding ``path``'s safetensors checkpoint, or ``None``. + + ``path`` is a :attr:`name_or_path`, so it is either a local directory or a hub id. + A hub id is resolved against the local HF cache only -- this never downloads, and + an uncached hub id simply yields ``None``. + + ``revision`` pins the lookup to the commit the model was actually loaded from. + Without it the cache lookup defaults to ``main``, which for a model loaded from a + pinned revision could read a *different* snapshot's names and filter the wrong + rules. ``None`` keeps the hub default, which is correct when the revision is + unknown because nothing else identifies the snapshot. + + A non-default ``cache_dir`` is *not* covered: ``transformers`` does not record it + on the model, so a model loaded from one resolves here against the default cache + and most likely finds nothing -- which leaves the guard inert and behavior + unchanged, the safe direction. + """ + directory = Path(path) + if directory.is_dir(): + return directory + for filename in (_SAFETENSORS_INDEX, _SAFETENSORS_SINGLE): + try: + cached = try_to_load_from_cache(path, filename, revision=revision) + except HFValidationError: # not a well-formed hub id either + return None + # Returns the file path on a hit, and a sentinel object (not a str) for a + # cached "this file does not exist" result. + if isinstance(cached, str): + return Path(cached).parent + return None + + +def _source_checkpoint_keys(model) -> set[str] | None: + """Tensor names of the checkpoint ``model`` was loaded from, or ``None``. + + ``None`` means "unknown" and leaves the reverse conversion exactly as it was: it is + returned for a model built from a config, a hub id absent from the local HF cache, + or a checkpoint layout this does not recognize. + """ + config = getattr(model, "config", None) + path = getattr(model, "name_or_path", None) or getattr(config, "_name_or_path", None) + if not path: + return None + try: + # ``transformers`` records the resolved commit on the config; a hub-id lookup + # must use it rather than falling back to ``main``. + directory = _checkpoint_dir(path, getattr(config, "_commit_hash", None)) + if directory is None: + return None + # ``indexed_weight_map`` is the repo's one index reader: ``weight_map`` for a + # sharded checkpoint, ``safe_open`` keys for a single file, ``{}`` for neither. + # An empty result means "nothing recorded", which is exactly our unknown. + return set(indexed_weight_map(directory)) or None + except Exception as exc: # unreadable or malformed checkpoint -> treat as unknown + # Deliberately no path and no exception text: SECURITY.md warns that logs may + # carry sensitive paths, and the type alone identifies the failure mode. + logger.debug("Cannot read source checkpoint keys: %s", type(exc).__name__) + return None + + +def _non_idempotent_renames(rules: list[RenameRule], source_keys: set[str] | None) -> set[int]: + """Indices of reverse renames that only ever invent names the source lacks. + + A reverse rename maps an in-memory name back to its source-checkpoint name, so it + must be a *no-op* on names that are already in source form -- the same invariant + :func:`_idempotent_pattern` enforces for prefix rules. A rule that breaks it may be + inventing a name rather than restoring one. + + That happens when a model declares a backwards-compatibility *load* rename as a + reversible ``WeightRenaming``. Reversing it rewrites the canonical name to a legacy + spelling the published checkpoint never used, and a deployment loader keyed on the + canonical name then cannot find the tensor. + + A rule is flagged only on positive evidence: it must rewrite at least one source key, + and every source key it rewrites must land on a name absent from the checkpoint. + + Flagging is not a verdict. A rule can be *needed* for one namespace and wrong for + another -- a mixed-namespace checkpoint whose sparse layers are stored under one + container name and whose dense and vision layers are already stored under the + in-memory one. :class:`_SourceAwareRenamer` therefore decides per key, and a flagged + rule still applies wherever it produces a name the checkpoint actually carries. + """ + if not source_keys: + return set() + suspect: set[int] = set() + for index, rule in enumerate(rules): + pattern = re.compile(_idempotent_pattern(rule.pattern, rule.repl)) + rewrites_a_source_key = False + every_result_absent = True + for key in source_keys: + mapped = _sub_scoped(pattern, rule.repl, key, rule.scope_prefixes) + if mapped == key: + continue + rewrites_a_source_key = True + if mapped in source_keys: + every_result_absent = False + break + if rewrites_a_source_key and every_result_absent: + suspect.add(index) + return suspect + + +class _SourceAwareRenamer: + """Apply reverse renames, preferring names the source checkpoint actually carries. + + The reverse of a rename is ambiguous: given a canonical in-memory name, both the + canonical and a legacy spelling are valid pre-images, and nothing inside the model + breaks the tie. The checkpoint the model was loaded from does, so a key whose + rewritten name is absent from it is resolved against the rules that do not invent + names (see :func:`_non_idempotent_renames`). + + Deciding per key rather than per rule matters: no rule is ever discarded, so a rule + that a mixed-namespace checkpoint needs for one set of layers keeps applying there + while the same rule stops rewriting layers already stored in their in-memory form. + + Without source evidence, or when neither candidate names a source tensor, the full + rule chain applies unchanged -- the behavior this class replaced. + """ + + def __init__(self, rename_rules: list[RenameRule], source_keys: set[str] | None): + self._full = _compile_rename_rules(rename_rules) + self._source_keys = source_keys or set() + suspect = _non_idempotent_renames(rename_rules, source_keys) + self._reduced = ( + _compile_rename_rules( + [rule for index, rule in enumerate(rename_rules) if index not in suspect] + ) + if suspect + else None + ) + + def _names_a_source_tensor(self, name: str) -> bool: + if name in self._source_keys: + return True + # A tensor the export adds -- a quantizer scale, or a bias the source omits -- + # has no counterpart in the checkpoint, so it is judged by the weight of the + # module it belongs to. That keeps every tensor of one module on the same + # decision instead of splitting the module across two namespaces. + parent, _, leaf = name.rpartition(".") + return bool(leaf) and f"{parent}.weight" in self._source_keys + + def __call__(self, key: str) -> str: + renamed = _apply_rename_rules(key, self._full) + if self._reduced is None or self._names_a_source_tensor(renamed): + return renamed + reduced = _apply_rename_rules(key, self._reduced) + return reduced if self._names_a_source_tensor(reduced) else renamed + + def _drop_shadowed_prefix_renames(model, rules: list[RenameRule]) -> list[RenameRule]: """Drop child-model reverse renames when the child namespace already exists. @@ -549,13 +718,16 @@ def _drop_shadowed_prefix_renames(model, rules: list[RenameRule]) -> list[Rename def _build_reverse_rules( model, -) -> tuple[list[SplitRule], list[_MergeRule], list[RenameRule], list[str]]: +) -> tuple[list[SplitRule], list[_MergeRule], list[RenameRule], list[str], set[str] | None]: """Derive reverse rules from the model's transformers conversion mapping. - Returns tensor-transform rules, rename rules, and fused expert leaves. The last is - used to guard against experts that were not pre-expanded. Returns empty lists when - no mapping applies (export unchanged). Uses transformers' own - ``reverse_transform()`` to get correctly-reversed name patterns, then translates them: + Returns tensor-transform rules, rename rules, fused expert leaves, and the key set + of the checkpoint the model was loaded from (``None`` when it cannot be read). The + expert leaves guard against experts that were not pre-expanded; the key set lets + :class:`_SourceAwareRenamer` resolve a rename against the names the checkpoint has. + Returns empty lists when no mapping applies (export unchanged). Uses transformers' + own ``reverse_transform()`` to get correctly-reversed name patterns, then translates + them: * ``WeightRenaming`` -> :class:`RenameRule` (carries scale siblings for free). * Expert ``WeightConverter`` (reverse contains ``SplitModulelist``): ModelOpt's @@ -580,7 +752,7 @@ def _build_reverse_rules( raise QuantConversionUnsupportedError(f"could not read conversion mapping: {exc}") from exc if not conversions: - return [], [], [], [] + return [], [], [], [], None try: from transformers.core_model_loading import ( @@ -660,7 +832,13 @@ def _build_reverse_rules( # substrings and are applied first. weight_renamings = _drop_shadowed_prefix_renames(model, weight_renamings) rename_rules = leaf_renamings + list(reversed(weight_renamings)) - return split_rules, merge_rules, rename_rules, expert_fused_leaves + return ( + split_rules, + merge_rules, + rename_rules, + expert_fused_leaves, + _source_checkpoint_keys(model), + ) # ModelOpt's export splits a fused ``gate_up_proj`` into these per-expert linears, diff --git a/modelopt/torch/export/unified_export_hf_streaming.py b/modelopt/torch/export/unified_export_hf_streaming.py index 26896da92f5..69cdc53117f 100644 --- a/modelopt/torch/export/unified_export_hf_streaming.py +++ b/modelopt/torch/export/unified_export_hf_streaming.py @@ -211,7 +211,7 @@ def _assert_no_split_rules(model: nn.Module) -> None: reversal cannot do. """ try: - split_rules, merge_rules, _, _ = _build_reverse_rules(model) + split_rules, merge_rules, _, _, _ = _build_reverse_rules(model) except Exception: return # build_reverse_name_mapper reports the failure with a warning if split_rules or merge_rules: diff --git a/tests/unit/torch/export/test_quant_aware_conversion.py b/tests/unit/torch/export/test_quant_aware_conversion.py index e2efe7c7ff6..1a798b2c1fd 100644 --- a/tests/unit/torch/export/test_quant_aware_conversion.py +++ b/tests/unit/torch/export/test_quant_aware_conversion.py @@ -29,8 +29,9 @@ import pytest import torch -from safetensors.torch import load_file +from safetensors.torch import load_file, save_file +import modelopt.torch.export.quant_aware_conversion as quant_aware_conversion import modelopt.torch.quantization as mtq from modelopt.torch.export.layerwise_export import LayerwiseExporter from modelopt.torch.export.quant_aware_conversion import ( @@ -38,6 +39,9 @@ RenameRule, SplitRule, _assert_experts_pre_expanded, + _build_reverse_rules, + _source_checkpoint_keys, + _SourceAwareRenamer, apply_reverse_rules, build_reverse_name_mapper, revert_quant_config_names, @@ -187,6 +191,297 @@ def test_idempotency_guard_leaves_disjoint_rules_alone(): assert list(out) == ["vision_model.radio_model.model.blocks.0.attn.weight"] +# A NemotronH-shaped source checkpoint: the canonical embedding name is plural, and the +# projector weights sit at the top level under their fused ``mlp1.`` names. +_SOURCE_KEYS = { + "backbone.embeddings.weight", + "backbone.layers.0.mixer.in_proj.weight", + "mlp1.0.weight", + "mlp1.1.weight", + "vision_projector.vision_final_layernorm.weight", +} + + +# A MiniMax-M3-shaped source checkpoint: the sparse text layers are stored under +# ``block_sparse_moe``, while the dense text layers and the vision tower are already +# stored under the in-memory ``mlp`` name. One rename rule serves both namespaces. +_MIXED_SOURCE_KEYS = { + "language_model.model.layers.0.mlp.down_proj.weight", + "language_model.model.layers.1.block_sparse_moe.gate.weight", + "language_model.model.layers.1.block_sparse_moe.experts.0.w1.weight", + "vision_tower.encoder.layers.0.mlp.fc1.weight", +} + +# The NemotronH reverse chain, in the order ``_build_reverse_rules`` emits it: the +# backwards-compatibility leaf rename first, then the ``model.`` -> ``backbone.`` prefix. +_NEMOTRON_H_REVERSE = [ + RenameRule(r"embeddings\.weight", "embedding.weight"), + RenameRule(r"^model\.", "backbone."), +] + + +def test_rename_inventing_a_name_absent_from_source_is_suppressed(): + """A BC *load* rename declared reversible would rewrite the canonical name. + + Reversing ``embedding.weight -> embeddings.weight`` turns the published + ``backbone.embeddings.weight`` into a legacy spelling the checkpoint never used, + which a deployment loader keyed on the canonical name cannot find. The prefix rule + in the same chain is still needed, so resolution has to span the whole chain. + """ + rename = _SourceAwareRenamer(_NEMOTRON_H_REVERSE, _SOURCE_KEYS) + assert rename("model.embeddings.weight") == "backbone.embeddings.weight" + + +def test_rename_restoring_a_source_name_is_applied(): + """A rename whose output is a real source name is left alone.""" + rule = RenameRule(r"^vision_projector\.mlp1\.norm\.", "mlp1.0.") + rename = _SourceAwareRenamer([rule], _SOURCE_KEYS) + assert rename("vision_projector.mlp1.norm.weight") == "mlp1.0.weight" + + +def test_rename_matching_no_source_key_is_applied(): + """No evidence means no suppression: a rule that never fires on source names runs.""" + rule = RenameRule(r"^sound_projector\.sound_encoder\.", "sound_encoder.") + rename = _SourceAwareRenamer([rule], _SOURCE_KEYS) + assert rename("sound_projector.sound_encoder.0.weight") == "sound_encoder.0.weight" + + +@pytest.mark.parametrize("source_keys", [None, set()]) +def test_guard_is_inert_when_source_is_unknown(source_keys): + """An unreadable or absent source checkpoint keeps the previous behaviour.""" + rename = _SourceAwareRenamer(_NEMOTRON_H_REVERSE, source_keys) + assert rename("model.embeddings.weight") == "backbone.embedding.weight" + + +# The same container rename shows up in two shapes depending on how the model declares it: +# one bare container rule, or a rule per anchored leaf. The per-key resolution must give +# the same answer either way, so every mixed-namespace assertion runs against both. +_MIXED_RULE_SHAPES = { + "bare-container": [RenameRule(r"\.mlp\.", ".block_sparse_moe.")], + "anchored-leaves": [ + RenameRule(r"\.mlp\.experts\.", ".block_sparse_moe.experts."), + RenameRule(r"\.mlp\.gate\.", ".block_sparse_moe.gate."), + ], +} + + +@pytest.mark.parametrize("shape", list(_MIXED_RULE_SHAPES)) +def test_mixed_namespace_rename_applies_only_where_the_source_needs_it(shape): + """One rule, two namespaces: suppression is per key, never per rule. + + A rule that rewrites only already-canonical source keys would be discarded by a + rule-level guard, taking the sparse-layer renames it is actually needed for with it. + """ + rename = _SourceAwareRenamer(_MIXED_RULE_SHAPES[shape], _MIXED_SOURCE_KEYS) + + # Sparse text layers: the rename is what restores their source names. + assert ( + rename("language_model.model.layers.1.mlp.experts.0.w1.weight") + == "language_model.model.layers.1.block_sparse_moe.experts.0.w1.weight" + ) + assert ( + rename("language_model.model.layers.1.mlp.gate.weight") + == "language_model.model.layers.1.block_sparse_moe.gate.weight" + ) + # Dense text layer and vision tower are already canonical: left untouched. + assert ( + rename("language_model.model.layers.0.mlp.down_proj.weight") + == "language_model.model.layers.0.mlp.down_proj.weight" + ) + assert ( + rename("vision_tower.encoder.layers.0.mlp.fc1.weight") + == "vision_tower.encoder.layers.0.mlp.fc1.weight" + ) + + +@pytest.mark.parametrize("shape", list(_MIXED_RULE_SHAPES)) +@pytest.mark.parametrize("suffix", [".weight", ".weight_scale", ".input_scale"]) +def test_mixed_namespace_scales_follow_their_weight(suffix, shape): + """A module's scale tensors must not be split across the two namespaces. + + Scale names never appear in the source checkpoint, so they are resolved by the + weight they belong to rather than judged absent on their own. + """ + rename = _SourceAwareRenamer(_MIXED_RULE_SHAPES[shape], _MIXED_SOURCE_KEYS) + sparse = "language_model.model.layers.1.mlp.experts.0.w1" + dense = "language_model.model.layers.0.mlp.down_proj" + assert ( + rename(sparse + suffix) + == "language_model.model.layers.1.block_sparse_moe.experts.0.w1" + suffix + ) + assert rename(dense + suffix) == dense + suffix + + +def test_mixed_namespace_config_names_track_the_weights(tmp_path): + """``exclude_modules`` / ``quantized_layers`` follow the same per-key decision.""" + pytest.importorskip("transformers.core_model_loading") + # Local import: transformers is an optional dependency for ModelOpt. + from transformers.core_model_loading import WeightRenaming + + (tmp_path / "model.safetensors.index.json").write_text( + json.dumps( + {"weight_map": dict.fromkeys(_MIXED_SOURCE_KEYS, "model-00001-of-00001.safetensors")} + ) + ) + model = types.SimpleNamespace( + name_or_path=str(tmp_path), + config=None, + _weight_conversions=[WeightRenaming(".block_sparse_moe.", ".mlp.")], + ) + mapper = build_reverse_name_mapper(model) + assert ( + mapper("language_model.model.layers.1.mlp.experts.0.w1") + == "language_model.model.layers.1.block_sparse_moe.experts.0.w1" + ) + assert ( + mapper("language_model.model.layers.0.mlp.down_proj") + == "language_model.model.layers.0.mlp.down_proj" + ) + + +def test_model_without_a_conversion_mapping_is_a_no_op(): + """The empty-mapping early return must stay shape-compatible with its callers. + + ``build_reverse_name_mapper`` unpacks the result, so a short tuple here would raise + ``ValueError`` for every model that declares no conversion mapping at all. + """ + model = types.SimpleNamespace(name_or_path=None, config=None, _weight_conversions=[]) + split_rules, merge_rules, rename_rules, expert_leaves, source_keys = _build_reverse_rules(model) + assert (split_rules, merge_rules, rename_rules, expert_leaves, source_keys) == ( + [], + [], + [], + [], + None, + ) + assert build_reverse_name_mapper(model) is None + state = {"model.embeddings.weight": torch.ones(2, 2)} + assert revert_weight_conversion_quant_aware(model, state) is state + + +def test_source_checkpoint_keys_reads_sharded_index(tmp_path): + """The guard's input comes from the checkpoint the model was loaded from.""" + (tmp_path / "model.safetensors.index.json").write_text( + json.dumps({"weight_map": dict.fromkeys(_SOURCE_KEYS, "model-00001-of-00001.safetensors")}) + ) + model = types.SimpleNamespace(name_or_path=str(tmp_path), config=None) + assert _source_checkpoint_keys(model) == _SOURCE_KEYS + assert _source_checkpoint_keys(types.SimpleNamespace(name_or_path="", config=None)) is None + + +def test_source_checkpoint_keys_reads_single_file(tmp_path): + """The single-file ``model.safetensors`` layout is read through ``safe_open``.""" + # A fresh tensor per key: ``save_file`` rejects names that share storage. + save_file({key: torch.ones(1) for key in _SOURCE_KEYS}, tmp_path / "model.safetensors") + model = types.SimpleNamespace(name_or_path=str(tmp_path), config=None) + assert _source_checkpoint_keys(model) == _SOURCE_KEYS + + +def test_source_checkpoint_keys_tolerates_a_malformed_index(tmp_path): + """An unreadable checkpoint is "unknown", not an export failure.""" + (tmp_path / "model.safetensors.index.json").write_text("{ not json") + model = types.SimpleNamespace(name_or_path=str(tmp_path), config=None) + assert _source_checkpoint_keys(model) is None + + (tmp_path / "model.safetensors.index.json").write_text(json.dumps({"metadata": {}})) + assert _source_checkpoint_keys(model) is None + + +def test_source_checkpoint_keys_resolves_a_cached_hub_id(tmp_path, monkeypatch): + """A hub-id ``name_or_path`` resolves through the local HF cache, offline. + + ``from_pretrained("nvidia/...")`` leaves ``name_or_path`` as the repo id, which is + not a directory. Without cache resolution the guard would silently do nothing for + the most common way ``hf_ptq.py`` is invoked. + """ + snapshot = tmp_path / "snapshots" / "abc123" + snapshot.mkdir(parents=True) + index = snapshot / "model.safetensors.index.json" + index.write_text( + json.dumps({"weight_map": dict.fromkeys(_SOURCE_KEYS, "model-00001-of-00001.safetensors")}) + ) + seen = {} + + def _fake_cache(repo_id, filename, revision=None): + seen["revision"] = revision + return str(index) if filename == "model.safetensors.index.json" else None + + monkeypatch.setattr(quant_aware_conversion, "try_to_load_from_cache", _fake_cache) + model = types.SimpleNamespace(name_or_path="nvidia/Some-Model", config=None) + assert _source_checkpoint_keys(model) == _SOURCE_KEYS + assert seen["revision"] is None + + +def test_cached_hub_id_lookup_uses_the_loaded_revision(tmp_path, monkeypatch): + """A model pinned to a revision must not read ``main``'s snapshot. + + ``try_to_load_from_cache`` defaults to ``main``. For a model loaded from a pinned + commit that would read a different cached checkpoint and filter the wrong rules, + so the commit ``transformers`` recorded on the config is passed through. + """ + index = tmp_path / "model.safetensors.index.json" + index.write_text( + json.dumps({"weight_map": dict.fromkeys(_SOURCE_KEYS, "model-00001-of-00001.safetensors")}) + ) + seen = {} + + def _fake_cache(repo_id, filename, revision=None): + seen["revision"] = revision + return str(index) if filename == "model.safetensors.index.json" else None + + monkeypatch.setattr(quant_aware_conversion, "try_to_load_from_cache", _fake_cache) + model = types.SimpleNamespace( + name_or_path="nvidia/Some-Model", + config=types.SimpleNamespace(_commit_hash="deadbeef", _name_or_path=None), + ) + assert _source_checkpoint_keys(model) == _SOURCE_KEYS + assert seen["revision"] == "deadbeef" + + +def test_source_aware_guard_applies_through_the_export_paths(tmp_path): + """The guard reaches both the tensor rename and the quant-config name mapping. + + Reproduces the NemotronH conversion mapping: a ``backbone.`` -> ``model.`` prefix + rename that genuinely restores source names, plus a backwards-compatibility + ``embedding.weight`` -> ``embeddings.weight`` load rename whose reverse invents a + name the published checkpoint never carried. Only the second must be suppressed. + """ + pytest.importorskip("transformers.core_model_loading") + # Local import: transformers is an optional dependency for ModelOpt. + from transformers.core_model_loading import WeightRenaming + + source_keys = {"backbone.embeddings.weight", "backbone.layers.0.mixer.in_proj.weight"} + (tmp_path / "model.safetensors.index.json").write_text( + json.dumps({"weight_map": dict.fromkeys(source_keys, "model-00001-of-00001.safetensors")}) + ) + model = types.SimpleNamespace( + name_or_path=str(tmp_path), + config=None, + _weight_conversions=[ + WeightRenaming("backbone.", "model."), + WeightRenaming("embedding.weight", "embeddings.weight"), + ], + ) + state = _nvfp4_linear("model.layers.0.mixer.in_proj", 8, 16) + state["model.embeddings.weight"] = torch.ones(8, 16) + + written = revert_weight_conversion_quant_aware(model, state) + # The ``backbone.`` prefix reverse still applies; the embedding name stays canonical. + assert "backbone.embeddings.weight" in written + assert "backbone.embedding.weight" not in written + assert "backbone.layers.0.mixer.in_proj.weight" in written + + config = { + "quantized_layers": {"model.layers.0.mixer.in_proj": {"quant_algo": "NVFP4"}}, + "exclude_modules": ["model.embeddings"], + } + revert_quant_config_names(config, build_reverse_name_mapper(model)) + assert config["quantized_layers"] == { + "backbone.layers.0.mixer.in_proj": {"quant_algo": "NVFP4"} + } + assert config["exclude_modules"] == ["backbone.embeddings"] + + def test_rename_carries_scale_siblings(): """A module rename rewrites weight + all scale siblings with identical values.""" sd = _nvfp4_linear("model.language_model.layers.10.mlp.experts.40.gate_proj", 8, 16) From 04ac0d50989c45bdb8389af68a2b4e7a91451221 Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Fri, 16 Oct 2026 03:32:16 +0530 Subject: [PATCH 22/35] Speed up CPU unit tests to stay under the 20 min CI limit (#2823) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Type of change: Bug fix, new tests (test cleanup) The `Unit tests / linux` job (serial, `timeout-minutes: 20`) was taking 13–19.5 min on main, and two recent main runs were cancelled at the limit. This cuts CPU unit-test time without dropping meaningful coverage. - **Bug fix:** `_source_checkpoint` returned `""` for models built with `from_config`, so `Path("")` resolved to the cwd and every export globbed the whole working tree (in CI, including the nox venv). `tests/unit/torch/export` went 718s → 15s locally. Regression test added. - **Multi-rank tests:** new session-scoped `cpu_dist_workers` fixture (two persistent gloo ranks) replaces a ~20s-on-CI spawn per test for autoquant (2), shapley, quantize indexer, `test_dist`, FSDP2 two-rank export. - **Moved to `tests/gpu`:** the 6 `ensure_local_checkpoint` spawn tests and `test_modeling_ptq_registration` (8 interpreters). - **Removed/deduped:** dropped the `test_presets` subprocess test (the `FutureWarning` test beside it covers it); `test_auto_quantize` matrix 60 → 9 cases (every model, format list and method still appears). - **`manual` (run on push/nightly):** slow AutoNAS-only cases and the two ONNX entropy-calibration cases (FastNAS representatives stay enabled). - **Guidelines:** AGENTS.md / CONTRIBUTING.md now say to check existing coverage first, avoid tests that restate the implementation or assert on mocks, and keep unit tests well under 5s. Clean CI-like venv (py3.12, torch 2.14 CPU, CPU-only): full `tests/unit` on the working tree was 1181s with local-only failures before, 217s and 4791 passed after (CI runners are slower, so expect roughly a 4–6 min saving, not these exact numbers). On this branch: affected unit dirs 3276 passed; moved `tests/gpu` tests run on CPU with gloo, 7 passed. New `_source_checkpoint` test fails without the fix. Not run: the moved tests on a GPU runner. The `linux` job now passes `--run-manual` on every non-PR event (push, nightly, dispatch) (nox `unit` session forwards extra args to pytest); all 17 `manual` unit tests pass locally (87s). Caveat: pinned ONNX entropy scales are no longer checked per-PR, only after merge. - Is this change backward compatible?: ✅ - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: ✅ - Did you update Changelog?: N/A - Did you get Claude approval on this PR?: ❌ 🤖 Generated with [Claude Code](https://claude.com/claude-code) * **Bug Fixes** * Checkpoint source detection now handles models with an empty source path correctly, avoiding an unintended fallback to the current directory. * **Developer Experience** * Unit-test commands can pass additional arguments to pytest. The Linux unit-test workflow runs manual tests for events other than pull requests; pull-request runs remain unchanged. * Contributor guidance now clarifies test coverage, runtime expectations, and where to place slower tests. * **Testing** * Added distributed checkpoint-resolution coverage and streamlined several CPU distributed tests. Some slower tests now run manually rather than by default. --------- Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Co-authored-by: Claude Sonnet 5.5 Signed-off-by: Rohan Joshi --- .github/workflows/unit_tests.yml | 4 +- CONTRIBUTING.md | 10 ++ modelopt/torch/export/unified_export_hf.py | 8 +- noxfile.py | 4 +- .../models/test_modeling_ptq_registration.py | 0 .../utils/test_hf_checkpoint_utils_dist.py | 107 ++++++++++++++++++ tests/unit/conftest.py | 20 ++++ .../test_calibrated_quantization.py | 4 + tests/unit/recipe/test_presets.py | 31 ----- .../export/test_fsdp2_parallel_export.py | 34 +++--- .../torch/export/test_unified_export_hf.py | 12 +- .../torch/nas/test_full_algorithms_vision.py | 2 +- tests/unit/torch/nas/test_nas.py | 16 ++- .../unit/torch/quantization/test_autoquant.py | 43 +++---- .../quantization/test_autoquant_shapley.py | 8 +- tests/unit/torch/quantization/test_dist.py | 6 +- .../torch/quantization/test_quantize_cpu.py | 6 +- .../torch/utils/test_hf_checkpoint_utils.py | 75 ------------ 18 files changed, 227 insertions(+), 163 deletions(-) rename tests/{unit => gpu}/torch/models/test_modeling_ptq_registration.py (100%) create mode 100644 tests/gpu/torch/utils/test_hf_checkpoint_utils_dist.py diff --git a/.github/workflows/unit_tests.yml b/.github/workflows/unit_tests.yml index b258db24b69..6fbca185d9f 100644 --- a/.github/workflows/unit_tests.yml +++ b/.github/workflows/unit_tests.yml @@ -67,7 +67,9 @@ jobs: env: COVERAGE_PROCESS_START: ${{ github.workspace }}/pyproject.toml COVERAGE_FILE: ${{ github.workspace }}/.coverage - run: pip install nox uv && nox -s "unit-3.12(torch_214, tf_latest)" + # Everything but PRs also runs the slow tests marked ``manual`` (e.g. AutoNAS). + PYTEST_EXTRA: ${{ github.event_name != 'pull_request' && '--run-manual' || '' }} + run: pip install nox uv && nox -s "unit-3.12(torch_214, tf_latest)" -- $PYTEST_EXTRA - name: Upload coverage reports to Codecov uses: codecov/codecov-action@v7 with: diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 3f344153850..d778f1bb10b 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -183,6 +183,16 @@ nox -s "unit-3.12(torch_211, tf_latest)" keeping CI/CD fast and lean. Default to one test per behavior: use `@pytest.mark.parametrize` instead of near-duplicate test functions, and don't add a per-branch test for a helper that a higher-level test already exercises. More tests is not better coverage — every checked-in test is CI time and maintenance forever. + Before adding a test, check whether an existing unit, GPU, or example test already exercises the path; if so, + extend it instead. Do not add tests that only restate the implementation, assert on mocks, pin log/warning + wording, or re-cover a branch of a helper whose caller is tested. A PR that adds many tests should be able to + say what regression each one catches. +- **Keep `tests/unit` fast — the `linux` job has a hard 20-minute limit** and runs the whole suite serially. + A unit test should take well under 5s (check with `pytest --durations=20`). Tests that spawn processes or + subprocesses (`spawn_multiprocess_job`, `subprocess.run` importing torch), train for many steps, or build + real-size models belong in `tests/gpu*`; batch several cases into one spawn rather than one spawn per case, + or reuse the session-scoped `cpu_dist_workers` fixture for multi-rank CPU tests. + Slow tests for features without active maintenance (e.g. AutoNAS) should be `@pytest.mark.manual` or nightly-only. - **Exercise the behavior a test claims to validate.** Mocks are useful for focused interface and wiring coverage, but replacing the implementation under test does not validate its real behavior. Include an end-to-end test that runs the actual implementation whenever the test claims backend or runtime behavior. For example, a test of diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index f6116f4746a..a1e474b80c8 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -1653,8 +1653,12 @@ def _source_checkpoint(model: nn.Module) -> str | None: tensors never reach ``exclude_modules`` because the two halves disagreed about where the checkpoint was. """ - return getattr(model, "_modelopt_source_checkpoint", None) or getattr( - getattr(model, "config", None), "_name_or_path", None + # ``or None``: a model built with ``from_config`` has ``_name_or_path == ""``, which would + # otherwise name the current directory. + return ( + getattr(model, "_modelopt_source_checkpoint", None) + or getattr(getattr(model, "config", None), "_name_or_path", None) + or None ) diff --git a/noxfile.py b/noxfile.py index afa96010ac4..46ffd1e1df6 100644 --- a/noxfile.py +++ b/noxfile.py @@ -69,7 +69,9 @@ def unit(session, torch_ver, tf_ver): tf_pins = TRANSFORMERS_VERSIONS[tf_ver] if tf_pins: session.install(*tf_pins) - session.run("python", "-m", "pytest", "tests/unit", *_cov_args(), env=_CPU_ONLY_ENV) + session.run( + "python", "-m", "pytest", "tests/unit", *_cov_args(), *session.posargs, env=_CPU_ONLY_ENV + ) @nox.session(python="3.12") diff --git a/tests/unit/torch/models/test_modeling_ptq_registration.py b/tests/gpu/torch/models/test_modeling_ptq_registration.py similarity index 100% rename from tests/unit/torch/models/test_modeling_ptq_registration.py rename to tests/gpu/torch/models/test_modeling_ptq_registration.py diff --git a/tests/gpu/torch/utils/test_hf_checkpoint_utils_dist.py b/tests/gpu/torch/utils/test_hf_checkpoint_utils_dist.py new file mode 100644 index 00000000000..72322ad50b4 --- /dev/null +++ b/tests/gpu/torch/utils/test_hf_checkpoint_utils_dist.py @@ -0,0 +1,107 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""Multi-rank tests for ``ensure_local_checkpoint`` (each case spawns a 2-process gloo job).""" + +import os +from functools import partial + +import pytest +import torch.distributed as dist +from _test_utils.torch.distributed.utils import spawn_multiprocess_job + +from modelopt.torch.utils.plugins import hf_checkpoint_utils + +pytest.importorskip("huggingface_hub") + + +class _FakeHfApi: + """``HfApi`` whose ``repo_info`` succeeds, so the Hub counts as reachable.""" + + def repo_info(self, repo_id): + pass + + +def _ensure_local_checkpoint_job(tmp_path, fail, group_ranks, rank, size): + utils = hf_checkpoint_utils + snapshot = tmp_path / "snapshot" + + def fake_snapshot_download(repo_id, **kwargs): + with (tmp_path / f"downloaded-by-rank{rank}").open("a") as marker: + marker.write(f"{repo_id}\n") # one line per call, so a repeat download shows + if fail: + raise OSError("hub unreachable") + snapshot.mkdir(exist_ok=True) + return str(snapshot) + + utils.snapshot_download = fake_snapshot_download + utils.HfApi = _FakeHfApi + os.environ.pop("HF_HUB_OFFLINE", None) # inherited from tests/unit/conftest.py + # new_group is collective over every rank, members or not. + group = None if group_ranks is None else dist.new_group(group_ranks, backend="gloo") + if group_ranks is not None and rank not in group_ranks: + return + if fail: + with pytest.raises(RuntimeError, match="hub unreachable"): + utils.ensure_local_checkpoint("org/model", group=group) + else: + result = utils.ensure_local_checkpoint("org/model", group=group) + assert result == ("org/model", str(snapshot)) + + +def _ensure_local_checkpoint_local_path_job(tmp_path, hidden_from_rank1, rank, size): + utils = hf_checkpoint_utils + + def fail(*args, **kwargs): + raise AssertionError("a local path must not be downloaded") + + utils.snapshot_download = fail + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir(exist_ok=True) + if hidden_from_rank1 and rank == 1: + # As if the directory sat on rank 0's node-local disk only. + isdir = os.path.isdir + utils.os.path.isdir = lambda path: path != str(checkpoint) and isdir(path) + with pytest.raises(RuntimeError, match="this rank cannot see"): + utils.ensure_local_checkpoint(str(checkpoint)) + else: + assert utils.ensure_local_checkpoint(str(checkpoint)) == (None, str(checkpoint)) + + +@pytest.mark.parametrize("hidden_from_rank1", [False, True]) +def test_ensure_local_checkpoint_lets_rank_0_resolve_a_local_path(tmp_path, hidden_from_rank1): + """Ranks that each checked a local path could disagree about it and strand one another in the + collective; rank 0 decides, and a rank that cannot see its answer raises instead.""" + job = partial(_ensure_local_checkpoint_local_path_job, tmp_path, hidden_from_rank1) + spawn_multiprocess_job(2, job, backend="gloo") + + +@pytest.mark.parametrize( + ("fail", "group_ranks", "downloader"), + [ + (False, None, "rank0"), + (True, None, "rank0"), + (False, [0, 1], "rank0"), + # Rank 0 of the group, not of the world: global rank 1. + (False, [1], "rank1"), + ], +) +def test_ensure_local_checkpoint_downloads_once_per_group(tmp_path, fail, group_ranks, downloader): + """Rank 0 of the group downloads once and every member gets its result -- including its + failure, so the waiting ranks raise instead of hanging.""" + job = partial(_ensure_local_checkpoint_job, tmp_path, fail, group_ranks) + spawn_multiprocess_job(2, job, backend="gloo") + assert [p.name for p in tmp_path.glob("downloaded-by-*")] == [f"downloaded-by-{downloader}"] + assert (tmp_path / f"downloaded-by-{downloader}").read_text().splitlines() == ["org/model"] diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index f397205f022..976d0e48206 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -15,6 +15,9 @@ import contextlib import os +import platform + +import pytest # Enforce no HuggingFace Hub network access for unit tests os.environ["HF_HUB_OFFLINE"] = "1" @@ -25,3 +28,20 @@ import huggingface_hub.constants as _hf_constants _hf_constants.HF_HUB_OFFLINE = True + + +@pytest.fixture(scope="session") +def cpu_dist_workers(): + """Two persistent gloo ranks shared by every multi-rank unit test. + + A fresh spawn per test re-imports torch and transformers in each child (~10-20s on CI), so the + workers start once. Jobs must leave the process group intact and undo any global state they set. + """ + if platform.system() == "Windows": + pytest.skip("Multi-rank tests are not run on Windows") + # Lazy: the onnx-only partial-install job has no torch. + from _test_utils.torch.distributed.utils import DistributedWorkerPool + + pool = DistributedWorkerPool(world_size=2, backend="gloo", teardown_fn=None) + yield pool + pool.shutdown() diff --git a/tests/unit/onnx/quantization/test_calibrated_quantization.py b/tests/unit/onnx/quantization/test_calibrated_quantization.py index e726f667d39..485099f9555 100644 --- a/tests/unit/onnx/quantization/test_calibrated_quantization.py +++ b/tests/unit/onnx/quantization/test_calibrated_quantization.py @@ -78,6 +78,8 @@ def _initializer_arrays(model: onnx.ModelProto) -> dict[str, np.ndarray]: 0.007997047156095505, [0.019685039296746254, 0.023622047156095505, 0.027559055015444756], id="int8-entropy", + # ~13s each; entropy calibration is also covered in tests/gpu/onnx + marks=pytest.mark.manual, ), pytest.param( "int8", @@ -94,6 +96,8 @@ def _initializer_arrays(model: onnx.ModelProto) -> dict[str, np.ndarray]: 0.028210056945681572, [0.06944013386964798, 0.08332816511392593, 0.09721619635820389], id="fp8-entropy", + # ~13s each; entropy calibration is also covered in tests/gpu/onnx + marks=pytest.mark.manual, ), pytest.param( "fp8", diff --git a/tests/unit/recipe/test_presets.py b/tests/unit/recipe/test_presets.py index f97b2d59bd4..3474cd51a43 100644 --- a/tests/unit/recipe/test_presets.py +++ b/tests/unit/recipe/test_presets.py @@ -22,9 +22,6 @@ """ import argparse -import subprocess -import sys -import textwrap import pytest @@ -202,31 +199,3 @@ def test_store_true_style_flag_stores_const_not_an_empty_list(): with pytest.warns(FutureWarning, match="--weight_only is deprecated"): args = parser.parse_args(["--weight_only"]) assert args.weight_only is True - - -def test_deprecation_reaches_stderr_under_the_real_default_filters(): - """The warning has to reach an actual CLI user, not just a test run. - - pytest enables every warning, so a category CPython suppresses looks healthy here and says - nothing in production. argparse invokes the action from its own module, so a - ``DeprecationWarning`` would be dropped by the default ``ignore::DeprecationWarning`` filter -- - the flag would keep working with nothing said. A subprocess is the only honest check: it uses - the interpreter's real filters rather than a reconstruction of them. - """ - script = textwrap.dedent( - """ - import argparse - from modelopt.recipe.presets import RecipeSupersededAction - - parser = argparse.ArgumentParser() - parser.add_argument("--qformat", action=RecipeSupersededAction) - parser.parse_args(["--qformat", "nvfp4"]) - """ - ) - proc = subprocess.run([sys.executable, "-c", script], capture_output=True, text=True) - - assert proc.returncode == 0, proc.stderr - assert "--qformat is deprecated" in proc.stderr, ( - "the deprecation is filtered out under Python's default filters, so a CLI user would " - f"never see it; stderr was: {proc.stderr!r}" - ) diff --git a/tests/unit/torch/export/test_fsdp2_parallel_export.py b/tests/unit/torch/export/test_fsdp2_parallel_export.py index 997a23020e5..c80f704bb6d 100644 --- a/tests/unit/torch/export/test_fsdp2_parallel_export.py +++ b/tests/unit/torch/export/test_fsdp2_parallel_export.py @@ -25,12 +25,10 @@ import copy import faulthandler import json -from functools import partial from pathlib import Path import pytest import torch -from _test_utils.torch.distributed.utils import spawn_multiprocess_job from safetensors.torch import load_file from torch.distributed.device_mesh import init_device_mesh from torch.distributed.fsdp import fully_shard @@ -156,18 +154,24 @@ def test_streaming_export_extra_state_dict_mtp(tmp_path): def _two_rank_streaming_export(rank, size, export_dir): - # A deadlocked rank exits with its stack instead of hanging the run on the join. + # A deadlocked rank exits with its stack instead of hanging the run on the join. Cancelled in + # ``finally`` because the worker is reused: a leftover timer would kill it in a later test. faulthandler.dump_traceback_later(45, exit=True) - # The unsharded reference is built and exported on every rank, so any collective its - # calibration or export issues has a partner. - expected, _ = _export_transformers_checkpoint(_tiny_quantized_llama(tie=False), torch.bfloat16) - # Calibrating after sharding leaves the root's parameters unsharded: FSDP2 skips the root's - # post-forward reshard. That is the state in which the owner's state_dict() used to reshard - # them on its rank alone, so the next gather window ran a collective the other rank skipped. - model = _tiny_quantized_llama(tie=False, mesh=init_device_mesh("cpu", (size,))) - _export_fsdp2_checkpoint_streaming(model, torch.bfloat16, export_dir=export_dir) - torch.distributed.barrier() - faulthandler.cancel_dump_traceback_later() + try: + # The unsharded reference is built and exported on every rank, so any collective its + # calibration or export issues has a partner. + expected, _ = _export_transformers_checkpoint( + _tiny_quantized_llama(tie=False), torch.bfloat16 + ) + # Calibrating after sharding leaves the root's parameters unsharded: FSDP2 skips the + # root's post-forward reshard. That is the state in which the owner's state_dict() used + # to reshard them on its rank alone, so the next gather window ran a collective the + # other rank skipped. + model = _tiny_quantized_llama(tie=False, mesh=init_device_mesh("cpu", (size,))) + _export_fsdp2_checkpoint_streaming(model, torch.bfloat16, export_dir=export_dir) + torch.distributed.barrier() + finally: + faulthandler.cancel_dump_traceback_later() if rank != 0: return exported = _load_all(Path(export_dir)) @@ -176,6 +180,6 @@ def _two_rank_streaming_export(rank, size, export_dir): assert torch.equal(exported[key].float(), tensor.float()), key -def test_streaming_export_two_cpu_ranks(tmp_path, skip_on_windows): +def test_streaming_export_two_cpu_ranks(tmp_path, cpu_dist_workers): """Two gloo ranks export what a single process does, without a rank waiting on the other.""" - spawn_multiprocess_job(2, partial(_two_rank_streaming_export, export_dir=str(tmp_path))) + cpu_dist_workers.run(_two_rank_streaming_export, export_dir=str(tmp_path)) diff --git a/tests/unit/torch/export/test_unified_export_hf.py b/tests/unit/torch/export/test_unified_export_hf.py index 9c9467ba17e..1c747186a1e 100644 --- a/tests/unit/torch/export/test_unified_export_hf.py +++ b/tests/unit/torch/export/test_unified_export_hf.py @@ -38,7 +38,11 @@ postprocess_state_dict, sync_tied_input_amax, ) -from modelopt.torch.export.unified_export_hf import _resolve_export_dtype, read_unplaced_weights +from modelopt.torch.export.unified_export_hf import ( + _resolve_export_dtype, + _source_checkpoint, + read_unplaced_weights, +) from modelopt.torch.quantization.nn import TensorQuantizer @@ -1075,3 +1079,9 @@ def test_union_degrades_to_the_recorded_set_without_the_loader(tmp_path, monkeyp carried = read_unplaced_weights(model) assert key in carried, "recorded keys must still carry when the structural pass is unavailable" + + +def test_source_checkpoint_is_none_for_a_model_built_from_config(): + """``from_config`` leaves ``_name_or_path`` empty; it must not resolve to the working directory, + or every export would copy the cwd as the "source checkpoint".""" + assert _source_checkpoint(SimpleNamespace(config=SimpleNamespace(_name_or_path=""))) is None diff --git a/tests/unit/torch/nas/test_full_algorithms_vision.py b/tests/unit/torch/nas/test_full_algorithms_vision.py index d217b79909f..f6c30c1da52 100644 --- a/tests/unit/torch/nas/test_full_algorithms_vision.py +++ b/tests/unit/torch/nas/test_full_algorithms_vision.py @@ -111,7 +111,7 @@ def run_search(): @pytest.mark.parametrize( ("flops_only", "bounded_latency"), [ - (False, True), # fast representative kept enabled for sanity + pytest.param(False, True, marks=pytest.mark.manual), pytest.param(True, True, marks=pytest.mark.manual), pytest.param(True, False, marks=pytest.mark.manual), pytest.param(False, False, marks=pytest.mark.manual), diff --git a/tests/unit/torch/nas/test_nas.py b/tests/unit/torch/nas/test_nas.py index 3d40b3a51d5..46b44a9ce3a 100644 --- a/tests/unit/torch/nas/test_nas.py +++ b/tests/unit/torch/nas/test_nas.py @@ -55,6 +55,7 @@ def get_data_loader(num_batches): yield torch.rand(2, 3, 56, 56), int(torch.randint(1000, (1,))) +@pytest.mark.manual # AutoNAS is not actively maintained; ~3s def test_automode(test_config): # Use model class to create a model inside mtn.convert model = mtn.convert((TinyMobileNetFeatures, (), {}), mode=[("autonas", test_config)]) @@ -118,6 +119,7 @@ def test_automode(test_config): assert not torch.allclose(out_dyn, out2) +@pytest.mark.manual # AutoNAS is not actively maintained; ~3s each @pytest.mark.parametrize("native_eval", [False, True]) def test_prep_for_eval(native_eval, test_config): """Test all four ways of calling prep_for_eval. @@ -288,6 +290,7 @@ def _get_current_bn_stats(): assert max_config == mtn.get_subnet_config(model) +@pytest.mark.manual # AutoNAS is not actively maintained; ~4s def test_set_auto_mode(test_config): """Test whether we can correctly set automode on/off.""" # setup model in autonas mode and training mode @@ -412,15 +415,24 @@ def _decorated_compare(): torch.randn(1, 16, 8, 8), ["autonas", "export_nas"], ), - ( + pytest.param( TinyMobileNetFeatures, (), "", False, torch.randn(1, 3, 64, 64), ["autonas", "export_nas"], + marks=pytest.mark.manual, # 5s; the InvertedResidual cases cover the same paths + ), + pytest.param( + TinyMobileNetFeatures, + (), + "", + False, + torch.randn(1, 3, 64, 64), + ["autonas"], + marks=pytest.mark.manual, ), - (TinyMobileNetFeatures, (), "", False, torch.randn(1, 3, 64, 64), ["autonas"]), (TinyMobileNetFeatures, (), "", False, torch.randn(1, 3, 64, 64), []), ], ) diff --git a/tests/unit/torch/quantization/test_autoquant.py b/tests/unit/torch/quantization/test_autoquant.py index f61cbf7b893..5ee55d81081 100644 --- a/tests/unit/torch/quantization/test_autoquant.py +++ b/tests/unit/torch/quantization/test_autoquant.py @@ -21,7 +21,6 @@ import pytest import torch -from _test_utils.torch.distributed.utils import spawn_multiprocess_job from _test_utils.torch.quantization.models import SimpleConv, SimpleConvLinear, SimpleLinear import modelopt.torch.opt as mto @@ -793,24 +792,29 @@ def test_active_moe_search_prefers_budget_lower_bound(): } +_BLOCKWISE_INT8 = [mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG, mtq.INT8_DEFAULT_CFG] +_AWQ_SMOOTHQUANT = [mtq.INT4_AWQ_CFG, mtq.INT8_SMOOTHQUANT_CFG] +_AWQ_CUSTOM = [mtq.INT4_AWQ_CFG, INT8_CUSTOM_QUANT_TEST_CFG] +_SMOOTHQUANT_ONLY = [mtq.INT8_SMOOTHQUANT_CFG] +_NONE_SMOOTHQUANT = [None, mtq.INT8_SMOOTHQUANT_CFG] + + +# A covering set rather than the full model x formats x method product: every model, format +# list and method appears at least once, and the methods only differ in how candidates are scored. @pytest.mark.parametrize( - "model_cls", - [SimpleConv, SimpleConvLinear, SimpleLinear, TransformerBlock], -) -@pytest.mark.parametrize( - ("search_formats", "min_bits", "search_bits"), + ("model_cls", "search_formats", "min_bits", "search_bits", "method"), [ - ([mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG, mtq.INT8_DEFAULT_CFG], 4.0, 6.0), - ([mtq.INT4_AWQ_CFG, mtq.INT8_SMOOTHQUANT_CFG], 4.0, 6.0), - ([mtq.INT4_AWQ_CFG, INT8_CUSTOM_QUANT_TEST_CFG], 4.0, 6.0), - ([mtq.INT8_SMOOTHQUANT_CFG], 8.0, 11.0), - ([None, mtq.INT8_SMOOTHQUANT_CFG], 8.0, 11.0), + (SimpleConv, _BLOCKWISE_INT8, 4.0, 6.0, "gradient"), + (SimpleConvLinear, _AWQ_SMOOTHQUANT, 4.0, 6.0, "gradient"), + (TransformerBlock, _NONE_SMOOTHQUANT, 8.0, 11.0, "gradient"), + (SimpleLinear, _AWQ_CUSTOM, 4.0, 6.0, "kl_div"), + (TransformerBlock, _SMOOTHQUANT_ONLY, 8.0, 11.0, "kl_div"), + (SimpleConvLinear, _BLOCKWISE_INT8, 4.0, 6.0, "kl_div"), + (SimpleConv, _AWQ_SMOOTHQUANT, 4.0, 6.0, "aumann_shapley"), + (SimpleLinear, _NONE_SMOOTHQUANT, 8.0, 11.0, "aumann_shapley"), + (TransformerBlock, _AWQ_CUSTOM, 4.0, 6.0, "aumann_shapley"), ], ) -@pytest.mark.parametrize( - "method", - ["gradient", "kl_div", "aumann_shapley"], -) def test_auto_quantize(model_cls, search_formats, min_bits, search_bits, method): model = model_cls() @@ -973,9 +977,8 @@ def _test_data_parallel_auto_quantize(rank, size): assert search_history["best"]["is_satisfied"] -def test_data_parallel_auto_quantize(skip_on_windows): - # 2 ranks fully exercise the cross-rank sync the test asserts; more just adds spawn overhead. - spawn_multiprocess_job(2, _test_data_parallel_auto_quantize, backend="gloo") +def test_data_parallel_auto_quantize(cpu_dist_workers): + cpu_dist_workers.run(_test_data_parallel_auto_quantize) def _test_data_parallel_moe_score_module(rank, size): @@ -1036,8 +1039,8 @@ def _test_data_parallel_moe_score_module(rank, size): assert scores == rank_zero_scores -def test_data_parallel_moe_score_module(skip_on_windows): - spawn_multiprocess_job(2, _test_data_parallel_moe_score_module, backend="gloo") +def test_data_parallel_moe_score_module(cpu_dist_workers): + cpu_dist_workers.run(_test_data_parallel_moe_score_module) def test_score_hparam_registration_preserves_order(): diff --git a/tests/unit/torch/quantization/test_autoquant_shapley.py b/tests/unit/torch/quantization/test_autoquant_shapley.py index b2a61fb4b1b..6752fcee65c 100644 --- a/tests/unit/torch/quantization/test_autoquant_shapley.py +++ b/tests/unit/torch/quantization/test_autoquant_shapley.py @@ -22,7 +22,6 @@ """ import copy -import functools import itertools import math from collections import namedtuple @@ -31,7 +30,6 @@ import pytest import torch import torch.nn.functional as F -from _test_utils.torch.distributed.utils import spawn_multiprocess_job import modelopt.torch.quantization as mtq import modelopt.torch.quantization.model_quant as _model_quant @@ -775,7 +773,7 @@ def _shapley_data_parallel(rank, size, baseline): assert got == pytest.approx(scores, rel=1e-3, abs=1e-9), f"{name}: {got} vs {scores}" -def test_data_parallel_aumann_shapley(skip_on_windows): +def test_data_parallel_aumann_shapley(cpu_dist_workers): """Data parallel aumann shapley.""" _model, single = _search( _Block(seed=0), method={"method": "aumann_shapley", "num_path_nodes": 2} @@ -785,9 +783,7 @@ def test_data_parallel_aumann_shapley(skip_on_windows): "f_corner": single["damage_model"]["f_corner"], "scores": {name: stat["scores"] for name, stat in single["candidate_stats"].items()}, } - spawn_multiprocess_job( - 2, functools.partial(_shapley_data_parallel, baseline=baseline), backend="gloo" - ) + cpu_dist_workers.run(_shapley_data_parallel, baseline=baseline) def test_anchor_ceiling_rejects_non_finite_measurements(): diff --git a/tests/unit/torch/quantization/test_dist.py b/tests/unit/torch/quantization/test_dist.py index 4130b8ce82a..0eb5329805e 100644 --- a/tests/unit/torch/quantization/test_dist.py +++ b/tests/unit/torch/quantization/test_dist.py @@ -15,7 +15,6 @@ import torch import torch.distributed as dist -from _test_utils.torch.distributed.utils import spawn_multiprocess_job from _test_utils.torch.quantization.models import SimpleLinear from torch.nn.parallel import DistributedDataParallel @@ -40,8 +39,7 @@ def forward_loop(model): amax = module.amax.clone() dist.all_reduce(amax, op=dist.ReduceOp.MAX) assert torch.allclose(amax, module.amax) - dist.destroy_process_group() -def test_data_parallel(skip_on_windows): - spawn_multiprocess_job(2, _test_data_parallel_helper, backend="gloo") +def test_data_parallel(cpu_dist_workers): + cpu_dist_workers.run(_test_data_parallel_helper) diff --git a/tests/unit/torch/quantization/test_quantize_cpu.py b/tests/unit/torch/quantization/test_quantize_cpu.py index 0c663875fb1..07d5fc309a9 100644 --- a/tests/unit/torch/quantization/test_quantize_cpu.py +++ b/tests/unit/torch/quantization/test_quantize_cpu.py @@ -21,7 +21,6 @@ import pytest import torch -from _test_utils.torch.distributed.utils import spawn_multiprocess_job from _test_utils.torch.quantization.models import SimpleConv, SimpleConvLinear, SimpleLinear from _test_utils.torch.quantization.quantize_common import ( INT4_AWQ_CLIP_CFG, @@ -722,11 +721,10 @@ def _test_indexer_check_covers_all_ranks(rank, size): mtq.quantize(_ToyStage([rank]), config, lambda m: None) finally: mtq.unregister(_ToyIndexer) - torch.distributed.destroy_process_group() -def test_indexer_k_check_covers_all_ranks(skip_on_windows): - spawn_multiprocess_job(2, _test_indexer_check_covers_all_ranks, backend="gloo") +def test_indexer_k_check_covers_all_ranks(cpu_dist_workers): + cpu_dist_workers.run(_test_indexer_check_covers_all_ranks) def test_indexer_q_patterns_need_the_query_quantizer(): diff --git a/tests/unit/torch/utils/test_hf_checkpoint_utils.py b/tests/unit/torch/utils/test_hf_checkpoint_utils.py index b436a9cb838..093e8e1cf66 100644 --- a/tests/unit/torch/utils/test_hf_checkpoint_utils.py +++ b/tests/unit/torch/utils/test_hf_checkpoint_utils.py @@ -16,18 +16,14 @@ """Tests for modelopt/torch/utils/plugins/hf_checkpoint_utils.py""" import json -import os import sys import warnings -from functools import partial from pathlib import Path from types import SimpleNamespace from unittest.mock import patch import pytest import torch -import torch.distributed as dist -from _test_utils.torch.distributed.utils import spawn_multiprocess_job from safetensors.torch import save_file from modelopt.torch.utils.plugins.hf_checkpoint_utils import ( @@ -977,74 +973,3 @@ def test_copy_non_model_files_follows_hub_snapshot_links_into_blobs(tmp_path): assert copied == ["tokenizer.json"] assert (tmp_path / "export" / "tokenizer.json").read_text() == "{}" - - -def _ensure_local_checkpoint_job(tmp_path, fail, group_ranks, rank, size): - utils = hf_checkpoint_utils - snapshot = tmp_path / "snapshot" - - def fake_snapshot_download(repo_id, **kwargs): - (tmp_path / f"downloaded-by-rank{rank}").write_text(repo_id) - if fail: - raise OSError("hub unreachable") - snapshot.mkdir(exist_ok=True) - return str(snapshot) - - utils.snapshot_download = fake_snapshot_download - utils.HfApi = _FakeHfApi - os.environ.pop("HF_HUB_OFFLINE", None) # inherited from tests/unit/conftest.py - # new_group is collective over every rank, members or not. - group = None if group_ranks is None else dist.new_group(group_ranks, backend="gloo") - if group_ranks is not None and rank not in group_ranks: - return - if fail: - with pytest.raises(RuntimeError, match="hub unreachable"): - utils.ensure_local_checkpoint("org/model", group=group) - else: - result = utils.ensure_local_checkpoint("org/model", group=group) - assert result == ("org/model", str(snapshot)) - - -def _ensure_local_checkpoint_local_path_job(tmp_path, hidden_from_rank1, rank, size): - utils = hf_checkpoint_utils - - def fail(*args, **kwargs): - raise AssertionError("a local path must not be downloaded") - - utils.snapshot_download = fail - checkpoint = tmp_path / "checkpoint" - checkpoint.mkdir(exist_ok=True) - if hidden_from_rank1 and rank == 1: - # As if the directory sat on rank 0's node-local disk only. - isdir = os.path.isdir - utils.os.path.isdir = lambda path: path != str(checkpoint) and isdir(path) - with pytest.raises(RuntimeError, match="this rank cannot see"): - utils.ensure_local_checkpoint(str(checkpoint)) - else: - assert utils.ensure_local_checkpoint(str(checkpoint)) == (None, str(checkpoint)) - - -@pytest.mark.parametrize("hidden_from_rank1", [False, True]) -def test_ensure_local_checkpoint_lets_rank_0_resolve_a_local_path(tmp_path, hidden_from_rank1): - """Ranks that each checked a local path could disagree about it and strand one another in the - collective; rank 0 decides, and a rank that cannot see its answer raises instead.""" - job = partial(_ensure_local_checkpoint_local_path_job, tmp_path, hidden_from_rank1) - spawn_multiprocess_job(2, job, backend="gloo") - - -@pytest.mark.parametrize( - ("fail", "group_ranks", "downloader"), - [ - (False, None, "rank0"), - (True, None, "rank0"), - (False, [0, 1], "rank0"), - # Rank 0 of the group, not of the world: global rank 1. - (False, [1], "rank1"), - ], -) -def test_ensure_local_checkpoint_downloads_once_per_group(tmp_path, fail, group_ranks, downloader): - """Rank 0 of the group downloads once and every member gets its result -- including its - failure, so the waiting ranks raise instead of hanging.""" - job = partial(_ensure_local_checkpoint_job, tmp_path, fail, group_ranks) - spawn_multiprocess_job(2, job, backend="gloo") - assert [p.name for p in tmp_path.glob("downloaded-by-*")] == [f"downloaded-by-{downloader}"] From 04d057a39fde768c12b222ecc5ed58df6b3d4728 Mon Sep 17 00:00:00 2001 From: mxinO <164952785+mxinO@users.noreply.github.com> Date: Thu, 15 Oct 2026 22:08:36 -0700 Subject: [PATCH 23/35] [1/2] Unified export: quantized-weight state capture (#2777) ### What does this PR do? Type of change: New feature and bug fix. Part 1 of 2: expose exact per-weight quantizer-state capture and metadata-only export specs. Capture is non-mutating and handles numbered grouped weights and supported transposed layouts. Callers can inspect deployment metadata without retaining tensor state. This also fixes shared format detection, grouped quantizer restoration and HF export metadata, including unquantized-layer exclusions in mixed configurations. Existing offline packing remains on its current path. Existing scale helpers are reused. The opaque snapshot also records logical weight shape, quantization axes, block sizes and packing layout, which a quantizer-buffer dictionary alone does not provide. It is detached in-memory state, not a new serialized configuration model. ### Usage For a converted and calibrated module: ```python from modelopt.torch.export import ( capture_quantized_weight_export_state, get_quantized_weight_export_spec, ) spec = get_quantized_weight_export_spec(module, "weight") state = capture_quantized_weight_export_state(module, "weight", cpu=False) ``` Capture returns `None` for an unquantized weight. The second PR adds consumers of these snapshots. Recapture state for each export when quantizer values may change. ### Testing Frozen-source container run: 169 focused CPU/GPU tests passed, no skips. Covers capture/spec, known scale values and axis/block metadata, detached-clone independence, required-calibration failures, ordinary offline export, GPU packing and TE grouped restoration. All changed-file pre-commit checks pass. This validates the extracted API and existing export paths, not training convergence. ### Before your PR is "*Ready for review*" - Backward compatible: Yes; additive APIs and shared export fixes. - New dependency or externally copied code: None. - Necessary tests: Capture/spec, exact grouped-weight resolution, detector, grouped restore and ordinary HF export regressions. - Changelog: Updated. - Claude approval: Not requested; independent agent review found no issues. ### Additional Information Split of #2251; the original PR is unchanged. Merge this first, then #2778 ([2/2]), which targets this branch. Megatron-Bridge and NeMo-RL integration require both. ## Summary by CodeRabbit * **New Features** * Added export specifications for supported quantized weights, including deployment metadata and non-mutating capture of quantizer state. * Expanded Hugging Face quantization metadata support for MXFP4, MXFP8, FP8 activation schemes, and block/group sizes, including mixed-format models and excluded layers. * Added support for the W4A16 MXFP4 format. * **Bug Fixes** * Preserved grouped weight calibration values during model restore. * Unsupported non-square FP8 block layouts are now rejected explicitly. --------- Signed-off-by: Meng Xin <164952785+mxinO@users.noreply.github.com> Signed-off-by: Rohan Joshi --- CHANGELOG.rst | 2 + modelopt/torch/export/__init__.py | 1 + modelopt/torch/export/convert_hf_config.py | 65 ++- modelopt/torch/export/quant_format.py | 1 + modelopt/torch/export/quant_utils.py | 261 +++++----- .../torch/export/quantized_weight_export.py | 444 ++++++++++++++++++ modelopt/torch/quantization/plugins/custom.py | 11 +- tests/gpu/torch/export/test_export.py | 122 ++++- .../plugins/test_transformer_engine.py | 20 + .../torch/export/test_get_quantization.py | 13 + .../export/test_quantized_weight_state.py | 353 ++++++++++++++ 11 files changed, 1168 insertions(+), 125 deletions(-) create mode 100644 modelopt/torch/export/quantized_weight_export.py create mode 100644 tests/unit/torch/export/test_quantized_weight_state.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 120b7015037..0dbdfcd82fa 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -13,6 +13,8 @@ Changelog *Quantization* +- Add metadata-only export specifications and non-mutating per-weight quantizer-state capture. +- HF export records canonical activation schemes, group/block sizes and weight-only MXFP4 metadata; mixed-format exports retain unquantized-layer exclusions, and non-square FP8 block layouts are rejected. - Add Q8_0 weight-only quantization with 32-value GGML blocks, packed unified HF and Megatron export, and a built-in ``q8_0`` PTQ recipe. - Add Hugging Face PTQ calibration and export support for Nemotron-H MTP modules, preserving calibrated expert input scales. - Backfill checkpoint aliases for three more published NVFP4 releases: ``moonshotai/Kimi-K2.7-Code``, ``google/gemma-4-26B-A4B-it`` and ``google/diffusiongemma-26B-A4B-it``. Each imports an existing general recipe wholesale rather than copying its body. diff --git a/modelopt/torch/export/__init__.py b/modelopt/torch/export/__init__.py index 874fa71b2b9..44d1ad2557e 100644 --- a/modelopt/torch/export/__init__.py +++ b/modelopt/torch/export/__init__.py @@ -38,6 +38,7 @@ from .moe_utils import * from .plugins import * from .quant_format import * +from .quantized_weight_export import * from .registry import * from .shard_cast_utils import * from .transformer_engine import * diff --git a/modelopt/torch/export/convert_hf_config.py b/modelopt/torch/export/convert_hf_config.py index d78410c661e..bbcc25a28da 100644 --- a/modelopt/torch/export/convert_hf_config.py +++ b/modelopt/torch/export/convert_hf_config.py @@ -43,9 +43,25 @@ def _quant_algo_to_group_config(quant_algo: str, group_size: int | None = None) } elif quant_algo == "FP8_PER_CHANNEL_PER_TOKEN": return { - "input_activations": {"dynamic": False, "num_bits": 8, "type": "float"}, + "input_activations": { + "dynamic": True, + "num_bits": 8, + "type": "float", + "strategy": "token", + }, "weights": {"dynamic": False, "num_bits": 8, "type": "float", "strategy": "channel"}, } + elif quant_algo == "FP8_PB_WO": + block_size = group_size or 128 + return { + "weights": { + "dynamic": False, + "num_bits": 8, + "type": "float", + "strategy": "block", + "block_structure": [block_size, block_size], + }, + } elif quant_algo == "NVFP4": gs = group_size or 16 return { @@ -115,13 +131,40 @@ def _quant_algo_to_group_config(quant_algo: str, group_size: int | None = None) gs = group_size or 32 return { "input_activations": { + "dynamic": True, + "num_bits": 8, + "type": "float", + "group_size": gs, + "strategy": "group", + }, + "weights": { "dynamic": False, "num_bits": 8, "type": "float", "group_size": gs, + "strategy": "group", }, - "weights": {"dynamic": False, "num_bits": 8, "type": "float", "group_size": gs}, } + elif quant_algo in ("MXFP4", "W4A16_MXFP4"): + gs = group_size or 32 + config = { + "weights": { + "dynamic": False, + "num_bits": 4, + "type": "float", + "group_size": gs, + "strategy": "group", + }, + } + if quant_algo == "MXFP4": + config["input_activations"] = { + "dynamic": True, + "num_bits": 4, + "type": "float", + "group_size": gs, + "strategy": "group", + } + return config elif quant_algo.lower() in GGML_FORMATS: ggml_format = GGML_FORMAT_REGISTRY[quant_algo.lower()] block_size, payload_bytes = ggml_format.block_size, ggml_format.block_bytes @@ -260,6 +303,19 @@ def convert_hf_quant_config_format(input_config: dict[str, Any]) -> dict[str, An if lora_rank is not None: config_group_details["lora_rank"] = lora_rank new_config["config_groups"] = {"group_0": config_group_details} + elif quant_algo_value in { + "FP8_PER_CHANNEL_PER_TOKEN", + "FP8_PB_WO", + "MXFP4", + "W4A16_MXFP4", + "MXFP8", + "W4A8_MXFP4_FP8", + "W4A8_NVFP4_FP8", + }: + group_size = original_quantization_details.get("group_size") + config_group_details = _quant_algo_to_group_config(quant_algo_value, group_size) + config_group_details["targets"] = ["Linear"] + new_config["config_groups"] = {"group_0": config_group_details} elif quant_algo_value == "MIXED_PRECISION": quantized_layers = original_quantization_details.get("quantized_layers", {}) @@ -292,6 +348,11 @@ def convert_hf_quant_config_format(input_config: dict[str, Any]) -> dict[str, An if quant_algo_value: new_config["quant_algo"] = quant_algo_value + group_size = original_quantization_details.get("group_size") + if group_size is not None: + # ModelOpt consumers use this top-level value to size block scales. + new_config["group_size"] = group_size + kv_cache_quant_algo = original_quantization_details.get("kv_cache_quant_algo") if kv_cache_quant_algo: if kv_cache_quant_algo == "FP8": diff --git a/modelopt/torch/export/quant_format.py b/modelopt/torch/export/quant_format.py index 8a520f972b4..8c503341760 100644 --- a/modelopt/torch/export/quant_format.py +++ b/modelopt/torch/export/quant_format.py @@ -31,6 +31,7 @@ QUANTIZATION_NVFP4_SVDQUANT = "nvfp4_svdquant" QUANTIZATION_W4A8_NVFP4_FP8 = "w4a8_nvfp4_fp8" QUANTIZATION_MXFP4 = "mxfp4" +QUANTIZATION_W4A16_MXFP4 = "w4a16_mxfp4" QUANTIZATION_MXFP8 = "mxfp8" QUANTIZATION_W4A8_MXFP4_FP8 = "w4a8_mxfp4_fp8" QUANTIZATION_W4A16_NVFP4 = "w4a16_nvfp4" diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py index 6450c1afc74..88092f689dc 100755 --- a/modelopt/torch/export/quant_utils.py +++ b/modelopt/torch/export/quant_utils.py @@ -42,7 +42,6 @@ QTensorWrapper, ) from modelopt.torch.quantization.utils import ( - QuantizerAttrNames, quantizer_attr_names, representative_weight_quantizer, weight_attr_names, @@ -74,6 +73,7 @@ QUANTIZATION_W4A8_AWQ, QUANTIZATION_W4A8_MXFP4_FP8, QUANTIZATION_W4A8_NVFP4_FP8, + QUANTIZATION_W4A16_MXFP4, QUANTIZATION_W4A16_NVFP4, ) @@ -225,7 +225,11 @@ def get_weight_scaling_factor(module: nn.Module, weight_name: str = "weight") -> weight_scaling_factor_2.to(weight.device), )[0] - if quantization_format in [QUANTIZATION_W4A8_MXFP4_FP8, QUANTIZATION_MXFP4]: + if quantization_format in [ + QUANTIZATION_W4A8_MXFP4_FP8, + QUANTIZATION_MXFP4, + QUANTIZATION_W4A16_MXFP4, + ]: return MXFP4QTensor.quantize(weight, block_size=weight_quantizer.block_sizes[-1])[ 1 ].reshape(*weight.shape[:-1], -1) @@ -466,134 +470,120 @@ def uses_iq_quantization(module) -> bool: return any(uses_iq_quantization(child) for _, child in module.named_children()) -def get_quantization_format(module) -> str | None: - """Gets the quantization string. - - Gets the quantization string by iterating through the module and its children. - The first non-None quantization string is returned. - """ - - def _get_quantization_from_layer(layer, quantizer_attr_names: QuantizerAttrNames): - # Singular form first, plural ModuleList fallback (fused-experts). - # Strip the "_weight_quantizer" suffix to recover the weight attr name. - weight_attr = quantizer_attr_names.weight_quantizer - weight_name = weight_attr[: -len("_weight_quantizer")].rstrip("_") or "weight" - weight_quantizer = representative_weight_quantizer(layer, weight_name) - input_quantizer = getattr(layer, quantizer_attr_names.input_quantizer, None) - - if weight_quantizer is None or not weight_quantizer.is_enabled: - return QUANTIZATION_NONE - - # Handle SequentialQuantizer - if isinstance(weight_quantizer, SequentialQuantizer): - assert ( - len(weight_quantizer) == 2 - and weight_quantizer[0].num_bits == 4 - and weight_quantizer[1].num_bits == (4, 3) - ), "Unsupported SequentialQuantizer configuration" - assert ( - weight_quantizer[0].block_sizes - and len(weight_quantizer[0].block_sizes) > 0 - and weight_quantizer[0].block_sizes[-1] > 0 - ), "Invalid block_sizes for SequentialQuantizer" - - return QUANTIZATION_W4A8_AWQ - - # Handle individual num_bits cases - if weight_quantizer.num_bits in GGML_FORMATS: - if weight_quantizer.backend != "ggml": - raise ValueError("GGML formats require the built-in 'ggml' quantization backend") - # Both exporters return before collecting input_scale and before the pre_quant_scale - # handling below, so an enabled activation quantizer would be dropped without a trace - # and the checkpoint would load as weight-only. Refuse instead. - if input_quantizer is not None and input_quantizer.is_enabled: - raise NotImplementedError( - "GGML export is weight-only, but this layer has an enabled input " - "quantizer. The GGML block payload carries no activation scale, so the " - "activation quantization would be silently lost." - ) - if input_quantizer is not None and hasattr(input_quantizer, "_pre_quant_scale"): - raise NotImplementedError( - "GGML export does not support an AWQ-style pre_quant_scale." - ) - return weight_quantizer.num_bits +def _get_quantization_from_quantizers( + layer: nn.Module, + weight_quantizer: TensorQuantizer | SequentialQuantizer | None, + input_quantizer: TensorQuantizer | SequentialQuantizer | None, +) -> str | None: + if weight_quantizer is None or not weight_quantizer.is_enabled: + return QUANTIZATION_NONE - if weight_quantizer.num_bits == 4: - assert len(weight_quantizer.block_sizes) > 0 and weight_quantizer.block_sizes[-1] > 0, ( - "Invalid block_sizes for INT4 quantizer" + if isinstance(weight_quantizer, SequentialQuantizer): + assert ( + len(weight_quantizer) == 2 + and weight_quantizer[0].num_bits == 4 + and weight_quantizer[1].num_bits == (4, 3) + ), "Unsupported SequentialQuantizer configuration" + assert ( + weight_quantizer[0].block_sizes + and len(weight_quantizer[0].block_sizes) > 0 + and weight_quantizer[0].block_sizes[-1] > 0 + ), "Invalid block_sizes for SequentialQuantizer" + return QUANTIZATION_W4A8_AWQ + + if weight_quantizer.num_bits in GGML_FORMATS: + if weight_quantizer.backend != "ggml": + raise ValueError("GGML formats require the built-in 'ggml' quantization backend") + # Both exporters return before collecting input_scale and before the pre_quant_scale + # handling below, so an enabled activation quantizer would be dropped without a trace + # and the checkpoint would load as weight-only. Refuse instead. + if input_quantizer is not None and input_quantizer.is_enabled: + raise NotImplementedError( + "GGML export is weight-only, but this layer has an enabled input " + "quantizer. The GGML block payload carries no activation scale, so the " + "activation quantization would be silently lost." ) - return QUANTIZATION_INT4_AWQ + if input_quantizer is not None and hasattr(input_quantizer, "_pre_quant_scale"): + raise NotImplementedError("GGML export does not support an AWQ-style pre_quant_scale.") + return weight_quantizer.num_bits - if weight_quantizer.num_bits == 8: - if input_quantizer is not None and input_quantizer.is_enabled: - return QUANTIZATION_INT8_SQ - else: - return QUANTIZATION_INT8_WO + if weight_quantizer.num_bits == 4: + assert len(weight_quantizer.block_sizes) > 0 and weight_quantizer.block_sizes[-1] > 0, ( + "Invalid block_sizes for INT4 quantizer" + ) + return QUANTIZATION_INT4_AWQ - if weight_quantizer.num_bits == (4, 3): - if weight_quantizer.block_sizes: - assert weight_quantizer.block_sizes[-1] > 0, "Invalid block_sizes for FP8 quantizer" - # Check if this is MXFP8 (dynamic block quantization with scale_bits (8, 0)) - block_sizes = getattr(weight_quantizer, "block_sizes") - if ( - isinstance(block_sizes, dict) - and block_sizes.get("type", "static") == "dynamic" - and block_sizes.get("scale_bits") == (8, 0) - ): - return QUANTIZATION_MXFP8 - if weight_quantizer.fake_quant: - return QUANTIZATION_FP8_PB_WO - else: - return QUANTIZATION_FP8_PB_REAL - if weight_quantizer.axis == 0: - return QUANTIZATION_FP8_PC_PT - return QUANTIZATION_FP8 + if weight_quantizer.num_bits == 8: + if input_quantizer is not None and input_quantizer.is_enabled: + return QUANTIZATION_INT8_SQ + return QUANTIZATION_INT8_WO - if weight_quantizer.num_bits == (2, 1): - # FP4 formats are all block quantization + if weight_quantizer.num_bits == (4, 3): + if weight_quantizer.block_sizes: + assert weight_quantizer.block_sizes[-1] > 0, "Invalid block_sizes for FP8 quantizer" block_sizes = getattr(weight_quantizer, "block_sizes") - scale_bits = block_sizes.get("scale_bits") - - if input_quantizer is not None and hasattr(weight_quantizer, "svdquant_lora_a"): - return QUANTIZATION_NVFP4_SVDQUANT - if input_quantizer is not None and hasattr(input_quantizer, "_pre_quant_scale"): - return QUANTIZATION_NVFP4_AWQ - if getattr(layer, "fused_with_prequant", False): - return QUANTIZATION_NVFP4_AWQ - if input_quantizer is None or not input_quantizer.is_enabled: - if scale_bits == (4, 3): - return QUANTIZATION_W4A16_NVFP4 - assert input_quantizer is not None, ( - f"input_quantizer is None for {quantizer_attr_names}" - ) - if ( - block_sizes.get("type", "static") == "dynamic" - and scale_bits == (8, 0) - and input_quantizer.is_enabled - and input_quantizer.num_bits == (4, 3) - and input_quantizer.block_sizes is None - ): - return QUANTIZATION_W4A8_MXFP4_FP8 if ( - block_sizes.get("type", "static") == "dynamic" - and scale_bits == (4, 3) - and input_quantizer.is_enabled - and input_quantizer.num_bits == (4, 3) - and input_quantizer.block_sizes is None + isinstance(block_sizes, dict) + and block_sizes.get("type", "static") == "dynamic" + and block_sizes.get("scale_bits") == (8, 0) ): - return QUANTIZATION_W4A8_NVFP4_FP8 + return QUANTIZATION_MXFP8 + if weight_quantizer.fake_quant: + return QUANTIZATION_FP8_PB_WO + return QUANTIZATION_FP8_PB_REAL + if weight_quantizer.axis == 0: + return QUANTIZATION_FP8_PC_PT + return QUANTIZATION_FP8 + + if weight_quantizer.num_bits == (2, 1): + block_sizes = getattr(weight_quantizer, "block_sizes") + scale_bits = block_sizes.get("scale_bits") + + if input_quantizer is not None and hasattr(weight_quantizer, "svdquant_lora_a"): + return QUANTIZATION_NVFP4_SVDQUANT + if input_quantizer is not None and hasattr(input_quantizer, "_pre_quant_scale"): + return QUANTIZATION_NVFP4_AWQ + if getattr(layer, "fused_with_prequant", False): + return QUANTIZATION_NVFP4_AWQ + if input_quantizer is None or not input_quantizer.is_enabled: if scale_bits == (4, 3): - return QUANTIZATION_NVFP4 - elif scale_bits == (8, 0): - return QUANTIZATION_MXFP4 + return QUANTIZATION_W4A16_NVFP4 + if scale_bits == (8, 0): + return QUANTIZATION_W4A16_MXFP4 + assert input_quantizer is not None, "input_quantizer is required for weight-activation FP4" + if ( + block_sizes.get("type", "static") == "dynamic" + and scale_bits == (8, 0) + and input_quantizer.is_enabled + and input_quantizer.num_bits == (4, 3) + and input_quantizer.block_sizes is None + ): + return QUANTIZATION_W4A8_MXFP4_FP8 + if ( + block_sizes.get("type", "static") == "dynamic" + and scale_bits == (4, 3) + and input_quantizer.is_enabled + and input_quantizer.num_bits == (4, 3) + and input_quantizer.block_sizes is None + ): + return QUANTIZATION_W4A8_NVFP4_FP8 + if scale_bits == (4, 3): + return QUANTIZATION_NVFP4 + if scale_bits == (8, 0): + return QUANTIZATION_MXFP4 + + raise NotImplementedError(f"Unsupported quantizer with num_bits: {weight_quantizer.num_bits}") - # Raise error for unsupported num_bits - raise NotImplementedError( - f"Unsupported quantizer with num_bits: {weight_quantizer.num_bits}" - ) +def get_quantization_format(module) -> str | None: + """Return the first enabled weight quantization format in a module tree.""" for weight_name in weight_attr_names(module): - quantization = _get_quantization_from_layer(module, quantizer_attr_names(weight_name)) + attr_names = quantizer_attr_names(weight_name) + quantization = _get_quantization_from_quantizers( + module, + representative_weight_quantizer(module, weight_name), + getattr(module, attr_names.input_quantizer, None), + ) if quantization != QUANTIZATION_NONE: return quantization @@ -707,6 +697,11 @@ def process_layer_quant_config(layer_config_dict): layer_config = {"quant_algo": "FP8"} elif v == "fp8_pc_pt": layer_config = {"quant_algo": "FP8_PER_CHANNEL_PER_TOKEN"} + elif v == "fp8_pb_wo": + layer_config = { + "quant_algo": "FP8_PB_WO", + "group_size": block_size_value, + } elif v == "int4_awq": layer_config = { "quant_algo": "W4A16_AWQ", @@ -766,6 +761,11 @@ def process_layer_quant_config(layer_config_dict): "quant_algo": "MXFP8", "group_size": block_size_value, } + elif v in ["mxfp4", "w4a16_mxfp4"]: + layer_config = { + "quant_algo": v.upper(), + "group_size": block_size_value, + } elif v in GGML_FORMATS: ggml_format = GGML_FORMAT_REGISTRY[v] block_size, payload_bytes = ggml_format.block_size, ggml_format.block_bytes @@ -794,6 +794,11 @@ def process_layer_quant_config(layer_config_dict): # If we have more than one quantization format, infer MIXED_PRECISION if len(quantization_formats) > 1: per_layer_config["quant_algo"] = "MIXED_PRECISION" + per_layer_config["exclude_modules"] = sorted( + _prefix_wildcard_summarize_exclude_modules( + exclude_modules, per_layer_config["quantized_layers"].keys() + ) + ) elif len(quantization_formats) == 1 and quantization_config is not None: per_layer_config.update(quantization_config) per_layer_config["exclude_modules"] = sorted( @@ -961,7 +966,11 @@ def to_quantized_weight( else weights_scaling_factor2, )[0]._quantized_data - if quantization in [QUANTIZATION_W4A8_MXFP4_FP8, QUANTIZATION_MXFP4]: + if quantization in [ + QUANTIZATION_W4A8_MXFP4_FP8, + QUANTIZATION_MXFP4, + QUANTIZATION_W4A16_MXFP4, + ]: return MXFP4QTensor.quantize(weight, block_size=block_size)[0]._quantized_data raise NotImplementedError(f"quantization format {quantization} not supported") @@ -1817,6 +1826,20 @@ def get_quant_config( for weight_name in weight_names: weight_block_size = get_weight_block_size(module, weight_name) + if quantization_format == QUANTIZATION_FP8_PB_WO: + # One group size can describe only square two-dimensional blocks. + weight_quantizer = representative_weight_quantizer(module, weight_name) + if weight_quantizer is not None and weight_quantizer.is_enabled: + block_axes = { + axis: size + for axis, size in weight_quantizer.block_sizes.items() + if isinstance(axis, int) + } + if block_axes != {-2: weight_block_size, -1: weight_block_size}: + raise NotImplementedError( + "Offline FP8 block export requires square blocks over " + f"the last two dimensions, got {block_axes} for {name}" + ) if weight_block_size > 0: block_size = weight_block_size break diff --git a/modelopt/torch/export/quantized_weight_export.py b/modelopt/torch/export/quantized_weight_export.py new file mode 100644 index 00000000000..1b1b3e5a5f2 --- /dev/null +++ b/modelopt/torch/export/quantized_weight_export.py @@ -0,0 +1,444 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""Capture state for functional quantized-weight export.""" + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from typing import Any + +import torch +import torch.nn as nn + +from modelopt.torch.quantization.qtensor import NVFP4QTensor, QTensorWrapper +from modelopt.torch.quantization.utils import quantizer_attr_names, representative_weight_quantizer + +from ..quantization.nn import SequentialQuantizer, TensorQuantizer +from .quant_format import ( + QUANTIZATION_FP8, + QUANTIZATION_FP8_PB_WO, + QUANTIZATION_FP8_PC_PT, + QUANTIZATION_MXFP4, + QUANTIZATION_MXFP8, + QUANTIZATION_NVFP4, + QUANTIZATION_W4A8_MXFP4_FP8, + QUANTIZATION_W4A8_NVFP4_FP8, + QUANTIZATION_W4A16_MXFP4, + QUANTIZATION_W4A16_NVFP4, +) +from .quant_utils import _get_quantization_from_quantizers, get_scaling_factor + +__all__ = [ + "capture_quantized_weight_export_state", + "get_quantized_weight_export_spec", +] + +_NVFP4_EXPORT_FORMATS = { + QUANTIZATION_NVFP4, + QUANTIZATION_W4A8_NVFP4_FP8, + QUANTIZATION_W4A16_NVFP4, +} +_MXFP4_EXPORT_FORMATS = {QUANTIZATION_MXFP4, QUANTIZATION_W4A8_MXFP4_FP8} +_FUNCTIONAL_WEIGHT_EXPORT_FORMATS = ( + _NVFP4_EXPORT_FORMATS + | _MXFP4_EXPORT_FORMATS + | { + QUANTIZATION_FP8, + QUANTIZATION_FP8_PC_PT, + QUANTIZATION_FP8_PB_WO, + QUANTIZATION_MXFP8, + } +) + + +class _UnsupportedQuantizedWeightExportFormatError(NotImplementedError): + pass + + +@dataclass(frozen=True) +class _ExportStateTensor: + name: str + value: torch.Tensor + axes: tuple[int, ...] = () + block_sizes: tuple[int, ...] = () + + +@dataclass(frozen=True) +class _QuantizedWeightExportState: + """Opaque state required to export one logical quantized weight.""" + + quantization_format: str + block_size: int + weight_shape: tuple[int, ...] + tensors: tuple[_ExportStateTensor, ...] + packing_permutation: tuple[int, ...] + input_quantizer_enabled: bool + static_nvfp4: bool = False + four_over_six: bool = False + + +@dataclass(frozen=True) +class _QuantizedWeightExportSpec: + """Stable deployment metadata for one logical quantized weight.""" + + quantization_format: str + block_size: int + input_quantizer_enabled: bool + + +def _same_storage(left: object, right: object) -> bool: + if left is right: + return True + if not isinstance(left, torch.Tensor) or not isinstance(right, torch.Tensor): + return False + if left.device.type == "meta" or right.device.type == "meta": + return False + return ( + left.device == right.device + and left.numel() == right.numel() + and left.untyped_storage().data_ptr() == right.untyped_storage().data_ptr() + and left.storage_offset() == right.storage_offset() + ) + + +def _resolve_weight_quantizer( + module: nn.Module, weight_name: str +) -> tuple[torch.Tensor, TensorQuantizer | SequentialQuantizer] | None: + weight = getattr(module, weight_name) + iter_weights = getattr(module, "iter_weights_for_calibration", None) + if iter_weights is not None: + for candidate, quantizer in iter_weights(): + if _same_storage(candidate, weight): + return candidate, quantizer + + quantizer = representative_weight_quantizer(module, weight_name) + if quantizer is None and weight_name.startswith("weight"): + quantizer = representative_weight_quantizer(module) + if quantizer is not None and weight_name.removeprefix("weight").isdigit(): + raise NotImplementedError( + f"Unable to resolve quantizer for numbered weight {weight_name!r}" + ) + if quantizer is None: + return None + # A representative quantizer cannot supply another expert's calibrated state. + if iter_weights is not None: + raise NotImplementedError( + f"Calibration iterator does not expose a matching full weight view for {weight_name!r}" + ) + return weight, quantizer + + +def _packing_permutation(weight: torch.Tensor, quantized_view: torch.Tensor) -> tuple[int, ...]: + ndim = weight.ndim + identity = tuple(range(ndim)) + if ndim >= 2: + transposed = weight.transpose(-1, -2) + if ( + tuple(quantized_view.shape) == tuple(transposed.shape) + and quantized_view.stride() == transposed.stride() + ): + return (*range(ndim - 2), ndim - 1, ndim - 2) + if ( + tuple(quantized_view.shape) == tuple(weight.shape) + and quantized_view.stride() == weight.stride() + ): + return identity + raise NotImplementedError( + "Unsupported quantized weight view " + f"shape/stride={tuple(quantized_view.shape)}/{quantized_view.stride()} for " + f"shape/stride={tuple(weight.shape)}/{weight.stride()}" + ) + + +def _state_tensor( + name: str, + value: torch.Tensor, + *, + axes: tuple[int, ...] = (), + block_sizes: tuple[int, ...] = (), + cpu: bool = True, +) -> _ExportStateTensor: + value = value.detach().cpu().clone() if cpu else value.detach().clone() + if value.numel() == 1 and not axes and not block_sizes: + return _ExportStateTensor(name, value.reshape(())) + if len(axes) != value.ndim or len(block_sizes) != value.ndim: + raise ValueError( + f"Explicit axis and block metadata is required for non-scalar {name}: " + f"shape={tuple(value.shape)}, axes={axes}, block_sizes={block_sizes}" + ) + return _ExportStateTensor(name, value, axes, block_sizes) + + +def _axis_scale_state( + name: str, + value: torch.Tensor, + packed_shape: tuple[int, ...], + axis: int | Sequence[int], + *, + cpu: bool = True, +) -> _ExportStateTensor: + axes = (axis,) if isinstance(axis, int) else tuple(axis) + axes = tuple(_normalize_weight_dim(dim, len(packed_shape)) for dim in axes) + expected_shape = tuple(packed_shape[dim] for dim in axes) + if value.numel() != torch.Size(expected_shape).numel(): + raise RuntimeError( + f"Scale {name!r} shape {tuple(value.shape)} does not match axes {axes} " + f"of packed weight shape {packed_shape}" + ) + return _state_tensor( + name, + value.reshape(expected_shape), + axes=axes, + block_sizes=(1,) * len(axes), + cpu=cpu, + ) + + +def _block_scale_state( + name: str, + value: torch.Tensor, + packed_shape: tuple[int, ...], + block_config: Mapping[int | str, Any], + *, + cpu: bool = True, +) -> _ExportStateTensor: + ndim = len(packed_shape) + block_sizes = [1] * ndim + for dim, block_size in block_config.items(): + if isinstance(dim, int) and block_size is not None: + block_sizes[_normalize_weight_dim(dim, ndim)] = int(block_size) + expected_shape = tuple( + (size + block_size - 1) // block_size for size, block_size in zip(packed_shape, block_sizes) + ) + if value.numel() != torch.Size(expected_shape).numel(): + raise RuntimeError( + f"Scale {name!r} shape {tuple(value.shape)} does not match block sizes " + f"{tuple(block_sizes)} of packed weight shape {packed_shape}" + ) + + expanded_shape = [] + expanded_axes = [] + expanded_block_sizes = [] + for axis, (size, block_size) in enumerate(zip(expected_shape, block_sizes)): + expanded_shape.append(size) + expanded_axes.append(axis) + expanded_block_sizes.append(block_size) + if block_size != 1: + expanded_shape.append(1) + expanded_axes.append(axis) + expanded_block_sizes.append(1) + if tuple(value.shape) == tuple(expanded_shape): + return _state_tensor( + name, + value, + axes=tuple(expanded_axes), + block_sizes=tuple(expanded_block_sizes), + cpu=cpu, + ) + return _state_tensor( + name, + value.reshape(expected_shape), + axes=tuple(range(ndim)), + block_sizes=tuple(block_sizes), + cpu=cpu, + ) + + +def _input_quantizer(module: nn.Module, weight_name: str): + quantizer = getattr(module, quantizer_attr_names(weight_name).input_quantizer, None) + if quantizer is None: + quantizer = getattr(module, "input_quantizer", None) + return quantizer + + +def _resolve_quantized_weight_export_inputs( + module: nn.Module, + weight_name: str, +) -> ( + tuple[ + torch.Tensor, + torch.Tensor, + TensorQuantizer, + TensorQuantizer | None, + str, + ] + | None +): + weight = getattr(module, weight_name) + if isinstance(weight, QTensorWrapper): + raise NotImplementedError("Functional export requires an uncompressed source weight") + + resolved = _resolve_weight_quantizer(module, weight_name) + if resolved is None: + return None + quantized_view, weight_quantizer = resolved + input_quantizer = _input_quantizer(module, weight_name) + if not weight_quantizer.is_enabled: + return None + quantization_format = _get_quantization_from_quantizers( + module, weight_quantizer, input_quantizer + ) + if quantization_format is None: + raise RuntimeError(f"Unable to resolve quantization format for {weight_name!r}") + if isinstance(weight_quantizer, SequentialQuantizer): + raise _UnsupportedQuantizedWeightExportFormatError( + f"Functional export does not support {quantization_format!r}" + ) + if quantization_format == QUANTIZATION_W4A16_MXFP4: + raise _UnsupportedQuantizedWeightExportFormatError( + "Functional export does not support weight-only MXFP4" + ) + if quantization_format == QUANTIZATION_FP8_PB_WO and weight_quantizer.block_sizes != { + -2: 128, + -1: 128, + }: + raise _UnsupportedQuantizedWeightExportFormatError( + "Functional export supports FP8 2D blockwise weight quantization only with " + "128x128 blocks" + ) + if quantization_format not in _FUNCTIONAL_WEIGHT_EXPORT_FORMATS: + raise _UnsupportedQuantizedWeightExportFormatError( + f"Functional export does not support {quantization_format!r}" + ) + return weight, quantized_view, weight_quantizer, input_quantizer, quantization_format + + +def capture_quantized_weight_export_state( + module: nn.Module, + weight_name: str = "weight", + *, + cpu: bool = True, +) -> _QuantizedWeightExportState | None: + """Capture detached state for one weight, or ``None`` when it is not quantized.""" + resolved = _resolve_quantized_weight_export_inputs(module, weight_name) + if resolved is None: + return None + weight, quantized_view, weight_quantizer, input_quantizer, quantization_format = resolved + + permutation = _packing_permutation(weight, quantized_view) + packed_shape = tuple(weight.shape[axis] for axis in permutation) + block_config = getattr(weight_quantizer, "block_sizes", None) or {} + block_size = int(block_config.get(-1, 0)) if isinstance(block_config, dict) else 0 + tensors = [] + static_nvfp4 = ( + quantization_format in _NVFP4_EXPORT_FORMATS + and NVFP4QTensor._is_static_quantizer(weight_quantizer) + ) + + if static_nvfp4: + if not block_size or packed_shape[-1] % block_size: + raise RuntimeError(f"Invalid static NVFP4 block size for {weight_name!r}") + per_block_amax = weight_quantizer.amax + global_amax = NVFP4QTensor._get_static_global_amax(weight_quantizer) + if per_block_amax is None or global_amax is None: + raise RuntimeError(f"Missing calibrated static NVFP4 state for {weight_name!r}") + block_shape = (*packed_shape[:-1], packed_shape[-1] // block_size) + tensors.append( + _state_tensor( + "weight_block_amax", + per_block_amax.reshape(block_shape), + axes=tuple(range(len(packed_shape))), + block_sizes=(1,) * (len(packed_shape) - 1) + (block_size,), + cpu=cpu, + ) + ) + tensors.append(_state_tensor("weight_global_amax", global_amax, cpu=cpu)) + elif quantization_format in _NVFP4_EXPORT_FORMATS: + weight_scale_2 = ( + weight_quantizer._amax.float() / 448.0 + if quantization_format == QUANTIZATION_W4A8_NVFP4_FP8 + else NVFP4QTensor.get_weights_scaling_factor_2_from_quantizer(weight_quantizer) + ) + tensors.append(_state_tensor("weight_scale_2", weight_scale_2, cpu=cpu)) + elif quantization_format == QUANTIZATION_FP8: + weight_scale = get_scaling_factor(weight_quantizer) + if weight_scale is None: + raise RuntimeError(f"Missing calibrated weight scale for {weight_name!r}") + tensors.append(_state_tensor("weight_scale", weight_scale, cpu=cpu)) + elif quantization_format == QUANTIZATION_FP8_PC_PT: + weight_scale = get_scaling_factor(weight_quantizer) + if weight_scale is None or weight_quantizer.axis is None: + raise RuntimeError(f"Missing calibrated per-channel scale for {weight_name!r}") + tensors.append( + _axis_scale_state( + "weight_scale", + weight_scale, + packed_shape, + weight_quantizer.axis, + cpu=cpu, + ) + ) + elif quantization_format == QUANTIZATION_FP8_PB_WO: + weight_scale = get_scaling_factor(weight_quantizer) + if weight_scale is None: + raise RuntimeError(f"Missing calibrated block scale for {weight_name!r}") + tensors.append( + _block_scale_state("weight_scale", weight_scale, packed_shape, block_config, cpu=cpu) + ) + elif quantization_format == QUANTIZATION_MXFP8: + cached_scale = getattr(weight_quantizer, "_scale", None) + if cached_scale is not None: + tensors.append( + _block_scale_state( + "weight_scale", cached_scale, packed_shape, block_config, cpu=cpu + ) + ) + elif quantization_format not in _MXFP4_EXPORT_FORMATS: + raise AssertionError(f"Unexpected functional export format {quantization_format!r}") + + if input_quantizer is not None and input_quantizer.is_enabled: + input_amax = input_quantizer.export_amax() + if input_amax is not None: + input_scale = ( + NVFP4QTensor.get_activation_scaling_factor(input_quantizer) + if quantization_format == QUANTIZATION_NVFP4 + else get_scaling_factor(input_quantizer) + ) + tensors.append(_state_tensor("input_scale", input_scale, cpu=cpu)) + + return _QuantizedWeightExportState( + quantization_format=quantization_format, + block_size=block_size, + weight_shape=tuple(weight.shape), + tensors=tuple(tensors), + packing_permutation=permutation, + input_quantizer_enabled=input_quantizer is not None and input_quantizer.is_enabled, + static_nvfp4=static_nvfp4, + four_over_six=bool(block_config.get("four_over_six", False)), + ) + + +def get_quantized_weight_export_spec( + module: nn.Module, + weight_name: str = "weight", +) -> _QuantizedWeightExportSpec | None: + """Return stable deployment metadata without retaining quantizer tensors.""" + resolved = _resolve_quantized_weight_export_inputs(module, weight_name) + if resolved is None: + return None + _, _, weight_quantizer, input_quantizer, quantization_format = resolved + block_config = getattr(weight_quantizer, "block_sizes", None) or {} + block_size = int(block_config.get(-1, 0)) if isinstance(block_config, dict) else 0 + return _QuantizedWeightExportSpec( + quantization_format, + block_size, + input_quantizer is not None and input_quantizer.is_enabled, + ) + + +def _normalize_weight_dim(weight_dim: int, ndim: int) -> int: + if not -ndim <= weight_dim < ndim: + raise IndexError(f"Weight dimension {weight_dim} is invalid for a rank-{ndim} weight") + return weight_dim % ndim diff --git a/modelopt/torch/quantization/plugins/custom.py b/modelopt/torch/quantization/plugins/custom.py index d16596908c9..54dad627457 100644 --- a/modelopt/torch/quantization/plugins/custom.py +++ b/modelopt/torch/quantization/plugins/custom.py @@ -173,10 +173,15 @@ def _has_complete_static_nvfp4_weight_state(quantizer, weight): _check_unsupported_states( quantizer if isinstance(quantizer, TensorQuantizer) else quantizer[0] ) + # Grouped linears validate each expert against its corresponding weight in their + # override. Resetting the container here would clear every expert amax before this + # single-weight compatibility path recalibrates only the representative quantizer. # Skip max_calibrate when saved static NVFP4 state is intact; else MSE scales get overwritten. - if _has_state( - self.weight_quantizer, "_amax" - ) and not _has_complete_static_nvfp4_weight_state(self.weight_quantizer, self.weight): + if ( + not isinstance(self.weight_quantizer, GroupedQuantizer) + and _has_state(self.weight_quantizer, "_amax") + and not _has_complete_static_nvfp4_weight_state(self.weight_quantizer, self.weight) + ): self.weight_quantizer.reset_amax() max_calibrate(self.weight_quantizer, lambda wq: wq(self.weight), distributed_sync=False) if _has_state(self.input_quantizer, "_pre_quant_scale"): diff --git a/tests/gpu/torch/export/test_export.py b/tests/gpu/torch/export/test_export.py index 6f90c9104c7..c25c135ee11 100644 --- a/tests/gpu/torch/export/test_export.py +++ b/tests/gpu/torch/export/test_export.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import copy import json from fnmatch import fnmatch @@ -32,7 +33,8 @@ partial_nvfp4_config, partial_w4a8_config, ) -from _test_utils.torch.transformers_models import get_tiny_qwen3_moe +from _test_utils.torch.transformers_models import get_tiny_llama, get_tiny_qwen3_moe +from safetensors.torch import load_file import modelopt.torch.quantization as mtq from modelopt.torch.export.quant_format import ( @@ -107,6 +109,7 @@ def test_get_quantization_format(config, expected): { "quant_algo": "MIXED_PRECISION", "kv_cache_quant_algo": None, + "exclude_modules": ["layer8"], "quantized_layers": { "layer1": {"quant_algo": "NVFP4", "group_size": 16}, "layer3": { @@ -509,6 +512,123 @@ def test_get_quant_config(config, expected): assert quant_config["quantization"] == expected +@pytest.mark.parametrize("activation", ["disabled", "mxfp4", "fp8", "mixed"]) +def test_mxfp4_offline_export_activation_schema(tmp_path, activation): + model = get_tiny_llama().to("cuda") + model.config.architectures = ["LlamaForCausalLM"] + config = copy.deepcopy(mtq.MXFP4_MLP_WEIGHT_ONLY_CFG) + if activation in ("mxfp4", "fp8", "mixed"): + config["quant_cfg"].append( + { + "quantizer_name": ( + "*layers.0.mlp*input_quantizer" + if activation == "mixed" + else "*mlp*input_quantizer" + ), + "cfg": ( + {"num_bits": (4, 3)} + if activation == "fp8" + else { + "num_bits": (2, 1), + "block_sizes": {-1: 32, "type": "dynamic", "scale_bits": (8, 0)}, + } + ), + } + ) + config["algorithm"] = "max" + inputs = torch.randint(model.config.vocab_size, (1, 8), device="cuda") + mtq.quantize(model, config, lambda m: m(inputs)) + export_hf_checkpoint(model, export_dir=tmp_path) + with open(tmp_path / "config.json") as file: + exported = json.load(file)["quantization_config"] + tensors = load_file(next(tmp_path.glob("*.safetensors"))) + assert tensors["model.layers.0.mlp.gate_proj.weight"].dtype == torch.uint8 + assert tensors["model.layers.0.mlp.gate_proj.weight"].shape == (32, 16) + assert tensors["model.layers.0.mlp.gate_proj.weight_scale"].shape == (32, 1) + + expected_algo = { + "disabled": "W4A16_MXFP4", + "mxfp4": "MXFP4", + "fp8": "W4A8_MXFP4_FP8", + "mixed": "MIXED_PRECISION", + }[activation] + assert exported["quant_algo"] == expected_algo + if activation == "mixed": + assert {cfg["quant_algo"] for cfg in exported["quantized_layers"].values()} == { + "MXFP4", + "W4A16_MXFP4", + } + assert len(exported["config_groups"]) == 2 + for group in exported["config_groups"].values(): + assert group["weights"]["num_bits"] == 4 + assert group["weights"]["group_size"] == 32 + has_input = activation in ("mxfp4", "fp8") or ( + activation == "mixed" and "layers.0.mlp." in group["targets"][0] + ) + assert ("input_activations" in group) == has_input + if has_input: + assert group["input_activations"]["num_bits"] == (8 if activation == "fp8" else 4) + + +@pytest.mark.parametrize("block_size", [64, 128, "mixed"]) +def test_fp8_block_offline_export_preserves_block_structure(tmp_path, block_size): + model = get_tiny_llama(hidden_size=128, intermediate_size=128).to("cuda") + model.config.architectures = ["LlamaForCausalLM"] + layer_blocks = [64, 128] if block_size == "mixed" else [block_size, block_size] + config = { + "algorithm": "max", + "quant_cfg": [{"quantizer_name": "*", "enable": False}], + } + for layer, size in enumerate(layer_blocks): + config["quant_cfg"].append( + { + "quantizer_name": f"*layers.{layer}.mlp*weight_quantizer", + "enable": True, + "cfg": {"num_bits": (4, 3), "block_sizes": {-2: size, -1: size}}, + } + ) + inputs = torch.randint(model.config.vocab_size, (1, 8), device="cuda") + mtq.quantize(model, config, lambda m: m(inputs)) + export_hf_checkpoint(model, export_dir=tmp_path) + with open(tmp_path / "config.json") as file: + exported = json.load(file)["quantization_config"] + tensors = load_file(next(tmp_path.glob("*.safetensors"))) + + assert exported["quant_algo"] == ("MIXED_PRECISION" if block_size == "mixed" else "FP8_PB_WO") + groups = exported["config_groups"].values() + assert {tuple(group["weights"]["block_structure"]) for group in groups} == { + (size, size) for size in layer_blocks + } + if block_size == "mixed": + for group in groups: + for target in group["targets"]: + layer = 0 if "layers.0." in target else 1 + assert group["weights"]["block_structure"] == [layer_blocks[layer]] * 2 + for layer, size in enumerate(layer_blocks): + name = f"model.layers.{layer}.mlp.gate_proj" + assert tensors[f"{name}.weight"].dtype == torch.float8_e4m3fn + assert tensors[f"{name}.weight_scale"].numel() == (128 // size) ** 2 + + +@pytest.mark.parametrize("block_sizes", [{-2: 64, -1: 128}, {-2: 128, -1: 64}, {-1: 64}]) +def test_offline_fp8_block_export_rejects_unsupported_layout(block_sizes): + model = torch.nn.Linear(128, 128, bias=False, device="cuda") + config = { + "algorithm": "max", + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + { + "quantizer_name": "*weight_quantizer", + "enable": True, + "cfg": {"num_bits": (4, 3), "block_sizes": block_sizes}, + }, + ], + } + mtq.quantize(model, config, lambda m: m(torch.randn(2, 128, device="cuda"))) + with pytest.raises(NotImplementedError, match="square blocks over the last two dimensions"): + get_quant_config(model) + + def test_qwen3_moe_nvfp4_experts_only_export_exclude_modules(tmp_path): """Test that NVFP4_EXPERTS_ONLY_CFG correctly excludes non-expert modules in HF export. diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_transformer_engine.py b/tests/gpu_megatron/torch/quantization/plugins/test_transformer_engine.py index 13ecc941c76..e98293c9f56 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_transformer_engine.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_transformer_engine.py @@ -118,6 +118,26 @@ def test_quantize(model_cls, config): quantize_model_and_forward(model, config, calib_data) +def test_grouped_post_restore_preserves_weight_amax(): + """Grouped restore must not clear calibrated state for non-representative experts.""" + model = TEGroupedLinear().cuda() + with torch.no_grad(): + for index in range(model.num_gemms): + getattr(model.net, f"weight{index}").fill_(0.25 + index) + + data = model.get_input().cuda() + mtq.quantize(model, copy.deepcopy(mtq.FP8_DEFAULT_CFG), lambda module: module(data)) + + quantizers = model.net.weight_quantizer + assert isinstance(quantizers, GroupedQuantizer) + before_restore = [quantizer.amax.clone() for quantizer in quantizers] + + model.net.modelopt_post_restore() + + for quantizer, expected in zip(quantizers, before_restore): + torch.testing.assert_close(quantizer.amax, expected) + + @pytest.mark.parametrize("share_weight_quantizer", [False, True]) def test_fold_weight_grouped_linear(share_weight_quantizer): model = TEGroupedLinear().cuda() diff --git a/tests/unit/torch/export/test_get_quantization.py b/tests/unit/torch/export/test_get_quantization.py index edf55bf4bea..f726f01f4e1 100644 --- a/tests/unit/torch/export/test_get_quantization.py +++ b/tests/unit/torch/export/test_get_quantization.py @@ -38,6 +38,7 @@ QUANTIZATION_NVFP4, QUANTIZATION_Q8_0, QUANTIZATION_W4A8_AWQ, + QUANTIZATION_W4A16_MXFP4, ) from modelopt.torch.export.quant_utils import ( _get_carried_over_module_names, @@ -50,6 +51,7 @@ seed_carried_over_exclusions, uses_iq_quantization, ) +from modelopt.torch.quantization.config import QuantizerAttributeConfig from modelopt.torch.quantization.nn import ( GroupedQuantizer, NVFP4StaticQuantizer, @@ -339,6 +341,17 @@ def test_get_quantization_format(config, expected): assert get_quantization_format(model) == expected +def test_mxfp4_weight_only_without_input_quantizer(): + module = torch.nn.Linear(32, 32, bias=False) + module.weight_quantizer = TensorQuantizer( + QuantizerAttributeConfig( + num_bits=(2, 1), + block_sizes={-1: 32, "type": "dynamic", "scale_bits": (8, 0)}, + ) + ) + assert get_quantization_format(module) == QUANTIZATION_W4A16_MXFP4 + + def test_nvfp4_static_quantizer_export(): """NVFP4StaticQuantizer: get_quantization_format returns NVFP4 and get_quant_config returns export config.""" model = ToyModel() diff --git a/tests/unit/torch/export/test_quantized_weight_state.py b/tests/unit/torch/export/test_quantized_weight_state.py new file mode 100644 index 00000000000..22ad3cac977 --- /dev/null +++ b/tests/unit/torch/export/test_quantized_weight_state.py @@ -0,0 +1,353 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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 copy + +import pytest +import torch +import torch.nn as nn + +import modelopt.torch.export.quantized_weight_export as weight_export +import modelopt.torch.quantization as mtq +from modelopt.torch.export.convert_hf_config import convert_hf_quant_config_format +from modelopt.torch.export.quantized_weight_export import ( + capture_quantized_weight_export_state, + get_quantized_weight_export_spec, +) +from modelopt.torch.quantization.config import QuantizerAttributeConfig +from modelopt.torch.quantization.nn import GroupedQuantizer, NVFP4StaticQuantizer, TensorQuantizer + + +def _fp8_linear() -> nn.Linear: + module = nn.Linear(4, 4, bias=False) + with torch.no_grad(): + module.weight.copy_(torch.arange(16, dtype=torch.float32).reshape(4, 4) / 8 - 1) + return mtq.quantize(module, copy.deepcopy(mtq.FP8_DEFAULT_CFG), lambda m: m(torch.ones(2, 4))) + + +def _static_w4a16_linear( + weight: torch.Tensor, + per_block_amax: torch.Tensor, + global_amax: torch.Tensor, +) -> nn.Linear: + module = nn.Linear(weight.shape[1], weight.shape[0], bias=False) + module.weight.data.copy_(weight) + cfg = QuantizerAttributeConfig( + num_bits=(2, 1), + block_sizes={-1: 16, "type": "static", "scale_bits": (4, 3)}, + ) + quantizer = NVFP4StaticQuantizer(quant_attribute_cfg=cfg) + quantizer.amax = per_block_amax.clone() + quantizer.global_amax = global_amax.clone() + module.weight_quantizer = quantizer + module.input_quantizer = TensorQuantizer() + module.input_quantizer.disable() + return module + + +class _SquareTransposedExperts(nn.Module): + def __init__(self): + super().__init__() + self.gate_up_proj = nn.Parameter(torch.arange(512, dtype=torch.float32).reshape(2, 16, 16)) + self.down_proj = nn.Parameter(torch.arange(512, dtype=torch.float32).reshape(2, 16, 16)) + + fp8_cfg = QuantizerAttributeConfig(num_bits=(4, 3)) + self.gate_up_proj_weight_quantizer = TensorQuantizer(fp8_cfg) + self.gate_up_proj_weight_quantizer._amax = torch.tensor(2.0) + self.gate_up_proj_input_quantizer = TensorQuantizer(fp8_cfg) + self.gate_up_proj_input_quantizer._amax = torch.tensor(1.0) + + nvfp4_cfg = QuantizerAttributeConfig( + num_bits=(2, 1), + block_sizes={-1: 16, "type": "static", "scale_bits": (4, 3)}, + ) + self.down_proj_weight_quantizer = NVFP4StaticQuantizer(quant_attribute_cfg=nvfp4_cfg) + self.down_proj_weight_quantizer.amax = torch.arange(32, dtype=torch.float32) + 1 + self.down_proj_weight_quantizer.global_amax = torch.tensor(32.0) + self.down_proj_input_quantizer = TensorQuantizer() + self.down_proj_input_quantizer.disable() + + def iter_weights_for_calibration(self): + for name in ("gate_up_proj", "down_proj"): + yield getattr(self, name).transpose(-1, -2), getattr(self, f"{name}_weight_quantizer") + + +class _GroupedWeights(nn.Module): + """Minimal TEGroupedLinear-style numbered-weight layout.""" + + def __init__(self): + super().__init__() + self.weight0 = nn.Parameter(torch.arange(16, dtype=torch.float32).reshape(4, 4)) + self.weight1 = nn.Parameter(torch.arange(16, 32, dtype=torch.float32).reshape(4, 4)) + cfg = QuantizerAttributeConfig(num_bits=(4, 3)) + self.quantizers = nn.ModuleList([TensorQuantizer(cfg), TensorQuantizer(cfg)]) + for index, quantizer in enumerate(self.quantizers, start=1): + quantizer._amax = torch.tensor(float(index)) + self.input_quantizer = TensorQuantizer() + self.input_quantizer.disable() + + def iter_weights_for_calibration(self): + yield self.weight0, self.quantizers[0] + yield self.weight1, self.quantizers[1] + + +def test_capture_does_not_modify_zero_amax(): + module = _fp8_linear() + module.weight_quantizer._amax.zero_() + module.input_quantizer._amax.zero_() + weight_amax = module.weight_quantizer._amax.clone() + input_amax = module.input_quantizer._amax.clone() + + capture_quantized_weight_export_state(module) + + torch.testing.assert_close(module.weight_quantizer._amax, weight_amax) + torch.testing.assert_close(module.input_quantizer._amax, input_amax) + + +_CAPTURE_FORMAT_CONFIGS = { + "fp8_pc_pt": {"num_bits": (4, 3), "axis": 0}, + "fp8_pb_wo": {"num_bits": (4, 3), "block_sizes": {-2: 128, -1: 128}}, + "nvfp4": { + "num_bits": (2, 1), + "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)}, + }, + "mxfp8": { + "num_bits": (4, 3), + "block_sizes": {-1: 32, "type": "dynamic", "scale_bits": (8, 0)}, + }, + "mxfp4": { + "num_bits": (2, 1), + "block_sizes": {-1: 32, "type": "dynamic", "scale_bits": (8, 0)}, + }, +} + + +def _capture_linear(quantization_format): + module = nn.Linear(256, 128, bias=False) + cfg = _CAPTURE_FORMAT_CONFIGS[quantization_format] + module.weight_quantizer = TensorQuantizer(QuantizerAttributeConfig(**cfg)) + module.input_quantizer = TensorQuantizer( + QuantizerAttributeConfig( + **({"num_bits": (4, 3), "type": "dynamic"} if "axis" in cfg else cfg) + ) + ) + if quantization_format == "fp8_pb_wo": + module.input_quantizer.disable() + return module + + +@pytest.mark.parametrize("quantization_format", _CAPTURE_FORMAT_CONFIGS) +def test_capture_preserves_format_state_and_detaches_tensors(quantization_format): + module = _capture_linear(quantization_format) + cfg = _CAPTURE_FORMAT_CONFIGS[quantization_format] + expected = {} + if quantization_format == "fp8_pc_pt": + scale = torch.arange(1, 129, dtype=torch.float32) + module.weight_quantizer._amax = (448 * scale).reshape(128, 1) + expected["weight_scale"] = (scale, (0,), (1,)) + elif quantization_format == "fp8_pb_wo": + scale = torch.tensor([[1.0, 2.0]]) + module.weight_quantizer._amax = 448 * scale + expected["weight_scale"] = (scale, (0, 1), (128, 128)) + elif quantization_format == "nvfp4": + module.weight_quantizer._amax = torch.tensor(2 * 6 * 448.0) + module.input_quantizer._amax = torch.tensor(3 * 6 * 448.0) + expected["weight_scale_2"] = (torch.tensor(2.0), (), ()) + expected["input_scale"] = (torch.tensor(3.0), (), ()) + elif quantization_format == "mxfp8": + scale = torch.full((128, 8), 127, dtype=torch.uint8) + module.weight_quantizer._scale = scale.clone() + expected["weight_scale"] = (scale, (0, 1), (1, 32)) + for quantizer in (module.weight_quantizer, module.input_quantizer): + if hasattr(quantizer, "_amax"): + quantizer._amax.requires_grad_() + + state = capture_quantized_weight_export_state(module) + + assert state.quantization_format == quantization_format + assert state.weight_shape == (128, 256) + assert state.packing_permutation == (0, 1) + assert state.block_size == (cfg.get("block_sizes") or {}).get(-1, 0) + assert state.input_quantizer_enabled is module.input_quantizer.is_enabled + with torch.no_grad(): + for quantizer in (module.weight_quantizer, module.input_quantizer): + for name in ("_amax", "_scale"): + if (value := getattr(quantizer, name, None)) is not None: + value.zero_() + assert {record.name for record in state.tensors} == set(expected) + for record in state.tensors: + value, axes, block_sizes = expected[record.name] + torch.testing.assert_close(record.value, value) + assert record.axes == axes + assert record.block_sizes == block_sizes + assert record.value.device.type == "cpu" + assert not record.value.requires_grad + + +@pytest.mark.parametrize("quantization_format", _CAPTURE_FORMAT_CONFIGS) +def test_capture_requires_calibration_only_for_static_scales(quantization_format): + module = _capture_linear(quantization_format) + + if quantization_format in {"mxfp8", "mxfp4"}: + state = capture_quantized_weight_export_state(module) + assert state.quantization_format == quantization_format + assert not state.tensors + else: + error = "amax" if quantization_format == "nvfp4" else "Missing calibrated" + with pytest.raises((RuntimeError, AssertionError), match=error): + capture_quantized_weight_export_state(module) + + +def test_capture_resolves_numbered_grouped_weights_by_storage(): + module = _GroupedWeights() + + state0 = capture_quantized_weight_export_state(module, "weight0") + state1 = capture_quantized_weight_export_state(module, "weight1") + + assert state0 is not None and state1 is not None + assert state0.quantization_format == state1.quantization_format == "fp8" + assert not state0.input_quantizer_enabled and not state1.input_quantizer_enabled + assert state0.tensors[0].value.item() != state1.tensors[0].value.item() + + +@pytest.mark.parametrize("has_iterator", [True, False]) +def test_capture_rejects_unresolved_numbered_weight(has_iterator): + module = _GroupedWeights() + module.weight_quantizer = GroupedQuantizer(*module.quantizers) + module.iter_weights_for_calibration = ( + (lambda: iter(((module.weight0, module.quantizers[0]),))) if has_iterator else None + ) + + with pytest.raises(NotImplementedError, match="weight1"): + capture_quantized_weight_export_state(module, "weight1") + + +def test_capture_rejects_full_fused_weight_with_per_expert_quantizers(): + module = _GroupedWeights() + module.weight = nn.Parameter(torch.stack((module.weight0, module.weight1))) + module.weight_quantizer = GroupedQuantizer(*module.quantizers) + module.iter_weights_for_calibration = lambda: ( + (module.weight[index], quantizer) for index, quantizer in enumerate(module.quantizers) + ) + + with pytest.raises(NotImplementedError, match="matching full weight view"): + capture_quantized_weight_export_state(module) + + +@pytest.mark.parametrize("parameter_name", ["bias", "layer_norm_weight"]) +def test_unquantized_sibling_has_no_export_state_or_spec(parameter_name): + module = _fp8_linear() + module.register_parameter(parameter_name, nn.Parameter(torch.ones(4))) + + assert capture_quantized_weight_export_state(module, parameter_name) is None + assert get_quantized_weight_export_spec(module, parameter_name) is None + + +def test_unquantized_weight_has_no_export_state_or_spec(): + module = nn.Linear(4, 4, bias=False) + + assert capture_quantized_weight_export_state(module) is None + assert get_quantized_weight_export_spec(module) is None + + +@pytest.mark.parametrize( + ("features", "quantizer_cfg", "error"), + [ + pytest.param(4, {"num_bits": 8}, "int8_wo", id="int8-weight-only"), + pytest.param( + 32, + { + "num_bits": (2, 1), + "block_sizes": {-1: 32, "type": "dynamic", "scale_bits": (8, 0)}, + }, + "weight-only MXFP4", + id="mxfp4-weight-only", + ), + pytest.param( + 128, + {"num_bits": (4, 3), "block_sizes": {-2: 64, -1: 128}}, + "only with 128x128 blocks", + id="fp8-rectangular-blocks", + ), + ], +) +def test_export_spec_rejects_unsupported_quantization(features, quantizer_cfg, error): + module = nn.Linear(features, features, bias=False) + module.weight_quantizer = TensorQuantizer(QuantizerAttributeConfig(**quantizer_cfg)) + module.input_quantizer = TensorQuantizer() + module.input_quantizer.disable() + + with pytest.raises(NotImplementedError, match=error): + get_quantized_weight_export_spec(module) + + +@pytest.mark.parametrize( + ("quantization_format", "block_size", "input_enabled"), + [("fp8", 0, True), ("fp8", 0, False), ("w4a16_nvfp4", 16, False)], +) +def test_export_spec_does_not_materialize_tensor_state( + quantization_format, block_size, input_enabled, monkeypatch +): + if quantization_format == "fp8": + module = _fp8_linear() + else: + weight = torch.arange(32, dtype=torch.float32).reshape(2, 16) / 32 + module = _static_w4a16_linear( + weight, weight.abs().amax(dim=1, keepdim=True), torch.tensor(1.0) + ) + monkeypatch.setattr( + weight_export, + "_state_tensor", + lambda *args, **kwargs: pytest.fail("export spec materialized tensor state"), + ) + + spec = get_quantized_weight_export_spec(module) + + assert spec is not None + assert spec.quantization_format == quantization_format + assert spec.block_size == block_size + if not input_enabled and module.input_quantizer.is_enabled: + module.input_quantizer.disable() + updated_spec = get_quantized_weight_export_spec(module) + assert updated_spec != spec + assert updated_spec.quantization_format == spec.quantization_format + assert updated_spec.block_size == spec.block_size + spec = updated_spec + assert spec.input_quantizer_enabled is input_enabled + + +def test_capture_uses_the_requested_weight_format_and_transposed_layout(): + module = _SquareTransposedExperts() + gate_state = capture_quantized_weight_export_state(module, "gate_up_proj") + down_state = capture_quantized_weight_export_state(module, "down_proj") + + assert gate_state.quantization_format == "fp8" + assert down_state.quantization_format == "w4a16_nvfp4" + assert gate_state.weight_shape == down_state.weight_shape == (2, 16, 16) + assert gate_state.packing_permutation == down_state.packing_permutation == (0, 2, 1) + + +def test_converted_hf_config_preserves_nvfp4_group_size(): + config = convert_hf_quant_config_format( + { + "quantization": { + "quant_algo": "W4A16_NVFP4", + "group_size": 32, + } + } + ) + + assert config["group_size"] == 32 + assert config["config_groups"]["group_0"]["weights"]["group_size"] == 32 From 0af87c1fa9d48fa79315153dd8358148b1cd4649 Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Fri, 16 Oct 2026 12:17:18 +0530 Subject: [PATCH 24/35] Sync GPTQ Hessians across data-parallel ranks (#2773) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Type of change: Bug fix GPTQ accumulated each linear's Hessian on the local rank only. Megatron PTQ shards calibration data across data-parallel ranks, so with DP > 1 every replicated module (attention, dense MLP, shared experts, and routed experts when EP < DP) ran GPTQ on 1/DP of the calibration set, and its replicas reached different quantized weights (the checkpoint keeps one of them). It affects every GPTQ recipe (NVFP4, INT4, GGML IQ*/Q8_0); max/MSE/AWQ already sync across DP. Fix: after the Hessian collection pass, `GPTQHelper.sync_hessian_across_dp()` sums the token-weighted Hessian and the token count over the module's `parallel_state.data_parallel_group` (the expert-DP group for routed experts), then renormalizes. It is a no-op when the group is uninitialized or has one rank, so single-process HF PTQ is unchanged. CPU-offloaded Hessians are moved to the GPU only for the all-reduce, one at a time. All ranks in a group build handles from `named_modules()` in the same order and reduce every handle, including experts that got no tokens, so collectives stay matched. `local_hessian` has the same gap but already warns that it is single-rank only; not changed here. Tracks OMNIML-6117. - New `tests/gpu_megatron/.../test_megatron.py::test_gptq_hessian_sync_across_dp` (2 GPUs): GPTQ at DP2 with different data per rank must produce the same Hessians as a run on the union of both ranks' data, and identical weights on both ranks. Fails with the fix reverted; passes with it (2× RTX 6000 Ada). Existing `test_gptq_mamba_hybrid`, `tests/gpu/torch/quantization/test_gptq.py` and `test_ggml_gptq_cuda.py` pass. - GLM-5.3-Flash mini model (4 layers, 16 experts), IQ2_XS layerwise GPTQ recipe, Megatron-Bridge `quantize.py` at EP4 on 4 GB300 GPUs, 128 × 4096-token samples: shared-expert Hessians now cover 524288 tokens (was 131072 = 1/4); routed experts unchanged (expert-DP size 1); total GPTQ time 145.0 s vs 145.2 s before. - Full GLM-5.3-Flash (91 layers, 288 experts) IQ2_XS PTQ at EP16 over 4 GB300 nodes, 512 × 4096-token samples: completes; shared-expert Hessians cover 2097152 tokens (was 131072 = 1/16); total GPTQ time 309 s vs 304 s before (+1.5%). Not run: TP > 1 or PP > 1 with GPTQ, HF FSDP2 multi-process GPTQ (its modules default to the world group as DP, which this now syncs, matching how amax is synced). - Is this change backward compatible?: ✅ - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: ✅ - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ - Did you get Claude approval on this PR?: ❌ Found while running GLM-5.3-Flash IQ2 PTQ with Megatron-Bridge (EP16 ⇒ DP16 for non-expert modules: dense MLPs and shared experts saw 32 of 512 samples each). 🤖 Generated with [Claude Code](https://claude.com/claude-code) * **Bug Fixes** * Fixed GPTQ quantization with Megatron data parallelism so all replicas use calibration data from across the data-parallel group. This prevents differences caused by replicas processing separate calibration batches and keeps their resulting quantized weights consistent when multiple data-parallel ranks are active. Single-rank behavior remains unchanged. Calibration data is combined across replicas before quantization. --------- Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 Signed-off-by: Rohan Joshi --- CHANGELOG.rst | 1 + modelopt/torch/quantization/model_calib.py | 1 + .../torch/quantization/utils/calib_utils.py | 20 ++++++++ .../quantization/plugins/test_megatron.py | 48 +++++++++++++++++++ tests/unit/torch/quantization/test_gptq.py | 28 ++++++++++- 5 files changed, 97 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 0dbdfcd82fa..b161b6e8415 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -175,6 +175,7 @@ Changelog - Fix YAML config I/O decoding with the locale codepage instead of UTF-8, which made a config containing any non-ASCII byte fail to load on a machine whose locale is not UTF-8 (notably Windows, where the default is cp1252). ``modelopt/recipe/loader.py``, the two ONNX autotune state files, the two transformers config readers, the distill config and the puzzletron profile now pass ``encoding="utf-8"`` explicitly. Only the YAML config paths are covered: these are the files most likely to carry non-ASCII text in comments, model names or paths, and the only ones read inside a user's process. - Hybrid (e.g. Nemotron-H) checkpoints saved by the ``examples/megatron_bridge`` scripts now record their layer spec in ``run_config.yaml`` in a form that reloads, so they can be converted to HuggingFace; a checkpoint saved by an earlier release still needs its ``model.hybrid_stack_spec`` block replaced by hand. - Fix ONNX preprocessing with custom TensorRT plugin libraries on TensorRT 11. +- Fix GPTQ with Megatron data parallelism: each data-parallel replica of a module now builds its Hessian from the whole calibration set instead of only its own shard, so all replicas reach the same quantized weights. 0.47.0 (2026-09-23) ^^^^^^^^^^^^^^^^^^^ diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index 77eeb9371b5..3741ed27883 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -2339,6 +2339,7 @@ def _make_gptq_handle(name, m): for handle in gptq_handles.values(): handle.cleanup() + handle.sync_hessian_across_dp() print_rank_0("Updating weights using GPTQ algorithm...") names = module_name_maps(model) diff --git a/modelopt/torch/quantization/utils/calib_utils.py b/modelopt/torch/quantization/utils/calib_utils.py index b1ade8e7809..f489bdd5d3e 100644 --- a/modelopt/torch/quantization/utils/calib_utils.py +++ b/modelopt/torch/quantization/utils/calib_utils.py @@ -41,6 +41,7 @@ import math import torch +import torch.distributed as dist from modelopt.torch.utils import print_rank_0 from modelopt.torch.utils.network import bind_forward_method, unpatch_forward_method @@ -170,6 +171,25 @@ def cleanup(self): """Unpatch the module's forward method.""" unpatch_forward_method(self.module, self.CACHE_NAME) + def sync_hessian_across_dp(self): + """Merge the Hessians of this module's data-parallel replicas, which saw different data.""" + parallel_state = getattr(self.module, "parallel_state", None) + group = parallel_state.data_parallel_group if parallel_state is not None else None + if group is None or not group.is_initialized() or group.world_size() == 1: + return + on_cuda = "nccl" in dist.get_backend(group.group) + device = ( + torch.device("cuda", torch.cuda.current_device()) if on_cuda else torch.device("cpu") + ) + n_samples = torch.tensor(self.n_samples, dtype=torch.float64, device=device) + hessian = self.hessian.to(device).mul_(self.n_samples) + dist.all_reduce(n_samples, group=group.group) + dist.all_reduce(hessian, group=group.group) + self.n_samples = int(n_samples.item()) + if self.n_samples: + hessian /= self.n_samples + self.hessian = hessian.to(self.hessian.device) + def free(self): """Release Hessian and working tensors to reclaim memory.""" self.hessian = None diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py index e2702f7cf30..177ce61fc4b 100644 --- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py +++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py @@ -97,6 +97,7 @@ ) from modelopt.torch.quantization.qtensor import QTensorWrapper from modelopt.torch.quantization.utils import is_quantized_linear +from modelopt.torch.quantization.utils.calib_utils import GPTQHelper from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector try: @@ -1524,6 +1525,53 @@ def _test_gptq_mamba_hybrid(rank, size): assert torch.isfinite(forward(model)).all() +def test_gptq_hessian_sync_across_dp(dist_workers_size_2): + """GPTQ at DP2 must use the Hessian of both ranks' calibration data and agree across ranks.""" + dist_workers_size_2.run(_test_gptq_hessian_sync_across_dp) + + +def _test_gptq_hessian_sync_across_dp(rank, size): + initialize_for_megatron(tensor_model_parallel_size=1, seed=SEED) + model = get_mcore_gpt_model(tensor_model_parallel_size=1, num_layers=1).cuda() + ref_model = get_mcore_gpt_model(tensor_model_parallel_size=1, num_layers=1).cuda() + ref_model.load_state_dict(model.state_dict()) + + torch.manual_seed(SEED + rank) + batch = get_batch(model, batch_size=2) + all_input_ids = [torch.empty_like(batch[0]) for _ in range(size)] + torch.distributed.all_gather(all_input_ids, batch[0]) + + def forward(model, input_ids_list): + for input_ids in input_ids_list: + model(input_ids, *batch[2:4], labels=batch[1], loss_mask=batch[4]) + + hessians = {} + update_weights = GPTQHelper.update_weights + + def capture_hessian(self, *args, **kwargs): + hessians[self.name] = self.hessian.clone() + return update_weights(self, *args, **kwargs) + + quant_cfg = copy.deepcopy(mtq.NVFP4_DEFAULT_CFG) + quant_cfg["algorithm"] = {"method": "gptq"} + with patch.object(GPTQHelper, "update_weights", capture_hessian): + mtq.quantize(model, quant_cfg, partial(forward, input_ids_list=[batch[0]])) + dp_hessians = dict(hessians) + # Every rank calibrates on the union of the data, so its Hessian is the full-data one. + mtq.quantize(ref_model, quant_cfg, partial(forward, input_ids_list=all_input_ids)) + + assert dp_hessians + for name, hessian in dp_hessians.items(): + # TF32 matmuls in the Hessian update leave ~1e-4 relative noise between the two runs. + assert (hessian - hessians[name]).norm() < 1e-3 * hessians[name].norm(), name + + for name, m in model.named_modules(): + if is_quantized_linear(m) and m.weight_quantizer.is_enabled: + weights = [torch.empty_like(m.weight) for _ in range(size)] + torch.distributed.all_gather(weights, m.weight.contiguous()) + assert all(torch.equal(w, weights[0]) for w in weights), f"{name} differs across ranks" + + def _auto_quantize_mamba_hybrid_cost_helper(rank, size, expert_model_parallel_size, result_path): initialize_for_megatron( tensor_model_parallel_size=1, diff --git a/tests/unit/torch/quantization/test_gptq.py b/tests/unit/torch/quantization/test_gptq.py index 345bd167112..1c55c089943 100644 --- a/tests/unit/torch/quantization/test_gptq.py +++ b/tests/unit/torch/quantization/test_gptq.py @@ -15,9 +15,13 @@ """CPU unit tests for GPTQ utilities.""" +from functools import partial + import torch +from _test_utils.torch.distributed.utils import spawn_multiprocess_job -from modelopt.torch.quantization.utils.calib_utils import update_hessian +from modelopt.torch.quantization.utils.calib_utils import GPTQHelper, update_hessian +from modelopt.torch.utils.distributed import ParallelState def test_update_hessian(): @@ -117,3 +121,25 @@ def test_update_hessian_zero_token_input_noops(): assert updated_hessian is hessian assert new_n_samples == n_samples torch.testing.assert_close(hessian, expected_hessian) + + +def _gptq_hessian_dp_sync_helper(rank, size, token_counts): + torch.manual_seed(0) + inputs = [torch.randn(n, 8) for n in token_counts] + linear = torch.nn.Linear(8, 4) + linear.parallel_state = ParallelState(data_parallel_group=None) + helper = GPTQHelper(linear, "linear") + helper.hessian, helper.n_samples = update_hessian(inputs[rank], helper.hessian, 0) + + helper.sync_hessian_across_dp() + + ref, n_ref = update_hessian(torch.cat(inputs), torch.zeros(8, 8), 0) + assert helper.n_samples == n_ref + torch.testing.assert_close(helper.hessian, ref) + + +def test_gptq_hessian_dp_sync(skip_on_windows): + """Token-weighted merge over a gloo (CPU) group, with one rank that saw no tokens.""" + spawn_multiprocess_job( + 2, partial(_gptq_hessian_dp_sync_helper, token_counts=(5, 0)), backend="gloo" + ) From eed13bb21e552b0c32e4d4d705b319931108ecfd Mon Sep 17 00:00:00 2001 From: yingguo-trt <244492186+yingguo-trt@users.noreply.github.com> Date: Fri, 16 Oct 2026 17:14:08 +0800 Subject: [PATCH 25/35] [6933164] fix: preserve Diffusers pipeline settings in component exports (#2813) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Type of change: Bug fix. Passing `components` to `export_hf_checkpoint` rebuilds `model_index.json` and drops pipeline settings such as Wan 2.2's `boundary_ratio` and SDXL's `force_zeros_for_empty_prompt`. Preserve public, non-component pipeline configuration while listing only the components actually exported. The default full-export path is unchanged. The regression test constructs tiny Wan and SDXL pipelines entirely from local components, saves and reloads the source, exports it, and verifies the exported configuration and reload. It covers Wan default/all/single-component export and SDXL all-component export, including a non-default SDXL setting. ### Usage ```python export_hf_checkpoint(pipe, export_dir="exported", components=["transformer_2"]) ``` The selected-component manifest now retains the pipeline settings needed when loading it with the omitted components supplied separately. Re-export affected checkpoints. ### Testing Focused test file: `tests/unit/torch/export/test_export_diffusers.py`. | Platform | Diffusers | Transformers | Result | | --- | --- | --- | --- | | Windows, CPU | 0.40.0 | 5.5.4 | 15 passed, 0 skipped | | Windows, CPU | 0.41.0 | 5.18.0 | 15 passed, 0 skipped | | Linux, CPU | 0.40.0 | 5.5.4 | 15 passed, 0 skipped | | Linux, CPU | 0.41.0 | 5.5.4 | 15 passed, 0 skipped | All runs used fresh Hugging Face caches and offline mode; caches remained empty. Linux runs also used containers with networking disabled. The Linux 0.41 environment includes its declared minimum `huggingface-hub==1.32.0`. - Negative control with the original exporter: the two explicit Wan exports and the SDXL export fail on missing configuration keys; default export passes. - The previous network-dependent fixture reproduces three setup errors with an empty offline cache. The replacement fixture eliminates that dependency. - Real B100 validation on ModelOpt 0.48.0rc0: full Wan2.2-A14B and SDXL component exports lose their pipeline setting before the exporter fix and retain it after the fix. These checks cover export metadata, not image/video quality. - Applicable pre-commit checks pass, including Ruff, mypy and Bandit. These are focused local/node results. The full upstream GitHub CI matrix is pending. ### Before your PR is "Ready for review" - Is this change backward compatible?: Yes. - Third-party code or new PIP dependencies?: None added. - Did you write necessary tests?: Yes, one parameterized regression test with offline fixtures. - Did you update Changelog?: Yes, under 0.48.0 for the requested release backport. - Did you get Claude approval on this PR?: Formal PR review pending; the offline-fixture issue identified in local review has been addressed. ### Additional Information Requested backport: `cherry-pick-0.48.0`. The PR targets `main`; release inclusion is pending maintainer cherry-pick. ## Summary by CodeRabbit * **New Features** * Added imatrix calibration support for GGML formats and a Nemotron 3.5 Super VL AutoQuantize recipe. * **Bug Fixes** * Fixed TensorRT 11 ONNX preprocessing and LayerNorm input quantization in affected presets and recipes. * Fixed Diffusers exports so pipeline settings, including Wan 2.2’s boundary ratio, are preserved when exporting all or selected components without an existing configuration file. * Existing exports are not updated automatically. Re-export affected checkpoints to a new directory to include the preserved settings. --------- Signed-off-by: yingguo-trt <244492186+yingguo-trt@users.noreply.github.com> Signed-off-by: Rohan Joshi --- CHANGELOG.rst | 1 + modelopt/torch/export/unified_export_hf.py | 16 ++- .../torch/export/test_export_diffusers.py | 127 ++++++++++++++++++ 3 files changed, 140 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index b161b6e8415..44221e73d49 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -131,6 +131,7 @@ Changelog - Fix unified HF export writing a weight name the source checkpoint does not use when a model declares a backwards-compatibility *load* rename as a reversible ``WeightRenaming``. Reversing such a rule rewrote the canonical name to a legacy spelling, and a deployment runtime keyed on the canonical name then failed to load the quantized checkpoint. Each exported name is now resolved against the names the source checkpoint actually carries, for both the exported tensors and the quantization config's module references. - Fix exports using the built-in Transformers NemotronH implementation retaining incompatible legacy remote-code config and model mappings. Re-export affected checkpoints before deployment. +- Fix ``export_hf_checkpoint`` dropping Diffusers pipeline settings, such as Wan 2.2 ``boundary_ratio``, from ``model_index.json`` when ``components`` is given. Re-export affected checkpoints to a new directory, since component exports do not overwrite an existing ``model_index.json``. - Fix INT8 real quantization (``mtq.compress``) of FP16 weights zeroing or clipping channels with a small amax: their scales underflowed in FP16. Scales of newly compressed INT8 weights are now stored in FP32. - Fix Hugging Face exports classifying registered Qwen VLM parameters as unplaced source weights, which overwrote quantized tensors with original weights and incorrectly excluded their modules from quantization. Re-export affected checkpoints. - Right-pad labeled calibration samples to avoid fully masked causal rows that can produce non-finite cuDNN attention gradients during AutoQuantize. Generation previews remove padding before decoding. diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index a1e474b80c8..f605ddaf636 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -1447,12 +1447,20 @@ def _export_diffusers_checkpoint( if not is_partial_export and not model_index_path.exists() and hasattr(pipe, "save_config"): pipe.save_config(export_dir) - # Last resort: synthesize a minimal model_index.json from exported components. + # Preserve pipeline settings while listing only the exported components. if not model_index_path.exists() and hasattr(pipe, "config") and pipe.config is not None: + pipeline_components = pipe.components model_index = { - "_class_name": type(pipe).__name__, - "_diffusers_version": diffusers.__version__, + key: value + for key, value in json.loads(pipe.to_json_string()).items() + if key not in pipeline_components and not key.startswith("_") } + model_index.update( + { + "_class_name": type(pipe).__name__, + "_diffusers_version": diffusers.__version__, + } + ) for name, comp in all_components.items(): module = type(comp).__module__ library = module.split(".")[0] @@ -1884,7 +1892,7 @@ def export_hf_checkpoint( export_dir: The target export path. save_modelopt_state: Whether to save the modelopt state_dict. components: Only used for diffusers pipelines. Optional list of component names - to export. If None, all quantized components are exported. + to export. If None, all components are exported. extra_state_dict: Extra state dictionary to add to the exported model. max_shard_size: Maximum size of each safetensors shard file. Defaults to "10GB". **kwargs: Runtime-specific post-processing options forwarded to diff --git a/tests/unit/torch/export/test_export_diffusers.py b/tests/unit/torch/export/test_export_diffusers.py index 753c81a4b0e..126550a978c 100644 --- a/tests/unit/torch/export/test_export_diffusers.py +++ b/tests/unit/torch/export/test_export_diffusers.py @@ -14,8 +14,11 @@ # limitations under the License. import copy +import inspect import json +from pathlib import Path +import numpy as np import pytest import torch import torch.nn as nn @@ -24,12 +27,23 @@ get_tiny_flux, get_tiny_flux2, get_tiny_unet, + get_tiny_wan22_transformer, + get_tiny_wan22_vae, ) pytest.importorskip("diffusers") +pytest.importorskip("transformers") +import diffusers from safetensors import safe_open from safetensors.torch import save_file +from transformers import ( + CLIPTextConfig, + CLIPTextModel, + CLIPTextModelWithProjection, + UMT5Config, + UMT5EncoderModel, +) import modelopt.torch.export.unified_export_hf as unified_export_hf import modelopt.torch.quantization as mtq @@ -88,6 +102,119 @@ def test_export_diffusers_models_non_quantized(tmp_path, model_factory): assert "quantization_config" not in config_data +@pytest.fixture +def tiny_pipeline(request, tmp_path, tiny_tokenizer): + """Build and reload a tiny pipeline without accessing the Hugging Face Hub.""" + if request.param == "wan22": + wan_pipeline = getattr(diffusers, "WanPipeline", None) + if ( + wan_pipeline is None + or "boundary_ratio" not in inspect.signature(wan_pipeline.__init__).parameters + ): + pytest.skip("Requires a Diffusers version with Wan 2.2 support") + text_encoder = UMT5EncoderModel( + UMT5Config( + vocab_size=len(tiny_tokenizer), + d_model=32, + d_kv=8, + d_ff=32, + num_layers=1, + num_heads=4, + ) + ) + pipe = wan_pipeline( + tokenizer=tiny_tokenizer, + text_encoder=text_encoder, + vae=get_tiny_wan22_vae(), + scheduler=diffusers.UniPCMultistepScheduler( + prediction_type="flow_prediction", use_flow_sigmas=True, flow_shift=3.0 + ), + transformer=get_tiny_wan22_transformer(), + transformer_2=get_tiny_wan22_transformer(), + boundary_ratio=0.875, + ) + settings = ("boundary_ratio", "expand_timesteps") + else: + text_config = CLIPTextConfig( + vocab_size=len(tiny_tokenizer), + hidden_size=32, + intermediate_size=32, + num_hidden_layers=1, + num_attention_heads=4, + projection_dim=32, + ) + pipe = diffusers.StableDiffusionXLPipeline( + vae=diffusers.AutoencoderKL(block_out_channels=(32,), norm_num_groups=8), + text_encoder=CLIPTextModel(text_config), + text_encoder_2=CLIPTextModelWithProjection(text_config), + tokenizer=tiny_tokenizer, + tokenizer_2=tiny_tokenizer, + unet=get_tiny_unet(in_channels=4, out_channels=4, cross_attention_dim=64), + scheduler=diffusers.DDIMScheduler(), + # Use a non-default value so reload cannot hide a missing setting. + force_zeros_for_empty_prompt=False, + add_watermarker=False, + ) + settings = ("force_zeros_for_empty_prompt",) + source_dir = tmp_path / "source_pipeline" + pipe.save_pretrained(source_dir) + return type(pipe).from_pretrained(source_dir), settings + + +@pytest.mark.parametrize( + ("tiny_pipeline", "selection"), + [("wan22", "default"), ("wan22", "all"), ("wan22", "transformer_2"), ("sdxl", "all")], + indirect=["tiny_pipeline"], +) +def test_export_diffusers_pipeline_config(tmp_path, tiny_pipeline, selection): + pipe, settings = tiny_pipeline + component_names = {name for name, component in pipe.components.items() if component is not None} + components = None if selection == "default" else sorted(component_names) + if selection == "transformer_2": + components = [selection] + expected_components = component_names if components is None else set(components) + export_dir = tmp_path / "export_pipeline" + + export_hf_checkpoint(pipe, export_dir=export_dir, components=components) + + model_index = _load_config(export_dir / "model_index.json") + for setting in settings: + assert model_index[setting] == pipe.config[setting] + assert set(model_index) & set(pipe.components) == expected_components + assert {path.name for path in export_dir.iterdir() if path.is_dir()} == expected_components + + # Loading the export must restore the settings; pass in the components it omits. + omitted = {name: pipe.components[name] for name in component_names - expected_components} + reloaded = type(pipe).from_pretrained(export_dir, **omitted) + for setting in settings: + assert reloaded.config[setting] == pipe.config[setting] + + +@pytest.mark.parametrize( + ("setting", "expected"), + [ + pytest.param(Path("settings") / "config.json", "settings/config.json", id="path"), + pytest.param(np.array([1, 2, 3]), [1, 2, 3], id="ndarray"), + ], +) +def test_export_diffusers_pipeline_config_serialization(tmp_path, setting, expected): + class ConfigPipeline(diffusers.DiffusionPipeline): + def __init__(self, unet, scheduler, setting=None): + super().__init__() + self.register_modules(unet=unet, scheduler=scheduler) + self.register_to_config(setting=setting) + + pipe = ConfigPipeline(get_tiny_unet(), diffusers.DDIMScheduler(), setting=setting) + export_dir = tmp_path / "export_pipeline" + + export_hf_checkpoint(pipe, export_dir=export_dir, components=["unet"]) + + model_index = _load_config(export_dir / "model_index.json") + assert model_index["setting"] == expected + assert set(model_index) & set(pipe.components) == {"unet"} + assert {path.name for path in export_dir.iterdir() if path.is_dir()} == {"unet"} + + def test_export_diffusers_unet_quantized_matches_llm_config(tmp_path, monkeypatch): model = get_tiny_unet() export_dir = tmp_path / "export_unet_quant" From 76825c9b7d3ed89cc6c47c9dc87fe64191cdb394 Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Fri, 16 Oct 2026 15:53:17 +0530 Subject: [PATCH 26/35] Fix post-merge unit CI and the gpu_megatron nightly timeout (#2832) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Type of change: Bug fix (CI / tests) Follow-up to #2823. Its auto-merge landed before the last two commits were pushed, so main is missing them and every push-triggered `Unit tests / linux` run now fails. - **noxfile.py:** with `--run-manual`, the `manual` unit tests run in a second pytest process (no coverage). `test_transformers_load_with_multi_thread` is documented as flaky and leaves BF16 default-dtype state, which broke later tests when `--run-manual` shared the process with the regular suite (`BFloat16 != float` in `test_accelerate`, `test_fused_experts`, `test_auto_quantize`). - **tests/gpu/.../test_hf_dflash.py:** fixes the nightly 2-GPU failure (3 nights in a row). With 4 layers and 2 draft layers `target_layer_ids == [1, 1]`, so the test's own precondition (target layers straddle the shard split) could never hold; 8 layers gives `[1, 5]`. - **tests/_test_utils/.../distributed/utils.py, tests/gpu_megatron/conftest.py, gpu_tests.yml:** stop one hung test from timing out the rest of the `gpu_megatron` nightly job (cancelled at the 60-min limit on 5 of the last 8 nights, with 18-28 pytest timeouts of 120 s each). When a job hangs or a failed rank leaves its peers blocked in a collective, the shared worker pool stayed stuck and every later test in the module waited out its own timeout. `DistributedWorkerPool` now prints the workers' stacks and kills them (job interrupted, or peers not done 30 s after a rank failed) and marks itself broken; the gpu_megatron fixtures respawn a fresh pool; job timeout 60 -> 90 min (a healthy run takes ~47 min). ### Testing - Reproduced the dtype pollution locally with `pytest tests/unit --run-manual` (1 failure); the `manual` tests alone pass (17 in 87s) and the regular suite passes without `--run-manual`. - DFlash test: fails with `AssertionError: [1, 1]` before, 14 speculative GPU tests pass on 2 GPUs after. - gpu_megatron: reproduced the cascade in the `nemo:26.08` CI image on 2 GPUs (RTX 6000 Ada); with the change `test_megatron.py` takes 5m42s (127 passed, 3 failed) instead of cascading into timeouts, and the 2 new pool tests pass. The first hang (nightly: `test_fp8_real_quantize` on the RTX Pro 6000 runner) did not reproduce, so it can still occur; it now costs ~2 min instead of ~30 and its worker stacks go to the log. The 3 failures here are `test_te_grouped_vs_sequential_default_loss` x2 and `test_auto_quantize_moe_ep` on 48 GB cards (likely memory pressure, not confirmed). - Not yet run end to end: the new nox flow in CI (will dispatch `Unit tests` on this branch). - Is this change backward compatible?: ✅ - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: N/A - Did you update Changelog?: N/A - Did you get Claude approval on this PR?: ❌ 🤖 Generated with [Claude Code](https://claude.com/claude-code) ## Summary by CodeRabbit * **Tests** * Improved reliability of distributed test runs by detecting failed or unresponsive workers, ending stalled jobs, and replacing unusable worker pools. * Added a way to run manual unit tests separately when requested, while keeping the standard test run’s coverage settings intact. * Expanded GPU test coverage for sharded hidden states across devices and generation output placement. --------- Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Co-authored-by: Claude Sonnet 5.5 Signed-off-by: Rohan Joshi --- noxfile.py | 13 ++- tests/_test_utils/torch/distributed/utils.py | 69 ++++++++++++++-- tests/gpu/conftest.py | 30 +++++-- .../speculative/plugins/test_hf_dflash.py | 16 ++-- .../utils/test_distributed_worker_pool.py | 81 +++++++++++++++++++ tests/gpu_megatron/conftest.py | 20 ++--- tests/unit/conftest.py | 22 +++-- 7 files changed, 211 insertions(+), 40 deletions(-) create mode 100644 tests/gpu/torch/utils/test_distributed_worker_pool.py diff --git a/noxfile.py b/noxfile.py index 46ffd1e1df6..edd0544e1cd 100644 --- a/noxfile.py +++ b/noxfile.py @@ -69,9 +69,20 @@ def unit(session, torch_ver, tf_ver): tf_pins = TRANSFORMERS_VERSIONS[tf_ver] if tf_pins: session.install(*tf_pins) + run_manual = "--run-manual" in session.posargs + pytest_args = [arg for arg in session.posargs if arg != "--run-manual"] session.run( - "python", "-m", "pytest", "tests/unit", *_cov_args(), *session.posargs, env=_CPU_ONLY_ENV + "python", "-m", "pytest", "tests/unit", *_cov_args(), *pytest_args, env=_CPU_ONLY_ENV ) + if run_manual: + # A separate process: some manual tests (e.g. multi-threaded from_pretrained) leak global + # state that breaks later tests. No coverage, so it cannot overwrite the main run's data. + session.run( + "python", + # Caller args first: the last ``-m`` wins, and it must be ours. + *("-m", "pytest", "tests/unit", *pytest_args, "--run-manual", "-m", "manual"), + env={**_CPU_ONLY_ENV, "COVERAGE_PROCESS_START": ""}, + ) @nox.session(python="3.12") diff --git a/tests/_test_utils/torch/distributed/utils.py b/tests/_test_utils/torch/distributed/utils.py index eae7c2ad38c..ea954efbcf5 100644 --- a/tests/_test_utils/torch/distributed/utils.py +++ b/tests/_test_utils/torch/distributed/utils.py @@ -13,8 +13,13 @@ # See the License for the specific language governing permissions and # limitations under the License. +import contextlib +import faulthandler import os +import queue +import signal import socket +import time import traceback import torch @@ -79,6 +84,18 @@ def spawn_multiprocess_job(size, job, backend="gloo"): assert not p.exitcode +def get_live_pool(pools, key, make_pool): + """Return ``pools[key]``, building it with ``make_pool`` if it is missing or was killed. + + Fixtures hand pools out through this, so a hang kills the pool for one test instead of + failing every later test that uses it. + """ + pool = pools.get(key) + if pool is None or pool.broken: + pools[key] = pool = make_pool() + return pool + + def default_worker_teardown(rank, world_size): """Minimal cleanup between tests in persistent workers.""" try: @@ -117,6 +134,10 @@ def _test_fn(rank, size): ... def __init__(self, world_size, backend="nccl", teardown_fn=default_worker_teardown): assert world_size > 0, "World size must be greater than 0" self.world_size = world_size + # Set once a job was interrupted: the workers are then killed, see ``run``. + self.broken = False + self.job_timeout = 600 + self.failure_grace = 30 # seconds the other ranks get to finish once one has failed ctx = mp.get_context("spawn") self._cmd_queues = [ctx.Queue() for _ in range(world_size)] self._result_queue = ctx.Queue() @@ -145,6 +166,9 @@ def __init__(self, world_size, backend="nccl", teardown_fn=default_worker_teardo @staticmethod def _worker_loop(rank, world_size, backend, port, cmd_queue, result_queue, teardown_fn): + if hasattr(signal, "SIGUSR1"): + # The parent signals this to get the stack of a hung job into the CI log. + faulthandler.register(signal.SIGUSR1, all_threads=True) os.environ["MASTER_ADDR"] = "localhost" os.environ["MASTER_PORT"] = str(port) os.environ["LOCAL_RANK"] = str(rank) @@ -186,19 +210,50 @@ def run(self, fn, *args, **kwargs): ``fn`` is called as ``fn(rank, world_size, *args, **kwargs)`` and must be picklable (top-level function or ``functools.partial`` of one). - """ - for q in self._cmd_queues: - q.put((fn, args, kwargs)) + If the wait is interrupted (e.g. by ``pytest-timeout``), or a rank fails and its peers do + not finish within ``failure_grace`` (they are typically blocked in a collective waiting for + it), the workers are still inside the job and every later ``run`` would hang until its own + timeout. The workers are killed instead and the pool is marked ``broken``; callers that + want to continue need a new pool. + """ + if self.broken: + raise RuntimeError("Worker pool was killed after an earlier job hung; create a new one") errors = [] - for _ in range(self.world_size): - status, rank, tb = self._result_queue.get(timeout=600) - if status == "error": - errors.append(f"--- Rank {rank} ---\n{tb}") + deadline = None # set when the first rank fails: its peers get ``failure_grace`` in total + try: + for q in self._cmd_queues: + q.put((fn, args, kwargs)) + for _ in range(self.world_size): + timeout = self.job_timeout if deadline is None else deadline - time.monotonic() + status, rank, tb = self._result_queue.get(timeout=max(timeout, 0)) + if status == "error": + errors.append(f"--- Rank {rank} ---\n{tb}") + deadline = deadline or time.monotonic() + self.failure_grace + except queue.Empty: + self._abort() + if not errors: + raise + errors.append("--- The other ranks did not finish after a rank failed; killed ---") + except BaseException: + self._abort() + raise if errors: raise RuntimeError("Worker(s) failed:\n" + "\n".join(errors)) + def _abort(self): + """Print every worker's stack (to stderr, for the CI log), then kill the workers.""" + self.broken = True + if hasattr(signal, "SIGUSR1"): + for p in self._processes: + with contextlib.suppress(OSError): + os.kill(p.pid, signal.SIGUSR1) + time.sleep(3) # let faulthandler write before the workers go away + for p in self._processes: + p.kill() + p.join(timeout=10) + def shutdown(self): """Signal all workers to exit and wait for them to finish.""" for q in self._cmd_queues: diff --git a/tests/gpu/conftest.py b/tests/gpu/conftest.py index b2b9344945b..40a272e8838 100644 --- a/tests/gpu/conftest.py +++ b/tests/gpu/conftest.py @@ -15,15 +15,29 @@ import pytest import torch -from _test_utils.torch.distributed.utils import DistributedWorkerPool, default_worker_teardown +from _test_utils.torch.distributed.utils import ( + DistributedWorkerPool, + default_worker_teardown, + get_live_pool, +) @pytest.fixture(scope="session") -def dist_workers(): - pool = DistributedWorkerPool( - world_size=torch.cuda.device_count(), - backend="nccl", - teardown_fn=default_worker_teardown, +def _dist_pools(): + pools = {} + yield pools + for pool in pools.values(): + pool.shutdown() + + +@pytest.fixture +def dist_workers(_dist_pools): + """The session-wide pool; replaced if an earlier test hung and killed it.""" + world_size = torch.cuda.device_count() + return get_live_pool( + _dist_pools, + world_size, + lambda: DistributedWorkerPool( + world_size=world_size, backend="nccl", teardown_fn=default_worker_teardown + ), ) - yield pool - pool.shutdown() diff --git a/tests/gpu/torch/speculative/plugins/test_hf_dflash.py b/tests/gpu/torch/speculative/plugins/test_hf_dflash.py index 0cb740c0e9f..69f0a117b63 100644 --- a/tests/gpu/torch/speculative/plugins/test_hf_dflash.py +++ b/tests/gpu/torch/speculative/plugins/test_hf_dflash.py @@ -287,20 +287,20 @@ def sharded_forward(*args, input_ids=None, **kwargs): model._base_model.forward = sharded_forward def test_generate_gathers_sharded_hidden_states(self): - model = get_tiny_llama(num_hidden_layers=4) + # 8 layers so the default target_layer_ids ([1, 5]) straddle the split; 4 gives [1, 1]. + model = get_tiny_llama(num_hidden_layers=8) mtsp.convert(model, [("dflash", get_dflash_config(block_size=4))]) model = model.to("cuda:1").eval() - # The fix is only exercised if the split actually separates two target layers. - assert ( - min(model.target_layer_ids) - < len(model._base_model.layers) // 2 - <= max(model.target_layer_ids) - ), model.target_layer_ids - self._shard_hidden_states(model, spare_device=torch.device("cuda:0")) input_ids = torch.tensor([[1, 2, 3, 4]], device="cuda:0") + # The fix is only exercised if the selected hidden states really sit on different devices. + with torch.no_grad(): + hidden = model._base_model(input_ids=input_ids, output_hidden_states=True).hidden_states + devices = {hidden[lid + 1].device for lid in model.target_layer_ids} + assert len(devices) > 1, (model.target_layer_ids, devices) + with torch.no_grad(): base_token, draft_tokens = model.pseudo_speculative_generate(input_ids, steps=2) diff --git a/tests/gpu/torch/utils/test_distributed_worker_pool.py b/tests/gpu/torch/utils/test_distributed_worker_pool.py new file mode 100644 index 00000000000..7ae3a966278 --- /dev/null +++ b/tests/gpu/torch/utils/test_distributed_worker_pool.py @@ -0,0 +1,81 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""A hung job must not leave a pool that stalls every later job.""" + +import queue +import time +from types import SimpleNamespace + +import pytest +import torch.distributed as dist +from _test_utils.torch.distributed.utils import DistributedWorkerPool, get_live_pool + + +def _hang(rank, size): + time.sleep(600) + + +def _rank_0_fails_and_rank_1_waits_for_it(rank, size): + if rank == 0: + raise ValueError("rank 0 failed") + dist.barrier() # never completes: its partner left + + +def test_a_hung_job_kills_the_pool_instead_of_stalling_later_runs(capfd): + pool = DistributedWorkerPool(world_size=1, backend="gloo", teardown_fn=None) + try: + pool.job_timeout = 1 + with pytest.raises(queue.Empty): + pool.run(_hang) + + assert pool.broken + assert not any(p.is_alive() for p in pool._processes) + # Fails at once rather than waiting out another timeout. + with pytest.raises(RuntimeError, match="killed after an earlier job hung"): + pool.run(_hang) + # The stuck worker's stack reached stderr, which is what makes the next CI hang diagnosable. + assert "_hang" in capfd.readouterr().err + finally: + pool.shutdown() + + +def test_a_failed_rank_does_not_make_its_peers_wait_out_the_job_timeout(): + pool = DistributedWorkerPool(world_size=2, backend="gloo", teardown_fn=None) + try: + pool.failure_grace = 2 + start = time.monotonic() + with pytest.raises(RuntimeError, match="rank 0 failed"): + pool.run(_rank_0_fails_and_rank_1_waits_for_it) + + assert time.monotonic() - start < 60 # not the 600 s job timeout + assert pool.broken + finally: + pool.shutdown() + + +def test_a_killed_pool_is_replaced_and_a_live_one_is_reused(): + pools = {} + made = [] + + def make(): + made.append(SimpleNamespace(broken=False)) + return made[-1] + + first = get_live_pool(pools, 2, make) + assert get_live_pool(pools, 2, make) is first + first.broken = True + second = get_live_pool(pools, 2, make) + assert second is not first and len(made) == 2 diff --git a/tests/gpu_megatron/conftest.py b/tests/gpu_megatron/conftest.py index 882b553d0da..5760b9fe22d 100644 --- a/tests/gpu_megatron/conftest.py +++ b/tests/gpu_megatron/conftest.py @@ -17,7 +17,7 @@ import pytest import torch from _test_utils.fs_utils import assert_unmodified_tree -from _test_utils.torch.distributed.utils import DistributedWorkerPool +from _test_utils.torch.distributed.utils import DistributedWorkerPool, get_live_pool from _test_utils.torch.transformers_models import get_tiny_tokenizer from megatron.core.parallel_state import destroy_model_parallel @@ -108,31 +108,31 @@ def _pool_cache(): def _get_pool(cache, world_size): - if world_size not in cache: - cache[world_size] = _make_pool(world_size) - return cache[world_size] + # A hung test kills its pool; a fresh one keeps a single hang from timing out every later + # test in the module (that cascade cost the nightly 20+ min). + return get_live_pool(cache, world_size, lambda: _make_pool(world_size)) -@pytest.fixture(scope="module") +@pytest.fixture def dist_workers(_pool_cache): - """Module-scoped pool with world_size=torch.cuda.device_count().""" + """Pool (cached per module) with world_size=torch.cuda.device_count().""" return _get_pool(_pool_cache, torch.cuda.device_count()) -@pytest.fixture(scope="module") +@pytest.fixture def dist_workers_size_1(_pool_cache): - """Module-scoped pool with world_size=1 for tests that require a single process.""" + """Pool (cached per module) with world_size=1 for tests that require a single process.""" return _get_pool(_pool_cache, 1) -@pytest.fixture(scope="module") +@pytest.fixture def dist_workers_size_2(_pool_cache): if torch.cuda.device_count() < 2: pytest.skip("Need at least 2 GPUs") return _get_pool(_pool_cache, 2) -@pytest.fixture(scope="module") +@pytest.fixture def dist_workers_size_4(_pool_cache): if torch.cuda.device_count() < 4: pytest.skip("Need at least 4 GPUs") diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 976d0e48206..c8ee193c624 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -31,17 +31,27 @@ @pytest.fixture(scope="session") -def cpu_dist_workers(): +def _cpu_pools(): + pools = {} + yield pools + for pool in pools.values(): + pool.shutdown() + + +@pytest.fixture +def cpu_dist_workers(_cpu_pools): """Two persistent gloo ranks shared by every multi-rank unit test. A fresh spawn per test re-imports torch and transformers in each child (~10-20s on CI), so the - workers start once. Jobs must leave the process group intact and undo any global state they set. + workers start once (a killed pool is rebuilt). Jobs must leave the process group intact and + undo any global state they set. """ if platform.system() == "Windows": pytest.skip("Multi-rank tests are not run on Windows") # Lazy: the onnx-only partial-install job has no torch. - from _test_utils.torch.distributed.utils import DistributedWorkerPool + from _test_utils.torch.distributed.utils import DistributedWorkerPool, get_live_pool - pool = DistributedWorkerPool(world_size=2, backend="gloo", teardown_fn=None) - yield pool - pool.shutdown() + # Replaced if an earlier test hung and killed it, so one hang fails one test. + return get_live_pool( + _cpu_pools, 2, lambda: DistributedWorkerPool(world_size=2, backend="gloo", teardown_fn=None) + ) From 5b480008bb1f372654f05b6bda627aaac50899da Mon Sep 17 00:00:00 2001 From: Ajinkya Rasane <131806219+ajrasane@users.noreply.github.com> Date: Fri, 16 Oct 2026 09:16:08 -0400 Subject: [PATCH 27/35] [6915426] Fix ONNX AutoTune benchmark initialization on TensorRT 11 (#2795) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Type of change: Bug fix ONNX AutoTune's TensorRT Python benchmark fails during initialization on TensorRT 11 because it accesses the removed `NetworkDefinitionCreationFlag.EXPLICIT_BATCH` enum. Create the network with `STRONGLY_TYPED` alone: explicit batching is already implicit in TensorRT 10 and 11, and strong typing remains enabled on both versions. Add a parametrized GPU regression test that constructs the real benchmark, verifies its network creation flags, and builds and measures tiny INT8 and FP8 Q/DQ MatMul candidates. Update the AutoTune guide's TensorRT Python example to use the same supported network flags, and document the integration-test fixture. The same constructor exists in ModelOpt 0.47, so this also includes a changelog entry. `BuilderFlag.DIRECT_IO` remains available in the tested versions and needs no change. ### Usage ```python from modelopt.onnx.quantization.autotune import TensorRTPyBenchmark benchmark = TensorRTPyBenchmark(warmup_runs=1, timing_runs=2) latency_ms = benchmark.run("model.onnx") ``` ### Testing - Before the fix, both new regression cases fail at the real constructor with the missing `EXPLICIT_BATCH` AttributeError. - TensorRT **11.1.0.106**, RTX 6000 Ada: `pytest tests/gpu/onnx/quantization/autotune/test_benchmark.py -q` — **36 passed**, including INT8/FP8 candidate builds and measurements, Python benchmark bytes/path/callable cases, and trtexec smoke tests. - TensorRT **10.13.3.9**, RTX 6000 Ada: `pytest tests/gpu/onnx/quantization/autotune/test_benchmark.py -k TestTensorRTPyBenchmark -q` — **5 passed**. - Re-ran the original minimal constructor reproduction on TensorRT 11: initialization succeeds. - Repository pre-commit checks on the changed files pass, including Ruff, mypy, and Bandit. - Verified the corrected guide's network creation on TensorRT 11.1.0.106 and 10.13.3.9; both initialize successfully with strong typing enabled. Full standalone AutoTune searches and quantization CLI accuracy evaluations were not run. Candidate measurements here validate the shared benchmark path, not model accuracy. ### Before your PR is "*Ready for review*" Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors) (e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(..., weights_only=False)`, `pickle`, etc.). - Is this change backward compatible?: ✅ - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: ✅ - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ - Did you get Claude approval on this PR?: N/A ### Additional Information [TensorRT 10-to-11 migration guide](https://docs.nvidia.com/deeplearning/tensorrt/latest/api/migration/tensorrt-10x-to-11x.html) ## Summary by CodeRabbit * **New Features** * Added imatrix calibration for GGML formats, including weighted calibration using per-column and per-expert statistics. Existing final codes remain unchanged, allowing calibration of weights skipped by GPTQ. * **Bug Fixes** * Fixed ONNX AutoTune benchmark initialization with TensorRT 11 while retaining strongly typed network support for TensorRT 10. * Fixed LayerNorm input quantization in presets and recipes that enable input quantizers. * **Documentation** * Updated the TensorRT Python API example to use strongly typed network configuration. * **Tests** * Added GPU integration coverage for INT8 and FP8 quantized MatMul benchmarks, including finite, positive latency checks. FP8 coverage is skipped on unsupported GPUs. --------- Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> Signed-off-by: Rohan Joshi --- CHANGELOG.rst | 1 + docs/source/guides/9_autotune.rst | 5 +-- .../onnx/quantization/autotune/benchmark.py | 3 +- .../quantization/autotune/test_benchmark.py | 42 +++++++++++++++++++ 4 files changed, 45 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 44221e73d49..52f6dc83d66 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -129,6 +129,7 @@ Changelog **Bug Fixes** +- Fix ONNX AutoTune benchmark initialization with TensorRT 11 while preserving strongly typed networks on TensorRT 10. - Fix unified HF export writing a weight name the source checkpoint does not use when a model declares a backwards-compatibility *load* rename as a reversible ``WeightRenaming``. Reversing such a rule rewrote the canonical name to a legacy spelling, and a deployment runtime keyed on the canonical name then failed to load the quantized checkpoint. Each exported name is now resolved against the names the source checkpoint actually carries, for both the exported tensors and the quantization config's module references. - Fix exports using the built-in Transformers NemotronH implementation retaining incompatible legacy remote-code config and model mappings. Re-export affected checkpoints before deployment. - Fix ``export_hf_checkpoint`` dropping Diffusers pipeline settings, such as Wan 2.2 ``boundary_ratio``, from ``model_index.json`` when ``components`` is given. Re-export affected checkpoints to a new directory, since component exports do not overwrite an existing ``model_index.json``. diff --git a/docs/source/guides/9_autotune.rst b/docs/source/guides/9_autotune.rst index 02f79395752..2e79e482692 100644 --- a/docs/source/guides/9_autotune.rst +++ b/docs/source/guides/9_autotune.rst @@ -569,10 +569,7 @@ Using TensorRT Python API # Create builder and logger logger = trt.Logger(trt.Logger.WARNING) builder = trt.Builder(logger) - network = builder.create_network( - 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH) - | 1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED) - ) + network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)) parser = trt.OnnxParser(network, logger) # Parse optimized ONNX model diff --git a/modelopt/onnx/quantization/autotune/benchmark.py b/modelopt/onnx/quantization/autotune/benchmark.py index ba5cf1142bf..6c62ee4e628 100644 --- a/modelopt/onnx/quantization/autotune/benchmark.py +++ b/modelopt/onnx/quantization/autotune/benchmark.py @@ -358,8 +358,7 @@ def __init__( trt.init_libnvinfer_plugins(self.trt_logger, "") self._plugin_registry = trt.get_plugin_registry() - self.network_flags = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH) - self.network_flags |= 1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED) + self.network_flags = 1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED) self._timing_cache = None self._load_timing_cache() self._shape_configs = {} diff --git a/tests/gpu/onnx/quantization/autotune/test_benchmark.py b/tests/gpu/onnx/quantization/autotune/test_benchmark.py index 925d45fffe6..012bf495ab6 100644 --- a/tests/gpu/onnx/quantization/autotune/test_benchmark.py +++ b/tests/gpu/onnx/quantization/autotune/test_benchmark.py @@ -31,9 +31,11 @@ from unittest.mock import MagicMock, patch import numpy as np +import onnx import pytest import torch from _test_utils.onnx.quantization.autotune.models import _create_simple_conv_onnx_model +from onnx import TensorProto, helper import modelopt.onnx.quantization.autotune.benchmark as bm from modelopt.onnx.quantization.autotune import TensorRTPyBenchmark, TrtExecBenchmark @@ -345,10 +347,50 @@ class TestTensorRTPyBenchmark: @pytest.fixture(autouse=True) def _require_tensorrt_and_torch_cuda(self): + """Skip benchmark integration tests when TensorRT or CUDA is unavailable.""" pytest.importorskip("tensorrt") if not bm.TORCH_CUDA_AVAILABLE: pytest.skip("CUDA not available") + @pytest.mark.parametrize("quant_type", [TensorProto.INT8, TensorProto.FLOAT8E4M3FN]) + def test_run_quantized_model(self, quant_type): + """Quantized candidates build and run with strongly typed networks.""" + if quant_type == TensorProto.FLOAT8E4M3FN and torch.cuda.get_device_capability() < (8, 9): + pytest.skip("FP8 requires SM89 or newer") + + nodes = [] + for name in ("input", "weight"): + nodes.extend( + [ + helper.make_node( + "QuantizeLinear", [name, "scale", "zero_point"], [f"{name}_q"] + ), + helper.make_node( + "DequantizeLinear", [f"{name}_q", "scale", "zero_point"], [f"{name}_dq"] + ), + ] + ) + nodes.append(helper.make_node("MatMul", ["input_dq", "weight_dq"], ["output"])) + graph = helper.make_graph( + nodes, + "quantized_matmul", + [helper.make_tensor_value_info("input", TensorProto.FLOAT, [16, 16])], + [helper.make_tensor_value_info("output", TensorProto.FLOAT, [16, 16])], + initializer=[ + helper.make_tensor("weight", TensorProto.FLOAT, [16, 16], [0.1] * 256), + helper.make_tensor("scale", TensorProto.FLOAT, [], [0.1]), + helper.make_tensor("zero_point", quant_type, [], [0]), + ], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 19)], ir_version=9) + onnx.checker.check_model(model) + + benchmark = TensorRTPyBenchmark(warmup_runs=1, timing_runs=2) + network = benchmark.builder.create_network(benchmark.network_flags) + assert network.get_flag(bm.trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED) + latency_ms = benchmark.run(model.SerializeToString()) + assert np.isfinite(latency_ms) and latency_ms > 0 + def test_run_with_bytes(self, simple_conv_model_bytes): """TensorRTPyBenchmark accepts model bytes and returns finite latency.""" benchmark = TensorRTPyBenchmark(warmup_runs=1, timing_runs=2) From 4b511684a6313723ad9e14a8014f4ed2edf810f5 Mon Sep 17 00:00:00 2001 From: Ajinkya Rasane <131806219+ajrasane@users.noreply.github.com> Date: Fri, 16 Oct 2026 12:46:07 -0400 Subject: [PATCH 28/35] [3/4] [6466805] Normalize canonical Dynamo quantized weight paths (#2684) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Type of change: new feature Part 3 of the four-PR Dynamo ONNX export stack replacing #2418. Adds a private canonical static-weight normalizer: initializer → optional constant reshape → format marker → optional restoring views/cast/transpose → one MatMul/Gemm weight input. It validates provenance and layout, materializes independent scales, and rejects shared weights, fanout, arbitrary views, or mismatched Q/DQ paths. Adds 484 source lines (1,175 total additions including tests). Final validation, metadata synchronization, and the packer integration hooks land together in part 4 so this slice stays below the source budget. No new helper behavior or changes to legacy packers are enabled here. ### Usage Internal prerequisite only. The normalizer is not a supported public export API; part 4 connects it to `get_onnx_bytes_and_metadata(..., dynamo_export=True)`. The normalized graph is an intermediate representation for the format packers and is not valid standalone ONNX. NVFP4 normalization removes the export-only Cast, so an fp16 Gemm can temporarily receive an fp32 weight. INT4 normalization removes restoring views and can leave blocked weights feeding MatMul directly. Part 4 runs the packers to restore the final layout and dtype, then performs final ONNX validation. ### Testing Ran the focused suite with CUDA hidden in the PyTorch 26.08 development container (PyTorch 2.14) on a source snapshot verified against the checkout: **93 passed**. ```bash CUDA_VISIBLE_DEVICES="" python -m pytest -o addopts= -p no:cacheprovider \ tests/unit/torch/quantization/test_tensor_quant_cpu.py \ tests/unit/torch/quantization/test_dynamo_onnx_export.py \ tests/unit/onnx/test_dynamo_adapter.py -q ``` Covers canonical paths; INT4 last-axis/full-block validation, prefix reshapes, restoring transpose and Cast rewiring; independent INT4/MXFP8 scales and packer name prefixes; NVFP4 Gemm transB toggling; unchanged FP8 attention Q/K/V and projected activations; unchanged INT8 convolution paths; and fail-closed static/dynamic weight-view errors. The original PR implementation fails 12 of the new axis/block and activation regression cases. Changed-file pre-commit checks and `git diff --check` passed. GPU and full multi-version CI were not rerun for this split. ### Before your PR is "*Ready for review*" Contributor guidelines: https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md. Commits are signed and signed off with `git commit -s -S`. Security Best Practices: https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors. - Is this change backward compatible?: ✅ - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A — neither added. - Did you write any new necessary tests?: ✅ - Did you update Changelog?: N/A — private prerequisite; part 4 contains the single user-facing feature entry. - Did you get Claude approval on this PR?: Pending — follow-up changes address the review; re-review is still required. ### Additional Information Merge order: capture → lowering → normalization → integration. This PR targets `ajrasane/dynamo-onnx-lowering` and must merge after part 2. Its tests call the normalizer without any dependency on later helper/packer changes. > 🤖 _Generated by Codex (AI agent)._ ### Replacement stack Merge in order; each PR targets the preceding branch. Retarget the next PR when its dependency merges. 1. [#2682: capture and CPU routing](https://github.com/NVIDIA/Model-Optimizer/pull/2682) 2. [#2683: private ONNXScript lowering](https://github.com/NVIDIA/Model-Optimizer/pull/2683) 3. [#2684: canonical weight normalization](https://github.com/NVIDIA/Model-Optimizer/pull/2684) 4. [#2685: helper and deployment integration](https://github.com/NVIDIA/Model-Optimizer/pull/2685) CI note: repository pull-request workflows target `main`, `release/*`, and `feature/*`. They do not automatically run for this stacked base. Full CI needs manual workflow dispatch or retargeting after the dependency merges; the independent local validation above passed. ## Summary by CodeRabbit * **New Features** * Added support for exporting quantized models through Dynamo for INT4, INT8, FP8, NVFP4, and MXFP8 formats. * Improved CPU quantization support, including per-tensor, per-axis, and block quantization. * **Bug Fixes** * Improved handling of quantized weights and scales during ONNX export, including supported reshapes and shared scales. * Preserved correct quantization behavior for FP8 attention activations and INT8 convolution weights. --------- Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> Signed-off-by: Rohan Joshi --- modelopt/onnx/export/_dynamo_adapter.py | 500 ++++++++++++++++ tests/unit/onnx/test_dynamo_adapter.py | 733 ++++++++++++++++++++++++ 2 files changed, 1233 insertions(+) create mode 100644 modelopt/onnx/export/_dynamo_adapter.py create mode 100644 tests/unit/onnx/test_dynamo_adapter.py diff --git a/modelopt/onnx/export/_dynamo_adapter.py b/modelopt/onnx/export/_dynamo_adapter.py new file mode 100644 index 00000000000..193ddc33d4e --- /dev/null +++ b/modelopt/onnx/export/_dynamo_adapter.py @@ -0,0 +1,500 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""Dynamo-only normalization and validation for quantized ONNX graphs.""" + +from collections.abc import Iterable + +import numpy as np +import onnx +from onnx import numpy_helper + +from modelopt.onnx.quantization.graph_indexing import ( + get_tensor_consumer_nodes, + get_tensor_producer_nodes, +) + +__all__ = [] + +_DYNAMO_INT4_NODE_PREFIX = "__modelopt_dynamo_int4__" +_DYNAMO_MXFP8_NODE_PREFIX = "__modelopt_dynamo_mxfp8__" + + +def _attribute(node: onnx.NodeProto, name: str): + attribute = next((item for item in node.attribute if item.name == name), None) + return onnx.helper.get_attribute_value(attribute) if attribute is not None else None + + +def _single_consumer(consumers, name: str, weight_name: str) -> onnx.NodeProto: + matches = consumers.get(name, []) + if len(matches) != 1: + raise NotImplementedError( + f"Unsupported Dynamo quantized weight topology for '{weight_name}': " + f"expected one consumer; found {len(matches)}." + ) + return matches[0] + + +def _constant_tensor(producers, name: str) -> onnx.TensorProto: + producer = producers.get(name) + if isinstance(producer, onnx.TensorProto): + return producer + if isinstance(producer, onnx.NodeProto) and producer.op_type == "Constant": + value = _attribute(producer, "value") + if isinstance(value, onnx.TensorProto): + return value + raise NotImplementedError("Dynamo ONNX export requires constant weight shapes and scales.") + + +def _replace_initializer(graph: onnx.GraphProto, tensor: onnx.TensorProto) -> None: + existing = next((item for item in graph.initializer if item.name == tensor.name), None) + if existing is None: + graph.initializer.append(tensor) + else: + existing.CopyFrom(tensor) + + +def _remove_nodes(graph: onnx.GraphProto, nodes: Iterable[onnx.NodeProto]) -> None: + output_names = {output for node in nodes for output in node.output} + for index in reversed(range(len(graph.node))): + if not output_names.isdisjoint(graph.node[index].output): + del graph.node[index] + + +def _path_after_marker( + consumers, value_name: str, weight_name: str +) -> tuple[list[onnx.NodeProto], onnx.NodeProto]: + path = [] + passthrough_order = {"Identity": 0, "Reshape": 1, "Cast": 1, "Transpose": 2} + seen = set() + previous_order = -1 + node = _single_consumer(consumers, value_name, weight_name) + while node.op_type in passthrough_order: + current_order = passthrough_order[node.op_type] + if node.op_type in seen or current_order < previous_order: + raise NotImplementedError( + f"Unsupported Dynamo quantized weight topology for '{weight_name}': " + f"unexpected {node.op_type} ordering." + ) + seen.add(node.op_type) + previous_order = current_order + path.append(node) + value_name = node.output[0] + node = _single_consumer(consumers, value_name, weight_name) + if node.op_type not in {"MatMul", "Gemm"} or len(node.input) < 2 or node.input[1] != value_name: + raise NotImplementedError( + f"Unsupported Dynamo quantized weight topology for '{weight_name}': " + "expected terminal MatMul/Gemm at input 1." + ) + return path, node + + +def _weight_source(graph, producers, consumers, marker: onnx.NodeProto): + initializers = {item.name: item for item in graph.initializer} + marker_input = marker.input[0] + producer = producers.get(marker_input) + reshape = ( + producer if isinstance(producer, onnx.NodeProto) and producer.op_type == "Reshape" else None + ) + weight_name = reshape.input[0] if reshape is not None else marker_input + if weight_name not in initializers: + raise NotImplementedError( + f"Unsupported Dynamo quantized weight '{marker_input}': " + "data input must be a static initializer." + ) + _single_consumer(consumers, weight_name, weight_name) + if reshape is not None: + _single_consumer(consumers, marker_input, weight_name) + _constant_tensor(producers, reshape.input[1]) + return initializers[weight_name], reshape + + +def _is_static_value(producers, name: str, cache: dict[str, bool]) -> bool: + stack: list[tuple[str, list[str] | None]] = [(name, None)] + visiting = set() + while stack: + current, inputs = stack.pop() + if current in cache: + continue + if inputs is not None: + cache[current] = all(cache[input_name] for input_name in inputs) + visiting.remove(current) + continue + producer = producers.get(current) + if isinstance(producer, onnx.TensorProto): + cache[current] = True + elif not isinstance(producer, onnx.NodeProto): + cache[current] = False + elif producer.op_type == "Constant": + cache[current] = True + elif current in visiting or not any(producer.input): + cache[current] = False + else: + inputs = [input_name for input_name in producer.input if input_name] + visiting.add(current) + stack.append((current, inputs)) + stack.extend((input_name, None) for input_name in reversed(inputs)) + return cache[name] + + +def _has_static_weight_source(producers, name: str, cache: dict[str, bool]) -> bool: + producer = producers.get(name) + if isinstance(producer, onnx.TensorProto): + return True + if ( + isinstance(producer, onnx.NodeProto) + and producer.op_type == "Reshape" + and isinstance(producers.get(producer.input[0]), onnx.TensorProto) + ): + return True + # Shape and index inputs can be dynamic; only input 0 carries the weight data. + data_name = name + seen = set() + while ( + isinstance(producer, onnx.NodeProto) + and producer.op_type + in { + "Identity", + "Reshape", + "Cast", + "Transpose", + "Slice", + "Gather", + "GatherElements", + "GatherND", + "Expand", + "Squeeze", + "Unsqueeze", + "Flatten", + } + and producer.input[0] not in seen + ): + data_name = producer.input[0] + seen.add(data_name) + producer = producers.get(data_name) + if _is_static_value(producers, data_name, cache): + raise NotImplementedError( + f"Unsupported Dynamo quantized weight '{name}': unsupported static or dynamic transformation " + "before the quantization marker." + ) + return False + + +def _validate_restoring_reshape( + path, weight: onnx.TensorProto, producers, format_name: str +) -> None: + reshape = next((node for node in path if node.op_type == "Reshape"), None) + if reshape is None: + return + restored_shape = [ + int(dim) for dim in numpy_helper.to_array(_constant_tensor(producers, reshape.input[1])) + ] + if restored_shape != list(weight.dims): + raise NotImplementedError( + f"Unsupported Dynamo {format_name} weight '{weight.name}': restoring Reshape does " + "not match the source initializer shape." + ) + + +def _fold_prefix_reshape( + graph: onnx.GraphProto, + marker: onnx.NodeProto, + weight: onnx.TensorProto, + reshape: onnx.NodeProto | None, + producers, +) -> None: + if reshape is None: + return + shape = [ + int(dim) for dim in numpy_helper.to_array(_constant_tensor(producers, reshape.input[1])) + ] + values = numpy_helper.to_array(weight).reshape(shape) + _replace_initializer(graph, numpy_helper.from_array(values, weight.name)) + marker.input[0] = weight.name + _remove_nodes(graph, [reshape]) + + +def _split_scale( + graph, producers, consumers, marker: onnx.NodeProto, *, materialize_initializer: bool +) -> None: + if len(marker.input) < 2: + return + scale_name = marker.input[1] + tensor = _constant_tensor(producers, scale_name) + if ( + materialize_initializer + and len(consumers.get(scale_name, [])) == 1 + and isinstance(producers.get(scale_name), onnx.TensorProto) + ): + return + array = np.asarray(numpy_helper.to_array(tensor)) + unique_name = f"{scale_name}_{marker.output[0]}" + marker.input[1] = unique_name + scale_tensor = numpy_helper.from_array(array, unique_name) + if materialize_initializer: + _replace_initializer(graph, scale_tensor) + else: + marker_index = next( + index for index, node in enumerate(graph.node) if node.output[0] == marker.output[0] + ) + graph.node.insert( + marker_index, + onnx.helper.make_node( + "Constant", [], [unique_name], name=f"{unique_name}_constant", value=scale_tensor + ), + ) + + +def _remove_identity(graph, path: list[onnx.NodeProto], marker_output: str) -> list[onnx.NodeProto]: + if not path or path[0].op_type != "Identity": + return path + identity = path.pop(0) + replacement = identity.output[0] + for node in graph.node: + for index, input_name in enumerate(node.input): + if input_name == replacement: + node.input[index] = marker_output + _remove_nodes(graph, [identity]) + return path + + +def _normalize_int4_path( + graph: onnx.GraphProto, + marker: onnx.NodeProto, + weight: onnx.TensorProto, + prefix_reshape: onnx.NodeProto | None, + path: list[onnx.NodeProto], + terminal: onnx.NodeProto, + producers, +) -> None: + original_shape = list(weight.dims) + block_size = int(_attribute(marker, "block_size") or 0) + if block_size <= 0 or not original_shape or original_shape[-1] % block_size: + raise NotImplementedError( + f"Unsupported Dynamo INT4 weight '{weight.name}': invalid block size {block_size}." + ) + + if prefix_reshape is not None: + blocked_shape = [ + int(dim) + for dim in numpy_helper.to_array(_constant_tensor(producers, prefix_reshape.input[1])) + ] + if not blocked_shape or blocked_shape[-1] != block_size: + raise NotImplementedError( + f"Unsupported Dynamo INT4 weight '{weight.name}': pre-Reshape must use " + f"block size {block_size}." + ) + else: + blocked_shape = [-1, block_size] + marker_rank = len(blocked_shape if prefix_reshape is not None else original_shape) + marker_axis = _attribute(marker, "axis") + if (1 if marker_axis is None else marker_axis) not in {-1, marker_rank - 1}: + raise NotImplementedError( + f"Unsupported Dynamo INT4 weight '{weight.name}': only last-axis blocks are supported." + ) + _validate_restoring_reshape(path, weight, producers, "INT4") + blocked_weight = numpy_helper.to_array(weight).reshape(blocked_shape) + _replace_initializer(graph, numpy_helper.from_array(blocked_weight, weight.name)) + marker.input[0] = weight.name + + path = _remove_identity(graph, path, marker.output[0]) + reshape = next((node for node in path if node.op_type == "Reshape"), None) + + block_axis = len(original_shape) - 1 + marker.attribute.append(onnx.helper.make_attribute("_target_shape", original_shape)) + transpose = next((node for node in path if node.op_type == "Transpose"), None) + if transpose is not None: + perm = list(_attribute(transpose, "perm") or []) + if sorted(perm) != list(range(len(original_shape))): + raise NotImplementedError( + f"Unsupported Dynamo INT4 weight '{weight.name}': invalid transpose permutation." + ) + marker.attribute.append(onnx.helper.make_attribute("_transpose_perm", perm)) + block_axis = perm.index(block_axis) + axis = next((attr for attr in marker.attribute if attr.name == "axis"), None) + if axis is None: + marker.attribute.append(onnx.helper.make_attribute("axis", block_axis)) + else: + axis.i = block_axis + + cast = next((node for node in path if node.op_type == "Cast"), None) + if cast is not None: + cast.input[0] = marker.output[0] + weight_output = cast.output[0] + else: + weight_output = marker.output[0] + terminal.input[1] = weight_output + _remove_nodes( + graph, + [node for node in (prefix_reshape, reshape, transpose) if node is not None], + ) + + scale = next(item for item in graph.initializer if item.name == marker.input[1]) + scale_values = numpy_helper.to_array(scale).reshape(-1, 1) + _replace_initializer(graph, numpy_helper.from_array(scale_values, scale.name)) + + +def _normalize_nvfp4_path( + graph: onnx.GraphProto, + marker: onnx.NodeProto, + weight: onnx.TensorProto, + prefix_reshape: onnx.NodeProto | None, + path: list[onnx.NodeProto], + terminal: onnx.NodeProto, + producers, +) -> None: + path = _remove_identity(graph, path, marker.output[0]) + removable = [] + prefix_nodes = [] + reshape = next((node for node in path if node.op_type == "Reshape"), None) + if prefix_reshape is None and reshape is not None: + raise NotImplementedError( + f"Unsupported Dynamo NVFP4 weight '{weight.name}': restoring Reshape requires a " + "pre-Reshape." + ) + if prefix_reshape is not None: + if reshape is None: + raise NotImplementedError( + f"Unsupported Dynamo NVFP4 weight '{weight.name}': missing restoring Reshape." + ) + _validate_restoring_reshape(path, weight, producers, "NVFP4") + marker.input[0] = weight.name + prefix_nodes.append(prefix_reshape) + removable.append(reshape) + path.remove(reshape) + cast = next((node for node in path if node.op_type == "Cast"), None) + if cast is not None: + removable.append(cast) + path.remove(cast) + + for node in removable: + output = node.output[0] + for candidate in graph.node: + for index, input_name in enumerate(candidate.input): + if input_name == output: + candidate.input[index] = marker.output[0] + _remove_nodes(graph, [*prefix_nodes, *removable]) + + if path and path[0].op_type == "Transpose" and terminal.op_type == "Gemm": + transpose = path[0] + if list(_attribute(transpose, "perm") or []) != [1, 0]: + raise NotImplementedError( + f"Unsupported Dynamo NVFP4 weight '{weight.name}': only a rank-2 transpose is supported." + ) + trans_b = next((item for item in terminal.attribute if item.name == "transB"), None) + if trans_b is None: + terminal.attribute.append(onnx.helper.make_attribute("transB", 1)) + else: + trans_b.i = 1 - trans_b.i + terminal.input[1] = marker.output[0] + _remove_nodes(graph, [transpose]) + + +def _validate_marker_domain(marker: onnx.NodeProto, expected_domain: str) -> None: + if marker.domain != expected_domain: + raise NotImplementedError( + f"Unsupported Dynamo quantization marker '{marker.op_type}': unexpected domain " + f"'{marker.domain}'." + ) + + +def normalize_dynamo_weight_paths(model: onnx.ModelProto) -> onnx.ModelProto: + """Normalize supported Dynamo static-weight paths for the legacy format packers.""" + graph = model.graph + producers = get_tensor_producer_nodes(graph, get_initializer_producers=True) + consumers = get_tensor_consumer_nodes(graph) + static_values: dict[str, bool] = {} + + for marker in list(graph.node): + kind = None + marker_output = marker.output[0] + if marker.op_type == "TRT_FP4QDQ": + _validate_marker_domain(marker, "trt") + kind = "nvfp4" + elif marker.op_type == "TRT_MXFP8DequantizeLinear" and _has_static_weight_source( + producers, marker.input[0], static_values + ): + _validate_marker_domain(marker, "trt") + kind = "mxfp8" + elif marker.op_type == "DequantizeLinear" and (_attribute(marker, "block_size") or 0) > 0: + producer = producers.get(marker.input[0]) + if ( + not isinstance(producer, onnx.NodeProto) + or producer.op_type != "TRT_FP4DynamicQuantize" + ): + if marker.domain != "trt": + _validate_marker_domain(marker, "") + kind = "int4" + elif marker.op_type in { + "TRT_FP8QuantizeLinear", + "QuantizeLinear", + } and _has_static_weight_source(producers, marker.input[0], static_values): + expected_domain = "trt" if marker.op_type == "TRT_FP8QuantizeLinear" else "" + _validate_marker_domain(marker, expected_domain) + dq = _single_consumer(consumers, marker.output[0], marker.input[0]) + expected_dq = ( + "TRT_FP8DequantizeLinear" + if marker.op_type == "TRT_FP8QuantizeLinear" + else "DequantizeLinear" + ) + if dq.op_type != expected_dq or dq.domain != marker.domain: + raise NotImplementedError( + f"Unsupported Dynamo quantized weight '{marker.input[0]}': expected a matching " + f"{expected_dq}." + ) + if list(dq.input[1:]) != list(marker.input[1:]): + raise NotImplementedError( + f"Unsupported Dynamo quantized weight '{marker.input[0]}': " + "Q/DQ scale and zero-point inputs must match." + ) + kind = "fp8" if marker.op_type.startswith("TRT_") else "int8" + marker_output = dq.output[0] + if kind is None: + continue + + if kind in {"fp8", "int8"}: + direct_consumers = consumers.get(marker_output, []) + if ( + len(direct_consumers) == 1 + and direct_consumers[0].op_type == "Conv" + and len(direct_consumers[0].input) > 1 + and direct_consumers[0].input[1] == marker_output + ): + continue + + weight, prefix_reshape = _weight_source(graph, producers, consumers, marker) + path, terminal = _path_after_marker(consumers, marker_output, weight.name) + if kind in {"fp8", "int8", "mxfp8"}: + _validate_restoring_reshape(path, weight, producers, kind.upper()) + _fold_prefix_reshape(graph, marker, weight, prefix_reshape, producers) + if kind in {"int4", "mxfp8"}: + _split_scale( + graph, + producers, + consumers, + marker, + materialize_initializer=kind == "int4", + ) + if kind == "int4": + if not marker.name.startswith(_DYNAMO_INT4_NODE_PREFIX): + marker.name = f"{_DYNAMO_INT4_NODE_PREFIX}{marker.name}" + _normalize_int4_path(graph, marker, weight, prefix_reshape, path, terminal, producers) + marker.domain = "" + elif kind == "mxfp8" and not marker.name.startswith(_DYNAMO_MXFP8_NODE_PREFIX): + marker.name = f"{_DYNAMO_MXFP8_NODE_PREFIX}{marker.name}" + elif kind == "nvfp4": + _normalize_nvfp4_path(graph, marker, weight, prefix_reshape, path, terminal, producers) + + return model diff --git a/tests/unit/onnx/test_dynamo_adapter.py b/tests/unit/onnx/test_dynamo_adapter.py new file mode 100644 index 00000000000..98bf9df004c --- /dev/null +++ b/tests/unit/onnx/test_dynamo_adapter.py @@ -0,0 +1,733 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""Behavior-level tests for the private Dynamo ONNX graph adapter.""" + +import numpy as np +import pytest +from onnx import TensorProto, helper, numpy_helper + +from modelopt.onnx.export._dynamo_adapter import normalize_dynamo_weight_paths + + +def _tensor(name, values, dtype=np.float32): + return numpy_helper.from_array(np.asarray(values, dtype=dtype), name) + + +def _model( + nodes, + *, + initializers=(), + inputs=(), + outputs=(), + value_info=(), + opset=23, + trt_opset=None, +): + graph = helper.make_graph( + nodes, + "dynamo_adapter", + list(inputs), + list(outputs), + list(initializers), + value_info=list(value_info), + ) + imports = [helper.make_opsetid("", opset)] + if trt_opset is not None: + imports.append(helper.make_opsetid("trt", trt_opset)) + return helper.make_model(graph, opset_imports=imports) + + +def _value(name, dtype=TensorProto.FLOAT, shape=(2, 4)): + return helper.make_tensor_value_info(name, dtype, list(shape)) + + +def _nvfp4_marker(weight="weight", output="weight_dq"): + return helper.make_node("TRT_FP4QDQ", [weight], [output], domain="trt", block_size=16) + + +@pytest.mark.parametrize("static_source", [False, True]) +def test_normalize_handles_deep_shared_provenance(static_source): + nodes = [] + source = "source" + for index in range(2000): + output = f"value_{index}" + nodes.append(helper.make_node("Add" if index % 2 else "Mul", [source, source], [output])) + source = output + for index in range(2): + nodes.extend( + [ + helper.make_node("QuantizeLinear", [source, "scale", "zero"], [f"q_{index}"]), + helper.make_node( + "DequantizeLinear", [f"q_{index}", "scale", "zero"], [f"dq_{index}"] + ), + helper.make_node("MatMul", ["other", f"dq_{index}"], [f"output_{index}"]), + ] + ) + model = _model( + nodes, + initializers=[ + _tensor("scale", 0.25), + _tensor("zero", 0, np.uint8), + *([_tensor("source", np.ones((4, 4)))] if static_source else []), + ], + inputs=[_value("other"), *([] if static_source else [_value("source", shape=(4, 4))])], + outputs=[_value(f"output_{index}") for index in range(2)], + ) + original = model.SerializeToString() + + if static_source: + with pytest.raises(NotImplementedError, match="before the quantization marker"): + normalize_dynamo_weight_paths(model) + else: + assert normalize_dynamo_weight_paths(model) is model + assert model.SerializeToString() == original + + +@pytest.mark.parametrize( + ("quantize_op", "dequantize_op", "domain"), + [ + ("TRT_FP8QuantizeLinear", "TRT_FP8DequantizeLinear", "trt"), + ("QuantizeLinear", "DequantizeLinear", ""), + ], +) +def test_normalize_accepts_canonical_qdq_weight_path(quantize_op, dequantize_op, domain): + weight = _tensor("weight", np.ones((4, 4))) + blocked_shape = _tensor("blocked_shape", [2, 2, 4], np.int64) + restored_shape = _tensor("restored_shape", [4, 4], np.int64) + scale = _tensor("scale", 0.25) + zero = _tensor("zero", 0, np.uint8) + model = _model( + [ + helper.make_node("Reshape", ["weight", "blocked_shape"], ["blocked"]), + helper.make_node( + quantize_op, ["blocked", "scale", "zero"], ["weight_q"], domain=domain + ), + helper.make_node( + dequantize_op, ["weight_q", "scale", "zero"], ["weight_dq"], domain=domain + ), + helper.make_node("Reshape", ["weight_dq", "restored_shape"], ["restored"]), + helper.make_node("MatMul", ["input", "restored"], ["output"]), + ], + initializers=[weight, blocked_shape, restored_shape, scale, zero], + inputs=[_value("input", shape=(2, 4))], + outputs=[_value("output", shape=(2, 4))], + trt_opset=1 if domain else None, + ) + + assert normalize_dynamo_weight_paths(model) is model + assert [node.op_type for node in model.graph.node] == [ + quantize_op, + dequantize_op, + "Reshape", + "MatMul", + ] + quantizer = model.graph.node[0] + assert quantizer.input[0] == "weight" + assert list(next(item for item in model.graph.initializer if item.name == "weight").dims) == [ + 2, + 2, + 4, + ] + + +def test_normalize_rejects_mismatched_qdq_parameters(): + model = _model( + [ + helper.make_node("QuantizeLinear", ["weight", "scale", "zero"], ["weight_q"]), + helper.make_node( + "DequantizeLinear", ["weight_q", "other_scale", "zero"], ["weight_dq"] + ), + helper.make_node("MatMul", ["input", "weight_dq"], ["output"]), + ], + initializers=[ + _tensor("weight", np.ones((4, 4))), + _tensor("scale", 0.25), + _tensor("other_scale", 0.5), + _tensor("zero", 0, np.uint8), + ], + inputs=[_value("input")], + outputs=[_value("output")], + ) + + with pytest.raises(NotImplementedError, match="scale and zero-point inputs must match"): + normalize_dynamo_weight_paths(model) + + +@pytest.mark.parametrize( + ("view_order", "trans_b"), + [("reshape_cast", None), ("cast_reshape", 0), ("cast_reshape", 1)], +) +def test_normalize_canonical_nvfp4_path_removes_export_only_views(view_order, trans_b): + weight = _tensor("weight", np.ones((2, 4))) + blocked_shape = _tensor("blocked_shape", [1, 2, 4], np.int64) + restored_shape = _tensor("restored_shape", [2, 4], np.int64) + views = { + "reshape_cast": [ + helper.make_node("Reshape", ["identity", "restored_shape"], ["restored"]), + helper.make_node("Cast", ["restored"], ["cast"], to=TensorProto.FLOAT16), + ], + "cast_reshape": [ + helper.make_node("Cast", ["identity"], ["cast"], to=TensorProto.FLOAT16), + helper.make_node("Reshape", ["cast", "restored_shape"], ["restored"]), + ], + }[view_order] + view_output = views[-1].output[0] + model = _model( + [ + helper.make_node("Reshape", ["weight", "blocked_shape"], ["blocked"]), + _nvfp4_marker("blocked", "marked"), + helper.make_node("Identity", ["marked"], ["identity"]), + *views, + helper.make_node("Transpose", [view_output], ["transposed"], perm=[1, 0]), + helper.make_node( + "Gemm", + ["input", "transposed"], + ["output"], + **({} if trans_b is None else {"transB": trans_b}), + ), + ], + initializers=[weight, blocked_shape, restored_shape], + inputs=[_value("input", TensorProto.FLOAT16, (1, 2 if trans_b else 4))], + outputs=[_value("output", TensorProto.FLOAT16, (1, 4 if trans_b else 2))], + trt_opset=1, + ) + + normalize_dynamo_weight_paths(model) + + assert [node.op_type for node in model.graph.node] == ["TRT_FP4QDQ", "Gemm"] + marker, gemm = model.graph.node + assert marker.input[0] == "weight" + assert gemm.input[1] == marker.output[0] + assert helper.get_attribute_value( + next(attr for attr in gemm.attribute if attr.name == "transB") + ) == 1 - (trans_b or 0) + + +@pytest.mark.parametrize( + ("view_op", "weight_source"), + [ + ("Slice", "initializer"), + ("Transpose", "initializer"), + ("Expand", "initializer"), + ("Gather", "initializer"), + ("GatherND", "initializer"), + ("GatherND", "constant"), + ], +) +def test_normalize_rejects_static_view_before_quantizer(view_op, weight_source): + if view_op == "Slice": + view = helper.make_node("Slice", ["weight"], ["view"]) + elif view_op == "Transpose": + view = helper.make_node("Transpose", ["weight"], ["view"], perm=[1, 0]) + elif view_op == "Expand": + view = helper.make_node("Expand", ["weight", "expanded_shape"], ["view"]) + else: + view = helper.make_node(view_op, ["weight", "runtime_indices"], ["view"]) + weight = _tensor("weight", np.ones((4, 4))) + weight_nodes = ( + [helper.make_node("Constant", [], ["weight"], value=weight)] + if weight_source == "constant" + else [] + ) + inputs = [_value("input")] + if view_op in {"Gather", "GatherND"}: + shape = (4, 1) if view_op == "GatherND" else (4,) + inputs.append(_value("runtime_indices", TensorProto.INT64, shape)) + model = _model( + [ + *weight_nodes, + view, + helper.make_node("QuantizeLinear", ["view", "scale", "zero"], ["weight_q"]), + helper.make_node("DequantizeLinear", ["weight_q", "scale", "zero"], ["weight_dq"]), + helper.make_node("MatMul", ["input", "weight_dq"], ["output"]), + ], + initializers=[ + *([weight] if weight_source == "initializer" else []), + _tensor("scale", 0.25), + _tensor("zero", 0, np.uint8), + _tensor("expanded_shape", [4, 4], np.int64), + ], + inputs=inputs, + outputs=[_value("output")], + ) + + with pytest.raises(NotImplementedError, match="before the quantization marker"): + normalize_dynamo_weight_paths(model) + + +def test_normalize_rejects_static_quantizer_without_matching_dequantizer(): + model = _model( + [ + helper.make_node("QuantizeLinear", ["weight", "scale", "zero"], ["weight_q"]), + helper.make_node("Identity", ["weight_q"], ["weight_dq"]), + helper.make_node("MatMul", ["input", "weight_dq"], ["output"]), + ], + initializers=[ + _tensor("weight", np.ones((4, 4))), + _tensor("scale", 0.25), + _tensor("zero", 0, np.uint8), + ], + inputs=[_value("input")], + outputs=[_value("output")], + ) + + with pytest.raises(NotImplementedError, match="matching DequantizeLinear"): + normalize_dynamo_weight_paths(model) + + +def test_normalize_rejects_wrong_marker_domain(): + model = _model( + [ + helper.make_node( + "QuantizeLinear", ["weight", "scale", "zero"], ["weight_q"], domain="custom" + ), + helper.make_node( + "DequantizeLinear", + ["weight_q", "scale", "zero"], + ["weight_dq"], + domain="custom", + ), + helper.make_node("MatMul", ["input", "weight_dq"], ["output"]), + ], + initializers=[ + _tensor("weight", np.ones((4, 4))), + _tensor("scale", 0.25), + _tensor("zero", 0, np.uint8), + ], + inputs=[_value("input")], + outputs=[_value("output")], + ) + + with pytest.raises(NotImplementedError, match="unexpected domain"): + normalize_dynamo_weight_paths(model) + + +def test_normalize_rejects_mismatched_post_marker_reshape(): + model = _model( + [ + helper.make_node("QuantizeLinear", ["weight", "scale", "zero"], ["weight_q"]), + helper.make_node("DequantizeLinear", ["weight_q", "scale", "zero"], ["weight_dq"]), + helper.make_node("Reshape", ["weight_dq", "wrong_shape"], ["restored"]), + helper.make_node("MatMul", ["input", "restored"], ["output"]), + ], + initializers=[ + _tensor("weight", np.ones((4, 4))), + _tensor("scale", 0.25), + _tensor("zero", 0, np.uint8), + _tensor("wrong_shape", [2, 8], np.int64), + ], + inputs=[_value("input")], + outputs=[_value("output")], + ) + + with pytest.raises(NotImplementedError, match="restoring Reshape does not match"): + normalize_dynamo_weight_paths(model) + + +def test_normalize_materializes_mxfp8_prefix_reshape(): + model = _model( + [ + helper.make_node("Reshape", ["linear.weight", "blocked_shape"], ["blocked"]), + helper.make_node( + "TRT_MXFP8DequantizeLinear", + ["blocked", "scale"], + ["weight_dq"], + domain="trt", + axis=-1, + block_size=32, + output_dtype=TensorProto.FLOAT, + ), + helper.make_node("Reshape", ["weight_dq", "restored_shape"], ["restored"]), + helper.make_node("MatMul", ["input", "restored"], ["output"]), + ], + initializers=[ + _tensor("linear.weight", np.ones((2, 64))), + _tensor("blocked_shape", [4, 32], np.int64), + _tensor("restored_shape", [2, 64], np.int64), + _tensor("scale", 1.0), + ], + inputs=[_value("input", shape=(1, 2))], + outputs=[_value("output", shape=(1, 64))], + trt_opset=1, + ) + + normalize_dynamo_weight_paths(model) + + marker = next(node for node in model.graph.node if node.op_type.startswith("TRT_MXFP8")) + assert marker.input[0] == "linear.weight" + weight = next(item for item in model.graph.initializer if item.name == "linear.weight") + assert list(weight.dims) == [4, 32] + + +def _unsupported_nvfp4_model(kind): + weight = _tensor("weight", np.ones((2, 4))) + marker = _nvfp4_marker() + common = { + "initializers": [weight], + "inputs": [_value("input")], + "outputs": [_value("output")], + "trt_opset": 1, + } + if kind == "shared weight": + nodes = [ + marker, + helper.make_node("Identity", ["weight"], ["other_weight_use"]), + helper.make_node("MatMul", ["input", "weight_dq"], ["output"]), + ] + elif kind == "fanout": + nodes = [ + marker, + helper.make_node("MatMul", ["input", "weight_dq"], ["output"]), + helper.make_node("Identity", ["weight_dq"], ["other_output"]), + ] + elif kind == "Slice": + nodes = [ + marker, + helper.make_node("Slice", ["weight_dq"], ["view"]), + helper.make_node("MatMul", ["input", "view"], ["output"]), + ] + else: + nodes = [marker, helper.make_node("Relu", ["weight_dq"], ["output"])] + return _model(nodes, **common) + + +@pytest.mark.parametrize("kind", ["shared weight", "fanout", "Slice", "non-compute"]) +def test_normalize_rejects_unsupported_weight_topologies(kind): + with pytest.raises(NotImplementedError, match="Unsupported Dynamo quantized weight topology"): + normalize_dynamo_weight_paths(_unsupported_nvfp4_model(kind)) + + +@pytest.mark.parametrize("domain", ["", "trt"]) +@pytest.mark.parametrize( + ("prefix_shape", "transpose", "axis"), + [ + (None, False, -1), + (None, False, 1), + (None, False, None), + ((2, 2, 2), False, 2), + ((2, 2, 2), True, -1), + ], +) +def test_normalize_int4_path_preserves_values_and_layout(prefix_shape, transpose, axis, domain): + values = np.arange(-4, 4, dtype=np.float32).reshape(2, 4) + scale_shape = (*prefix_shape[:-1], 1) if prefix_shape else (2, 2) + initializers = [_tensor("weight", values), _tensor("scale", np.ones(scale_shape))] + nodes = [] + if prefix_shape: + initializers.extend( + [ + _tensor("blocked_shape", prefix_shape, np.int64), + _tensor("restored_shape", values.shape, np.int64), + ] + ) + nodes.append(helper.make_node("Reshape", ["weight", "blocked_shape"], ["blocked"])) + nodes.append( + helper.make_node( + "DequantizeLinear", + ["blocked" if prefix_shape else "weight", "scale"], + ["weight_dq"], + name="weight_dequantizer", + domain=domain, + block_size=2, + **({} if axis is None else {"axis": axis}), + ) + ) + output = "weight_dq" + if prefix_shape: + nodes.extend( + [ + helper.make_node("Reshape", [output, "restored_shape"], ["restored"]), + helper.make_node("Cast", ["restored"], ["cast"], to=TensorProto.FLOAT16), + ] + ) + output = "cast" + if transpose: + nodes.append(helper.make_node("Transpose", [output], ["transposed"], perm=[1, 0])) + output = "transposed" + nodes.append(helper.make_node("MatMul", ["input", output], ["output"])) + dtype = TensorProto.FLOAT16 if prefix_shape else TensorProto.FLOAT + model = _model( + nodes, + initializers=initializers, + inputs=[_value("input", dtype, (1, 4 if transpose else 2))], + outputs=[_value("output", dtype, (1, 2 if transpose else 4))], + trt_opset=1 if domain else None, + ) + + normalize_dynamo_weight_paths(model) + + assert [node.op_type for node in model.graph.node] == [ + "DequantizeLinear", + *(["Cast"] if prefix_shape else []), + "MatMul", + ] + marker = model.graph.node[0] + assert marker.name == "__modelopt_dynamo_int4__weight_dequantizer" + assert marker.domain == "" + assert marker.input[0] == "weight" + attributes = {attr.name: helper.get_attribute_value(attr) for attr in marker.attribute} + assert attributes["_target_shape"] == [2, 4] + assert attributes["axis"] == (0 if transpose else 1) + if transpose: + assert attributes["_transpose_perm"] == [1, 0] + else: + assert "_transpose_perm" not in attributes + tensors = {tensor.name: tensor for tensor in model.graph.initializer} + np.testing.assert_array_equal( + numpy_helper.to_array(tensors["weight"]), values.reshape(prefix_shape or (4, 2)) + ) + np.testing.assert_array_equal(numpy_helper.to_array(tensors[marker.input[1]]), np.ones((4, 1))) + if prefix_shape: + assert model.graph.node[1].input[0] == marker.output[0] + assert model.graph.node[-1].input[1] == ("cast" if prefix_shape else marker.output[0]) + + +@pytest.mark.parametrize( + ("weight_shape", "prefix_shape", "axis", "message"), + [ + ((2, 4), None, 0, "axis"), + ((2, 4), (2, 2, 2), 0, "axis"), + ((2, 4), (2, 2, 2), 1, "axis"), + ((2, 4), (2, 2, 2), None, "axis"), + ((2, 3), None, -1, "block size"), + ((2, 3), (3, 2), -1, "block size"), + ], +) +def test_normalize_rejects_unsupported_int4_blocks(weight_shape, prefix_shape, axis, message): + initializers = [_tensor("weight", np.ones(weight_shape)), _tensor("scale", 1.0)] + nodes = [] + if prefix_shape: + initializers.append(_tensor("blocked_shape", prefix_shape, np.int64)) + nodes.append(helper.make_node("Reshape", ["weight", "blocked_shape"], ["blocked"])) + nodes.extend( + [ + helper.make_node( + "DequantizeLinear", + ["blocked" if prefix_shape else "weight", "scale"], + ["weight_dq"], + block_size=2, + **({} if axis is None else {"axis": axis}), + ), + helper.make_node("MatMul", ["input", "weight_dq"], ["output"]), + ] + ) + model = _model( + nodes, + initializers=initializers, + inputs=[_value("input", shape=(1, weight_shape[0]))], + outputs=[_value("output", shape=(1, weight_shape[1]))], + ) + + with pytest.raises(NotImplementedError, match=message): + normalize_dynamo_weight_paths(model) + + +@pytest.mark.parametrize("format_name", ["int4", "nvfp4"]) +def test_normalize_rejects_unmatched_restoring_reshape(format_name): + weight = _tensor("weight", np.ones((2, 4))) + scale = _tensor("scale", np.ones((1, 2))) + blocked_shape = _tensor("blocked_shape", [2, 2, 2], np.int64) + wrong_shape = _tensor("wrong_shape", [4, 2], np.int64) + model = _model( + [ + helper.make_node("Reshape", ["weight", "blocked_shape"], ["blocked"]), + ( + helper.make_node( + "DequantizeLinear", + ["blocked", "scale"], + ["weight_dq"], + axis=-1, + block_size=2, + ) + if format_name == "int4" + else _nvfp4_marker("blocked") + ), + helper.make_node("Reshape", ["weight_dq", "wrong_shape"], ["restored"]), + helper.make_node("MatMul", ["input", "restored"], ["output"]), + ], + initializers=[weight, scale, blocked_shape, wrong_shape], + inputs=[_value("input")], + outputs=[_value("output")], + trt_opset=1 if format_name == "nvfp4" else None, + ) + + with pytest.raises(NotImplementedError, match="restoring Reshape does not match"): + normalize_dynamo_weight_paths(model) + + +def test_normalize_rejects_unpaired_nvfp4_restoring_reshape(): + model = _model( + [ + _nvfp4_marker(), + helper.make_node("Reshape", ["weight_dq", "view_shape"], ["view"]), + helper.make_node("MatMul", ["input", "view"], ["output"]), + ], + initializers=[ + _tensor("weight", np.ones((2, 4))), + _tensor("view_shape", [4, 2], np.int64), + ], + inputs=[_value("input")], + outputs=[_value("output")], + trt_opset=1, + ) + + with pytest.raises(NotImplementedError, match="restoring Reshape requires a pre-Reshape"): + normalize_dynamo_weight_paths(model) + + +@pytest.mark.parametrize("format_name", ["int4", "mxfp8"]) +def test_normalize_splits_shared_scale_without_mutating_source(format_name): + marker_op = "DequantizeLinear" if format_name == "int4" else "TRT_MXFP8DequantizeLinear" + source_scale = np.asarray([[0.25, 0.5], [0.75, 1.0]], dtype=np.float32) + initializers = [ + _tensor("weight_a", np.ones((2, 64))), + _tensor("weight_b", np.ones((2, 64))), + _tensor("shared_scale", source_scale), + ] + nodes = [] + for suffix in ("a", "b"): + nodes.extend( + [ + helper.make_node( + marker_op, + [f"weight_{suffix}", "shared_scale"], + [f"weight_{suffix}_dq"], + name=f"dequantize_{suffix}", + domain="" if format_name == "int4" else "trt", + axis=-1, + block_size=32, + ), + helper.make_node( + "MatMul", [f"input_{suffix}", f"weight_{suffix}_dq"], [f"output_{suffix}"] + ), + ] + ) + model = _model( + nodes, + initializers=initializers, + inputs=[_value("input_a", shape=(1, 2)), _value("input_b", shape=(1, 2))], + outputs=[_value("output_a", shape=(1, 64)), _value("output_b", shape=(1, 64))], + trt_opset=1 if format_name == "mxfp8" else None, + ) + + normalize_dynamo_weight_paths(model) + + dequantizers = [node for node in model.graph.node if node.op_type == marker_op] + assert [node.name for node in dequantizers] == [ + f"__modelopt_dynamo_{format_name}__dequantize_{suffix}" for suffix in ("a", "b") + ] + split_names = [node.input[1] for node in dequantizers] + assert len(set(split_names)) == 2 + assert "shared_scale" not in split_names + tensors = {tensor.name: tensor for tensor in model.graph.initializer} + np.testing.assert_array_equal(numpy_helper.to_array(tensors["shared_scale"]), source_scale) + constants = { + node.output[0]: helper.get_attribute_value(node.attribute[0]) + for node in model.graph.node + if node.op_type == "Constant" + } + for name in split_names: + if format_name == "int4": + np.testing.assert_array_equal( + numpy_helper.to_array(tensors[name]), source_scale.reshape(-1, 1) + ) + else: + assert name not in tensors + np.testing.assert_array_equal(numpy_helper.to_array(constants[name]), source_scale) + + +@pytest.mark.parametrize("domain", ["", "trt"]) +@pytest.mark.parametrize("operand", ["query", "key", "value", "projected_key"]) +def test_normalize_leaves_fp8_attention_activation_qdq_intact(domain, operand): + transposed = operand in {"key", "projected_key"} + activation_shape = (2, 8, 4) if operand == "value" else (2, 4, 8) + other_shape = (2, 8, 4) if operand == "query" else (2, 4, 8) + initializers = [ + _tensor("scale", 0.25), + helper.make_tensor("zero", TensorProto.FLOAT8E4M3FN, [], [0]), + ] + nodes = [] + quantizer_input = "activation" + if operand == "projected_key": + initializers.append(_tensor("projection_weight", np.ones((8, 8)))) + nodes.extend( + [ + helper.make_node("MatMul", ["activation", "projection_weight"], ["projected"]), + helper.make_node("Identity", ["projected"], ["prepared"]), + ] + ) + quantizer_input = "prepared" + prefix = "TRT_FP8" if domain else "" + nodes.extend( + [ + helper.make_node( + f"{prefix}QuantizeLinear", + [quantizer_input, "scale", "zero"], + ["activation_q"], + domain=domain, + ), + helper.make_node( + f"{prefix}DequantizeLinear", + ["activation_q", "scale", "zero"], + ["activation_dq"], + domain=domain, + ), + ] + ) + activation_output = "activation_dq" + if transposed: + nodes.append( + helper.make_node("Transpose", [activation_output], ["transposed"], perm=[0, 2, 1]) + ) + activation_output = "transposed" + operands = [activation_output, "other"] if operand == "query" else ["other", activation_output] + nodes.append(helper.make_node("MatMul", operands, ["output"])) + model = _model( + nodes, + initializers=initializers, + inputs=[_value("activation", shape=activation_shape), _value("other", shape=other_shape)], + outputs=[_value("output", shape=(2, 4, 4))], + trt_opset=1 if domain else None, + ) + original = model.SerializeToString() + + assert normalize_dynamo_weight_paths(model) is model + assert model.SerializeToString() == original + + +def test_normalize_leaves_standard_int8_conv_weight_qdq_intact(): + model = _model( + [ + helper.make_node("QuantizeLinear", ["weight", "scale", "zero"], ["weight_q"]), + helper.make_node("DequantizeLinear", ["weight_q", "scale", "zero"], ["weight_dq"]), + helper.make_node("Conv", ["input", "weight_dq"], ["output"]), + ], + initializers=[ + _tensor("weight", np.ones((4, 3, 1, 1))), + _tensor("scale", np.ones(4)), + _tensor("zero", np.zeros(4), np.int8), + ], + inputs=[_value("input", shape=(1, 3, 2, 2))], + outputs=[_value("output", shape=(1, 4, 2, 2))], + ) + + normalize_dynamo_weight_paths(model) + + assert [node.op_type for node in model.graph.node] == [ + "QuantizeLinear", + "DequantizeLinear", + "Conv", + ] From 7ea7a972e7fe964b80eaeacab2d6678cdde647ae Mon Sep 17 00:00:00 2001 From: Ajinkya Rasane <131806219+ajrasane@users.noreply.github.com> Date: Fri, 16 Oct 2026 14:47:25 -0400 Subject: [PATCH 29/35] [6914448] Keep Qwen vision blocks in high precision for vLLM (#2822) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Type of change: Bug fix The Qwen3-VL and Qwen3.5 FP8 vision recipes previously quantized vision-block Linears whose 3D attention inputs are flattened by vLLM's FlashInfer FP8 GEMM, causing vision encoder initialization to fail. Narrow the shared visual quantizer rules to primary and deepstack merger Linears, whose inputs are 2D. Update both recipe descriptions, the recipe guide, the 0.48.0 changelog entry, and the quantizer-selection and checkpoint-export tests. This changes existing recipe output: vision blocks now remain in high precision and affected checkpoints must be regenerated. ### Usage ```python from modelopt.recipe import load_recipe recipe = load_recipe("model_type/qwen3_5/ptq/fp8_vision-kv_none") quant_cfg = recipe.quantize.model_dump() ``` ### Testing - Verified the four Qwen3-VL/Qwen3.5 recipe-selection cases fail with the previous quantizer rules and pass with this change (4 passed). - Ran the four GPU calibration and unified-HF-export cases for vision-only and joint FP8 recipes (4 passed). - Exported a Qwen3.5 vision-language checkpoint after four image-calibration samples on a GB300. The export marks vision blocks as excluded and retains FP8 merger scales. Unmodified vLLM 0.31.0 started successfully, returned HTTP 200 from `/health`, and answered an image prompt correctly. - Ran all applicable pre-commit hooks on the changed files, including recipe validation. ### Before your PR is "*Ready for review*" Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors) (e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(..., weights_only=False)`, `pickle`, etc.). - Is this change backward compatible?: ❌ — rerunning either existing recipe now leaves vision blocks in high precision; previously exported checkpoints are unchanged. - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: ✅ - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ ### Additional Information This avoids the failing 3D FP8 path by reducing vision quantization coverage. It does not change vLLM's FP8 kernel or repair already exported checkpoints. End-to-end serving used four calibration images; a full-size calibration run with the changed recipe was not repeated. ## Summary by CodeRabbit * **Bug Fixes** * Updated Qwen3-VL and Qwen3.5 FP8 vision recipes to quantize primary and deepstack vision-merger layers while keeping vision blocks and patch embeddings in high precision for vLLM compatibility. * Recipes that also enable language-model FP8 quantization retain that option while keeping vision blocks and patch embeddings in high precision. * Calibration checks now account for vision-merger input quantizers. * Re-quantize affected checkpoints before deployment to apply the updated settings. * **Documentation** * Clarified which model components are quantized by the FP8 vision recipes. --------- Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> Signed-off-by: Rohan Joshi --- CHANGELOG.rst | 1 + examples/hf_ptq/hf_ptq.py | 6 +++++- .../qwen3_5/ptq/fp8_vision-kv_none.yaml | 7 +++---- .../qwen3_5/ptq/fp8_vision_lm-kv_fp8_cast.yaml | 10 +++++----- .../qwen3_vl/ptq/fp8_vision-kv_none.yaml | 7 +++---- .../qwen3_vl/ptq/fp8_vision_lm-kv_fp8_cast.yaml | 10 +++++----- .../qwen3_vl/ptq/vision_fp8.quant_cfg.yaml | 10 ++++------ modelopt_recipes/ptq.md | 5 +++-- .../export/test_qwen_vision_recipe_export.py | 2 +- tests/unit/recipe/test_qwen_vision_recipe.py | 16 ++++------------ 10 files changed, 34 insertions(+), 40 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 52f6dc83d66..1cc5fc19482 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -130,6 +130,7 @@ Changelog **Bug Fixes** - Fix ONNX AutoTune benchmark initialization with TensorRT 11 while preserving strongly typed networks on TensorRT 10. +- Keep vision-block Linears in high precision in the ``qwen3_vl`` and ``qwen3_5`` FP8 vision recipes so exported checkpoints avoid a vLLM vision-attention shape failure. The recipes still quantize primary and deepstack merger Linears; re-quantize affected checkpoints before deployment. - Fix unified HF export writing a weight name the source checkpoint does not use when a model declares a backwards-compatibility *load* rename as a reversible ``WeightRenaming``. Reversing such a rule rewrote the canonical name to a legacy spelling, and a deployment runtime keyed on the canonical name then failed to load the quantized checkpoint. Each exported name is now resolved against the names the source checkpoint actually carries, for both the exported tensors and the quantization config's module references. - Fix exports using the built-in Transformers NemotronH implementation retaining incompatible legacy remote-code config and model mappings. Re-export affected checkpoints before deployment. - Fix ``export_hf_checkpoint`` dropping Diffusers pipeline settings, such as Wan 2.2 ``boundary_ratio``, from ``model_index.json`` when ``components`` is given. Re-export affected checkpoints to a new directory, since component exports do not overwrite an existing ``model_index.json``. diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index e75c7a4d70c..d7a14a68192 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -286,11 +286,15 @@ def _validate_recipe_calibration(args: argparse.Namespace, recipe) -> None: if not isinstance(recipe, ModelOptPTQRecipe) or args.calib_with_images: return + visual_input_names = { + "*visual.*input_quantizer", + "*visual.*merger*.input_quantizer", + } visual_input_entry = next( ( entry for entry in reversed(recipe.quantize.quant_cfg) - if entry.quantizer_name == "*visual.*input_quantizer" + if entry.quantizer_name in visual_input_names ), None, ) diff --git a/modelopt_recipes/model_type/qwen3_5/ptq/fp8_vision-kv_none.yaml b/modelopt_recipes/model_type/qwen3_5/ptq/fp8_vision-kv_none.yaml index d6ed27b32d9..8854b376af9 100644 --- a/modelopt_recipes/model_type/qwen3_5/ptq/fp8_vision-kv_none.yaml +++ b/modelopt_recipes/model_type/qwen3_5/ptq/fp8_vision-kv_none.yaml @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -# Dense Qwen3.5 Vision Encoder FP8 with a high-precision language model and KV cache. +# Dense Qwen3.5 vision-merger FP8 with a high-precision language model and KV cache. # modelopt-schema: modelopt.recipe.config.ModelOptPTQRecipe imports: @@ -22,9 +22,8 @@ imports: metadata: description: >- - FP8 quantization of dense Qwen3.5 Vision Encoder Linear layers, including any merger Linears; - the language model, KV cache, patch embedding, and vision-attention operands remain in high - precision. + FP8 quantization of dense Qwen3.5 vision-merger Linears; the vision blocks, language model, + KV cache, and patch embedding remain in high precision for vLLM compatibility. quantize: algorithm: max diff --git a/modelopt_recipes/model_type/qwen3_5/ptq/fp8_vision_lm-kv_fp8_cast.yaml b/modelopt_recipes/model_type/qwen3_5/ptq/fp8_vision_lm-kv_fp8_cast.yaml index 173d32b286a..7ccf42855a0 100644 --- a/modelopt_recipes/model_type/qwen3_5/ptq/fp8_vision_lm-kv_fp8_cast.yaml +++ b/modelopt_recipes/model_type/qwen3_5/ptq/fp8_vision_lm-kv_fp8_cast.yaml @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -# Dense Qwen3.5 Vision Encoder and language model FP8 with FP8 KV-cache cast. +# Dense Qwen3.5 vision-merger and language-model FP8 with FP8 KV-cache cast. # modelopt-schema: modelopt.recipe.config.ModelOptPTQRecipe imports: @@ -25,9 +25,9 @@ imports: metadata: description: >- - W8A8 FP8 quantization of dense Qwen3.5 Vision Encoder and language-model Linear layers, - including any merger Linears, with FP8 KV-cache cast; patch embedding and vision-attention - operands remain in high precision. + W8A8 FP8 quantization of dense Qwen3.5 vision-merger and language-model Linears, with FP8 + KV-cache cast; vision blocks and patch embedding remain in high precision for vLLM + compatibility. quantize: algorithm: max @@ -36,6 +36,6 @@ quantize: - $import: w8a8_fp8_fp8 - $import: kv_fp8_cast - $import: default_disabled_quantizers - # Re-enable visual Linears after the standard vision exclusion. The imported snippet also + # Re-enable merger Linears after the standard vision exclusion. The imported snippet also # keeps vision-attention BMM quantizers disabled as a final defense-in-depth rule. - $import: vision_fp8 diff --git a/modelopt_recipes/model_type/qwen3_vl/ptq/fp8_vision-kv_none.yaml b/modelopt_recipes/model_type/qwen3_vl/ptq/fp8_vision-kv_none.yaml index 9fc2e4ee590..f80aad5962c 100644 --- a/modelopt_recipes/model_type/qwen3_vl/ptq/fp8_vision-kv_none.yaml +++ b/modelopt_recipes/model_type/qwen3_vl/ptq/fp8_vision-kv_none.yaml @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -# Qwen3-VL Vision Encoder FP8 with a high-precision language model and KV cache. +# Qwen3-VL vision-merger FP8 with a high-precision language model and KV cache. # modelopt-schema: modelopt.recipe.config.ModelOptPTQRecipe imports: @@ -22,9 +22,8 @@ imports: metadata: description: >- - FP8 quantization of Qwen3-VL Vision Encoder Linear layers, including primary and deepstack - merger Linears; the language model, KV cache, patch embedding, and vision-attention operands - remain in high precision. + FP8 quantization of Qwen3-VL primary and deepstack vision-merger Linears; the vision blocks, + language model, KV cache, and patch embedding remain in high precision for vLLM compatibility. quantize: algorithm: max diff --git a/modelopt_recipes/model_type/qwen3_vl/ptq/fp8_vision_lm-kv_fp8_cast.yaml b/modelopt_recipes/model_type/qwen3_vl/ptq/fp8_vision_lm-kv_fp8_cast.yaml index fbd449d2d42..131859ab2c6 100644 --- a/modelopt_recipes/model_type/qwen3_vl/ptq/fp8_vision_lm-kv_fp8_cast.yaml +++ b/modelopt_recipes/model_type/qwen3_vl/ptq/fp8_vision_lm-kv_fp8_cast.yaml @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -# Qwen3-VL Vision Encoder and language model FP8 with FP8 KV-cache cast. +# Qwen3-VL vision-merger and language-model FP8 with FP8 KV-cache cast. # modelopt-schema: modelopt.recipe.config.ModelOptPTQRecipe imports: @@ -25,9 +25,9 @@ imports: metadata: description: >- - W8A8 FP8 quantization of Qwen3-VL Vision Encoder and language-model Linear layers, including - primary and deepstack merger Linears, with FP8 KV-cache cast; patch embedding and - vision-attention operands remain in high precision. + W8A8 FP8 quantization of Qwen3-VL primary and deepstack vision-merger and language-model + Linears, with FP8 KV-cache cast; vision blocks and patch embedding remain in high precision + for vLLM compatibility. quantize: algorithm: max @@ -36,6 +36,6 @@ quantize: - $import: w8a8_fp8_fp8 - $import: kv_fp8_cast - $import: default_disabled_quantizers - # Re-enable visual Linears after the standard vision exclusion. The imported snippet also + # Re-enable merger Linears after the standard vision exclusion. The imported snippet also # keeps vision-attention BMM quantizers disabled as a final defense-in-depth rule. - $import: vision_fp8 diff --git a/modelopt_recipes/model_type/qwen3_vl/ptq/vision_fp8.quant_cfg.yaml b/modelopt_recipes/model_type/qwen3_vl/ptq/vision_fp8.quant_cfg.yaml index 9486d5a7fa2..a650aa895a8 100644 --- a/modelopt_recipes/model_type/qwen3_vl/ptq/vision_fp8.quant_cfg.yaml +++ b/modelopt_recipes/model_type/qwen3_vl/ptq/vision_fp8.quant_cfg.yaml @@ -13,23 +13,21 @@ # See the License for the specific language governing permissions and # limitations under the License. -# FP8 Linear weights and inputs in the Qwen visual branch. Patch embedding -# and vision-attention operands remain in high precision. +# Keep vision blocks in high precision: vLLM's FP8 GEMM flattens the 3D +# attention inputs expected by the vision encoder. Merger Linears use 2D inputs. # modelopt-schema: modelopt.torch.quantization.config.QuantizerCfgListConfig imports: fp8: configs/numerics/fp8 --- - - quantizer_name: '*visual.*weight_quantizer' + - quantizer_name: '*visual.*merger*.weight_quantizer' parent_class: 'nn.Linear' cfg: $import: fp8 - - quantizer_name: '*visual.*input_quantizer' + - quantizer_name: '*visual.*merger*.input_quantizer' parent_class: 'nn.Linear' cfg: $import: fp8 - - quantizer_name: '*visual.*patch_embed*' - enable: false # Keep this defense-in-depth rule last when the snippet follows a KV-cache unit. - quantizer_name: '*visual.*_bmm_quantizer' enable: false diff --git a/modelopt_recipes/ptq.md b/modelopt_recipes/ptq.md index d0db470af53..af3c10a117d 100644 --- a/modelopt_recipes/ptq.md +++ b/modelopt_recipes/ptq.md @@ -334,8 +334,9 @@ linear layers and MSE-calibrated NVFP4 to routed experts, with expert remain unquantized. **`qwen3_vl/ptq`** and **`qwen3_5/ptq`** provide FP8 recipes for the `visual` branch of -validated Qwen3-VL and dense Qwen3.5 checkpoints. `fp8_vision-kv_none` enables only Vision -Encoder `nn.Linear` weights and inputs, including the primary merger and any deepstack mergers. +validated Qwen3-VL and dense Qwen3.5 checkpoints. `fp8_vision-kv_none` enables FP8 only on the +primary and any deepstack vision-merger `nn.Linear` weights and inputs. Vision blocks remain in +high precision because vLLM's FP8 GEMM flattens their 3D attention inputs. `fp8_vision_lm-kv_fp8_cast` combines that visual configuration with the standard W8A8 FP8 model and FP8 KV-cache-cast units. Both keep patch embedding and vision-attention BMM operands in high precision. The shared visual snippet lives under `qwen3_vl`; thin wrappers remain discoverable diff --git a/tests/gpu/torch/export/test_qwen_vision_recipe_export.py b/tests/gpu/torch/export/test_qwen_vision_recipe_export.py index 9c1a804ea1c..389375428b1 100644 --- a/tests/gpu/torch/export/test_qwen_vision_recipe_export.py +++ b/tests/gpu/torch/export/test_qwen_vision_recipe_export.py @@ -92,7 +92,7 @@ def calibration_forward(_model): scale_names = {name for name in tensor_names if name.endswith(("input_scale", "weight_scale"))} assert scale_names assert any(name.startswith("model.visual.") for name in scale_names) - assert any(".attn.qkv.weight_scale" in name for name in scale_names) + assert not any(".blocks." in name for name in scale_names) assert any(".merger.linear_fc1.weight_scale" in name for name in scale_names) if model_type == "qwen3_vl": assert any( diff --git a/tests/unit/recipe/test_qwen_vision_recipe.py b/tests/unit/recipe/test_qwen_vision_recipe.py index aeae38749bf..264e3b73239 100644 --- a/tests/unit/recipe/test_qwen_vision_recipe.py +++ b/tests/unit/recipe/test_qwen_vision_recipe.py @@ -52,14 +52,6 @@ def test_qwen_vision_recipes_select_expected_quantizers( enabled = {name for name, module in modules.items() if getattr(module, "is_enabled", False)} expected_visual_suffixes = { - "blocks.0.attn.qkv.weight_quantizer", - "blocks.0.attn.qkv.input_quantizer", - "blocks.0.attn.proj.weight_quantizer", - "blocks.0.attn.proj.input_quantizer", - "blocks.0.mlp.linear_fc1.weight_quantizer", - "blocks.0.mlp.linear_fc1.input_quantizer", - "blocks.0.mlp.linear_fc2.weight_quantizer", - "blocks.0.mlp.linear_fc2.input_quantizer", "merger.linear_fc1.weight_quantizer", "merger.linear_fc1.input_quantizer", "merger.linear_fc2.weight_quantizer", @@ -74,11 +66,11 @@ def test_qwen_vision_recipes_select_expected_quantizers( "deepstack_merger_list.0.linear_fc2.input_quantizer", } ) - assert {f"model.visual.{suffix}" for suffix in expected_visual_suffixes} <= enabled - assert not any("patch_embed" in name for name in enabled) + visual_quantizers = {name for name in enabled if name.startswith("model.visual.")} + assert visual_quantizers == {f"model.visual.{suffix}" for suffix in expected_visual_suffixes} - vision_weight = modules["model.visual.blocks.0.attn.qkv.weight_quantizer"] - vision_input = modules["model.visual.blocks.0.attn.qkv.input_quantizer"] + vision_weight = modules["model.visual.merger.linear_fc1.weight_quantizer"] + vision_input = modules["model.visual.merger.linear_fc1.input_quantizer"] assert vision_weight.num_bits == (4, 3) assert vision_input.num_bits == (4, 3) From 725cb773b31928cf8bb0b6c006d95e0d85cd8823 Mon Sep 17 00:00:00 2001 From: Gwena Cunha <4861122+gcunhase@users.noreply.github.com> Date: Fri, 16 Oct 2026 15:29:27 -0400 Subject: [PATCH 30/35] Fix concat elimination crashing when a Concat input is a constant (#2838) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Type of change: Bug fix `get_concat_eliminated_tensors` groups every input of a quantized `Concat` so the inputs share one scaling factor. It also grouped constant (initializer) inputs. The calibrators only collect ranges for activation tensors, so applying the group raised `KeyError: '_ReduceMin'` in `_compute_data_minmax_calibrator`. A ViT hits this: constant folding leaves the class token (`/Expand_output_0`) as an initializer that feeds `Concat`, and when the Concat output is quantized the default `concat_elimination` pass puts it in a group. With `--autotune` the crash came after the whole search had finished. The fix skips `gs.Constant` inputs when building the group. The constant keeps its own scale, so TensorRT may keep one Q/DQ pair at that Concat that concat elimination would have removed. ```sh python -m modelopt.onnx.quantization --onnx_path model.onnx --quantize_mode fp8 --autotune default ``` - New unit test `tests/unit/onnx/quantization/test_concat_elimination.py` - Real model: ViT-tiny (timm), FP8 `--autotune default`, H100, TensorRT 10.16.2.11. Before: crash after 2.0 h. After: completes in 2.0 h (ModelOpt keeps the FP16 model, since FP8 measured 0.895x). ```python import timm import torch import modelopt.onnx.quantization as moq model = timm.create_model("vit_tiny_patch16_224", pretrained=False).eval() torch.onnx.export( model, torch.randn(1, 3, 224, 224), "vit_tiny.onnx", input_names=["input"], opset_version=17 ) moq.quantize("vit_tiny.onnx", quantize_mode="fp8", autotune=True) ``` - Is this change backward compatible?: ✅ - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: ✅ - Did you update Changelog?: ✅ - Did you get Claude approval on this PR?: ❌ (not run yet) None. 🤖 Generated with [Claude Code](https://claude.com/claude-code) * **Bug Fixes** * ONNX calibration no longer crashes when a concatenation includes a constant input. * Constant inputs are excluded from shared-scale grouping, while eligible activation inputs continue to be grouped. * **Documentation** * Updated the changelog with release highlights, including GGML imatrix calibration, a Nemotron 3.5 Super VL AutoQuantize recipe, and external-draft speculative decoding. --------- Signed-off-by: gcunhase <4861122+gcunhase@users.noreply.github.com> Co-authored-by: Claude Sonnet 5.5 Signed-off-by: Rohan Joshi --- CHANGELOG.rst | 1 + modelopt/onnx/quantization/qdq_graph.py | 4 +++ tests/_test_utils/onnx/lib_test_models.py | 32 +++++++++++++++++ .../quantization/test_concat_elimination.py | 34 +++++++++++++++++++ 4 files changed, 71 insertions(+) create mode 100644 tests/unit/onnx/quantization/test_concat_elimination.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 1cc5fc19482..d598e93bee6 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -179,6 +179,7 @@ Changelog - Hybrid (e.g. Nemotron-H) checkpoints saved by the ``examples/megatron_bridge`` scripts now record their layer spec in ``run_config.yaml`` in a form that reloads, so they can be converted to HuggingFace; a checkpoint saved by an earlier release still needs its ``model.hybrid_stack_spec`` block replaced by hand. - Fix ONNX preprocessing with custom TensorRT plugin libraries on TensorRT 11. - Fix GPTQ with Megatron data parallelism: each data-parallel replica of a module now builds its Hessian from the whole calibration set instead of only its own shard, so all replicas reach the same quantized weights. +- Fix ONNX quantization with the default ``concat_elimination`` pass crashing in calibration with ``KeyError: '_ReduceMin'`` when a ``Concat`` input is a constant, such as the class token of a ViT. 0.47.0 (2026-09-23) ^^^^^^^^^^^^^^^^^^^ diff --git a/modelopt/onnx/quantization/qdq_graph.py b/modelopt/onnx/quantization/qdq_graph.py index 47ee6ef1464..c712b7e385c 100755 --- a/modelopt/onnx/quantization/qdq_graph.py +++ b/modelopt/onnx/quantization/qdq_graph.py @@ -842,6 +842,10 @@ def union(x, y): # Find quantized inputs to concat quantized_inputs = [] for input_tensor in node.inputs: + # A constant (initializer) input is not an activation: the calibrator collects no range + # for it, so it cannot share a scaling factor with the other tensors of the group. + if isinstance(input_tensor, gs.Constant): + continue input_name = input_tensor.name input_consumers = input_name_to_nodes[input_name] if any(consumer.name in nodes_to_quantize for consumer in input_consumers): diff --git a/tests/_test_utils/onnx/lib_test_models.py b/tests/_test_utils/onnx/lib_test_models.py index 1b669317e8a..57a0621c5da 100644 --- a/tests/_test_utils/onnx/lib_test_models.py +++ b/tests/_test_utils/onnx/lib_test_models.py @@ -1344,3 +1344,35 @@ def build_matmul_1xn_model(): onnx.checker.check_model(model_inferred) return model_inferred + + +# x ----- relu -------------+ +# | +# cls (initializer, or --- +--- concat --- matmul --- out +# relu(y) if not constant_input) +# +# ``cls`` mimics a ViT class token: a constant Concat input that has no producer node. +def build_concat_with_constant_input_model(constant_input=True): + x = helper.make_tensor_value_info("x", onnx.TensorProto.FLOAT, [1, 4, 8]) + y = helper.make_tensor_value_info("y", onnx.TensorProto.FLOAT, [1, 1, 8]) + out = helper.make_tensor_value_info("out", onnx.TensorProto.FLOAT, [1, 5, 8]) + nodes = [helper.make_node("Relu", ["x"], ["relu_out"], name="relu")] + inputs = [x] + initializers = [ + helper.make_tensor("weight", onnx.TensorProto.FLOAT, [8, 8], np.ones(64, dtype=np.float32)) + ] + if constant_input: + cls = helper.make_tensor( + "cls", onnx.TensorProto.FLOAT, [1, 1, 8], np.ones(8, dtype=np.float32) + ) + initializers.append(cls) + else: + inputs.append(y) + nodes.append(helper.make_node("Relu", ["y"], ["cls"], name="relu_cls")) + nodes += [ + helper.make_node("Concat", ["cls", "relu_out"], ["concat_out"], name="concat", axis=1), + helper.make_node("MatMul", ["concat_out", "weight"], ["out"], name="matmul"), + ] + graph = helper.make_graph(nodes, "concat_with_constant_input", inputs, [out], initializers) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + return onnx.shape_inference.infer_shapes(model) diff --git a/tests/unit/onnx/quantization/test_concat_elimination.py b/tests/unit/onnx/quantization/test_concat_elimination.py new file mode 100644 index 00000000000..9edf9712a79 --- /dev/null +++ b/tests/unit/onnx/quantization/test_concat_elimination.py @@ -0,0 +1,34 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +from _test_utils.onnx.lib_test_models import build_concat_with_constant_input_model + +from modelopt.onnx.quantization.qdq_graph import get_concat_eliminated_tensors + + +def test_concat_elimination_groups_activation_inputs(): + model = build_concat_with_constant_input_model(constant_input=False) + groups = get_concat_eliminated_tensors(model, ["concat", "matmul", "relu", "relu_cls"]) + assert groups["concat_out"] == {"cls", "relu_out", "concat_out"} + assert groups["cls"] == groups["concat_out"] + + +def test_concat_elimination_skips_constant_inputs(): + """A constant Concat input has no calibration range, so it must not join the scale group.""" + model = build_concat_with_constant_input_model(constant_input=True) + groups = get_concat_eliminated_tensors(model, ["concat", "matmul", "relu"]) + assert "cls" not in groups + assert groups["concat_out"] == {"relu_out", "concat_out"} + assert all("cls" not in group for group in groups.values()) From 7cdde516fe2c0a3d5d32652a10b692e4944acb4e Mon Sep 17 00:00:00 2001 From: Frida Hou <201670829+Fridah-nv@users.noreply.github.com> Date: Fri, 16 Oct 2026 12:33:00 -0700 Subject: [PATCH 31/35] Fix fused-experts export for W4A8 AWQ (SequentialQuantizer) weight quantizers (#2807) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Type of change: Bug fix `hf_ptq.py --recipe general/auto_quantize/w4a8_awq_beta_fp8_at_6p0bits` on a fused-MoE model (Qwen3-30B-A3B-Instruct-2507) fails at export with `RuntimeError: Changing shape when setting amax is not allowed.` (NVBug 6926543 / OMNIML-6125). **Root cause.** `_export_fused_experts` assumed a single `TensorQuantizer` per expert weight, but block-wise INT4 weights break that in two ways: 1. W4A8 weights use a `SequentialQuantizer`: an INT4 block stage (per-block `_amax`) followed by a per-tensor FP8 stage. It has no `_amax` of its own, so every expert looked uncalibrated; the weight-derived fallback then wrote a scalar amax into both stages (`SequentialQuantizer.amax` forwards to each stage), which the INT4 stage's block-shaped amax rejects. The failure is deterministic for any W4A8 fused-MoE model, not only experts that saw no calibration tokens. For fused experts the FP8 stage is never calibrated, so the fallback is the normal path for it. 2. For any INT4 block quantizer (plain `INT4_AWQ` too), the dim-0 slice into gate/up left the stage's `_amax_shape_for_export` (recorded from the fused weight at calibration) stale, so `export_amax()` reshaped the sliced amax back to the fused row count (`size of tensor a (256) must match tensor b (512)`). **Fix** (`modelopt/torch/export/moe_utils.py`): handle each `TensorQuantizer` stage separately for the calibrated check, the dim-0 slice (also updating `_amax_shape_for_export`) and the fallback. The fallback yields the layout export expects (per-block `(rows, blocks)` for INT4 stages, the existing amax's shape otherwise) and raises on layouts it cannot group instead of writing a wrong shape. Float-block formats (NVFP4) keep the previous scalar path. ### Usage ```bash python examples/hf_ptq/hf_ptq.py --model Qwen/Qwen3-30B-A3B-Instruct-2507 --dataset cnn_dailymail \ --recipe general/auto_quantize/w4a8_awq_beta_fp8_at_6p0bits --batch_size 8 --calib_size 512 \ --export_path --trust_remote_code ``` ### Testing - **Repro:** the ticket's exact traceback (`moe_utils.py:124`) on the `0.48.0rc0` tag with the ticket's pins (transformers 5.12.1, kernels 0.12.3, compressed-tensors 0.19.0), Qwen3-30B-A3B-Instruct-2507, 1 GPU (GB300; the ticket used GH200, both aarch64). - **Full 30B end-to-end with the fix** (ticket command, calib 512, batch 8): export completes, no traceback, 23 GB checkpoint with `MIXED_PRECISION` `hf_quant_config.json` (29 W4A8_AWQ + 19 FP8 expert layers). - **Unit tests** (`tests/unit/torch/quantization/plugins/test_fused_experts.py`): W4A8 `SequentialQuantizer` with calibrated / zero / missing amax; plain per-tensor / per-channel / per-block zero amax; refusal of ungroupable layouts. All 7 fail on `origin/main`'s `moe_utils.py` and pass with the fix; the file passes 59/59. - **Existing behavior unchanged**: tiny random-weight Qwen3-MoE (4 layers, 32 experts), each recipe quantized once and exported with `origin/main`'s and the fixed `moe_utils`, checkpoints compared tensor by tensor: | Recipe | `origin/main` | This PR | |---|---|---| | `NVFP4_DEFAULT_CFG` | exports | byte-identical (1623 tensors, same `hf_quant_config`) | | `NVFP4_AWQ_LITE_CFG` | exports | byte-identical (1627 tensors) | | `FP8_DEFAULT_CFG` | exports | byte-identical (1223 tensors) | | `INT4_AWQ_CFG` | fails (`tensor a (256) vs b (512)`) | exports | | `W4A8_AWQ_BETA_CFG` | fails (`Changing shape when setting amax`) | exports | - **W4A8 scales**: exported packed-weight / scale shapes are consistent, no NaN or non-positive scales, gate and up `weight_scale_2` identical for every expert. The FP8-stage amax derived by the fallback equals what a dense Linear gets from calibration (`max|W|` of the final weight; checked on `q_proj`/`o_proj`). - **Not run:** serving the exported checkpoints in vLLM / TRT-LLM, accuracy evaluation, GH200 hardware, multi-GPU. `tests/unit/torch/export/test_fsdp2_parallel_export.py::test_streaming_export_two_cpu_ranks` fails with and without this change (same on plain `origin/main`). ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: ✅ - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A (bug in the unreleased 0.48.0 cycle) - Did you get Claude approval on this PR?: ❌ (two local `/code-review` rounds done; findings addressed, one postponed below) ### Additional Information - NVBug 6926543 is keyed `Committed_ModelOpt_0.48.0`: needs the `cherry-pick-0.48.0` label. - Postponed (pre-existing, separate change): for non-gated experts and `down_proj`, the live quantizer is mutated in place (`delattr(_amax)`), so an exception mid-export leaves it modified. ## Summary by CodeRabbit * **Bug Fixes** * Improved export reliability for quantized mixture-of-experts models using sequential weight quantization. * Missing or zero calibration scales now use weight-derived fallbacks while preserving supported per-tensor, per-channel, per-row, and per-block layouts. * Fused gate and up projections use shared or split scales as appropriate, with scale slicing handled independently for each quantizer stage. * Unsupported layouts and incompatible block sizes report errors; scale dimensions that cannot be evenly split trigger a warning instead. --------- Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Co-authored-by: Claude Sonnet 5.5 Signed-off-by: Rohan Joshi --- modelopt/torch/export/moe_utils.py | 162 ++++++++++++------ .../plugins/test_fused_experts.py | 156 ++++++++++++++++- 2 files changed, 262 insertions(+), 56 deletions(-) diff --git a/modelopt/torch/export/moe_utils.py b/modelopt/torch/export/moe_utils.py index 734302690f1..2affd111b02 100644 --- a/modelopt/torch/export/moe_utils.py +++ b/modelopt/torch/export/moe_utils.py @@ -45,6 +45,66 @@ def _delete_fused_moe_source_attrs(module: nn.Module) -> None: delattr(module, attr) +def _tensor_quantizers(quantizer) -> list: + """The ``TensorQuantizer`` stages of ``quantizer`` (several for W4A8's SequentialQuantizer). + + A SequentialQuantizer has no ``_amax`` of its own and forwards ``amax`` writes to every + stage, so each stage must be checked and set on its own. + """ + # Deferred: ``export/__init__`` star-imports this module, so keep quantization out of its + # import-time dependencies. + from modelopt.torch.quantization.nn import SequentialQuantizer + + return list(quantizer) if isinstance(quantizer, SequentialQuantizer) else [quantizer] + + +def _is_uncalibrated(quantizer) -> bool: + amax = getattr(quantizer, "_amax", None) + return getattr(quantizer, "is_enabled", False) and (amax is None or bool(torch.all(amax == 0))) + + +def _fallback_amax(weight: torch.Tensor, quantizer) -> torch.Tensor: + """Weight-derived amax for an uncalibrated expert, in the shape export expects. + + The amax setter rejects shape changes, and a block/channel-wise quantizer (e.g. the INT4 + stage of W4A8 AWQ) keeps a zero ``_amax`` of that shape, so a scalar fallback would raise. + A missing ``_amax`` on an INT4 block quantizer still needs per-block ``(rows, blocks)``. + """ + like = getattr(quantizer, "_amax", None) + block_sizes = getattr(quantizer, "block_sizes", None) + block = block_sizes.get(-1) if block_sizes and block_sizes.get("type") == "static" else None + # Integer static-block stages (INT4 of W4A8) need a per-block amax. Float-block formats + # (NVFP4) keep a per-tensor amax plus separate block scales, so they take the scalar path. + if like is None and isinstance(quantizer.num_bits, int) and block: + if weight.shape[-1] % block: + raise ValueError( + f"Weight of shape {tuple(weight.shape)} is not divisible into blocks of {block}." + ) + like = weight.new_empty(weight.shape[0], weight.shape[-1] // block) + if like is None or like.numel() == 1: + amax = torch.maximum(weight.amax(), -weight.amin()).float() + return amax if like is None else amax.view_as(like) + # Block-quant calibration records the flat (rows * blocks, 1); also accept (rows, blocks) and + # a per-row (rows,). Refuse any other layout (e.g. per input channel) rather than guess its + # grouping. + rows = weight.shape[0] + if like.dim() == 1 and like.shape[0] == rows: + blocks = 1 + elif like.dim() == 2 and like.shape[-1] == 1 and like.shape[0] % rows == 0: + blocks = like.shape[0] // rows + elif like.dim() == 2 and like.shape[0] == rows: + blocks = like.shape[1] + else: + blocks = 0 + if not blocks or weight.shape[-1] % blocks: + raise ValueError( + f"Cannot derive amax of shape {tuple(like.shape)} from a weight of shape " + f"{tuple(weight.shape)}." + ) + groups = weight.reshape(rows, blocks, -1) + return torch.maximum(groups.amax(dim=-1), -groups.amin(dim=-1)).float().reshape(like.shape) + + def _export_fused_experts( module: nn.Module, dtype: torch.dtype, @@ -111,17 +171,18 @@ def _export_fused_experts( # mismatched weight_scale_2 and garbled MoE output at inference. # Non-gated experts have no gate/up fusion, so this shared-amax step is # skipped — their single up_proj uses the generic per-projection fallback. - first_proj_q = first_proj_weight_quantizers[idx] - if ( - is_gated - and getattr(first_proj_q, "is_enabled", False) - and ( - not hasattr(first_proj_q, "_amax") - or first_proj_q._amax is None - or torch.all(first_proj_q._amax == 0) - ) - ): - first_proj_q.amax = first_proj[idx].abs().amax().to(torch.float32) + uncalibrated_stages = ( + [ + q + for q in _tensor_quantizers(first_proj_weight_quantizers[idx]) + if _is_uncalibrated(q) + ] + if is_gated + else [] + ) + for q in uncalibrated_stages: + q.amax = _fallback_amax(first_proj[idx], q) + if uncalibrated_stages: warnings.warn( f"Expert {idx} gate_up_proj weight quantizer was not calibrated " f"(amax missing or zero). Using fused-tensor amax as fallback " @@ -163,53 +224,46 @@ def _export_fused_experts( copy.deepcopy(w_quantizer_src) if uses_first_proj_quantizers else w_quantizer_src ) - # For per-channel amax (dim >= 1), proportionally slice dim-0 - # to match the split weight. - if ( - hasattr(w_quantizer, "_amax") - and w_quantizer._amax is not None - and w_quantizer._amax.dim() >= 1 - ): - amax = w_quantizer._amax - # Per-block _amax (NVFP4 static) collapses the row axis we want - # to slice on; restore it so dim-0 slicing splits gate/up. - if amax.numel() != fused_total and amax.numel() % fused_total == 0: - amax = amax.contiguous().view(fused_total, amax.numel() // fused_total) - amax_dim0 = amax.shape[0] - if fused_total % amax_dim0 == 0: - slice_start = fused_start * amax_dim0 // fused_total - slice_end = (fused_start + weight_slice.shape[0]) * amax_dim0 // fused_total - sliced = amax[slice_start:slice_end].contiguous() - # The amax setter refuses shape changes; drop _amax first. - if hasattr(w_quantizer, "_amax"): - delattr(w_quantizer, "_amax") - w_quantizer.amax = sliced - else: + for stage in _tensor_quantizers(w_quantizer): + # For per-channel amax (dim >= 1), proportionally slice dim-0 + # to match the split weight. + if getattr(stage, "_amax", None) is not None and stage._amax.numel() > 1: + amax = stage._amax + # Per-block _amax (NVFP4 static) collapses the row axis we want + # to slice on; restore it so dim-0 slicing splits gate/up. + if amax.numel() != fused_total and amax.numel() % fused_total == 0: + amax = amax.contiguous().view(fused_total, amax.numel() // fused_total) + amax_dim0 = amax.shape[0] + if fused_total % amax_dim0 == 0: + slice_start = fused_start * amax_dim0 // fused_total + slice_end = (fused_start + weight_slice.shape[0]) * amax_dim0 // fused_total + sliced = amax[slice_start:slice_end].contiguous() + # The amax setter refuses shape changes; drop _amax first. + delattr(stage, "_amax") + stage.amax = sliced + # Block-quant calibration recorded the fused amax's row count here, and + # export_amax() reshapes to it, undoing the slice. + export_shape = getattr(stage, "_amax_shape_for_export", None) + if export_shape is not None and export_shape[0] == amax_dim0: + stage._amax_shape_for_export = (sliced.shape[0], *export_shape[1:]) + else: + warnings.warn( + f"Expert {idx} {proj_name}: fused amax dim0 ({amax_dim0}) does not " + f"evenly divide fused_total ({fused_total}). Skipping amax slicing, " + f"which may produce incorrect quantization scales.", + stacklevel=2, + ) + + # If the weight quantizer was never calibrated, compute amax from weights. + if _is_uncalibrated(stage): + stage.amax = _fallback_amax(weight_slice, stage) warnings.warn( - f"Expert {idx} {proj_name}: fused amax dim0 ({amax_dim0}) does not " - f"evenly divide fused_total ({fused_total}). Skipping amax slicing, " - f"which may produce incorrect quantization scales.", + f"Expert {idx} {proj_name} weight quantizer was not calibrated " + f"(amax missing or zero). Using weight-derived amax as fallback. " + f"Consider using more calibration data to activate all experts.", stacklevel=2, ) - # If the weight quantizer was never calibrated, compute amax from weights. - if ( - hasattr(w_quantizer, "is_enabled") - and w_quantizer.is_enabled - and ( - not hasattr(w_quantizer, "_amax") - or w_quantizer._amax is None - or torch.all(w_quantizer._amax == 0) - ) - ): - w_quantizer.amax = weight_slice.abs().amax().to(torch.float32) - warnings.warn( - f"Expert {idx} {proj_name} weight quantizer was not calibrated " - f"(amax missing or zero). Using weight-derived amax as fallback. " - f"Consider using more calibration data to activate all experts.", - stacklevel=2, - ) - wrapper = nn.Module() wrapper.weight = nn.Parameter(weight_slice.contiguous(), requires_grad=False) wrapper.weight_quantizer = w_quantizer diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py index 6e7375dd695..959471cf13e 100644 --- a/tests/unit/torch/quantization/plugins/test_fused_experts.py +++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py @@ -15,6 +15,8 @@ """Tests for _QuantFusedExperts: generic fused MoE quantization and export.""" +import warnings + import pytest import torch import torch.nn as nn @@ -27,12 +29,12 @@ import modelopt.torch.quantization as mtq import modelopt.torch.quantization.nn.modules.tensor_quantizer as tensor_quantizer_module from modelopt.torch.export.model_utils import _release_exported_tensors -from modelopt.torch.export.moe_utils import _export_fused_experts +from modelopt.torch.export.moe_utils import _export_fused_experts, _fallback_amax from modelopt.torch.export.quant_utils import get_quant_config, get_quantization_format from modelopt.torch.quantization.config import QuantizerAttributeConfig from modelopt.torch.quantization.conversion import _normalize_fused_experts_quantizer_name from modelopt.torch.quantization.model_calib import local_hessian_calibrate -from modelopt.torch.quantization.nn import QuantModuleRegistry, TensorQuantizer +from modelopt.torch.quantization.nn import QuantModuleRegistry, SequentialQuantizer, TensorQuantizer from modelopt.torch.quantization.plugins.huggingface import ( _fused_experts_wrapper_class, _is_fused_experts_module, @@ -696,6 +698,156 @@ def _spy_export(wrapper, dtype, **_kwargs): if QuantModuleRegistry.get(expert_type) is not None: QuantModuleRegistry.unregister(expert_type) + @pytest.mark.parametrize("amax_state", ["calibrated", "zero", "missing"]) + def test_w4a8_sequential_weight_quantizer_exports_split_block_amax( + self, monkeypatch, amax_state + ): + """W4A8 AWQ weights use a SequentialQuantizer: an INT4 block stage (flat per-block + ``_amax`` of shape ``(rows * blocks, 1)``, as block-quant calibration records it) + followed by a per-tensor FP8 stage. Export must treat each stage on its own: slice the + block amax into the gate/up halves (updating the stale ``_amax_shape_for_export``), and + for experts that saw no calibration tokens derive a per-block fallback instead of + forcing a scalar into the block stage (``Changing shape when setting amax``). + """ + block = 16 + experts = _SyntheticFusedExperts() + expert_type = type(experts) + if QuantModuleRegistry.get(expert_type) is None: + QuantModuleRegistry.register({expert_type: "test.SyntheticFusedExperts"})( + _QuantFusedExperts + ) + try: + converted = QuantModuleRegistry.convert(experts) + + def _w4a8_quantizer(rows, fused): + int4 = TensorQuantizer( + QuantizerAttributeConfig(num_bits=4, block_sizes={-1: block, "type": "static"}) + ) + fp8 = TensorQuantizer(QuantizerAttributeConfig(num_bits=(4, 3))) + # gate_up_proj reduces over HIDDEN_DIM, down_proj over INTERMEDIATE_DIM. + blocks = (HIDDEN_DIM if fused else INTERMEDIATE_DIM) // block + int4._amax_shape_for_export = (rows, -1) + if amax_state != "missing": + # Row r of gate_up carries the marker r + 1, so a wrong slice is visible. + marker = torch.arange(1, rows + 1).float().repeat_interleave(blocks)[:, None] + int4.amax = ( + marker.clone() if amax_state == "calibrated" else torch.zeros_like(marker) + ) + fp8.amax = torch.tensor(2.0 if amax_state == "calibrated" else 0.0) + return SequentialQuantizer(int4, fp8) + + for idx in range(NUM_EXPERTS): + converted.gate_up_proj_weight_quantizers[idx] = _w4a8_quantizer( + 2 * INTERMEDIATE_DIM, fused=True + ) + converted.down_proj_weight_quantizers[idx] = _w4a8_quantizer( + HIDDEN_DIM, fused=False + ) + + seen = [] + + def _spy_export(wrapper, dtype, **_kwargs): + int4, fp8 = wrapper.weight_quantizer + seen.append( + ( + wrapper.weight.shape[0], + int4.export_amax().detach().clone(), + fp8._amax.detach().clone(), + ) + ) + + monkeypatch.setattr( + "modelopt.torch.export.unified_export_hf._export_quantized_weight", + _spy_export, + ) + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + _export_fused_experts(converted, torch.float16) + fallback_warned = any("was not calibrated" in str(w.message) for w in caught) + assert fallback_warned == (amax_state != "calibrated") + + assert len(seen) == 3 * NUM_EXPERTS + for rows, int4_amax, fp8_amax in seen: + assert int4_amax.shape[0] == rows # sliced to the projection, not the fused rows + assert torch.all(int4_amax > 0) + assert torch.all(fp8_amax > 0) + # Export order per expert is gate, up, down; gate and up must share the FP8 stage. + for gate, up in zip(seen[0::3], seen[1::3]): + assert torch.equal(gate[2], up[2]) + if amax_state == "calibrated": + gate_up_rows = {int4_amax[0, 0].item() for rows, int4_amax, _ in seen if rows == 16} + # gate starts at row marker 1, up at INTERMEDIATE_DIM + 1. + assert gate_up_rows == {1.0, float(INTERMEDIATE_DIM + 1)} + finally: + if QuantModuleRegistry.get(expert_type) is not None: + QuantModuleRegistry.unregister(expert_type) + + @pytest.mark.parametrize( + ("amax_layout", "expected_gate_shape", "expected_down_shape"), + [ + ("per_tensor", (1,), (1,)), + ("per_channel", (INTERMEDIATE_DIM, 1), (HIDDEN_DIM, 1)), + ("per_row_1d", (INTERMEDIATE_DIM,), (HIDDEN_DIM,)), + ("per_block", (INTERMEDIATE_DIM, 2), (HIDDEN_DIM, 2)), + # What block-quant calibration records: (rows * blocks, 1). + ("flat_block", (INTERMEDIATE_DIM, 2), (HIDDEN_DIM, 2)), + ], + ) + def test_uncalibrated_zero_amax_fallback_keeps_layout( + self, monkeypatch, amax_layout, expected_gate_shape, expected_down_shape + ): + """The weight-derived fallback must reproduce the layout of the quantizer's zero amax.""" + experts = _SyntheticFusedExperts() + expert_type = type(experts) + if QuantModuleRegistry.get(expert_type) is None: + QuantModuleRegistry.register({expert_type: "test.SyntheticFusedExperts"})( + _QuantFusedExperts + ) + try: + converted = QuantModuleRegistry.convert(experts) + layouts = { + "per_tensor": lambda rows: (1,), + "per_channel": lambda rows: (rows, 1), + "per_row_1d": lambda rows: (rows,), + "per_block": lambda rows: (rows, 2), + "flat_block": lambda rows: (rows * 2, 1), + } + for idx in range(NUM_EXPERTS): + for wq, rows in ( + (converted.gate_up_proj_weight_quantizers[idx], 2 * INTERMEDIATE_DIM), + (converted.down_proj_weight_quantizers[idx], HIDDEN_DIM), + ): + wq._disabled = False + wq.amax = torch.zeros(*layouts[amax_layout](rows)) + + seen = [] + + def _spy_export(wrapper, dtype, **_kwargs): + seen.append(wrapper.weight_quantizer._amax.detach().clone()) + + monkeypatch.setattr( + "modelopt.torch.export.unified_export_hf._export_quantized_weight", + _spy_export, + ) + with pytest.warns(UserWarning, match="was not calibrated"): + _export_fused_experts(converted, torch.float16) + + assert all(torch.all(a > 0) for a in seen) + # Export order per expert is gate, up, down. + for gate, up, down in zip(seen[0::3], seen[1::3], seen[2::3]): + assert gate.shape == up.shape == expected_gate_shape + assert down.shape == expected_down_shape + finally: + if QuantModuleRegistry.get(expert_type) is not None: + QuantModuleRegistry.unregister(expert_type) + + def test_fallback_amax_refuses_layouts_it_cannot_group(self): + quantizer = TensorQuantizer(QuantizerAttributeConfig(num_bits=8, axis=1)) + quantizer.amax = torch.zeros(1, HIDDEN_DIM) # per input channel: not a (rows, blocks) grid + with pytest.raises(ValueError, match="Cannot derive amax"): + _fallback_amax(torch.randn(INTERMEDIATE_DIM, HIDDEN_DIM), quantizer) + # --------------------------------------------------------------------------- # Tests for tied-experts dedup in _export_fused_experts From 2f1f95776c45f14f9657f9bcc98008fcb75bdc4c Mon Sep 17 00:00:00 2001 From: yingguo-trt <244492186+yingguo-trt@users.noreply.github.com> Date: Sat, 17 Oct 2026 03:50:03 +0800 Subject: [PATCH 32/35] [https://nvbugspro.nvidia.com/bug/6918844] Fix LSQ FP8 block-scale lower bound (#2772) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Type of change: Bug fix Fix the lower bound used by `StaticBlockScaleQuantizer._block_scale_from_amax()` for FP8-quantized LSQ block scales. The previous bound could inflate small but representable scales, distort quantized outputs, and block gradients through the clamped scale-to-amax path. `scaled_e4m3` maps `per_tensor_scale` to the E4M3 maximum, 448. The minimum positive encoded value therefore corresponds to a physical scale of `(2**-9 / 448) * per_tensor_scale`; the previous bound omitted the division by 448. This change applies that inverse mapping. The existing GPU regression now checks exact agreement with static quantization and analytical amax gradients for small scales in tied and default post-only learnable configurations. Existing CPU assertions are updated to the corrected bound. ### Usage No API or recipe changes are required. The correction applies when LSQ computes FP8-quantized block scales; the unquantized-scale path is unchanged. ### Testing Standalone validation on NVIDIA B300, verified from the archived pytest JUnit results: | Test file | Result | | --- | --- | | `tests/unit/torch/quantization/test_lsq.py` | 48 passed | | `tests/gpu/torch/quantization/test_lsq_cuda.py` | 7 passed | | `tests/unit/recipe/test_lsq_recipes.py` | 6 passed | All 61 tests passed with no skips or errors. Recipe tests check configuration/loading, not training convergence. - **Before/after control:** The same GPU regression body fails both parameter cases with the old lower-bound formula (**2 failed**) and passes with the fix (**2 passed**), including the analytical gradient assertions. - **Formatting/static checks:** Applicable pre-commit hooks and `git diff --check` passed. - **Internal Jenkins validation:** [Build #173](http://dlswqa-nas:18880/job/modelopt-quant-module/173/) compared baseline `2c4aed219` with fixed commit `0be8b9401`. Its frozen tied-LSQ forward diagnostic changed from a zeroed small block to 14 nonzero values matching the static reference. This diagnostic is report-only (`OBSERVED`); the overall build remains **FAILURE / INVALID** because a separate static NVFP4 small-scale check fails on both revisions. The standalone validation wrapper also reported failure because it misparsed the intentionally failing old-formula JUnit results; the pytest outcomes above were checked directly from the raw reports. The implementation and test assertions validated at `0be8b9401` are unchanged. Subsequent commits only update explanatory comments, test docstrings, and CHANGELOG placement. Full-model LSQ training convergence was not evaluated. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ No API/configuration changes; affected numerical behavior is corrected. - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`?: N/A — no copied code or new dependencies. - Did you write any new necessary tests?: ✅ Extended the existing GPU regression and updated existing CPU assertions. - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ Added under 0.48.0 Bug Fixes. - Did you get Claude approval on this PR?: Pending. ### Additional Information Related issue: [NVBug 6918844](https://nvbugspro.nvidia.com/bug/6918844). ## Summary by CodeRabbit * **Bug Fixes** * Corrected the lower bound for FP8 block scales so small, representable scales are not clamped so aggressively that quantized values become zero. * Preserved scale gradients for these small values when using LSQ, improving quantization behavior at the lower end of the supported range. This helps ensure low-magnitude values remain represented and their scales can continue to be adjusted during quantization. --------- Signed-off-by: yingguo-trt <244492186+yingguo-trt@users.noreply.github.com> Signed-off-by: Rohan Joshi Co-authored-by: Rohan Joshi Signed-off-by: Rohan Joshi --- CHANGELOG.rst | 1 + .../nn/modules/tensor_quantizer.py | 6 ++- tests/gpu/torch/quantization/test_lsq_cuda.py | 47 +++++++++++++------ tests/unit/torch/quantization/test_lsq.py | 7 ++- 4 files changed, 42 insertions(+), 19 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index d598e93bee6..ccdcb86260f 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -129,6 +129,7 @@ Changelog **Bug Fixes** +- Fix LSQ clamping small, representable FP8 block scales too aggressively, which could zero quantized values and suppress scale gradients. - Fix ONNX AutoTune benchmark initialization with TensorRT 11 while preserving strongly typed networks on TensorRT 10. - Keep vision-block Linears in high precision in the ``qwen3_vl`` and ``qwen3_5`` FP8 vision recipes so exported checkpoints avoid a vLLM vision-attention shape failure. The recipes still quantize primary and deepstack merger Linears; re-quantize affected checkpoints before deployment. - Fix unified HF export writing a weight name the source checkpoint does not use when a model declares a backwards-compatibility *load* rename as a reversible ``WeightRenaming``. Reversing such a rule rewrote the canonical name to a legacy spelling, and a deployment runtime keyed on the canonical name then failed to load the quantized checkpoint. Each exported name is now resolved against the names the source checkpoint actually carries, for both the exported tensors and the quantization config's module references. diff --git a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py index 1893ed3e800..454d18bdaf6 100644 --- a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py +++ b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py @@ -63,7 +63,7 @@ static_blockwise_fp4_fake_quant, ) from ...utils import is_torch_export_mode -from ...utils.numeric_utils import fp8_max_for_normalization +from ...utils.numeric_utils import E4M3_MAX, fp8_max_for_normalization from ..functional import normalized_hadamard_transform # torch.finfo(...).tiny gives the smallest normal E4M3 value; scale clamping needs @@ -1730,7 +1730,9 @@ def _block_scale_from_amax(self, amax: torch.Tensor, quantize: bool) -> torch.Te """Compute the per-block scale from a per-block amax, optionally FP8-quantizing it.""" if quantize: per_tensor_scale = _amax_to_scale(self.global_amax, self._quant_max_bound) - min_value = _FP8_E4M3_MIN_POSITIVE * per_tensor_scale.view(-1) + # scaled_e4m3 maps per_tensor_scale to E4M3_MAX (448); apply the inverse + # mapping to express the minimum encoded FP8 scale in physical scale units. + min_value = (_FP8_E4M3_MIN_POSITIVE / E4M3_MAX) * per_tensor_scale.view(-1) scale = _amax_to_scale(amax, self._quant_max_bound, min_value=min_value) return scaled_e4m3(scale, per_tensor_scale, None, 4, 3) return _amax_to_scale(amax, self._quant_max_bound, min_value=1e-8) diff --git a/tests/gpu/torch/quantization/test_lsq_cuda.py b/tests/gpu/torch/quantization/test_lsq_cuda.py index 4e9890737ec..465be40467e 100644 --- a/tests/gpu/torch/quantization/test_lsq_cuda.py +++ b/tests/gpu/torch/quantization/test_lsq_cuda.py @@ -152,9 +152,19 @@ def test_lsq_quantize_e2e(config): assert out.shape == (2, 64) -def test_lsq_fp4_fake_quantize_differentiable(): - """Test that _fake_quantize in FP4 LSQ mode is differentiable.""" +@pytest.mark.parametrize("tied_amax", [True, False], ids=["tied", "default_dual"]) +def test_lsq_fp4_fake_quantize_differentiable(tied_amax): + """Representable small scales retain static outputs and analytical LSQ gradients.""" device = torch.device("cuda") + scales = torch.tensor([2.0**-5, 1.0, 448.0], device=device) + codes = torch.tensor( + [0.625, -0.625, 1, -1, 1.5, -1.5, 2, -2, 3, -3, 4, -4, 6, -6, 0, 0], device=device + ) + x = scales[:, None] * codes + rounded = codes.clone() + rounded[:2] = torch.tensor([0.5, -0.5], device=device) + expected = scales[:, None] * rounded + tq = TensorQuantizer() tq._num_bits = (2, 1) tq._unsigned = False @@ -162,25 +172,32 @@ def test_lsq_fp4_fake_quantize_differentiable(): tq._disabled = False tq._block_sizes = {-1: 16, "type": "static", "scale_bits": (4, 3)} tq._pass_through_bwd = True - tq.register_buffer("_amax", torch.ones(4, device=device)) + tq.register_buffer("_amax", 6.0 * scales) tq.to(device) sbsq = StaticBlockScaleQuantizer.from_tensor_quantizer( - tq, global_amax=torch.tensor(1.0, device=device) + tq, global_amax=torch.tensor(6.0 * 448.0, device=device) ) - # global_amax=1.0 with NVFP4 _quant_max_bound=6.0 yields per_tensor_scale = 1/6. - sbsq.amax = torch.ones(4, device=device) * 3.0 - sbsq.enable_lsq( - quantize_scales=True, - learnable_amax=["post"], + static_output = sbsq(x) + torch.testing.assert_close(static_output, expected, rtol=0, atol=0) + if tied_amax: + sbsq.enable_lsq(learnable_amax=["pre", "post"], tied_amax=True) + else: + sbsq.enable_lsq() + + out = sbsq(x) + torch.testing.assert_close(out, static_output, rtol=0, atol=0) + (out[0, 0] + out[1, 0]).backward() + # At x/scale=0.625, Q(x/scale)=0.5 and d(scale)/d(amax)=1/6. + # Tied: (0.5 - 0.625)/6; default dual: 0.5/6 with fixed pre scale. + gradient = -1.0 / 48.0 if tied_amax else 1.0 / 12.0 + torch.testing.assert_close( + sbsq.amax_post.grad, + torch.tensor([gradient, gradient, 0.0], device=device), + rtol=1e-6, + atol=1e-7, ) - x = torch.randn(4, 16, device=device) - out = sbsq._fake_quantize(x) - assert out.shape == x.shape - out.sum().backward() - assert sbsq._amax_post.grad is not None - @pytest.mark.parametrize("tied_amax", [False, True]) def test_lsq_fp4_small_global_amax(tied_amax): diff --git a/tests/unit/torch/quantization/test_lsq.py b/tests/unit/torch/quantization/test_lsq.py index 1a1136d6ece..623b92342f1 100644 --- a/tests/unit/torch/quantization/test_lsq.py +++ b/tests/unit/torch/quantization/test_lsq.py @@ -40,6 +40,7 @@ _amax_to_scale, ) from modelopt.torch.quantization.tensor_quant import int_cast_ste +from modelopt.torch.quantization.utils.numeric_utils import E4M3_MAX from modelopt.torch.quantization.utils.shared_input import SharedWeightGlobalAmaxState from modelopt.torch.utils import to_empty_if_meta_device @@ -426,6 +427,7 @@ def spy_block_scale(amax, quantize): assert quantize_flags == [True, False] def test_skip_pre_scale_quantization_uses_raw_scale_floor(self, monkeypatch): + """Keep the raw pre-scale floor unchanged while normalizing the FP8 post-scale floor.""" q = self._make_lsq_quantizer() q._quantize_scales = True q._quantize_pre_scale = False @@ -446,7 +448,7 @@ def fake_amax_to_scale(amax, maxbound, min_value=1e-8): out = q._fake_quantize(torch.randn(4, 16)) assert out.shape == (4, 16) - assert torch.equal(min_values[0], torch.tensor([_FP8_E4M3_MIN_POSITIVE])) + assert torch.equal(min_values[0], torch.tensor([_FP8_E4M3_MIN_POSITIVE / E4M3_MAX])) assert min_values[1] == 1e-8 @@ -484,6 +486,7 @@ def _make_tied_lsq_group(self, global_amax=3.0, n_members=3): return members def test_block_scale_tracks_shared_update(self): + """Check each member's FP8 scale floor against the updated shared global amax.""" members = self._make_tied_lsq_group(global_amax=3.0) new_value = 5.0 # Mutate the shared buffer in place, mimicking export unification. @@ -493,7 +496,7 @@ def test_block_scale_tracks_shared_update(self): expected = _amax_to_scale(torch.tensor(new_value), member._quant_max_bound) scale = member._block_scale_from_amax(member.amax_post, quantize=True) assert scale.shape == member.amax_post.shape - assert torch.all(scale >= _FP8_E4M3_MIN_POSITIVE * expected) + assert torch.all(scale >= (_FP8_E4M3_MIN_POSITIVE / E4M3_MAX) * expected) def test_members_produce_identical_output_after_shared_update(self): members = self._make_tied_lsq_group(global_amax=3.0) From a95a86a719278ebb1ac32f551a09bd8c9c2e1e5c Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Sat, 17 Oct 2026 03:29:17 +0530 Subject: [PATCH 33/35] Load distributed object payloads with weights_only (#2837) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Type of change: Bug fix `broadcast`/`allgather` in `modelopt/torch/utils/distributed.py` and the Megatron expert-parallel (EP) export gather handled peer-supplied buffers with unrestricted pickle (`torch.load(..., weights_only=False)` and `torch.distributed.gather_object`), so a compromised rank could execute code on the others. - `_deserialize` and the EP export load use `torch.load(..., weights_only=True)`, passed explicitly so `TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD=1` cannot re-enable pickle. - The EP export gathers raw bytes with a new `gather_bytes` (root-only, chunked tensor gather) instead of `gather_object`, which unpickles the outer object before any restricted load. - `Hparam`/`TracedHp` (broadcast by the Megatron pruning path) are registered as safe globals, and `Hparam` pickles its default importance estimator as a marker instead of a bound method (which needs `getattr`). - Not covered: `DistributedProcessGroup.get_dist_syncd_obj` uses `all_gather_object`; left for a follow-up. **Why not `utils/serialization.py::safe_load`, JSON or safetensors?** `safe_load` deliberately honors `TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD` so users can trust a local checkpoint; trusting a local file must not also trust every peer, so peer payloads pin `weights_only=True`. Payloads are small heterogeneous Python objects (bools, dicts, tensors, hparams), which JSON/safetensors cannot carry without a larger redesign. ### Usage No API change. Custom classes sent through `broadcast`/`allgather` must be registered with `torch.serialization.add_safe_globals`. ### Testing - `tests/unit/torch/utils/test_distributed.py`: a `__reduce__` gadget is rejected by `_deserialize`, with and without `TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD=1` (fails when the fix is reverted); a `TracedHp` subclass (`DepthHparam`) round-trips through `_serialize`/`_deserialize`. - `tests/gpu/torch/utils/test_distributed.py`, run locally on 2 GPUs (8 passed): real `broadcast`/`allgather` round trip (incl. `TracedHp`); `gather_bytes` (uneven sizes, empty sender, small chunks); `gather_state_dicts` (the EP export's gather-and-load) accepts ordinary state dicts on the root only and rejects a gadget sent from a non-root rank without running it, with and without the env override. - The `gather_*` tests need 2 ranks and use `need_2_gpus`, so they skip on CI's single-GPU `gpu` runner; checked locally with `CUDA_VISIBLE_DEVICES=0` (4 passed, 4 skipped). - `tests/unit/torch/{nas,opt}`: all pass locally except 4 `opt/plugins` collection errors that also occur without this change (broken package metadata in the local env). - CI: `gpu_megatron` (real Megatron EP export and NAS/prune paths, which the local env cannot import), `gpu_vllm`, `linux` unit and example jobs passed on 1661191391. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ❌ objects of unregistered classes sent via `broadcast`/`allgather` now fail to load; register them with `add_safe_globals`. - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: ✅ - Did you update Changelog?: ✅ - Did you get Claude approval on this PR?: ❌ ## Summary by CodeRabbit * **New Features** * Distributed workflows can gather byte payloads and state dictionaries at a destination rank. * **Security** * Distributed data transfers use restricted deserialization to help prevent unsafe object execution, including when a PyTorch environment setting would otherwise allow unrestricted loading. * Expert-parallel exports transfer data without loading pickled objects. * **Compatibility** * Distributed broadcasts and allgathers reject unregistered custom classes. Register classes with PyTorch’s safe-globals mechanism before using them in these operations. * Hyperparameters and traced hyperparameters are registered for safe PyTorch checkpoint loading. --------- Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Co-authored-by: Claude Sonnet 5.5 Signed-off-by: Rohan Joshi --- CHANGELOG.rst | 2 + .../torch/export/unified_export_megatron.py | 25 ++---- modelopt/torch/nas/hparams/concat.py | 3 + modelopt/torch/nas/hparams/container.py | 5 ++ modelopt/torch/nas/plugins/megatron.py | 3 + modelopt/torch/nas/traced_hp.py | 5 ++ modelopt/torch/opt/hparam.py | 16 ++++ modelopt/torch/utils/distributed.py | 64 ++++++++++++- tests/gpu/torch/utils/test_distributed.py | 89 ++++++++++++++++++- tests/unit/torch/utils/test_distributed.py | 27 +++++- 10 files changed, 213 insertions(+), 26 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index ccdcb86260f..9ed9ca6fb13 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -68,6 +68,7 @@ Changelog **Backward Breaking Changes** +- ``modelopt.torch.utils.distributed.broadcast``/``allgather`` now reject objects of unregistered custom classes. Register them with ``torch.serialization.add_safe_globals``. - Bump minimum transformers version to ``5.5`` instead of ``4.57``; transformers 4.x is no longer supported. Upgrade with ``pip install -U "nvidia-modelopt[hf]"``. - ``modelopt.torch.distill.plugins.megatron.TopKLogitsKLLoss`` (``logit_kl_topk`` in ``DistillationConfig``) is renamed to ``TopLogitsKLLoss``, keeping the old name as a deprecated alias. It now normalizes both distributions over the full vocabulary instead of re-normalizing over the Top-K entries, and always appends a "ghost" token holding the probability mass outside the Top-K to both student and teacher (matching Megatron-LM's offline cached-logits KD loss). Loss values change for existing ``logit_kl_topk`` runs. - ``LogitsAndIntermediatesLossBalancer`` (Megatron distillation plugin) no longer rescales the distillation loss to the magnitude of the LM loss when the LM loss is included. The total is now the fixed convex combination ``(1 - alpha) * lm_loss + alpha * kd_loss`` with ``DistillationConfig.kd_loss_alpha`` (in [0, 1]), matching Megatron-LM's offline cached-logits KD. The default ``kd_loss_alpha=1.0`` skips the LM loss, as before. ``DistillationConfig.skip_lm_loss`` and ``kd_loss_scale`` are removed and raise ``ValueError`` if passed; set ``kd_loss_alpha`` instead. In ``examples/megatron_bridge/distill.py``, ``--kd_loss_alpha`` replaces ``--no_skip_lm_loss`` and ``--kd_loss_scale``. @@ -129,6 +130,7 @@ Changelog **Bug Fixes** +- Harden the distributed object broadcast/allgather helpers and the Megatron expert-parallel export gather so a compromised peer rank can no longer execute code through a crafted pickle payload: payloads load with ``weights_only=True`` and the export gathers raw bytes instead of pickled objects. - Fix LSQ clamping small, representable FP8 block scales too aggressively, which could zero quantized values and suppress scale gradients. - Fix ONNX AutoTune benchmark initialization with TensorRT 11 while preserving strongly typed networks on TensorRT 10. - Keep vision-block Linears in high precision in the ``qwen3_vl`` and ``qwen3_5`` FP8 vision recipes so exported checkpoints avoid a vLLM vision-attention shape failure. The recipes still quantize primary and deepstack merger Linears; re-quantize affected checkpoints before deployment. diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py index 1d6a0fc99d7..502c8f863de 100644 --- a/modelopt/torch/export/unified_export_megatron.py +++ b/modelopt/torch/export/unified_export_megatron.py @@ -18,7 +18,6 @@ """Code that export quantized Megatron Core models for deployment.""" -import io import json import os import re @@ -38,6 +37,7 @@ from modelopt.torch.quantization.ggml import GGML_FORMAT_REGISTRY from modelopt.torch.quantization.nn.modules.tensor_quantizer import GroupedQuantizer from modelopt.torch.utils import import_plugin, warn_rank_0 +from modelopt.torch.utils.distributed import gather_state_dicts from modelopt.torch.utils.plugins.hf_checkpoint_utils import ( copy_hf_ckpt_remote_code, copy_non_safetensor_files_from_ckpt, @@ -1709,29 +1709,14 @@ def _grouped_mlp_slicing( self._record_layer_quant_config(expert_prefix, seen_qformat, seen_block_size) if ep_size > 1: - # all_gather_object pickles trip on quantized uint8 tensors whose - # UntypedStorage has no dtype attr -- round-trip through torch.save bytes. - _buf = io.BytesIO() - torch.save(local_expert_state, _buf) - local_bytes = _buf.getvalue() - del _buf # Gather to EP rank 0 only, which writes the shards: holding every expert on every EP # rank multiplies host memory by EP (a full GLM-5.3-Flash export OOMs at EP4). + # Quantized uint8 tensors trip all_gather_object pickles, so go through torch.save bytes. ep_group = get_expert_model_parallel_group() - gathered_bytes: list | None = [None] * ep_size if ep_rank == 0 else None - torch.distributed.gather_object( - local_bytes, - gathered_bytes, - dst=torch.distributed.get_global_rank(ep_group, 0), - group=ep_group, - ) - del local_bytes - for b in gathered_bytes or (): - # weights_only=False: our own torch.save output from a sibling EP rank - # in this job's collective, not user-supplied. - s_loaded = torch.load(io.BytesIO(b), map_location="cpu", weights_only=False) + for s_loaded in gather_state_dicts( + local_expert_state, torch.distributed.get_global_rank(ep_group, 0), ep_group + ): self._state_dict.update(s_loaded) - del gathered_bytes else: self._state_dict.update(local_expert_state) diff --git a/modelopt/torch/nas/hparams/concat.py b/modelopt/torch/nas/hparams/concat.py index 31274ab052e..9234939e1ed 100644 --- a/modelopt/torch/nas/hparams/concat.py +++ b/modelopt/torch/nas/hparams/concat.py @@ -317,3 +317,6 @@ def build_concat_hp(inputs: list[TracedHp]): concat_hp._importance_estimators = None return concat_hp + + +torch.serialization.add_safe_globals([ConcatTracedHp]) diff --git a/modelopt/torch/nas/hparams/container.py b/modelopt/torch/nas/hparams/container.py index dc7ee17f2f2..61548b29852 100644 --- a/modelopt/torch/nas/hparams/container.py +++ b/modelopt/torch/nas/hparams/container.py @@ -17,6 +17,8 @@ from collections.abc import Callable +import torch + from modelopt.torch.trace import Symbol, SymDepth from ..traced_hp import TracedHp, TracedHpRegistry @@ -46,3 +48,6 @@ def _resolve_dependencies( self.choices = list(choices & set(self.choices)) return super()._resolve_dependencies(sym, get_hp) + + +torch.serialization.add_safe_globals([DepthHparam]) diff --git a/modelopt/torch/nas/plugins/megatron.py b/modelopt/torch/nas/plugins/megatron.py index c2d31b6f3de..9daa81f0cf4 100644 --- a/modelopt/torch/nas/plugins/megatron.py +++ b/modelopt/torch/nas/plugins/megatron.py @@ -1100,6 +1100,9 @@ def _resolve_dependencies( raise NotImplementedError("MambaDInnerHp does not support `_resolve_dependencies`!") +torch.serialization.add_safe_globals([NumAttentionHeadsHp, MambaNumHeadsHp, MambaDInnerHp]) + + class _DynamicExtendedRMSNorm(DynamicModule): """An ExtendedRMSNorm (GroupNorm) layer with dynamic hyperparams. diff --git a/modelopt/torch/nas/traced_hp.py b/modelopt/torch/nas/traced_hp.py index 2f424bef0ce..0a29842f80e 100644 --- a/modelopt/torch/nas/traced_hp.py +++ b/modelopt/torch/nas/traced_hp.py @@ -17,6 +17,8 @@ from collections.abc import Callable +import torch + from modelopt.torch.opt.hparam import Hparam from modelopt.torch.trace import Symbol @@ -147,3 +149,6 @@ def _resolve_dependencies( self.choices = [self.original] return mapping + + +torch.serialization.add_safe_globals([TracedHp]) diff --git a/modelopt/torch/opt/hparam.py b/modelopt/torch/opt/hparam.py index c36e2ec84e4..df537c68352 100644 --- a/modelopt/torch/opt/hparam.py +++ b/modelopt/torch/opt/hparam.py @@ -82,6 +82,19 @@ def __init__(self, choices: Sequence[HPType], original: HPType | None = None) -> # an optional order to enforce for active_slice self._slice_order: torch.LongTensor | None = None + def __getstate__(self): + # A bound method pickles via ``getattr``, which ``weights_only`` loading rightly rejects. + state = self.__dict__.copy() + # Custom estimators stay as-is (deepcopy relies on this) and are rejected by weights_only loads. + if state.get("_importance_estimators") == [self._default_get_importance]: + state["_importance_estimators"] = "default" + return state + + def __setstate__(self, state): + self.__dict__.update(state) + if state.get("_importance_estimators") == "default": + self._importance_estimators = [self._default_get_importance] + def __iter__(self) -> Generator[HPType, None, None]: """Iterate over choices.""" yield from self.choices @@ -284,3 +297,6 @@ def __iand__(self, hp: "Hparam"): self._importance_estimators = None return self + + +torch.serialization.add_safe_globals([Hparam]) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index c3cea6acaf4..181bcff632d 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -21,7 +21,7 @@ import sys import time import traceback -from collections.abc import Callable +from collections.abc import Callable, Iterator from contextlib import suppress from datetime import timedelta from typing import Any @@ -38,6 +38,8 @@ "backend", "barrier", "fsdp2_wrap", + "gather_bytes", + "gather_state_dicts", "is_available", "is_fsdp2_model", "is_initialized", @@ -108,9 +110,9 @@ def _deserialize(tensor: torch.Tensor, size: int | None = None) -> Any: buffer = tensor.numpy().tobytes() if size is not None: buffer = buffer[:size] - # Security NOTE: weights_only=False is used here on internally-generated buffer, not on untrusted user input - obj = torch.load(io.BytesIO(buffer), weights_only=False) - return obj + # Explicit so TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD cannot re-enable pickle for peer payloads. Custom + # classes sent through broadcast/allgather must be registered with add_safe_globals. + return torch.load(io.BytesIO(buffer), weights_only=True) def _broadcast(tensor: torch.Tensor, src: int = 0, group=None) -> None: @@ -144,6 +146,60 @@ def broadcast(obj: Any, src: int = 0, group=None) -> Any: return obj +def gather_bytes( + data: bytes, dst: int, group=None, chunk_bytes: int = 1 << 28 +) -> list[bytearray] | None: + """Gathers byte strings to global rank ``dst`` only, without unpickling anything on the way. + + Unlike ``torch.distributed.gather_object`` it never runs pickle, and it moves ``chunk_bytes`` per + rank at a time to bound device memory. Returns the per-rank bytes on ``dst`` and ``None`` elsewhere. + """ + n = size(group) + device = torch.device("cuda" if torch.distributed.get_backend(group) == "nccl" else "cpu") + local = ( + torch.frombuffer(bytearray(data), dtype=torch.uint8) + if data + else torch.empty(0, dtype=torch.uint8) + ) + sizes_t = [torch.zeros(1, dtype=torch.long, device=device) for _ in range(n)] + torch.distributed.all_gather(sizes_t, torch.tensor([len(data)], device=device), group=group) + sizes = [int(t.item()) for t in sizes_t] + + is_dst = torch.distributed.get_rank() == dst + out = [bytearray(s) for s in sizes] if is_dst else None + for off in range(0, max(sizes), chunk_bytes): + length = min(chunk_bytes, max(sizes) - off) + buf = torch.zeros(length, dtype=torch.uint8, device=device) + piece = local[off : off + length] + buf[: piece.numel()] = piece.to(device) + bufs = [torch.empty_like(buf) for _ in range(n)] if is_dst else None + torch.distributed.gather(buf, bufs, dst=dst, group=group) + if bufs is not None and out is not None: + for i, b in enumerate(bufs): + k = min(length, max(0, sizes[i] - off)) + out[i][off : off + k] = b[:k].cpu().numpy().tobytes() + return out + + +def gather_state_dicts(state: dict, dst: int, group=None) -> Iterator[dict]: + """Gathers ``state`` from every rank of ``group`` to global rank ``dst``, loaded with ``weights_only``. + + Yields one rank's dict at a time on ``dst`` (to bound host memory) and nothing elsewhere. The collective + runs eagerly, so every rank must call this before iterating. + """ + buf = io.BytesIO() + torch.save(state, buf) + gathered = gather_bytes(buf.getvalue(), dst, group) + del buf + + def _load() -> Iterator[dict]: + while gathered: + # Explicit weights_only so TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD cannot re-enable pickle for peers. + yield torch.load(io.BytesIO(gathered.pop(0)), map_location="cpu", weights_only=True) + + return _load() + + def _allgather(tensors: list[torch.Tensor], tensor: torch.Tensor, group=None) -> None: if backend() == "torch": torch.distributed.all_gather(tensors, tensor, group) diff --git a/tests/gpu/torch/utils/test_distributed.py b/tests/gpu/torch/utils/test_distributed.py index 4091c7465dc..d1581e1ca03 100644 --- a/tests/gpu/torch/utils/test_distributed.py +++ b/tests/gpu/torch/utils/test_distributed.py @@ -15,6 +15,8 @@ """GPU/distributed tests for ``modelopt.torch.utils.distributed``.""" +import os +import pickle from functools import partial import pytest @@ -25,7 +27,14 @@ from torch.distributed.checkpoint.state_dict import StateDictOptions, set_model_state_dict from torch.distributed.tensor import DTensor -from modelopt.torch.utils.distributed import fsdp2_wrap +from modelopt.torch.nas.traced_hp import TracedHp +from modelopt.torch.utils.distributed import ( + allgather, + broadcast, + fsdp2_wrap, + gather_bytes, + gather_state_dicts, +) VOCAB_SIZE = 32 N_EXPERTS = 4 @@ -129,3 +138,81 @@ def test_fsdp2_wrap_moves_ignored_params_to_device(dist_workers, cpu_offload): dist_workers.run( partial(_test_fsdp2_wrap_moves_ignored_params_to_device, cpu_offload=cpu_offload) ) + + +def _test_broadcast_allgather_objects(rank, size): + hp = TracedHp([1, 2, 4]) if rank == 0 else None + out = broadcast([True, {"latency": 1.5}, hp], src=0) + assert out[0] is True + assert out[1] == {"latency": 1.5} + assert out[2].choices == [1, 2, 4] + assert allgather({"rank": rank}) == [{"rank": r} for r in range(size)] + + +def test_broadcast_allgather_objects(dist_workers): + dist_workers.run(_test_broadcast_allgather_objects) + + +def _test_gather_bytes(rank, size): + data = bytes([rank + 1]) * (5 + 3 * rank) if rank else b"" + out = gather_bytes(data, dst=1, group=None, chunk_bytes=4) + if rank == 1: + assert out == [ + bytearray(b"" if r == 0 else bytes([r + 1]) * (5 + 3 * r)) for r in range(size) + ] + else: + assert out is None + + +def test_gather_bytes(need_2_gpus, dist_workers): + dist_workers.run(_test_gather_bytes) + + +class _Gadget: + def __init__(self, marker): + self.marker = marker + + def __reduce__(self): + return (os.mkdir, (self.marker,)) + + +def _test_gather_state_dicts(rank, size): + out = list(gather_state_dicts({f"w{rank}": torch.full((2,), float(rank))}, dst=0)) + if rank == 0: + assert [list(d) for d in out] == [[f"w{r}"] for r in range(size)] + assert torch.equal(out[1]["w1"], torch.ones(2)) + else: + assert out == [] + + +def test_gather_state_dicts(need_2_gpus, dist_workers): + dist_workers.run(_test_gather_state_dicts) + + +def _test_gather_state_dicts_rejects_gadget(rank, size, marker, force_no_weights_only): + if force_no_weights_only: + os.environ["TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD"] = "1" + # A non-root peer sends a pickle gadget; the root must refuse it without running it. + state = {"w": _Gadget(marker)} if rank == 1 else {"w": torch.zeros(1)} + gathered = gather_state_dicts(state, dst=0) + if rank == 0: + with pytest.raises(pickle.UnpicklingError): + list(gathered) + assert not os.path.exists(marker) + else: + assert list(gathered) == [] + + +@pytest.mark.parametrize("force_no_weights_only", [False, True]) +def test_gather_state_dicts_rejects_gadget( + need_2_gpus, dist_workers, tmp_path, force_no_weights_only +): + marker = str(tmp_path / "gadget_ran") + dist_workers.run( + partial( + _test_gather_state_dicts_rejects_gadget, + marker=marker, + force_no_weights_only=force_no_weights_only, + ) + ) + assert not os.path.exists(marker) diff --git a/tests/unit/torch/utils/test_distributed.py b/tests/unit/torch/utils/test_distributed.py index 7b630f93d50..f88377e2a5f 100644 --- a/tests/unit/torch/utils/test_distributed.py +++ b/tests/unit/torch/utils/test_distributed.py @@ -15,11 +15,15 @@ """Pure-function tests for ``modelopt.torch.utils.distributed``.""" +import os +import pickle + import pytest import torch import torch.nn as nn -from modelopt.torch.utils.distributed import _off_dtype_params +from modelopt.torch.nas.hparams.container import DepthHparam +from modelopt.torch.utils.distributed import _deserialize, _off_dtype_params, _serialize def _model(*sizes_and_dtypes) -> nn.Module: @@ -61,3 +65,24 @@ def test_off_dtype_params_dominant_is_by_numel(params, expected): off = _off_dtype_params(model) kept = {n for n, p in model.named_parameters() if p not in off} assert kept == {expected} + + +class _Gadget: + def __reduce__(self): + return (os.system, ("exit 1",)) + + +@pytest.mark.parametrize("force_no_weights_only", [False, True]) +def test_deserialize_rejects_pickle_gadget(monkeypatch, force_no_weights_only): + # The override that trusts local checkpoints must not also trust peer payloads. + if force_no_weights_only: + monkeypatch.setenv("TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD", "1") + with pytest.raises(pickle.UnpicklingError): + _deserialize(_serialize(_Gadget())) + + +def test_deserialize_traced_hp_subclass(): + # The allowlist matches exact classes, so each TracedHp subclass must be registered too. + hp = _deserialize(_serialize(DepthHparam([1, 2]))) + assert type(hp) is DepthHparam + assert hp.choices == [1, 2] From 2a52aa90b5e63d0e937df8647c65f60725ff4a15 Mon Sep 17 00:00:00 2001 From: Chad Voegele Date: Fri, 16 Oct 2026 17:29:27 -0500 Subject: [PATCH 34/35] Fix scale-aware PTQ for native FP8 experts (#2820) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Chad's Agent Type of change: Bug fix. Native Transformers `FP8Experts` store fused FP8 weights with inverse scales. The generic wrapper exposed raw values to weight calibration and packing, while native FP8 dispatch bypassed ModelOpt's fake quantization. Dequantize one expert projection at a time for forward, calibration, and export, keeping scale multiplication in FP32 until the output cast. Export removes consumed inverse/activation scales. Native packed FP4 inputs are rejected. Explicit `unpack_weight()` materializes floating parameters before folding, matching FP8Linear's interface. Folding is inherited from the fused wrapper. Native forward still reads scale attributes after unpacking, so export removes them. Ordinary PTQ/export keeps per-expert dequantization without unpacking entire modules. Raw FP8 folding retains FP8Linear's existing unsupported behavior; shared folding code is unchanged. Independent of recipe PR #2817; analogous to #1962. FP8Linear is unchanged. Use a recipe matching native expert names; M2.7 validation also needs Call `experts.unpack_weight()` explicitly before permanent or temporary folding. - Current revision: 99 focused fused-expert/FP8Linear plugin tests pass on both Transformers 5.5.0 and 5.18.0. Gated/non-gated temporary-fold tests verify parameter identity, dtype, exact values, quantizer state, and forward results after normal and exceptional exits. The new tests fail on the previous head without explicit unpacking. - Six real native FP8Experts GPU cases pass across gated/non-gated layouts and FP32/FP16/BF16 inputs, with NVFP4 calibration/export and independent decoding. - Fresh full MiniMax M2.7 PTQ/export passes with #2817: 512 samples, batch 8, 64/64 calibration batches, 14 shards (134.52 GB versus 230.12 GB source payload). All 47,616 expert projections have correct packed/block-scale shapes. All 191,087 tensor headers and shard extents agree with the index; no stale inverse/activation scales or declaration mismatches; reverse naming succeeds. - 27 crops (layers 0/30/61, experts 0/127/255, gate/up/down) have 9.28–9.53% relative L2 error and norm ratios 0.997–1.003 against the **scaled FP8 source**. Independent E2M1 and ModelOpt NVFP4 decoders agree within FP32 rounding. Both generation previews complete. - Controlled environment: four GB300 GPUs, aarch64, `nvcr.io/nvidia/pytorch:26.08-py3`, Transformers 5.18.0, compressed-tensors 0.19.0, kernels 0.17.0, datasets 5.0.1; original ModelOpt 0.48.0rc0 source plus backend and recipe patches. Wrapper ASTs/export code match this PR. Model revision `d494266a4affc0d2995ba1fa35c8481cbd84294b`; CNN/DailyMail revision `96df5e686bee6baa90b8bee7c28b81fa3fa6223d`, config 3.0.0. The weight-quantizer guard remained active. - Fresh calibration: 26m10s; export: 150.1s; peak memory: 62.75–65.32 GB/GPU. Tokenizer/vocabulary/merges/chat-template files are byte-identical; generation settings differ only in the recorded Transformers version. Native config normalization and the loader-generated expert-parallel scale-sharding plan are accounted for. - Changed-file pre-commit hooks and `git diff --check` pass. Current repository CI: 32 successful checks, 1 skipped, 19 GPU/example jobs queued without assigned runners; no failures. Earlier revision: full-model smoke/full512 and 54 CI checks passed; the full512 results above are from the fresh rerun. - Backward compatible: ✅ - Copied code/new dependencies: N/A - Necessary tests: ✅ - Changelog: ✅ - Claude approval: N/A Full-model PTQ uses a four-GPU device map. FSDP2, offload, serving, and benchmark accuracy are untested. Excluded attention projections export as BF16, routers as FP32, and the LM head as BF16. --------- Signed-off-by: Chad Voegele Signed-off-by: Rohan Joshi --- CHANGELOG.rst | 1 + modelopt/torch/export/moe_utils.py | 19 +- .../torch/quantization/plugins/huggingface.py | 73 ++++- .../torch/quantization/test_fp8_experts.py | 81 ++++++ .../plugins/test_fused_experts.py | 259 +++++++++++++++++- 5 files changed, 421 insertions(+), 12 deletions(-) create mode 100644 tests/gpu/torch/quantization/test_fp8_experts.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 9ed9ca6fb13..38f8f6f5c50 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -183,6 +183,7 @@ Changelog - Fix ONNX preprocessing with custom TensorRT plugin libraries on TensorRT 11. - Fix GPTQ with Megatron data parallelism: each data-parallel replica of a module now builds its Hessian from the whole calibration set instead of only its own shard, so all replicas reach the same quantized weights. - Fix ONNX quantization with the default ``concat_elimination`` pass crashing in calibration with ``KeyError: '_ReduceMin'`` when a ``Concat`` input is a constant, such as the class token of a ViT. +- Fix PTQ of native Transformers FP8 experts to apply source weight scales during calibration and export. Re-quantize affected checkpoints from the original FP8 source. 0.47.0 (2026-09-23) ^^^^^^^^^^^^^^^^^^^ diff --git a/modelopt/torch/export/moe_utils.py b/modelopt/torch/export/moe_utils.py index 2affd111b02..02660142736 100644 --- a/modelopt/torch/export/moe_utils.py +++ b/modelopt/torch/export/moe_utils.py @@ -40,6 +40,11 @@ def _delete_fused_moe_source_attrs(module: nn.Module) -> None: "down_proj", "down_proj_weight_quantizers", "down_proj_input_quantizer", + f"{first_proj_attr}_scale_inv", + "down_proj_scale_inv", + "gate_up_proj_activation_scale", + "up_proj_activation_scale", + "down_proj_activation_scale", ): if hasattr(module, attr): delattr(module, attr) @@ -160,6 +165,8 @@ def _export_fused_experts( for idx in range(n): expert = nn.Module() + first_weight = module.get_expert_weight(first_proj_attr, idx, dtype) + down_weight = module.get_expert_weight("down_proj", idx, dtype) # If the gate_up source quantizer was never calibrated (rare expert # that received no calibration tokens), derive its amax once from the @@ -181,7 +188,7 @@ def _export_fused_experts( else [] ) for q in uncalibrated_stages: - q.amax = _fallback_amax(first_proj[idx], q) + q.amax = _fallback_amax(first_weight, q) if uncalibrated_stages: warnings.warn( f"Expert {idx} gate_up_proj weight quantizer was not calibrated " @@ -193,16 +200,16 @@ def _export_fused_experts( if is_gated: projections = [ - ("gate_proj", first_proj[idx, :expert_dim, :], 0, fused_dim0, True), - ("up_proj", first_proj[idx, expert_dim:, :], expert_dim, fused_dim0, True), - ("down_proj", down[idx], 0, down.shape[1], False), + ("gate_proj", first_weight[:expert_dim, :], 0, fused_dim0, True), + ("up_proj", first_weight[expert_dim:, :], expert_dim, fused_dim0, True), + ("down_proj", down_weight, 0, down.shape[1], False), ] else: # Non-gated: the single up_proj maps 1:1 to its weight quantizer, so it # is exported whole (no dim-0 split, no shared gate/up weight_scale_2). projections = [ - ("up_proj", first_proj[idx], 0, fused_dim0, True), - ("down_proj", down[idx], 0, down.shape[1], False), + ("up_proj", first_weight, 0, fused_dim0, True), + ("down_proj", down_weight, 0, down.shape[1], False), ] for ( diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index 4110a95e5bd..6a6c5d18bfc 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -28,7 +28,7 @@ import transformers from torch import Tensor from torch.nn.functional import linear -from transformers.integrations.finegrained_fp8 import FP8Linear +from transformers.integrations.finegrained_fp8 import FP8Experts, FP8Linear from transformers.models.t5.modeling_t5 import T5Attention from modelopt.torch.kernels.common.attention import IS_AVAILABLE as TRITON_FA_AVAILABLE @@ -868,12 +868,16 @@ def _quantized_linear(input, weight, bias=None): if self._down_proj_linear: idx = self._current_expert_idx input = self.down_proj_input_quantizer(input) - weight = self.down_proj_weight_quantizers[idx](weight) + weight = self.down_proj_weight_quantizers[idx]( + self.get_expert_weight("down_proj", idx, input.dtype) + ) else: idx = self._get_expert_idx_from_first_proj(weight) self._current_expert_idx = idx input = getattr(self, self._first_proj_input_quantizer_attr)(input) - weight = getattr(self, self._first_proj_weight_quantizers_attr)[idx](weight) + weight = getattr(self, self._first_proj_weight_quantizers_attr)[idx]( + self.get_expert_weight(self._first_proj_attr, idx, input.dtype) + ) self._down_proj_linear = not self._down_proj_linear return _orig_linear(input, weight, bias) @@ -885,6 +889,10 @@ def forward(self, *args, **kwargs): self._down_proj_linear = False return super().forward(*args, **kwargs) + def get_expert_weight(self, name, idx, dtype=None): + """Return one expert's projection in the calibration/export weight domain.""" + return getattr(self, name)[idx] + def iter_weights_for_calibration(self): """Yield ``(weight_slice, quantizer)`` per-expert pairs. @@ -901,7 +909,7 @@ def iter_weights_for_calibration(self): if weight is None or quantizers is None: continue for idx, q in enumerate(quantizers): - yield weight[idx], q + yield self.get_expert_weight(weight_name, idx), q def iter_weight_quantizers_for_calibration(self): """Yield the per-expert weight quantizers without slicing the fused 3-D weight. @@ -947,6 +955,58 @@ class _QuantNonGatedFusedExperts(_QuantFusedExperts): _is_gated = False +class _QuantFP8Experts(_QuantFusedExperts): + """Scale-aware PTQ for native Transformers FP8 experts, dequantized per expert.""" + + @property + def _first_proj_attr(self): + return "gate_up_proj" if self.has_gate else "up_proj" + + @property + def _is_gated(self): + return self.has_gate + + def _setup(self): + if getattr(self.config, "expert_dtype", "fp8") != "fp8": + raise NotImplementedError("Only FP8 E4M3 native expert weights are supported") + super()._setup() + + def get_expert_weight(self, name, idx, dtype=None): + weight = getattr(self, name)[idx] + if weight.element_size() > 1: + return weight if dtype is None else weight.to(dtype) + dtype = dtype or torch.get_default_dtype() + scales = getattr(self, f"{name}_scale_inv")[idx].float() + rows, cols = self.block_size or (weight.shape[0] // scales.shape[0], weight.shape[1]) + scales = scales.repeat_interleave(rows, 0).repeat_interleave(cols, 1) + return (weight.float() * scales[: weight.shape[0], : weight.shape[1]]).to(dtype) + + def linear(self, input, weight, weight_scale_inv, activation_scale=None): + # Preserve the FP8 view's storage offset for the per-expert quantizer lookup. + return torch.nn.functional.linear(input, weight) + + def _validate_fold_weight(self): + if any( + getattr(self, name).element_size() == 1 for name in (self._first_proj_attr, "down_proj") + ): + raise ValueError("Call unpack_weight() before folding native FP8 expert weights.") + super()._validate_fold_weight() + + def fold_weight(self, keep_attrs: bool = False): + self._validate_fold_weight() + super().fold_weight(keep_attrs) + + def unpack_weight(self): + """Materialize floating weights before folding; retain scales for native forward.""" + for name in (self._first_proj_attr, "down_proj"): + weight = getattr(self, name) + if weight.element_size() == 1: + unpacked = torch.empty_like(weight, dtype=torch.get_default_dtype()) + for idx in range(self.num_experts): + unpacked[idx].copy_(self.get_expert_weight(name, idx)) + setattr(self, name, nn.Parameter(unpacked, requires_grad=False)) + + def _get_fused_experts_quantizer_attr_names(module): """Return quantizer attribute names for a converted fused-experts module.""" first_proj_attr = getattr(module, "_first_proj_attr", "gate_up_proj") @@ -1199,6 +1259,9 @@ def unpack_weight(self): except ImportError: pass +if FP8Experts not in QuantModuleRegistry: + QuantModuleRegistry.register({FP8Experts: "hf.FP8Experts"})(_QuantFP8Experts) + if FP8Linear not in QuantModuleRegistry: QuantModuleRegistry.register({FP8Linear: "hf.FP8Linear"})(_QuantFP8Linear) diff --git a/tests/gpu/torch/quantization/test_fp8_experts.py b/tests/gpu/torch/quantization/test_fp8_experts.py new file mode 100644 index 00000000000..f33e80e7e12 --- /dev/null +++ b/tests/gpu/torch/quantization/test_fp8_experts.py @@ -0,0 +1,81 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""Native FP8 expert PTQ and packed NVFP4 export.""" + +import pytest +import torch +from transformers import PretrainedConfig +from transformers.integrations.finegrained_fp8 import FP8Experts + +import modelopt.torch.quantization as mtq +from modelopt.torch.export.moe_utils import _export_fused_experts + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("has_gate", [True, False]) +def test_fp8_experts_nvfp4_export(has_gate, dtype): + torch.manual_seed(42) + config = PretrainedConfig( + hidden_size=256, intermediate_size=256, num_experts=2, hidden_act="silu" + ) + experts = FP8Experts(config, block_size=(128, 128), has_gate=has_gate).cuda() + first = "gate_up_proj" if has_gate else "up_proj" + expected = {} + for name in (first, "down_proj"): + weight = getattr(experts, name) + scales = getattr(experts, f"{name}_scale_inv") + with torch.no_grad(): + weight.copy_(torch.randn(weight.shape, device="cuda").to(weight.dtype)) + scales.copy_( + torch.linspace(0.01, 0.08, scales.numel(), device="cuda").reshape(scales.shape) + ) + expected[name] = weight.float() * scales.repeat_interleave(128, 1).repeat_interleave(128, 2) + inputs = torch.randn(2, 256, device="cuda", dtype=dtype) + indices = torch.tensor([[0], [1]], device="cuda") + routing = torch.ones(2, 1, device="cuda") + mtq.quantize(experts, mtq.NVFP4_DEFAULT_CFG, lambda m: m(inputs, indices, routing)) + assert torch.isfinite(experts(inputs, indices, routing)).all() + _export_fused_experts(experts, torch.float32) + levels = torch.tensor( + [0, 0.5, 1, 1.5, 2, 3, 4, 6, 0, -0.5, -1, -1.5, -2, -3, -4, -6], device="cuda" + ) + for idx in range(2): + exported = getattr(experts, str(idx)) + projections = [("down_proj", expected["down_proj"][idx])] + if has_gate: + projections += [ + ("gate_proj", expected[first][idx, :256]), + ("up_proj", expected[first][idx, 256:]), + ] + torch.testing.assert_close( + exported.gate_proj.weight_scale_2, exported.up_proj.weight_scale_2 + ) + else: + projections += [("up_proj", expected[first][idx])] + for name, reference in projections: + proj = getattr(exported, name) + packed = proj.weight + codes = torch.stack([packed & 15, packed >> 4], -1).reshape(reference.shape).long() + decoded = ( + levels[codes] + * proj.weight_scale.float().repeat_interleave(16, -1) + * proj.weight_scale_2 + ) + relative_l2 = torch.linalg.vector_norm(decoded - reference) / torch.linalg.vector_norm( + reference + ) + assert relative_l2 < 0.15 + assert not any("scale_inv" in key for key in experts.state_dict()) diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py index 959471cf13e..4f224e3b7c0 100644 --- a/tests/unit/torch/quantization/plugins/test_fused_experts.py +++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -15,7 +15,9 @@ """Tests for _QuantFusedExperts: generic fused MoE quantization and export.""" +import copy import warnings +from contextlib import nullcontext import pytest import torch @@ -25,6 +27,8 @@ pytest.importorskip("transformers") from _test_utils.torch.quantization.tied_modules import tie_fused_experts_3d_params +from transformers import PretrainedConfig +from transformers.integrations.finegrained_fp8 import FP8Experts import modelopt.torch.quantization as mtq import modelopt.torch.quantization.nn.modules.tensor_quantizer as tensor_quantizer_module @@ -47,6 +51,71 @@ ) from modelopt.torch.quantization.utils.core_utils import weight_attr_names + +def test_native_fp8_experts_dequantized_before_nvfp4_export(): + """Calibrate and pack scaled FP8 gate/up and down weights, preserving their BF16 values.""" + generator = torch.Generator().manual_seed(42) + experts = FP8Experts( + PretrainedConfig(hidden_size=256, intermediate_size=256, num_experts=2, hidden_act="silu"), + block_size=(128, 128), + ) + reference = {} + for name, scale_range in (("gate_up_proj", (0.01, 0.08)), ("down_proj", (0.15, 0.3))): + weight = getattr(experts, name) + scales = getattr(experts, f"{name}_scale_inv") + assert weight.dtype == torch.float8_e4m3fn + with torch.no_grad(): + weight.copy_(torch.randn(weight.shape, generator=generator).to(weight.dtype)) + scales.copy_(torch.linspace(*scale_range, scales.numel()).reshape(scales.shape)) + # Build only the reference here; ModelOpt must consume the untouched FP8 source. + blocks = weight.float().reshape(2, weight.shape[1] // 128, 128, 2, 128) + reference[name] = (blocks * scales[:, :, None, :, None]).reshape(weight.shape).bfloat16() + + registered_before = QuantModuleRegistry.get(FP8Experts) + try: + mtq.quantize(experts, mtq.NVFP4_DEFAULT_CFG) + for name, expected in reference.items(): + assert getattr(experts, name).dtype == torch.float8_e4m3fn + quantizers = getattr(experts, f"{name}_weight_quantizers") + for idx, quantizer in enumerate(quantizers): + assert quantizer.is_enabled and quantizer.num_bits == (2, 1) + torch.testing.assert_close( + quantizer.amax.float(), + expected[idx].float().abs().amax(), + rtol=0.004, + atol=0, + ) + + _export_fused_experts(experts, torch.bfloat16) + levels = torch.tensor([0, 0.5, 1, 1.5, 2, 3, 4, 6, 0, -0.5, -1, -1.5, -2, -3, -4, -6]) + for idx in range(2): + exported = getattr(experts, str(idx)) + torch.testing.assert_close( + exported.gate_proj.weight_scale_2, exported.up_proj.weight_scale_2 + ) + for name, expected in ( + ("gate_proj", reference["gate_up_proj"][idx, :256]), + ("up_proj", reference["gate_up_proj"][idx, 256:]), + ("down_proj", reference["down_proj"][idx]), + ): + projection = getattr(exported, name) + packed = projection.weight + assert packed.dtype == torch.uint8 + codes = torch.stack((packed & 15, packed >> 4), dim=-1).reshape(expected.shape) + decoded = ( + levels[codes.long()] + * projection.weight_scale.float().repeat_interleave(16, dim=-1) + * projection.weight_scale_2.float() + ) + relative_l2 = torch.linalg.vector_norm( + decoded - expected.float() + ) / torch.linalg.vector_norm(expected.float()) + assert relative_l2 < 0.15, f"Expert {idx} {name}: relative L2 {relative_l2}" + finally: + if registered_before is None and QuantModuleRegistry.get(FP8Experts) is not None: + QuantModuleRegistry.unregister(FP8Experts) + + # --------------------------------------------------------------------------- # Synthetic fused expert module matching the HF transformers 5.0+ pattern # --------------------------------------------------------------------------- @@ -1726,3 +1795,191 @@ def test_quantize_collects_amax(self): finally: if registered_before is None and QuantModuleRegistry.get(experts_type) is not None: QuantModuleRegistry.unregister(experts_type) + + +@pytest.mark.parametrize("activation_scheme", ["dynamic", "static"]) +@pytest.mark.parametrize("fold", [True, False]) +@pytest.mark.parametrize("has_gate", [True, False]) +@pytest.mark.parametrize("block_size", [(128, 128), (64, 128), None]) +def test_native_fp8_experts_scales(has_gate, block_size, fold, activation_scheme): + torch.manual_seed(42) + config = PretrainedConfig( + hidden_size=256, intermediate_size=256, num_experts=2, hidden_act="silu" + ) + experts = FP8Experts( + config, block_size=block_size, has_gate=has_gate, activation_scheme=activation_scheme + ) + first = "gate_up_proj" if has_gate else "up_proj" + expected = {} + for name in (first, "down_proj"): + weight = getattr(experts, name) + scales = getattr(experts, f"{name}_scale_inv") + with torch.no_grad(): + weight.copy_(torch.randn(weight.shape).to(weight.dtype)) + scales.copy_(torch.linspace(1e-8, 0.08, scales.numel()).reshape(scales.shape)) + rows, cols = block_size or (weight.shape[1] // scales.shape[1], weight.shape[2]) + expanded = scales.float().repeat_interleave(rows, 1).repeat_interleave(cols, 2) + expected[name] = weight.float() * expanded[:, : weight.shape[1], : weight.shape[2]] + + reference = copy.deepcopy(experts) + for name, weight in expected.items(): + setattr(reference, name, nn.Parameter(weight, requires_grad=False)) + mtq.quantize( + experts, {"quant_cfg": [{"quantizer_name": "*", "enable": False}], "algorithm": None} + ) + for (weight, _), expected_weight in zip( + experts.iter_weights_for_calibration(), + [w for name in (first, "down_proj") for w in expected[name]], + strict=True, + ): + torch.testing.assert_close(weight.float(), expected_weight) + + for name in (first, "down_proj"): + for idx in range(2): + for dtype in (torch.float16, torch.bfloat16): + torch.testing.assert_close( + experts.get_expert_weight(name, idx, dtype), + expected[name][idx].to(dtype), + rtol=0, + atol=0, + ) + + inputs = torch.randn(2, 256) + indices = torch.tensor([[0], [1]]) + routing = torch.ones(2, 1) + torch.testing.assert_close( + experts(inputs, indices, routing), reference(inputs, indices, routing) + ) + for _, quantizer in experts.iter_weights_for_calibration(): + quantizer.enable() + mtq.calibrate(experts, algorithm="max", forward_loop=None) + for weight, quantizer in experts.iter_weights_for_calibration(): + torch.testing.assert_close(quantizer.amax, weight.abs().amax()) + quantizer.disable() + if fold: + experts.unpack_weight() + experts.fold_weight() + torch.testing.assert_close( + experts(inputs.bfloat16(), indices, routing), + copy.deepcopy(reference).bfloat16()(inputs.bfloat16(), indices, routing), + ) + torch.testing.assert_close( + experts(inputs, indices, routing), reference(inputs, indices, routing) + ) + _export_fused_experts(experts, torch.float32) + for idx in range(2): + exported = getattr(experts, str(idx)) + if has_gate: + torch.testing.assert_close(exported.gate_proj.weight, expected[first][idx, :256]) + torch.testing.assert_close(exported.up_proj.weight, expected[first][idx, 256:]) + else: + torch.testing.assert_close(exported.up_proj.weight, expected[first][idx]) + torch.testing.assert_close(exported.down_proj.weight, expected["down_proj"][idx]) + assert not any("scale_inv" in key or "activation_scale" in key for key in experts.state_dict()) + + +@pytest.mark.parametrize("has_gate", [True, False]) +@pytest.mark.parametrize("unpacked_projection", [None, "first", "down_proj"]) +@pytest.mark.parametrize("fold_api", ["direct", "permanent", "temporary"]) +def test_native_fp8_experts_reject_packed_fold(has_gate, unpacked_projection, fold_api): + torch.manual_seed(42) + experts = FP8Experts( + PretrainedConfig(hidden_size=32, intermediate_size=32, num_experts=2, hidden_act="silu"), + block_size=None, + has_gate=has_gate, + ) + first = "gate_up_proj" if has_gate else "up_proj" + for name in (first, "down_proj"): + weight = getattr(experts, name) + with torch.no_grad(): + weight.copy_(torch.randn(weight.shape).to(weight.dtype)) + getattr(experts, f"{name}_scale_inv").fill_(0.1) + # Put a foldable module first to check model-wide validation precedes all mutation. + model = nn.Sequential(nn.Linear(32, 32), experts) + mtq.quantize(model, mtq.INT8_DEFAULT_CFG, forward_loop=None) + if unpacked_projection is not None: + name = first if unpacked_projection == "first" else unpacked_projection + weight = torch.stack([experts.get_expert_weight(name, idx) for idx in range(2)]) + setattr(experts, name, nn.Parameter(weight, requires_grad=False)) + params = dict(model.named_parameters()) + state = copy.deepcopy(model.state_dict()) + quantizers = [q for q in model.modules() if isinstance(q, TensorQuantizer)] + enabled = [q.is_enabled for q in quantizers] + inputs = torch.randn(2, 32) + indices = torch.tensor([[0], [1]]) + routing = torch.ones(2, 1) + expected = experts(inputs, indices, routing) + + with pytest.raises(ValueError, match=r"Call unpack_weight\(\) before folding"): + if fold_api == "direct": + experts.fold_weight() + elif fold_api == "permanent": + mtq.fold_weight(model) + else: + with mtq.temporarily_fold_weights(model): + pytest.fail("Packed FP8 weights must be rejected before entering the context") + + assert [q.is_enabled for q in quantizers] == enabled + for name, param in model.named_parameters(): + assert param is params[name] + assert param.dtype == state[name].dtype + # Compare FP8 payloads in FP32 because CPU close/equality kernels do not support FP8. + torch.testing.assert_close( + {name: value.float() for name, value in model.state_dict().items()}, + {name: value.float() for name, value in state.items()}, + rtol=0, + atol=0, + ) + torch.testing.assert_close(experts(inputs, indices, routing), expected) + + +@pytest.mark.parametrize("has_gate", [True, False]) +@pytest.mark.parametrize("exception", [True, False]) +@pytest.mark.parametrize("snapshot_device", [None, "cpu"]) +def test_native_fp8_experts_unpacked_temporary_fold(has_gate, exception, snapshot_device): + torch.manual_seed(42) + config = PretrainedConfig( + hidden_size=32, intermediate_size=32, num_experts=2, hidden_act="silu" + ) + experts = FP8Experts(config, block_size=None, has_gate=has_gate) + first = "gate_up_proj" if has_gate else "up_proj" + for name in (first, "down_proj"): + weight = getattr(experts, name) + scales = getattr(experts, f"{name}_scale_inv") + with torch.no_grad(): + weight.copy_(torch.randn(weight.shape).to(weight.dtype)) + scales.copy_(torch.linspace(0.01, 0.08, scales.numel()).reshape(scales.shape)) + mtq.quantize( + experts, + { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + {"quantizer_name": "*weight_quantizer", "cfg": {"num_bits": 4, "axis": None}}, + ], + "algorithm": "max", + }, + ) + experts.unpack_weight() + params = {name: getattr(experts, name) for name in (first, "down_proj")} + state = copy.deepcopy(experts.state_dict()) + quantizers = [q for _, q in experts.iter_weights_for_calibration()] + inputs = torch.randn(2, 32) + indices = torch.tensor([[0], [1]]) + routing = torch.ones(2, 1) + expected = experts(inputs, indices, routing) + expected_exit = pytest.raises(RuntimeError, match="fold exit") if exception else nullcontext() + with expected_exit, mtq.temporarily_fold_weights(experts, snapshot_device=snapshot_device): + assert all(not q.is_enabled for q in quantizers) + for name, param in params.items(): + assert getattr(experts, name) is param + assert param.dtype == torch.get_default_dtype() + assert not torch.equal(param, state[name]) + torch.testing.assert_close(experts(inputs, indices, routing), expected) + if exception: + raise RuntimeError("fold exit") + for name, param in params.items(): + assert getattr(experts, name) is param + assert param.dtype == torch.get_default_dtype() + assert all(q.is_enabled for q in quantizers) + torch.testing.assert_close(experts.state_dict(), state, rtol=0, atol=0) + torch.testing.assert_close(experts(inputs, indices, routing), expected) From ab34c9fab30cbebc2017018ffc8d94d658cd8404 Mon Sep 17 00:00:00 2001 From: Rohan Joshi Date: Fri, 16 Oct 2026 23:09:47 +0000 Subject: [PATCH 35/35] Backport the fold_weight validation hook from #2656 for #2820 #2820 overrides QuantModule._validate_fold_weight() and relies on mtq.fold_weight() validating every module before folding any. Both come from #2656 (mergeable LoRA QAD), which is not on release/0.48.0, so the #2820 native FP8 expert unit tests fail with "'super' object has no attribute '_validate_fold_weight'" and partial folds. Backport only those two pieces from #2656: the no-op base hook in QuantModule and the validation pass in mtq.fold_weight(). The LoRA feature itself is not included. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Rohan Joshi --- modelopt/torch/quantization/model_quant.py | 3 +++ modelopt/torch/quantization/nn/modules/quant_module.py | 3 +++ 2 files changed, 6 insertions(+) diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index 5b53526aba6..f40df709a04 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -1078,6 +1078,9 @@ def fold_weight(model: nn.Module, keep_attrs: bool = False): Any weight-quantizer rotation is folded into the weights and disabled so subsequent forwards do not re-rotate the already-folded weights. """ + for module in model.modules(): + if isinstance(module, QuantModule): + module._validate_fold_weight() for name, module in model.named_modules(): if isinstance(module, QuantModule): module.fold_weight(keep_attrs) diff --git a/modelopt/torch/quantization/nn/modules/quant_module.py b/modelopt/torch/quantization/nn/modules/quant_module.py index 522f2d0db09..8b1d8a4d349 100644 --- a/modelopt/torch/quantization/nn/modules/quant_module.py +++ b/modelopt/torch/quantization/nn/modules/quant_module.py @@ -168,6 +168,9 @@ def _fold_weight_quantizer( if hasattr(quantizer, attr_name): delattr(quantizer, attr_name) + def _validate_fold_weight(self): + """Validate that weight folding is supported before any module is modified.""" + def fold_weight(self, keep_attrs: bool = False): """Bake each fake-quant weight quantizer into its weight for faster eval.