From 3c4446699073793c4de49ded783a1af0865842cf Mon Sep 17 00:00:00 2001 From: Chad Voegele Date: Fri, 16 Oct 2026 21:16:09 -0500 Subject: [PATCH 01/43] Fix MiniMax M2.7 expert selectors and source dtype documentation (#2817) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Chad's Agent **Draft: full-model export correctness is blocked by native FP8 expert scaling.** Type of change: Bug fix, documentation, tests. The MiniMax M2.7 expert-only recipe selects `block_sparse_moe.experts` but misses the modern Transformers `mlp.experts` layout, leaving no weight quantizers enabled. Add the modern patterns while retaining legacy support. Regression tests cover linear and fused experts, expert W4A4 settings, exclusions, and unchanged parameters. Correct the source-dtype description. Safetensors headers at `MiniMaxAI/MiniMax-M2.7` revision `d494266a4affc0d2995ba1fa35c8481cbd84294b` show all 248 attention projection weights stored as FP8 E4M3, all 62 router gate weights and correction biases as FP32, and `lm_head.weight` as BF16. The config's BF16 default is not every tensor's dtype. These are checkpoint storage dtypes; excluding modules from this recipe does not guarantee their runtime/export dtype. ### Root cause: legacy remote code versus native Transformers This is not a naming change introduced by Transformers 5.18. Native MiniMax-M2 support landed on January 9, 2026 in [Transformers PR #42028](https://github.com/huggingface/transformers/pull/42028) ([commit `7a2bf25f36`](https://github.com/huggingface/transformers/commit/7a2bf25f36264f3f5409d2743722e4f6b3504e86)) and shipped in v5.0.0, already using `mlp.experts` and fused expert tensors. The original production setup recorded Transformers 4.57.1 and the legacy remote-code implementation; this reproduction uses native Transformers 5.18.0. The environment's exact upgrade date is unknown. The vendor checkpoint stores separate expert projections: ```text model.layers..block_sparse_moe.experts..w1.weight # gate model.layers..block_sparse_moe.experts..w3.weight # up model.layers..block_sparse_moe.experts..w2.weight # down ``` Native Transformers renames `block_sparse_moe` to `mlp`, stacks experts, and concatenates gate/up projections: ```text model.layers..mlp.experts.gate_up_proj # stacked w1+w3 model.layers..mlp.experts.down_proj # stacked w2 ``` Recipe selectors match the loaded module/quantizer names, not the checkpoint keys. Therefore the legacy-only recipe misses native expert quantizers. Both selector families are retained to support both loading implementations; the nonmatching family is a no-op. Two distinct issues are involved: the recipe's missing native selectors cause the reported guard failure; enabling them exposes a separate native `FP8Experts` inverse-scale omission in the tested ModelOpt calibration/export path. The published `nvidia/MiniMax-M2.7-NVFP4` checkpoint is not affected by this issue: it was produced with an older Transformers version that did not surface it. ### Usage Existing `--recipe modelopt_recipes/models/MiniMaxAI/MiniMax-M2.7/ptq/nvfp4_experts_only-kv_fp8_cast.yaml` commands need no changes. ### Testing - Original recipe with the new regression tests: 2 modern-layout failures, 2 legacy-layout passes. - Metadata-only full-model check with ModelOpt 0.48.0rc0 and Transformers 5.18: fixed recipe enables 31,744 fused expert weight quantizers and 124 expert input quantizers; original recipe enables zero expert weight quantizers. The 124 KV quantizers are unchanged. - Fixed recipe after rebase: `pytest tests/unit/recipe` — 509 passed. - `pre-commit run --files ` and `git diff --check` passed. - Full-model GB300 original recipe: reproduced the enabled-weight guard failure with the exact command and pinned environment. - Fixed recipe: completed all 64 calibration batches (512 samples, batch size 8, 10m27s), post-calibration preview, and 14-shard export; `hf_ptq.py` exited 0. All 47,616 expert projection weights have packed NVFP4 storage and FP8 block scales. - **Export validation failed:** nine projection samples from expert 0 in layers 0/30/61 have relative L2 error 1,989–4,474 against the dequantized FP8 source. ModelOpt's own NVFP4 dequantizer confirms that they instead approximate raw FP8 stored values (about 9% error); multiplying by the original FP8 inverse scales restores that approximately 9% error. Native expert inverse scales were omitted. The exporter also warns that fused `gate_up_proj_scale_inv` prevents reverse weight conversion. - All GitHub CI checks passed (27 successful, 2 skipped). ### Before your PR is "*Ready for review*" - Backward compatible: ✅ - Copied code or new PIP dependency: N/A - Necessary tests: ✅ - Changelog: N/A; unreleased recipe correction. - Claude approval: N/A; not requested. ### Additional Information Full-model validation uses PyTorch 26.08, ModelOpt 0.48.0rc0, Transformers 5.18.0, compressed-tensors 0.19.0, and kernels 0.17.0. Model and CNN/DailyMail revisions are pinned; only the recipe changes between runs. No weight-quantization check bypass is used. The recipe selectors are necessary but insufficient for a valid export with this environment. ModelOpt's generic fused-expert wrapper handles `FP8Experts` structurally, but its calibration/export weight path needs the native FP8 inverse scales. Completing that backend support changes library code; it cannot be demonstrated by changing only this recipe. Keep this PR in draft until that blocker is resolved. The exported checkpoint is not ready for serving or accuracy evaluation. ## Summary by CodeRabbit * **New Features** * The MiniMax-M2.7 quantization recipe now supports expert layers in both `block_sparse_moe` and `mlp` layouts, including fused experts. * NVFP4 quantization is now configured for experts in either layout. Attention weights retain their source FP8 format, router weights remain FP32, and LM-head weights remain BF16. * **Documentation** * Updated the recipe description to clarify supported expert layouts and the formats retained for attention, router, and LM-head weights. Signed-off-by: Chad Voegele Signed-off-by: Rohan Joshi --- .../ptq/nvfp4_experts_only-kv_fp8_cast.yaml | 8 +- modelopt_recipes/ptq.md | 5 +- tests/unit/recipe/test_minimax_m2_7_recipe.py | 82 +++++++++++++++++++ 3 files changed, 93 insertions(+), 2 deletions(-) create mode 100644 tests/unit/recipe/test_minimax_m2_7_recipe.py diff --git a/modelopt_recipes/models/MiniMaxAI/MiniMax-M2.7/ptq/nvfp4_experts_only-kv_fp8_cast.yaml b/modelopt_recipes/models/MiniMaxAI/MiniMax-M2.7/ptq/nvfp4_experts_only-kv_fp8_cast.yaml index 1cf8b612d84..cdabe12919b 100644 --- a/modelopt_recipes/models/MiniMaxAI/MiniMax-M2.7/ptq/nvfp4_experts_only-kv_fp8_cast.yaml +++ b/modelopt_recipes/models/MiniMaxAI/MiniMax-M2.7/ptq/nvfp4_experts_only-kv_fp8_cast.yaml @@ -26,7 +26,9 @@ imports: metadata: description: >- MiniMax M2.7 expert-only NVFP4 weights and activations with max calibration; - FP8 KV-cache cast mode; attention projections, router, and LM head stay BF16. + FP8 KV-cache cast mode; attention projections, router, and LM head are + excluded from additional quantization. The source stores attention weights + in FP8, router weights in FP32, and LM-head weights in BF16. Reproduces the published nvidia/MiniMax-M2.7-NVFP4 checkpoint from MiniMaxAI/MiniMax-M2.7. @@ -40,5 +42,9 @@ quantize: cfg: {$import: nvfp4} - quantizer_name: '*block_sparse_moe.experts.*input_quantizer' cfg: {$import: nvfp4} + - quantizer_name: '*mlp.experts.*weight_quantizer' + cfg: {$import: nvfp4} + - quantizer_name: '*mlp.experts.*input_quantizer' + cfg: {$import: nvfp4} - $import: kv_fp8_cast - $import: default_disabled_quantizers diff --git a/modelopt_recipes/ptq.md b/modelopt_recipes/ptq.md index af3c10a117d..1d4b754bfb6 100644 --- a/modelopt_recipes/ptq.md +++ b/modelopt_recipes/ptq.md @@ -476,7 +476,10 @@ checkpoint's** quant config verbatim: and the KV cache use FP8. It uses max calibration. - **`models/MiniMaxAI/MiniMax-M2.7/ptq/nvfp4_experts_only-kv_fp8_cast`** reproduces `nvidia/MiniMax-M2.7-NVFP4` with max-calibrated NVFP4 W4A4 experts (block size - 16) and FP8 KV-cache cast mode; attention projections, router, and LM head stay BF16. + 16) and FP8 KV-cache cast mode, supporting both `block_sparse_moe.experts` and + `mlp.experts` layouts, including fused experts. Attention projections, router, + and LM head are excluded from additional quantization: the source checkpoint + stores attention weights in FP8, router weights in FP32, and LM-head weights in BF16. - **`models/Qwen/Qwen3.8-27B/ptq/nvfp4_w4a4_mlp_fp8_attn_local_hessian`** reproduces NVFP4 W4A4 mixed precision quantization recipe used for `nvidia/Qwen3.8-27B-NVFP4`: MLP projections and `lm_head` use NVFP4 W4A4, self-attention and diff --git a/tests/unit/recipe/test_minimax_m2_7_recipe.py b/tests/unit/recipe/test_minimax_m2_7_recipe.py new file mode 100644 index 00000000000..fd54dcd7888 --- /dev/null +++ b/tests/unit/recipe/test_minimax_m2_7_recipe.py @@ -0,0 +1,82 @@ +# 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 pytest +import torch +import torch.nn as nn + +import modelopt.torch.quantization as mtq +from modelopt.recipe import load_recipe +from modelopt.torch.quantization.nn import TensorQuantizer +from modelopt.torch.quantization.plugins.huggingface import register_fused_experts_on_the_fly + + +class _FusedExperts(nn.Module): + def __init__(self): + super().__init__() + self.num_experts = 2 + self.intermediate_dim = 32 + self.gate_up_proj = nn.Parameter(torch.randn(2, 64, 32)) + self.down_proj = nn.Parameter(torch.randn(2, 32, 32)) + + def forward(self, hidden_states): + return hidden_states + + +@pytest.mark.parametrize("layout", ["block_sparse_moe", "mlp"]) +@pytest.mark.parametrize("fused", [False, True]) +def test_minimax_m2_7_expert_selection(layout, fused): + model = nn.Module() + model.model = nn.Module() + layer = nn.Module() + model.model.layers = nn.ModuleList([layer]) + moe = nn.Module() + layer.add_module(layout, moe) + moe.experts = ( + _FusedExperts() + if fused + else nn.ModuleList( + [nn.ModuleDict({p: nn.Linear(32, 32) for p in ("gate_proj", "up_proj", "down_proj")})] + ) + ) + moe.gate = nn.Linear(32, 2) + moe.shared_experts = nn.Linear(32, 32) + layer.self_attn = nn.ModuleDict( + {p: nn.Linear(32, 32) for p in ("q_proj", "k_proj", "v_proj", "o_proj")} + ) + model.lm_head = nn.Linear(32, 32) + preserved = {name: p.detach().clone() for name, p in model.named_parameters()} + register_fused_experts_on_the_fly(model) + + config = load_recipe( + "models/MiniMaxAI/MiniMax-M2.7/ptq/nvfp4_experts_only-kv_fp8_cast" + ).quantize.model_dump() + assert config["algorithm"]["method"] == "max" + config["algorithm"] = None + mtq.quantize(model, config) + + enabled = { + name: q + for name, q in model.named_modules() + if isinstance(q, TensorQuantizer) and q.is_enabled + } + assert len(enabled) == 6 # Three W4A4 projections, or four fused weights and two inputs. + for name, q in enabled.items(): + assert f".{layout}.experts." in name + assert q.num_bits == (2, 1) + assert q.block_sizes[-1] == 16 + assert q.block_sizes["type"] == "dynamic" + for name, p in model.named_parameters(): + torch.testing.assert_close(p, preserved[name], rtol=0, atol=0) From 6060bce04c94c54eba13b047c505a6829714cf64 Mon Sep 17 00:00:00 2001 From: LinCanNerd <64252217+LinCanNerd@users.noreply.github.com> Date: Sat, 17 Oct 2026 09:36:21 +0200 Subject: [PATCH 02/43] Fix import modelopt.torch without the [hf] extra (#2671) 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 `import modelopt.torch` fails on a base install (`pip install nvidia-modelopt`, no extras), although the [installation docs](https://nvidia.github.io/Model-Optimizer/getting_started/_installation_for_Linux.html) state that the base install provides the `modelopt.torch` requirements. Reproduced with the released 0.47.0 in a fresh venv containing only `torch` and `nvidia-modelopt==0.47.0`: ``` >>> import modelopt.torch.quantization ModuleNotFoundError: No module named 'requests' # after `pip install requests`: ModuleNotFoundError: No module named 'huggingface_hub' ``` Root cause: two modules that `modelopt.torch` always imports have module-scope imports of packages that are not base dependencies: - `modelopt/torch/utils/dataset_utils.py` imports `requests` (since #997). It is only used in `download_hf_dataset_as_jsonl`, which already imports `datasets` and `huggingface_hub` locally. `requests` is not declared anywhere; it was only present transitively. - `modelopt/torch/opt/plugins/huggingface.py` imports `huggingface_hub` (since #975). `modelopt/torch/opt/__init__.py` imports this plugin under `import_plugin("huggingface")`, but `opt/plugins/__init__.py` imports it unconditionally, and `unified_export_hf.py` imports from it directly. It is only used in `_get_modelopt_state_path`. The fix moves each import into the only function that uses it. CI did not catch this because every test environment installs extras that bring both packages in, so the new `tests/unit/torch/test_import.py` imports `modelopt.torch` in a subprocess with the `[hf]` extra's modules (plus `requests`) made unavailable. Out of scope: `modelopt.torch.export` (imports `huggingface_hub` deliberately via `utils/plugins/hf_checkpoint_utils.py`, #2427) and `modelopt.torch.puzzletron` (has its own `[puzzletron]` extra) still need their optional dependencies. ### Usage No API changes. ### Testing - Fresh venv with only CPU `torch` and this branch installed (`pip install -e .`, base dependencies only): `import modelopt.torch` and `modelopt.torch.{opt,quantization,sparsity,distill,nas,prune,speculative,peft,trace,utils,fastgen,models}` all import; before the change every one failed with the `ModuleNotFoundError` above. - `tests/unit/torch/test_import.py` fails on `main` (`import of requests halted; None in sys.modules`) and passes with this change. - `pytest tests/unit/torch/utils/test_dataset_utils.py tests/unit/torch/opt tests/unit/torch/test_import.py`: 122 passed (with `datasets` installed, `test_dataset_utils.py`: 70 passed). - `pre-commit run` on the changed files: all hooks pass. Found while setting up ModelOpt with pip-only dependencies on a Jetson AGX Thor. ### 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 🤖 Generated with [Claude Code](https://claude.com/claude-code) ## Summary by CodeRabbit * **Bug Fixes** * Fixed an issue that could prevent `modelopt.torch` from importing when optional Hugging Face dependencies were not installed. You can now import the package without those dependencies; they are needed only when using features that require them, such as downloading datasets or looking up model state in the Hugging Face cache. Dataset download and retry behavior remains unchanged. Signed-off-by: LinCanNerd Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 (1M context) Co-authored-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Signed-off-by: Rohan Joshi --- CHANGELOG.rst | 1 + modelopt/torch/opt/plugins/huggingface.py | 6 ++-- modelopt/torch/utils/dataset_utils.py | 2 +- tests/unit/torch/test_import.py | 41 +++++++++++++++++++++++ 4 files changed, 47 insertions(+), 3 deletions(-) create mode 100644 tests/unit/torch/test_import.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index ef059b41228..cf4d4172b2a 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -132,6 +132,7 @@ Changelog **Bug Fixes** +- Fix ``import modelopt.torch`` failing with ``ModuleNotFoundError`` for ``requests`` or ``huggingface_hub`` when ``nvidia-modelopt`` is installed without the ``[hf]`` extra. - 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. diff --git a/modelopt/torch/opt/plugins/huggingface.py b/modelopt/torch/opt/plugins/huggingface.py index 19b2e357631..afbe4a3fde9 100644 --- a/modelopt/torch/opt/plugins/huggingface.py +++ b/modelopt/torch/opt/plugins/huggingface.py @@ -23,8 +23,6 @@ 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 import print_rank_0 @@ -65,6 +63,10 @@ def _get_modelopt_state_path(model_name_or_path: str) -> str: """ if os.path.isdir(model_name_or_path): return os.path.join(model_name_or_path, _MODELOPT_STATE_SAVE_NAME) + # Optional [hf] dependency: this module is imported even when huggingface_hub is not installed. + from huggingface_hub import try_to_load_from_cache + from huggingface_hub.errors import HFValidationError + try: return try_to_load_from_cache(model_name_or_path, _MODELOPT_STATE_SAVE_NAME) or "" except HFValidationError: diff --git a/modelopt/torch/utils/dataset_utils.py b/modelopt/torch/utils/dataset_utils.py index f882a6e7364..76260c216e3 100644 --- a/modelopt/torch/utils/dataset_utils.py +++ b/modelopt/torch/utils/dataset_utils.py @@ -25,7 +25,6 @@ from pathlib import Path from typing import TYPE_CHECKING, Any, Literal -import requests import torch from torch.utils.data import DataLoader from torch.utils.data.distributed import DistributedSampler @@ -1318,6 +1317,7 @@ def download_hf_dataset_as_jsonl( Returns: List of paths to downloaded JSONL files. """ + import requests from datasets import load_dataset from huggingface_hub.utils import build_hf_headers diff --git a/tests/unit/torch/test_import.py b/tests/unit/torch/test_import.py new file mode 100644 index 00000000000..3ceb933c8b3 --- /dev/null +++ b/tests/unit/torch/test_import.py @@ -0,0 +1,41 @@ +# 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 subprocess +import sys + +# Top-level modules of the optional ``[hf]`` extra, plus ``requests`` which it only pulls in +# transitively. A base ``pip install nvidia-modelopt`` provides none of them. +_HF_EXTRA_MODULES = [ + "accelerate", + "datasets", + "deepspeed", + "diffusers", + "httpx", + "huggingface_hub", + "nltk", + "peft", + "requests", + "sentencepiece", + "tiktoken", + "transformers", + "wonderwords", +] + + +def test_modelopt_torch_imports_without_hf_extra(): + # A None entry in sys.modules makes importing that module raise ModuleNotFoundError. + code = f"import sys; sys.modules.update(dict.fromkeys({_HF_EXTRA_MODULES!r})); import modelopt.torch" + subprocess.run([sys.executable, "-c", code], check=True) From 1c7c7b46b9a4749d125d7dc93c45458af5761c8b Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Tue, 20 Oct 2026 20:12:01 +0530 Subject: [PATCH 03/43] Pin "leptonai<0.47" to fix tools/launcher CI job (#2892) ### What does this PR do? Type of change: Fix CI failure ## Summary by CodeRabbit * **Bug Fixes** * Improved compatibility with the launcher by constraining the Lepton AI runtime dependency to supported versions. Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Signed-off-by: Rohan Joshi --- tools/launcher/pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/tools/launcher/pyproject.toml b/tools/launcher/pyproject.toml index 8ba825776d4..02adf5991fe 100644 --- a/tools/launcher/pyproject.toml +++ b/tools/launcher/pyproject.toml @@ -4,6 +4,7 @@ version = "0.1.0" description = "ModelOpt Launcher — submit quantization, training, and evaluation jobs to Slurm clusters" requires-python = ">=3.10" dependencies = [ + "leptonai<0.47", # 0.47 removed leptonai.api.v1, which nemo-run (<=0.11.1) imports "nemo-run>=0.8.0", "pyyaml", ] From a7572357af8c78f0752e673dfa80f982af7e28ac Mon Sep 17 00:00:00 2001 From: Wei-Ming Chen <17592131+meenchen@users.noreply.github.com> Date: Tue, 20 Oct 2026 21:47:46 -0700 Subject: [PATCH 04/43] Fix layerwise MoE source-weight accounting (NVBug 6954041) (#2876) ### What does this PR do? Type of change: Bug fix Fix layerwise MoE export copying already-quantized fused expert tensors back from the source checkpoint as BF16 weights. Those copies incorrectly add the routed expert container to `exclude_modules`, causing vLLM to select an unquantized MoE backend and fail loading `w2_input_scale`. Snapshot complete source-key accounting at exporter binding, before per-layer packing removes the original fused parameters. Reuse the snapshot, including an empty result, for that source checkpoint. Only names are retained; genuinely unmodeled auxiliary weights still carry over. Whole-model export is unchanged. ### Usage No API changes. Re-export affected checkpoints from the original source using the existing layerwise recipe. ### Testing - 186 focused CPU tests passed across model loading, reverse conversion, unified export, and layerwise calibration/resume (Transformers 5.18.0). - Real tiny native Qwen3.5-MoE VLM max calibration and checkpoint export with `general/ptq/nvfp4_experts_only-kv_fp8_layerwise_export`, compared with whole-model export, with and without an auxiliary source weight: passed. Packed expert weights/scales remain present, original fused BF16 weights are absent, and expert containers are not excluded. The auxiliary weight remains present and excluded. - Extended existing `test_moe_export_matches` with source provenance and the public export API. Both cases passed in a local one-layer CPU adaptation; both failed on the duplicate fused expert tensors when the snapshot lookup was reverted. The committed GPU test retains its original two-layer model and quantization configuration. - Meta-model binding smoke: 40 decoder layers, 256 experts/layer, 636 source keys; all parameters remained on meta and the retained source-key snapshot was empty. Full binding took 2.99 s and peaked at 12.48 MiB of Python-tracked allocations. This is a metadata probe, not a full-size PTQ/GPU memory benchmark. - Applicable pre-commit hooks and `git diff --check`: passed. - Not run: unadapted CUDA tests, full Qwen3.5-35B-A3B PTQ, or vLLM serving. GPU launcher MCP tools are unavailable in this session. ### Before your PR is "*Ready for review*" Contributor guidelines and security coding practices reviewed. Commit `b139a70eb` is signed and signed off; its GPG signature was verified before publishing. - 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; neither. - Did you write any new necessary tests?: Yes; extended the existing MoE export comparison to cover empty accounting and genuine auxiliary carry-over. - Did you update Changelog?: N/A; layerwise export is new in the unreleased 0.48 cycle. - Did you get Claude approval on this PR?: No; pending review. ### Additional Information NVBug 6954041. Release backport: 0.48.0. ## Summary by CodeRabbit * **Bug Fixes** * Layerwise exports now preserve eligible auxiliary weights from the source checkpoint, keeping exported results consistent with baseline exports. Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com> Signed-off-by: Rohan Joshi --- modelopt/torch/export/layerwise_export.py | 8 +++++ modelopt/torch/export/unified_export_hf.py | 8 +++++ .../gpu/torch/export/test_layerwise_export.py | 32 +++++++++++++++---- 3 files changed, 41 insertions(+), 7 deletions(-) diff --git a/modelopt/torch/export/layerwise_export.py b/modelopt/torch/export/layerwise_export.py index 92f31011ea3..2ff3abe966f 100644 --- a/modelopt/torch/export/layerwise_export.py +++ b/modelopt/torch/export/layerwise_export.py @@ -52,9 +52,11 @@ _fuse_shared_input_modules, _prepare_moe_inputs, _resolve_export_dtype, + _source_checkpoint, _write_hf_export_config, collect_shared_input_modules, save_non_weight_artifacts, + unplaced_keys_for, ) from .unified_export_hf_streaming import _assert_no_split_rules @@ -181,6 +183,7 @@ def __init__( self._dtype = dtype self._bound = False self._finalized = False + self._source_key_snapshot: tuple[str, list[str]] | None = None self._announced_on: list[nn.Module] = [] # A VLM calibrates its language model but exports the whole thing, so announce on # both: whichever of the two mtq.quantize is handed will find this exporter. @@ -272,6 +275,11 @@ def bind(self, calibrated_layers: list[nn.Module]) -> None: "match the original HF hub checkpoint." ) + # Packing removes fused source parameters; account for them while the tree is intact. + source = _source_checkpoint(model) + if source and Path(source).is_dir(): + self._source_key_snapshot = (str(source), unplaced_keys_for(model, source)) + self._bound = True def export_layer( diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index f605ddaf636..646ba21d714 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -1750,6 +1750,14 @@ def unplaced_keys_for(model: nn.Module, ckpt: "str | Path") -> list[str]: module scope. Where those are absent the recorded set stands alone -- degraded, but no worse than before this function existed. """ + # Local import breaks the cycle: layerwise_export imports this module. + from .layerwise_export import LAYERWISE_EXPORTER_ATTR + + exporter = getattr(model, LAYERWISE_EXPORTER_ATTR, None) + snapshot = getattr(exporter, "_source_key_snapshot", None) + if snapshot is not None and snapshot[0] == str(ckpt): + return list(snapshot[1]) + recorded = getattr(model, "_modelopt_unplaced_source_keys", None) or [] structural: list[str] = [] try: diff --git a/tests/gpu/torch/export/test_layerwise_export.py b/tests/gpu/torch/export/test_layerwise_export.py index 84244a8fec1..e442fe11c39 100644 --- a/tests/gpu/torch/export/test_layerwise_export.py +++ b/tests/gpu/torch/export/test_layerwise_export.py @@ -27,7 +27,7 @@ get_tiny_llama, get_tiny_qwen3_moe, ) -from safetensors.torch import load_file +from safetensors.torch import load_file, save_file import modelopt.torch.quantization as mtq from modelopt.torch.export.layerwise_export import ( @@ -610,20 +610,38 @@ def test_tied_embeddings_are_refused(tmp_path, make_cfg): mtq.quantize(model, cfg, _calib) -def test_moe_export_matches(tmp_path): - """MoE layers take a different path: fused expert inputs and gate/up amax sync.""" +@pytest.mark.parametrize("carry_auxiliary_weight", [False, True]) +def test_moe_export_matches(tmp_path, carry_auxiliary_weight): + """Packing fused experts must not reclassify their source tensors as unplaced.""" + source = tmp_path / "source" + source.mkdir() + baseline, layerwise = _build_moe_model(), _build_moe_model() + source_state = {k: v.cpu().contiguous() for k, v in baseline.state_dict().items()} + auxiliary_key = "mtp.eh_proj.weight" + if carry_auxiliary_weight: + source_state[auxiliary_key] = torch.ones(2, 2, dtype=torch.bfloat16) + save_file(source_state, source / "model.safetensors") + for model in (baseline, layerwise): + model._modelopt_source_checkpoint = str(source) + # Transformers can hide auxiliary weights from its unexpected_keys report. + model._modelopt_unplaced_source_keys = [] + baseline_dir = tmp_path / "baseline" base = _nvfp4_cfg() base["algorithm"] = {"method": "max", "layerwise": {"enable": True}} - export_hf_checkpoint(mtq.quantize(_build_moe_model(), base, _calib), export_dir=baseline_dir) + export_hf_checkpoint(mtq.quantize(baseline, base, _calib), export_dir=baseline_dir) export_dir = tmp_path / "fused" - _layerwise_quantize( - _build_moe_model(), _layerwise_cfg(export_dir, tmp_path / "ckpt", base=_nvfp4_cfg()) + mtq.quantize( + layerwise, _layerwise_cfg(export_dir, tmp_path / "ckpt", base=_nvfp4_cfg()), _calib ) + export_hf_checkpoint(layerwise, export_dir=export_dir) - _assert_same_checkpoint(_load_checkpoint(baseline_dir), _load_checkpoint(export_dir)) + exported = _load_checkpoint(export_dir) + _assert_same_checkpoint(_load_checkpoint(baseline_dir), exported) _assert_same_quant_config(baseline_dir, export_dir) + if carry_auxiliary_weight: + assert torch.equal(exported[auxiliary_key], source_state[auxiliary_key]) def test_export_consumes_the_model_without_affecting_the_checkpoint(tmp_path): From 2101318b92397f814ea37c8c3ead30085267a317 Mon Sep 17 00:00:00 2001 From: noeyy-mino <174223378+noeyy-mino@users.noreply.github.com> Date: Wed, 21 Oct 2026 17:29:50 +0800 Subject: [PATCH 05/43] [NVBug 6966551] Fix spelling errors and duplicated words in docs and docstrings. (#2910) 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 : https://nvbugs/6966551 Fix misspellings and duplicated words found by LLM Doc Check on the published Model Optimizer documentation, plus a few more found by the typos pre-commit hook in the same files. Only comments, docstrings, and docs are changed; there is no functional change. - Wrong API/module names: mnt.search() -> mtn.search(), ``fasnas`` -> ``fastnas`` - Misspellings: infiinite, longs, ans, incompatibitly, wihin, rigoruos, Op top, steo, accomodates, Irreggular, fintuned, dont, splitted, distrubted, dimention, heterogenuous, homogenous - Duplicated words: "the the", "process process", "layout layout", "loss loss" ### Usage N/A ### Testing N/A ### Before your PR is "*Ready for review*" - Is this change backward compatible?: N/A - 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?: ✅ ### Additional Information N/A ## Summary by CodeRabbit * **Documentation** * Corrected spelling, grammar, duplicated words, and wording across user guides, examples, and API documentation. * Clarified descriptions for quantization, distillation, configuration, NAS, optimization, serialization, and distributed processing. * Corrected references to the `mtn.search()` API in relevant guidance and comments. * Updated technical comments and docstrings for greater accuracy and consistency. * **Tests** * Updated test comments for consistency and readability. --------- Signed-off-by: Noey Yang <174223378+noeyy-mino@users.noreply.github.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Co-authored-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com> Signed-off-by: Rohan Joshi --- docs/source/guides/7_nas.rst | 2 +- docs/source/guides/_pytorch_quantization.rst | 2 +- examples/hf_ptq/README.md | 2 +- examples/puzzletron/main.py | 2 +- modelopt/onnx/op_types.py | 2 +- modelopt/torch/distill/layerwise_distillation_model.py | 2 +- modelopt/torch/distill/loss_balancers.py | 4 ++-- modelopt/torch/export/trtllm/model_config.py | 4 ++-- modelopt/torch/models/specs.py | 2 +- modelopt/torch/nas/autonas.py | 2 +- modelopt/torch/opt/__init__.py | 2 +- modelopt/torch/opt/conversion.py | 4 ++-- modelopt/torch/opt/dynamic.py | 2 +- modelopt/torch/opt/mode.py | 2 +- modelopt/torch/prune/fastnas.py | 2 +- modelopt/torch/puzzletron/puzzletron_nas_plugin.py | 2 +- modelopt/torch/speculative/plugins/megatron_eagle.py | 8 ++++---- modelopt/torch/speculative/utils.py | 2 +- modelopt/torch/utils/network.py | 6 +++--- modelopt/torch/utils/serialization.py | 2 +- tests/gpu/torch/puzzletron/test_nas_search.py | 2 +- tests/unit/onnx/test_onnx_utils.py | 2 +- 22 files changed, 30 insertions(+), 30 deletions(-) diff --git a/docs/source/guides/7_nas.rst b/docs/source/guides/7_nas.rst index 14566bfcf95..c86dc3cc198 100644 --- a/docs/source/guides/7_nas.rst +++ b/docs/source/guides/7_nas.rst @@ -326,7 +326,7 @@ The conversion process ---------------------- ModelOpt will automatically generate a search space for you from your custom PyTorch model. -This is a one time process process performed during pruning and NAS. Once a model is converted, you +This is a one time process performed during pruning and NAS. Once a model is converted, you can save and restore it for downstream tasks like training, inference, and fine-tuning. To help you better understand how the search space is derived from your model, we go through the diff --git a/docs/source/guides/_pytorch_quantization.rst b/docs/source/guides/_pytorch_quantization.rst index f8f12b068ba..1e7a2b05c09 100644 --- a/docs/source/guides/_pytorch_quantization.rst +++ b/docs/source/guides/_pytorch_quantization.rst @@ -142,7 +142,7 @@ quantization and weight & activation quantization). ``effective_bits`` constrain You may specify a ``effective_bits`` constraint such as 4.8 for mixed precision quantization using :attr:`NVFP4_DEFAULT_CFG` & :attr:`FP8_DEFAULT_CFG`. ``AutoQuantize`` will automatically quantize highly sensitive layers in :attr:`FP8_DEFAULT_CFG` while keeping less sensitive layers in :attr:`NVFP4_DEFAULT_CFG` (and even skip quantization for any extremely sensitive layers) so that -the the final mixed precision quantized model has an effective quantized bits of 4.8. +the final mixed precision quantized model has an effective quantized bits of 4.8. This model would give a better accuracy than the model quantized with vanilla :attr:`NVFP4_DEFAULT_CFG` since the more aggressive :attr:`NVFP4_DEFAULT_CFG` quantization was not applied for the highly sensitive layers. diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index 7e3c7d8fc44..107ce82c6a7 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -363,7 +363,7 @@ weight-only and weight & activation quantization) — the effective number of bi You may specify an `effective_bits` target such as 5.4 for mixed precision quantization using `NVFP4_DEFAULT_CFG` & `FP8_DEFAULT_CFG`. `AutoQuantize` will automatically quantize highly sensitive layers in `FP8_DEFAULT_CFG` while keeping less sensitive layers in `NVFP4_DEFAULT_CFG` (and even skip quantization for any extremely sensitive layers) so that -the the final mixed precision quantized model has an effective quantized bits of 5.4. This model would give a better accuracy than the model quantized with vanilla `NVFP4_DEFAULT_CFG` configuration since the more aggressive `NVFP4_DEFAULT_CFG` quantization was not applied for the highly sensitive layers. +the final mixed precision quantized model has an effective quantized bits of 5.4. This model would give a better accuracy than the model quantized with vanilla `NVFP4_DEFAULT_CFG` configuration since the more aggressive `NVFP4_DEFAULT_CFG` quantization was not applied for the highly sensitive layers. Here is an example usage for `AutoQuantize` algorithm (Please see [auto_quantize](https://nvidia.github.io/Model-Optimizer/reference/generated/modelopt.torch.quantization.model_quant.html#modelopt.torch.quantization.model_quant.auto_quantize) API for more details): diff --git a/examples/puzzletron/main.py b/examples/puzzletron/main.py index f093e5b7e68..09bae07de9d 100644 --- a/examples/puzzletron/main.py +++ b/examples/puzzletron/main.py @@ -156,7 +156,7 @@ def run_mip_only(hydra_config_path: str): mtpz.mip.run_mip_sweep(hydra_cfg) else: # mip_and_realize_models (distributed processing) - # TODO: How to make it part of mnt.search() api, similarly to run_full_puzzletron() API + # TODO: How to make it part of mtn.search() api, similarly to run_full_puzzletron() API mtpz.tools.mprint("Puzzletron Progress 7/8: running MIP and realizing models (multi-gpu)") mtpz.mip.launch_mip_and_realize_model(hydra_cfg) diff --git a/modelopt/onnx/op_types.py b/modelopt/onnx/op_types.py index 637c0ad7a45..aa2ac4ef374 100644 --- a/modelopt/onnx/op_types.py +++ b/modelopt/onnx/op_types.py @@ -175,7 +175,7 @@ def is_non_reshape_copy_op(op_type: str): def is_irregular_mem_access_op(op_type: str): - """Returns whether the given op type is of Irreggular mem access category or not.""" + """Returns whether the given op type is of irregular mem access category or not.""" return op_type in [ "Gather", "GatherElements", diff --git a/modelopt/torch/distill/layerwise_distillation_model.py b/modelopt/torch/distill/layerwise_distillation_model.py index 4a98f5060c5..a8d52b8d988 100644 --- a/modelopt/torch/distill/layerwise_distillation_model.py +++ b/modelopt/torch/distill/layerwise_distillation_model.py @@ -34,7 +34,7 @@ class LayerwiseDistillationModel(DistillationModel): """Meta-model wrapper to support layerwise-enabled knowledge-distillation learning. The LayerwiseDistillationModel is a subclass of the DistillationModel that injects teacher inputs - into the corresponding student layers. This accomodates the case where the student model is the + into the corresponding student layers. This accommodates the case where the student model is the teacher with specific submodules replaced, which now need to be trained to mimic the original submodule in the teacher. """ diff --git a/modelopt/torch/distill/loss_balancers.py b/modelopt/torch/distill/loss_balancers.py index 0ceaca3e3c2..0f199d860e7 100644 --- a/modelopt/torch/distill/loss_balancers.py +++ b/modelopt/torch/distill/loss_balancers.py @@ -66,7 +66,7 @@ def forward(self, loss: dict[str, torch.Tensor]) -> torch.Tensor: ``{"student_loss": torch.tensor(...), "MSELoss_0": torch.tensor(...), "MSELoss_1": torch.tensor(...)}``. Returns: - The total loss after balancing student and kd loss loss components. + The total loss after balancing student and kd loss components. """ raise NotImplementedError @@ -116,7 +116,7 @@ def forward(self, loss: dict[str, torch.Tensor]) -> torch.Tensor: loss: The loss dict to aggregate. Returns: - The total loss after balancing student and kd loss loss components. + The total loss after balancing student and kd loss components. """ kd_loss_dict = loss.copy() diff --git a/modelopt/torch/export/trtllm/model_config.py b/modelopt/torch/export/trtllm/model_config.py index c3d630f9fc9..1bee58298f7 100755 --- a/modelopt/torch/export/trtllm/model_config.py +++ b/modelopt/torch/export/trtllm/model_config.py @@ -108,7 +108,7 @@ class LinearConfig: # For methods like W4A8 AWQ, we have two quantizers for weights # For W4A8, the first quantizer is for INT4 quantization and the second quantizer is for FP8 quantization - # `weight_scaling_factor_2` is the scaling factor the the second FP8 quantizer + # `weight_scaling_factor_2` is the scaling factor for the second FP8 quantizer weights_scaling_factor_2: torch.Tensor = None prequant_scaling_factor: torch.Tensor = None @@ -319,7 +319,7 @@ class RelativeAttentionTableConfig: class AttentionConfig: """The attention layer config.""" - # QKV can either be stored as splitted (for easier postprocessing) + # QKV can either be stored as split (for easier postprocessing) # or merged (for TRT LLM export) qkv: QKVConfig | LinearConfig = None dense: LinearConfig = None diff --git a/modelopt/torch/models/specs.py b/modelopt/torch/models/specs.py index aa9110223ee..0a203780b7f 100644 --- a/modelopt/torch/models/specs.py +++ b/modelopt/torch/models/specs.py @@ -351,7 +351,7 @@ def hf_model_type(model) -> str | None: def match_moe_block(module: "nn.Module", model_type: str | None = None) -> "MoESpec | None": - """Return the MoE layout layout for ``module``, resolved by model type. + """Return the MoE layout for ``module``, resolved by model type. ``model_type`` (the root ``model.config.model_type``) is a strict filter: only that model's own spec is consulted, and an unregistered model type resolves to diff --git a/modelopt/torch/nas/autonas.py b/modelopt/torch/nas/autonas.py index 00393657c97..2cdf6349eb4 100644 --- a/modelopt/torch/nas/autonas.py +++ b/modelopt/torch/nas/autonas.py @@ -476,7 +476,7 @@ class RandomSearcher(IterativeSearcher): """An iterative searcher that samples subnets randomly.""" def sample(self) -> dict[str, Any]: - """Random sample new subset during each steo.""" + """Random sample new subnet during each step.""" return self._sample() diff --git a/modelopt/torch/opt/__init__.py b/modelopt/torch/opt/__init__.py index 950be0b986f..c46bf64687c 100644 --- a/modelopt/torch/opt/__init__.py +++ b/modelopt/torch/opt/__init__.py @@ -22,7 +22,7 @@ * export a model back to a regular pytorch model after optimization; * save, restore, and manage checkpoints from which the model modifications can be resumed. -Please refer to each individual sub-module to learn more about the various concepts wihin +Please refer to each individual sub-module to learn more about the various concepts within ``modelopt.torch.opt`` and how to use them to implement a model optimization algorithm. """ diff --git a/modelopt/torch/opt/conversion.py b/modelopt/torch/opt/conversion.py index bb3fe0ce4a5..a480b7d940c 100644 --- a/modelopt/torch/opt/conversion.py +++ b/modelopt/torch/opt/conversion.py @@ -19,8 +19,8 @@ (mode) of the algorithm. This module provides the state manager, which is a standardized interface (class) to record and store state information in the model. -Op top of the state manager, this module provides utilities to save a history of these modifications -("modelopt state dict") and restoring a unmodified model to the state indicated in the state dict. +On top of the state manager, this module provides utilities to save a history of these modifications +("modelopt state dict") and restoring an unmodified model to the state indicated in the state dict. """ import copy diff --git a/modelopt/torch/opt/dynamic.py b/modelopt/torch/opt/dynamic.py index 6c0e9dd3c5c..18baa449e9b 100644 --- a/modelopt/torch/opt/dynamic.py +++ b/modelopt/torch/opt/dynamic.py @@ -608,7 +608,7 @@ def convert(cls, module: nn.Module, **setup_kwargs: Any) -> "DynamicModule": Patching is achieved by updating the ``__class__`` attribute of the module to its dynamic counterpart. The dynamic counterpart is a subclass of the original class, hence, we ensure the module is fully compatible with the original class. Simultaneously, we can inject the - corresponding dynamic behavior in a standardized and rigoruos fashion. + corresponding dynamic behavior in a standardized and rigorous fashion. """ def bind_forward_method_if_needed(self): diff --git a/modelopt/torch/opt/mode.py b/modelopt/torch/opt/mode.py index d9ac1bf07bf..4ccac5473b0 100644 --- a/modelopt/torch/opt/mode.py +++ b/modelopt/torch/opt/mode.py @@ -94,7 +94,7 @@ def next_prohibited_modes(self) -> set[str] | None: """Modes that should not be applied after this mode. This is used to specify that certain modes should not be applied after this mode. - A ``None`` value indicates that incompatibitly check is skipped for this mode. + A ``None`` value indicates that the incompatibility check is skipped for this mode. This is useful if there are many allowed modes but only a few prohibited modes. Returns: diff --git a/modelopt/torch/prune/fastnas.py b/modelopt/torch/prune/fastnas.py index 4852efdad2b..2385f5b19b2 100644 --- a/modelopt/torch/prune/fastnas.py +++ b/modelopt/torch/prune/fastnas.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Module implementing ``fasnas`` pruning algorithm for search.""" +"""Module implementing ``fastnas`` pruning algorithm for search.""" import warnings from collections.abc import Callable diff --git a/modelopt/torch/puzzletron/puzzletron_nas_plugin.py b/modelopt/torch/puzzletron/puzzletron_nas_plugin.py index 253674f97af..1295fcf1f4f 100644 --- a/modelopt/torch/puzzletron/puzzletron_nas_plugin.py +++ b/modelopt/torch/puzzletron/puzzletron_nas_plugin.py @@ -105,7 +105,7 @@ def convert_puzzletron_model(model: nn.Module, config: PuzzletronConfig) -> Conv 2. Score the pruning activations. 3. Prune the model and save pruned checkpoints - The output of this step will be used by mnt.search() to perform the NAS search. + The output of this step will be used by mtn.search() to perform the NAS search. """ # Required for mtn.search() to read NAS configuration model.hydra_config_dir = config.hydra_config_dir diff --git a/modelopt/torch/speculative/plugins/megatron_eagle.py b/modelopt/torch/speculative/plugins/megatron_eagle.py index 22c337ed41d..65e68d8f280 100644 --- a/modelopt/torch/speculative/plugins/megatron_eagle.py +++ b/modelopt/torch/speculative/plugins/megatron_eagle.py @@ -368,11 +368,11 @@ def sharded_state_dict( ) if num_layers > 1: - # For multi-layer EAGLE, we use heterogenuous distrubted checkpoint + # For multi-layer EAGLE, we use heterogeneous distributed checkpoint sharded_prefix = state_dict_prefix sharded_pp_offset = [] else: - # For single EAGLE, we use homogenous distrubted checkpoint for backward compatibility + # For single EAGLE, we use homogeneous distributed checkpoint for backward compatibility sharded_prefix = layer_prefix sharded_pp_offset = [ (0, global_layer_offset, num_layers) @@ -592,7 +592,7 @@ def forward( """Forward function.""" # NOTE: Even if sequence_parallel is used, the rotary_seq_len must be in the original # length. Since we get the seq_len from hidden_states.shape[0], we need to - # multiply the the tp back. + # multiply the tp back. # Similarly, if context parallel is used, the rotary_seq_len must also be # multiplied by context parallel size. rotary_seq_len = hidden_states.shape[0] @@ -830,7 +830,7 @@ def _get_eagle_module_inputs( ttt_step: int = 0, ): """Getting EAGLE module inputs.""" - # gather_from_sequence_parallel_region gathers from the first dimention + # gather_from_sequence_parallel_region gathers from the first dimension # so we need to transpose input_ids first # [b,s] -> [s,b] input_ids = input_ids.clone().transpose(0, 1).contiguous() diff --git a/modelopt/torch/speculative/utils.py b/modelopt/torch/speculative/utils.py index 4db1f4b0f76..bf33761eb74 100644 --- a/modelopt/torch/speculative/utils.py +++ b/modelopt/torch/speculative/utils.py @@ -82,7 +82,7 @@ def calibrate_frequent_vocab(tokenizer, text, target_vocab_size, output_file=Non def get_default_attention_mask_and_position_ids(input_ids: torch.Tensor): - """Compute default attention_mask ans position_ids given input_ids.""" + """Compute default attention_mask and position_ids given input_ids.""" seq_len = input_ids.shape[-1] position_ids = torch.arange(seq_len, dtype=torch.long, device=input_ids.device) diff --git a/modelopt/torch/utils/network.py b/modelopt/torch/utils/network.py index d9038faae86..101c6518bd0 100644 --- a/modelopt/torch/utils/network.py +++ b/modelopt/torch/utils/network.py @@ -278,7 +278,7 @@ def standardize_named_model_args( Returns: A tuple (args_normalized, args_with_default) where args_normalized is a dictionary of ordered model args where the key represents a unique - serialized string based on the the argument's name in the function signature and the + serialized string based on the argument's name in the function signature and the value contains the actual argument, args_with_default is a set indicating whether the argument was retrieved from the default value in the function signature of the model's ``forward()`` method or whether the @@ -520,7 +520,7 @@ def run_forward_loop( Args: model: The model with which we run forward. data_loader: An iterator with data samples. - max_iters: Number of batches to run; by default it is infiinite or until ``data_loader`` + max_iters: Number of batches to run; by default it is infinite or until ``data_loader`` is exhausted. collect_func: A ``Callable`` that takes a batch of data from the ``data_loader`` as input and returns the input to ``model.forward()`` such that the return value @@ -585,7 +585,7 @@ def create_param_grad_clear_hook(param): """Create a hook to clear gradients for a parameter. The hook will be fired after the gradient is accumulated for the parameter. - Important: For this to work, ``accum_grad`` should be kept alive as longs as this utility is needed. + Important: For this to work, ``accum_grad`` should be kept alive as long as this utility is needed. """ # For methods such as AutoQuantize, involving backward - # We want to clear parameter gradients as soon as they are computed to save memory diff --git a/modelopt/torch/utils/serialization.py b/modelopt/torch/utils/serialization.py index dc880b86b80..6ccae6ee8fd 100644 --- a/modelopt/torch/utils/serialization.py +++ b/modelopt/torch/utils/serialization.py @@ -56,7 +56,7 @@ def safe_save(obj: Any, f: str | os.PathLike | BinaryIO, **kwargs) -> None: def safe_load(f: str | os.PathLike | BinaryIO | bytes, **kwargs) -> Any: """Load a checkpoint securely using ``weights_only=True`` by default. - NOTE: We dont set default ``weights_only`` (interpret as True for torch>=2.6) so you can override it with + NOTE: We don't set default ``weights_only`` (interpret as True for torch>=2.6) so you can override it with ``export TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD=1`` if you see ``pickle.UnpicklingError`` and trust the checkpoint. """ if isinstance(f, (bytes, bytearray)): diff --git a/tests/gpu/torch/puzzletron/test_nas_search.py b/tests/gpu/torch/puzzletron/test_nas_search.py index ae0dff47be5..2cdc3e26d44 100644 --- a/tests/gpu/torch/puzzletron/test_nas_search.py +++ b/tests/gpu/torch/puzzletron/test_nas_search.py @@ -66,7 +66,7 @@ def _test_nas_search_multiprocess_job( ) # - # Run the mnt.search() step + # Run the mtn.search() step # mtn.search( converted_model, diff --git a/tests/unit/onnx/test_onnx_utils.py b/tests/unit/onnx/test_onnx_utils.py index 590a9f24186..209e9c46036 100644 --- a/tests/unit/onnx/test_onnx_utils.py +++ b/tests/unit/onnx/test_onnx_utils.py @@ -304,7 +304,7 @@ def test_random_onnx_weights(): # Removed model weights should be greater than 18 MB assert original_model_size - len(onnx_model_wo_weights) > 18e6 - # After assigning random weights, model size should be slightly greater than the the original + # After assigning random weights, model size should be slightly greater than the original # size due to some extra metadata onnx_model_randomized = randomize_weights_onnx_bytes(onnx_model_wo_weights) assert len(onnx_model_randomized) > original_model_size From c603b63407ee61e5b377e2194f12f9a16184e8cd Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Wed, 21 Oct 2026 15:30:04 +0530 Subject: [PATCH 06/43] Leave MTP layers unquantized at any depth in the default PTQ exclusions (#2911) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Type of change: Bug fix The default PTQ exclusions (`configs/ptq/units/default_disabled_quantizers`, imported by the `general/ptq/*` recipes and `mtq.*_CFG` presets) disable multi-token-prediction layers with `mtp.*`, which only matches an MTP block at the model root (HF checkpoints, Megatron LLMs). Megatron vision-language models nest it under the language model (`language_model.mtp.layers.0...`), so their MTP layers were quantized. For Qwen3.5/3.6 VLMs quantized with `examples/megatron_bridge/quantize.py` and a general recipe, `export_quantized_megatron_to_hf.py --vllm_fakequant_export` then refuses the checkpoint ("MTP quantization is not supported by vLLM fakequant export/reload"). The pattern is now `*mtp.*`, matching the model-specific recipes that already use `*mtp*`. Models whose MTP block was nested and previously quantized by a general recipe now keep it in BF16, which is what the exclusion intends. Recipes that quantize MTP on purpose (e.g. the Nemotron 3.5 Super ones) do not import this unit after their MTP entries, so they are unaffected. No API change. - New `tests/unit/recipe/test_presets.py::test_default_exclusions_disable_mtp_at_any_depth`: with `general/ptq/fp8_default-kv_fp8`, MTP quantizers at the root and under `language_model.` are disabled and decoder quantizers stay enabled; fails without the change. - `tests/unit/recipe`: 492 passed. - Found while validating vLLM fake-quant export of Qwen3.6-35B-A3B (#2889): the default `nvfp4_default-kv_fp8` checkpoint was refused on `language_model.mtp.layers.0.eh_proj.input_quantizer`. - Is this change backward compatible?: ✅ (nested MTP layers that a general recipe used to quantize now stay in BF16, as the exclusion intends) - 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?: ❌ Related to #2889 (vLLM fake-quant export through Megatron-Bridge). 🤖 Generated with [Claude Code](https://claude.com/claude-code) * **Bug Fixes** * Corrected default post-training quantization recipes for Megatron vision-language models so MTP layers remain unquantized consistently, including nested language-model layers. * Preserved quantization behavior for unrelated decoder projection layers. * Added regression coverage to verify both root-level and nested MTP exclusions in the default FP8 recipe. * **Documentation** * Added a changelog entry for the 0.49.0 bug fix. 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 + .../units/default_disabled_quantizers.yaml | 2 +- tests/unit/recipe/test_presets.py | 20 +++++++++++++++++++ 3 files changed, 22 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index cf4d4172b2a..b54c82ce477 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -187,6 +187,7 @@ Changelog - 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. +- The default PTQ recipes now leave multi-token-prediction (MTP) layers unquantized in Megatron vision-language models too (``language_model.mtp.*``), as they already did at the model root. 0.47.0 (2026-09-23) ^^^^^^^^^^^^^^^^^^^ diff --git a/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml b/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml index cac4f6d80aa..cd5097dbded 100644 --- a/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml +++ b/modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml @@ -46,7 +46,7 @@ enable: false - quantizer_name: '*router*' enable: false - - quantizer_name: 'mtp.*' + - quantizer_name: '*mtp.*' enable: false - quantizer_name: 'output.*' enable: false diff --git a/tests/unit/recipe/test_presets.py b/tests/unit/recipe/test_presets.py index 3474cd51a43..167aa46e5a0 100644 --- a/tests/unit/recipe/test_presets.py +++ b/tests/unit/recipe/test_presets.py @@ -24,6 +24,7 @@ import argparse import pytest +import torch import modelopt.torch.quantization as mtq from modelopt.recipe import load_recipe, presets @@ -95,6 +96,25 @@ def test_w4a16_nvfp4_preset_disables_vllm_marlin_incompatible_projections(): } <= disabled_quantizers +def test_default_exclusions_disable_mtp_at_any_depth(): + # Megatron VLMs nest the MTP block under ``language_model.``; HF and Megatron LLMs keep it at the root. + model = torch.nn.Module() + for path in ("mtp", "language_model.mtp", "language_model.decoder"): + parent = model + for part in path.split("."): + if not hasattr(parent, part): + parent.add_module(part, torch.nn.Module()) + parent = getattr(parent, part) + parent.add_module("proj", torch.nn.Linear(16, 16)) + + quant_cfg = load_recipe("general/ptq/fp8_default-kv_fp8").quantize.model_dump() + mtq.quantize(model, {**quant_cfg, "algorithm": None}) + + assert model.language_model.decoder.proj.input_quantizer.is_enabled + assert not model.mtp.proj.input_quantizer.is_enabled + assert not model.language_model.mtp.proj.input_quantizer.is_enabled + + def test_local_hessian_layerwise_is_scoped_to_the_model_recipe(): # Layerwise calibration needs identifiable decoder layers, so it belongs to the model # recipe the published numbers come from -- not to the shared preset, which also backs From d7713ab94de0ed4d57e323d81d2536111b0726a8 Mon Sep 17 00:00:00 2001 From: kaix-nv Date: Wed, 21 Oct 2026 12:51:55 -0700 Subject: [PATCH 07/43] Fix llm_sparsity finetune saving a sharded state dict under FSDP2 (#2919) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Type of change: Bug fix `export_hf_ckpt.py` cannot restore the checkpoint written by the weight-sparsity `finetune.py`. Every parameter fails with: ``` While copying the parameter named "model.embed_tokens.weight" ... an exception occurred : ('aten.copy_.default got mixed torch.Tensor and DTensor, need to convert all torch.Tensor to DTensor before calling distributed operators!',) ``` **Root cause.** Transformers 5 builds the Trainer's FSDP plugin with `fsdp_config.get("version", 2)`, so `launch_finetune.sh`'s `--fsdp 'full_shard auto_wrap'` now runs FSDP2. Under FSDP2, `model.state_dict()` returns sharded DTensors. `get_model_state_dict()` in `finetune.py` only moved them to CPU, so `finetuned_modelopt_state.pth` held rank 0's DTensor shards. Restoring those into the unsharded model in `export_hf_ckpt.py` fails. This affects every `--sparsity_fmt`, not only `sparse_magnitude`. **Fix.** Gather the full state dict with `torch.distributed.checkpoint.state_dict.get_model_state_dict(..., StateDictOptions(full_state_dict=True, cpu_offload=True))`, the same pattern as `examples/speculative_decoding/eagle_utils.py`. The function is looked up at call time, so ModelOpt's `DynamicModule`-aware patch from `modelopt/torch/opt/_hooks.py` applies. On a sparsified FSDP2 model, the patched and unpatched gathers gave identical tensors. The existing finetune example test now also runs `export_hf_ckpt.py` on the finetuned checkpoint. Before this PR, the test never restored that checkpoint, so CI missed the bug. Checkpoints written by the old `finetune.py` under FSDP2 hold only rank 0's shards and cannot be recovered. Users need to re-run finetuning. No API change. The documented flow now works end to end: ```bash bash launch_finetune.sh ... --output_dir $FT_DIR python export_hf_ckpt.py ... --modelopt_restore_path $FT_DIR/finetuned_modelopt_state.pth ``` Run locally on 2x RTX A6000 with torch 2.9.1, transformers 5.12.1 (the NVBug version), and accelerate 1.15.0: - `pytest tests/examples/llm_sparsity/weight_sparsity/test_llama_sparsify.py -k finetune`, covering `test_llama_sparsity_finetune` (1 GPU) and `test_llama_sparsity_finetune_multi_gpu` (2 GPU). Both pass with the fix. With only `finetune.py` reverted to `main`, both fail at the new export step with the NVBug's `got mixed torch.Tensor and DTensor` error. - In a standalone FSDP2 repro (tiny Llama, `sparse_magnitude`, 2 ranks), the restored weights keep 2:4 sparsity (at most 2 nonzeros per group of 4). - `pre-commit run` on the changed files passes. Not run: the full Llama-3.1-8B flow on GB200 from the NVBug; peak host memory at 8B scale (rank 0 holds the full CPU state dict, about 16 GB in bf16, as with FSDP1's full state dict); FSDP1 (`fsdp_config.version=1`); and non-FSDP runs. 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?: ✅ (extended the existing finetune example test with the export step) - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ - Did you get Claude approval on this PR?: ❌ NVBug 6927612 (committed for ModelOpt 0.48.0, so this needs a cherry-pick to the 0.48.0 release branch). 🤖 Generated with [Claude Code](https://claude.com/claude-code) * **Bug Fixes** * Fixed sparsity fine-tuning checkpoint saving for compatibility with FSDP2. * Enabled saved checkpoints to be restored and exported as unsharded Hugging Face checkpoints. * **Tests** * Added coverage validating checkpoint export and restoration, including model, data type, sequence length, and output settings. Signed-off-by: Kai Xu Signed-off-by: Rohan Joshi --- CHANGELOG.rst | 1 + examples/llm_sparsity/weight_sparsity/finetune.py | 13 ++++++++----- .../weight_sparsity/test_llama_sparsify.py | 13 +++++++++++++ 3 files changed, 22 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index b54c82ce477..30b384a11dd 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -188,6 +188,7 @@ Changelog - 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. - The default PTQ recipes now leave multi-token-prediction (MTP) layers unquantized in Megatron vision-language models too (``language_model.mtp.*``), as they already did at the model root. +- Fix ``examples/llm_sparsity/weight_sparsity/finetune.py`` saving a sharded state dict under FSDP2, the default FSDP version in Transformers 5, which made ``export_hf_ckpt.py`` fail to restore it. Re-run finetuning to produce a loadable ``finetuned_modelopt_state.pth``. 0.47.0 (2026-09-23) ^^^^^^^^^^^^^^^^^^^ diff --git a/examples/llm_sparsity/weight_sparsity/finetune.py b/examples/llm_sparsity/weight_sparsity/finetune.py index 6eb199adc50..892e08577af 100644 --- a/examples/llm_sparsity/weight_sparsity/finetune.py +++ b/examples/llm_sparsity/weight_sparsity/finetune.py @@ -30,6 +30,7 @@ from dataclasses import dataclass, field import torch +import torch.distributed.checkpoint.state_dict as dcp_state_dict import transformers import utils from torch.utils.data import Dataset @@ -206,11 +207,13 @@ def preprocess( def get_model_state_dict(trainer: transformers.Trainer): - """Collects the state dict.""" - state_dict = trainer.model.state_dict() - cpu_state_dict = {key: value.cpu() for key, value in state_dict.items()} - del state_dict - return cpu_state_dict + """Collects the full (unsharded) state dict on rank 0's CPU; must be called on all ranks.""" + # Under FSDP2 (the transformers 5.x default), model.state_dict() holds sharded DTensors that + # cannot be loaded into an unsharded model, so gather full tensors instead. + return dcp_state_dict.get_model_state_dict( + trainer.model, + options=dcp_state_dict.StateDictOptions(full_state_dict=True, cpu_offload=True), + ) class SupervisedDataset(Dataset): diff --git a/tests/examples/llm_sparsity/weight_sparsity/test_llama_sparsify.py b/tests/examples/llm_sparsity/weight_sparsity/test_llama_sparsify.py index 6277b0fa5b4..f6a7e31d7ec 100644 --- a/tests/examples/llm_sparsity/weight_sparsity/test_llama_sparsify.py +++ b/tests/examples/llm_sparsity/weight_sparsity/test_llama_sparsify.py @@ -112,6 +112,19 @@ def _test_llama_sparsity_finetune(tiny_llama_path, tmp_path, data_path, sparsity max_length=64, ) + # Finally restore the finetuned (FSDP-saved) state into an unsharded model and export it + run_example_command( + extend_cmd_parts( + ["python", "export_hf_ckpt.py"], + model_name_or_path=tiny_llama_path, + modelopt_restore_path=finetune_output / "finetuned_modelopt_state.pth", + dtype=dtype, + model_max_length=64, + output_dir=tmp_path / "export_output", + ), + "llm_sparsity/weight_sparsity", + ) + @pytest.mark.parametrize( ("sparsity_fmt", "dtype"), From bbb07c38d69914c6b57572e245fa39a3bcff0589 Mon Sep 17 00:00:00 2001 From: Zhehao Hu <69692533+Zhehao-Hu@users.noreply.github.com> Date: Wed, 21 Oct 2026 20:24:12 -0700 Subject: [PATCH 08/43] Fix attention quantizer leakage in diffusers HF exports (#2904) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Type of change: ? Bug fix Exporting Qwen-Image FP8 with --quantize-mha leaves internal {q,k,v}_bmm_quantizer._amax buffers in the HF checkpoint, including buffers from disabled quantizers. TRT-LLM rejects these unexpected keys. Hide attention quantizers during HF export by default, except for FLUX.1, which retains enabled Q/K/V quantizers for compatibility with the static FP8 attention path in [NVIDIA/TensorRT-LLM#18396](https://github.com/NVIDIA/TensorRT-LLM/pull/18396). An explicit option overrides this default. Disabled attention quantizers are always omitted, and quantizers are restored on the live model after saving. Infer KV-cache metadata only for Transformers components. Diffusion components, including the native LTX-2 transformer, no longer emit KV-cache quantization metadata. Default exports omit attention quantizer state for Qwen-Image and other diffusion components, while FLUX.1 retains enabled Q/K/V quantizers without an additional flag: ```python export_hf_checkpoint(pipe, export_dir="./hf_ckpt") ``` Explicitly retain enabled Q/K/V quantizers for a compatible runtime: ```python export_hf_checkpoint( pipe, export_dir="./hf_ckpt", export_attention_quantizers=True, ) ``` Set export_attention_quantizers=False to omit attention quantizers, including for FLUX.1. Equivalent diffusers CLI option: ```sh --extra-param export_attention_quantizers=true --extra-param export_attention_quantizers=false ``` - 93 focused export and quantization-config tests passed. - Regression cases run real FP8 calibration and sharded safetensors export with tiny Qwen-Image and FLUX models. They verify clean Qwen exports, FLUX’s default retention of enabled Q/K/V scales, explicit overrides, omission of disabled buffers, and restoration of live quantizers. - Additional coverage verifies that the native diffusion save path used by LTX-2 omits KV-cache metadata while a real Transformers model retains it. This exercises the native export interface, not a full LTX-2 model. - Regression tests reproduce the original buffer leak. The FLUX-default and native-diffusion metadata checks also fail against the implementation before the review fixes. - Pre-commit checks passed. - One optional ONNX example suite was skipped during earlier validation because onnx_graphsurgeon was absent. - Full B100/TRT-LLM generation and full LTX-2 deployment were not run; the test host has no usable GPU. 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?: ❌ - Consumers of legacy attention `_amax` keys must enable `export_attention_quantizers=True`. - 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 - Not yet run Fixes NVBug 6927340. - **Breaking Changes** - Hugging Face diffusion exports omit attention quantizer data by default. - FLUX.1 exports retain enabled Q/K/V quantizer values for runtime compatibility. - KV-cache metadata is excluded from native diffusion exports. - **New Features** - Added the `export_attention_quantizers` option to retain or omit attention quantizers. - Disabled quantizers remain excluded regardless of this option. - Transformers exports continue to include KV-cache metadata. - **Documentation** - Added guidance on attention quantizer and KV-cache export behavior. Signed-off-by: Zhehao Hu Signed-off-by: Rohan Joshi --- CHANGELOG.rst | 1 + examples/diffusers/README.md | 11 ++ examples/diffusers/quantization/quantize.py | 6 +- modelopt/torch/export/diffusers_utils.py | 21 ++- modelopt/torch/export/quant_utils.py | 6 +- modelopt/torch/export/unified_export_hf.py | 38 ++++- .../torch/export/test_export_diffusers.py | 134 +++++++++++++++++- 7 files changed, 208 insertions(+), 9 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 30b384a11dd..dd2d156d601 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -111,6 +111,7 @@ Changelog - ``get_te_hybrid_stack_spec`` was removed from ``modelopt.torch.nas.plugins.megatron``; it had no use outside tests. Use ``modelopt.torch.utils.plugins.megatron_layer_specs.te_hybrid_stack_spec_sequential_mlp`` for the SequentialMLP layout, or ``megatron.core.models.hybrid.hybrid_layer_specs.hybrid_stack_spec`` for grouped GEMM. - Unified HuggingFace export now fails with ``NotImplementedError`` when it meets an MoE block whose expert projection names it does not know, instead of assuming Mixtral's ``w1``/``w2``/``w3``. If you hit this, register a ``ModelSpec`` for the model under ``modelopt/torch/models/``. Every MoE architecture ModelOpt exported correctly before this change is registered, so no supported model regresses. - ``--recipe`` (and ``modelopt.recipe.load_recipe``) now resolve a recipe path **filesystem-first**: a recipe of the same relative path in the current working directory takes precedence over the shipped built-in of that name, matching how recipe ``$import`` paths already resolve. Previously the built-in won. +- Diffusion HF exports now omit attention quantizer buffers by default, except for FLUX.1, which retains enabled Q/K/V ``_amax`` keys for runtime compatibility; other consumers of those keys can opt in with ``export_attention_quantizers=True`` (CLI: ``--extra-param export_attention_quantizers=true``). Disabled attention quantizers remain excluded, and diffusion components, including LTX-2, no longer emit KV-cache quantization metadata. **Deprecations** diff --git a/examples/diffusers/README.md b/examples/diffusers/README.md index 8d6e61e9644..364aa0754ef 100644 --- a/examples/diffusers/README.md +++ b/examples/diffusers/README.md @@ -124,6 +124,17 @@ python quantize.py \ --hf-ckpt-dir ./hf_ckpt ``` +HF exports omit attention quantizer state by default, including when calibration uses +`--quantize-mha`, except for FLUX.1. Qwen-Image's TRT-LLM loader requires this clean +checkpoint. FLUX.1 retains enabled Q/K/V `_amax` keys by default for compatibility with +the static FP8 attention path in [TensorRT-LLM PR #18396](https://github.com/NVIDIA/TensorRT-LLM/pull/18396). +Use `--extra-param export_attention_quantizers=true` to retain enabled Q/K/V quantizers +explicitly, or `export_attention_quantizers=false` to omit them, including for FLUX.1. +Disabled quantizers are always omitted. The equivalent Python option is +`export_hf_checkpoint(pipe, export_dir="./hf_ckpt", export_attention_quantizers=True)`. +Diffusion attention quantization, including native LTX-2 components, is not exported +as KV-cache quantization. Transformers components retain their KV-cache metadata. + #### Wan 2.2 VAE NVFP4 (Conv3D Implicit GEMM) The Wan 2.2 VAE (`AutoencoderKLWan`, shared between the 5B and 14B pipelines) is built from 3D convolutions. When quantizing the VAE with NVFP4, the `Conv3d` layers are automatically dispatched through a custom BF16 WMMA implicit-GEMM kernel with fused FP4 activation quantization. Requires SM80+ (Ampere or newer). See [`modelopt/torch/kernels/quantization/conv/README.md`](../../modelopt/torch/kernels/quantization/conv/README.md) for kernel details. diff --git a/examples/diffusers/quantization/quantize.py b/examples/diffusers/quantization/quantize.py index 502801157fe..73ce6808df3 100644 --- a/examples/diffusers/quantization/quantize.py +++ b/examples/diffusers/quantization/quantize.py @@ -389,7 +389,11 @@ def export_hf_ckpt(self, pipe: Any, model_config: ModelConfig | None = None) -> self.logger.info(f"Merging base safetensors from {merged_path} for LTX2 export") kwargs["merged_base_safetensor_path"] = merged_path if model_config: - for key in ("enable_swizzle_layout", "enable_layerwise_quant_metadata"): + for key in ( + "enable_swizzle_layout", + "enable_layerwise_quant_metadata", + "export_attention_quantizers", + ): val = model_config.extra_params.get(key) if val is not None: normalized = str(val).strip().lower() diff --git a/modelopt/torch/export/diffusers_utils.py b/modelopt/torch/export/diffusers_utils.py index 0309c83e1f6..2892e16a02e 100644 --- a/modelopt/torch/export/diffusers_utils.py +++ b/modelopt/torch/export/diffusers_utils.py @@ -717,7 +717,7 @@ def get_diffusion_components( @contextmanager -def hide_quantizers_from_state_dict(model: nn.Module): +def hide_quantizers_from_state_dict(model: nn.Module, export_attention_quantizers: bool = False): """Context manager that temporarily removes quantizer modules from the model. This allows save_pretrained to save the model without quantizer buffers like _amax. @@ -725,6 +725,8 @@ def hide_quantizers_from_state_dict(model: nn.Module): Args: model: The model with quantizers to temporarily hide. + export_attention_quantizers: Keep enabled Q/K/V BMM quantizers for runtimes + that consume their legacy ``*_bmm_quantizer._amax`` keys. Yields: None - the model can be saved within the context. @@ -739,8 +741,23 @@ def hide_quantizers_from_state_dict(model: nn.Module): # first since we mutate the module tree while iterating. for name, module in list(model.named_modules()): backup = {} - for attr in ["weight_quantizer", "input_quantizer", "output_quantizer"]: + for attr in [ + "weight_quantizer", + "input_quantizer", + "output_quantizer", + "q_bmm_quantizer", + "k_bmm_quantizer", + "v_bmm_quantizer", + "softmax_quantizer", + "bmm2_output_quantizer", + ]: child = getattr(module, attr, None) + if ( + export_attention_quantizers + and attr in ("q_bmm_quantizer", "k_bmm_quantizer", "v_bmm_quantizer") + and getattr(child, "is_enabled", False) + ): + continue if isinstance(child, nn.Module): backup[attr] = child delattr(module, attr) diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py index 88092f689dc..91372bbd951 100755 --- a/modelopt/torch/export/quant_utils.py +++ b/modelopt/torch/export/quant_utils.py @@ -1758,6 +1758,8 @@ def _get_unquantized_moe_router_names(model: nn.Module) -> list[str]: def get_quant_config( model: nn.Module, is_modelopt_qlora: bool = False, + *, + include_kv_cache: bool = True, ) -> dict[str, Any]: """Generate quantization config for a model. @@ -1767,6 +1769,8 @@ def get_quant_config( Args: model: The PyTorch model to make config for. is_modelopt_qlora: Whether the model is a modelopt-trained QLoRA model. + include_kv_cache: Infer KV-cache metadata from attention quantizers. Disable + for diffusion models, whose attention quantizers do not represent a cache. Returns: Dictionary containing the quantization configuration @@ -1883,7 +1887,7 @@ def get_quant_config( for quantizer in (k_quantizer, v_quantizer, output_quantizer) if is_kv_boundary and quantizer is not None and quantizer.is_enabled ] - if enabled_kv_quantizers and is_language_model_module: + if include_kv_cache and enabled_kv_quantizers and is_language_model_module: module_kv_quant = get_kv_cache_dtype(module) if module_kv_quant != QUANTIZATION_NONE: kv_cache_formats.add(module_kv_quant) diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 646ba21d714..0ee1f1cd2eb 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -64,6 +64,13 @@ except ImportError: HAS_DIFFUSERS = False +try: + from transformers import PreTrainedModel + + _TRANSFORMERS_MODEL_TYPES: tuple[type, ...] = (PreTrainedModel,) +except ImportError: + _TRANSFORMERS_MODEL_TYPES = () + from modelopt.torch.opt.conversion import ModeloptStateManager, modelopt_state from modelopt.torch.opt.plugins.huggingface import _MODELOPT_STATE_SAVE_NAME @@ -1242,6 +1249,7 @@ def _export_diffusers_checkpoint( export_dir: Path, components: list[str] | None, max_shard_size: int | str = "10GB", + export_attention_quantizers: bool | None = None, **kwargs, ) -> None: """Internal: Export diffusion(-like) model/pipeline checkpoint. @@ -1259,6 +1267,10 @@ def _export_diffusers_checkpoint( max_shard_size: Maximum size of each shard file. If the model exceeds this size, it will be sharded into multiple files and a .safetensors.index.json will be created. Use smaller values like "5GB" or "2GB" to force sharding. + export_attention_quantizers: Preserve enabled Q/K/V BMM quantizer state for + runtimes that consume legacy attention ``_amax`` keys, such as static FP8 FLUX. + None preserves these keys for FLUX.1 and omits them for other diffusion + components. Transformers components always retain them for KV-cache scales. **kwargs: Runtime-specific post-processing options forwarded to :func:`_postprocess_safetensors`. See its docstring for details. """ @@ -1303,6 +1315,14 @@ def _export_diffusers_checkpoint( # Infer dtype if not provided component_dtype = dtype if dtype is not None else infer_dtype_from_model(component) + is_transformers_component = isinstance(component, _TRANSFORMERS_MODEL_TYPES) + # TRT-LLM's static FP8 FLUX.1 path consumes the legacy Q/K/V amax keys. + # Transformers components need these keys for their declared KV-cache scheme. + keep_attention_quantizers = is_transformers_component or ( + isinstance(component, diffusers.FluxTransformer2DModel) + if export_attention_quantizers is None + else export_attention_quantizers + ) if is_quantized: # Fuse QKV linears that share the same input (unify amax values) @@ -1329,7 +1349,12 @@ def _export_diffusers_checkpoint( # promoted buffers are always removed, even if save / post-process / config # update raises (keeps the live module reusable for a repeated export). try: - quant_config = get_quant_config(component, is_modelopt_qlora=False) + quant_config = get_quant_config( + component, + is_modelopt_qlora=False, + # Native diffusion components (including LTX-2) have no KV cache. + include_kv_cache=is_transformers_component, + ) if quant_config: quantization_details = quant_config.get("quantization", {}) # Record the SVDQuant low-rank size so consumers know the LoRA shape. @@ -1345,12 +1370,12 @@ def _export_diffusers_checkpoint( # - diffusers ModelMixin.save_pretrained does NOT accept state_dict parameter # - for non-diffusers modules (e.g., LTX-2 transformer), fall back to torch.save if hasattr(component, "save_pretrained"): - with hide_quantizers_from_state_dict(component): + with hide_quantizers_from_state_dict(component, keep_attention_quantizers): component.save_pretrained( component_export_dir, max_shard_size=max_shard_size ) else: - with hide_quantizers_from_state_dict(component): + with hide_quantizers_from_state_dict(component, keep_attention_quantizers): _save_component_state_dict_safetensors(component, component_export_dir) # Post-process — merge, metadata, padding, swizzle @@ -1905,7 +1930,12 @@ def export_hf_checkpoint( max_shard_size: Maximum size of each safetensors shard file. Defaults to "10GB". **kwargs: Runtime-specific post-processing options forwarded to :func:`_postprocess_safetensors` for diffusion model exports. - See its docstring for supported keys. + See its docstring for supported keys. Diffusion exports also accept + ``export_attention_quantizers=True`` to retain enabled Q/K/V BMM quantizer + state for runtimes using legacy ``*_bmm_quantizer._amax`` keys, or False + to omit it. When unspecified, FLUX.1 retains enabled Q/K/V quantizers for + runtime compatibility; other diffusion components omit them. + Transformers components always retain them for KV-cache scales. """ # Weights the model never loaded (MTP head, auxiliary tower, ...) are copied straight from the # source so the exported checkpoint is the complete model. Merged here, ahead of the path diff --git a/tests/unit/torch/export/test_export_diffusers.py b/tests/unit/torch/export/test_export_diffusers.py index 126550a978c..eedeae677e1 100644 --- a/tests/unit/torch/export/test_export_diffusers.py +++ b/tests/unit/torch/export/test_export_diffusers.py @@ -17,6 +17,7 @@ import inspect import json from pathlib import Path +from types import SimpleNamespace import numpy as np import pytest @@ -26,6 +27,7 @@ get_tiny_dit, get_tiny_flux, get_tiny_flux2, + get_tiny_qwen_image_transformer, get_tiny_unet, get_tiny_wan22_transformer, get_tiny_wan22_vae, @@ -36,11 +38,13 @@ import diffusers from safetensors import safe_open -from safetensors.torch import save_file +from safetensors.torch import load_file, save_file from transformers import ( CLIPTextConfig, CLIPTextModel, CLIPTextModelWithProjection, + LlamaConfig, + LlamaModel, UMT5Config, UMT5EncoderModel, ) @@ -53,6 +57,7 @@ hide_quantizers_from_state_dict, ) from modelopt.torch.export.unified_export_hf import _postprocess_safetensors, export_hf_checkpoint +from modelopt.torch.quantization.nn import TensorQuantizer def _load_config(config_path): @@ -281,6 +286,133 @@ def test_flux2_dummy_inputs_shape(): assert "guidance" in inputs +@pytest.mark.parametrize( + ("model_factory", "export_attention_quantizers", "expect_attention_quantizers"), + [ + (get_tiny_qwen_image_transformer, None, False), + (get_tiny_flux, None, True), + (get_tiny_flux, False, False), + (get_tiny_flux, True, True), + ], +) +def test_export_diffusers_attention_quantizers( + tmp_path, model_factory, export_attention_quantizers, expect_attention_quantizers +): + model = model_factory().eval() + quant_cfg = copy.deepcopy(mtq.FP8_DEFAULT_CFG) + quant_cfg["quant_cfg"].append( + { + "quantizer_name": "*[qkv]_bmm_quantizer", + "cfg": {"num_bits": (4, 3), "axis": None}, + "enable": True, + } + ) + + def calibrate(m): + inputs = generate_diffusion_dummy_inputs(m, torch.device("cpu"), torch.float32) + assert inputs is not None + m(**inputs) + + mtq.quantize(model, quant_cfg, calibrate) + # The example disables sensitive blocks after calibration, leaving their amax buffers. + mtq.disable_quantizer(model, "transformer_blocks.0.*") + quantizers = { + name: module + for name, module in model.named_modules() + if isinstance(module, TensorQuantizer) + } + attention_amax = { + f"{name}._amax": module.amax.clone() + for name, module in quantizers.items() + if name.endswith(("q_bmm_quantizer", "k_bmm_quantizer", "v_bmm_quantizer")) + and module.is_enabled + } + assert attention_amax + assert any(not module.is_enabled and hasattr(module, "_amax") for module in quantizers.values()) + + kwargs = ( + {"export_attention_quantizers": export_attention_quantizers} + if export_attention_quantizers is not None + else {} + ) + export_hf_checkpoint(model, export_dir=tmp_path, max_shard_size="1KB", **kwargs) + + exported = {} + for shard in tmp_path.glob("*.safetensors"): + exported.update(load_file(str(shard))) + expected_attention = attention_amax if expect_attention_quantizers else {} + assert {key for key in exported if "quantizer" in key} == set(expected_attention) + for key, value in expected_attention.items(): + torch.testing.assert_close(exported[key], value) + assert any(value.dtype == torch.float8_e4m3fn for value in exported.values()) + assert any(key.endswith(".weight_scale") for key in exported) + config = _load_config(tmp_path / "config.json")["quantization_config"] + assert not any(key.startswith("kv_cache") for key in config) + for name, module in quantizers.items(): + assert model.get_submodule(name) is module + for key, value in attention_amax.items(): + torch.testing.assert_close(model.state_dict()[key], value) + + +@pytest.mark.parametrize("is_transformers_component", [False, True]) +@pytest.mark.parametrize("export_attention_quantizers", [None, False, True]) +def test_export_diffusion_component_kv_metadata( + tmp_path, is_transformers_component, export_attention_quantizers +): + if is_transformers_component: + model = LlamaModel( + LlamaConfig( + vocab_size=32, + hidden_size=16, + intermediate_size=32, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=2, + ) + ).eval() + + def calibrate(m): + m(input_ids=torch.ones(1, 4, dtype=torch.long)) + + else: + # Exercise the native nn.Module save path used by LTX-2, with real attention. + model = nn.Sequential(get_tiny_flux()).eval() + + def calibrate(m): + inputs = generate_diffusion_dummy_inputs(m[0], torch.device("cpu"), torch.float32) + m[0](**inputs) + + quant_cfg = copy.deepcopy(mtq.FP8_DEFAULT_CFG) + quant_cfg["quant_cfg"].append( + {"quantizer_name": "*[kv]_bmm_quantizer", "cfg": {"num_bits": (4, 3), "axis": None}} + ) + mtq.quantize(model, quant_cfg, calibrate) + kv_amax = { + f"{name}._amax": module.amax.clone() + for name, module in model.named_modules() + if name.endswith(("k_bmm_quantizer", "v_bmm_quantizer")) and module.is_enabled + } + assert kv_amax + # LTX-2 exposes its native transformer through the stage-1 ledger. + pipe = SimpleNamespace(stage_1_model_ledger=SimpleNamespace(transformer=lambda: model)) + unified_export_hf._export_diffusers_checkpoint( + pipe, None, tmp_path, None, export_attention_quantizers=export_attention_quantizers + ) + + config = _load_config(tmp_path / "config.json")["quantization_config"] + assert ("kv_cache_scheme" in config) == is_transformers_component + if is_transformers_component: + assert config["kv_cache_scheme"] == {"dynamic": False, "num_bits": 8, "type": "float"} + else: + assert not any(key.startswith("kv_cache") for key in config) + + exported = load_file(str(tmp_path / "model.safetensors")) + expected_kv = kv_amax if is_transformers_component or export_attention_quantizers else {} + assert {key for key in exported if "quantizer" in key} == set(expected_kv) + for key, value in expected_kv.items(): + torch.testing.assert_close(exported[key], value) + + def test_svdquant_diffusers_export_promotes_clean_keys(): """Fast CPU check of the diffusers SVDQuant export promotion. From f3e9435a602f621d310696ca921831999c3d9feb Mon Sep 17 00:00:00 2001 From: noeyy-mino <174223378+noeyy-mino@users.noreply.github.com> Date: Thu, 22 Oct 2026 15:04:22 +0800 Subject: [PATCH 09/43] [NVBug 6972128] Fix broken links (#2937) 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: https://nvbugs/6972128, documentation. Fix 13 broken links (stale heading anchors, wrong relative paths, a renamed config directory, and references to docs files that never existed) and 3 passages that contradicted the list or command next to them. Text-only change. ### Usage N/A ### Testing N/A ### Before your PR is "*Ready for review*" - Is this change backward compatible?: N/A - 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?: ✅ ### Additional Information N/A ## Summary by CodeRabbit * **Documentation** * Updated links in the ONNX quantization, QAT/QAD, sparsity, autotune, pruning, and compression guides to point to current documentation and configuration references. * Clarified that the QAD sample was tested on a Linux RTX 5090 and currently encounters an out-of-memory issue; updated the dataset-preprocessing script name. * Updated the preset guide to describe three preset types and reposition guidance on reusable snippets and schema validation. --------- Signed-off-by: Noey Yang <174223378+noeyy-mino@users.noreply.github.com> Signed-off-by: noeyy-mino <174223378+noeyy-mino@users.noreply.github.com> Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com> Signed-off-by: Rohan Joshi --- README.md | 2 +- examples/llm_qat/README.md | 2 +- examples/llm_sparsity/weight_sparsity/README.md | 2 +- examples/onnx_ptq/autotune/README.md | 8 ++++---- examples/pruning/cifar_resnet.ipynb | 8 ++++---- examples/windows/diffusers/qad_example/README.md | 4 ++-- modelopt/torch/puzzletron/anymodel/README.md | 4 ++-- modelopt_recipes/configs/ptq/presets/README.md | 12 ++++++------ 8 files changed, 21 insertions(+), 21 deletions(-) diff --git a/README.md b/README.md index dc838483e88..532bd1669e2 100644 --- a/README.md +++ b/README.md @@ -139,7 +139,7 @@ more fine-grained control on installed dependencies or for alternative docker im |------------|----------------| | LLM / VLM Quantization | [View Support Matrix](./examples/hf_ptq/README.md#support-matrix) | | Diffusers Quantization | [View Support Matrix](./examples/diffusers/README.md#support-matrix) | -| ONNX Quantization | [View Support Matrix](./examples/torch_onnx/README.md#onnx-export-supported-llm-models) | +| ONNX Quantization | [View Support Matrix](./examples/torch_onnx/README.md#supported-models) | | Windows Quantization | [View Support Matrix](./examples/windows/README.md#support-matrix) | | Quantization Aware Training | [View Support Matrix](./examples/llm_qat/README.md#support-matrix) | | Pruning | [View Support Matrix](./examples/pruning/README.md#support-matrix) | diff --git a/examples/llm_qat/README.md b/examples/llm_qat/README.md index cd88ae30498..0925d4dd1d9 100644 --- a/examples/llm_qat/README.md +++ b/examples/llm_qat/README.md @@ -39,7 +39,7 @@ The Qwen3-8B example below requires a minimum of **2 x 80GB GPUs**. ## Run End-to-End QAT/QAD Example All arguments can be set via YAML, CLI, or both (CLI overrides YAML). See -[ARGUMENTS.md](ARGUMENTS.md), `--help`, and [Configuration](#advanced-configuration). +[ARGUMENTS.md](ARGUMENTS.md), `--help`, and [Advanced Topics](#advanced-topics). ### Quantization Recipes diff --git a/examples/llm_sparsity/weight_sparsity/README.md b/examples/llm_sparsity/weight_sparsity/README.md index 42bac030ce9..2330595a270 100644 --- a/examples/llm_sparsity/weight_sparsity/README.md +++ b/examples/llm_sparsity/weight_sparsity/README.md @@ -133,7 +133,7 @@ The output should be similar to the following: ROUGE scores: {'rouge1': 42.174, 'rouge2': 19.2724, 'rougeL': 28.6989, 'rougeLsum': 39.1394} ``` -Please refer to [link](../llm_eval/README.md#Evaluation-scripts-for-LLM-tasks) for more details of how to evaluate the sparsified models on other benchmarks, such as MMLU and HumanEval. +Please refer to [link](../../llm_eval/README.md#evaluation-scripts-for-llm-tasks) for more details of how to evaluate the sparsified models on other benchmarks, such as MMLU and HumanEval. ## Export and Serve a Hugging Face Checkpoint diff --git a/examples/onnx_ptq/autotune/README.md b/examples/onnx_ptq/autotune/README.md index d8a86b8bae0..14cc9772af7 100644 --- a/examples/onnx_ptq/autotune/README.md +++ b/examples/onnx_ptq/autotune/README.md @@ -2,7 +2,7 @@ This example demonstrates automated Q/DQ (Quantize/Dequantize) node placement optimization for ONNX models using TensorRT performance measurements. -> **Warning:** This example uses the direct Autotune entry point for lower-level Q/DQ placement experiments. If you are starting from an unquantized ONNX model and care about **accuracy**, please use the ONNX PTQ workflow with `--autotune` enabled and with representative calibration data. See [../README#optimize-qdq-node-placement-with-autotune](../README#optimize-qdq-node-placement-with-autotune). +> **Warning:** This example uses the direct Autotune entry point for lower-level Q/DQ placement experiments. If you are starting from an unquantized ONNX model and care about **accuracy**, please use the ONNX PTQ workflow with `--autotune` enabled and with representative calibration data. See [../README.md#optimize-qdq-node-placement-with-autotune](../README.md#optimize-qdq-node-placement-with-autotune). ## Table of Contents @@ -265,7 +265,7 @@ All examples above use the command-line interface. For **low-level programmatic **See the API Reference documentation for low-level usage:** -- [`docs/source/reference/2_qdq_placement.rst`](../../docs/source/reference/2_qdq_placement.rst) +- [`modelopt.onnx.quantization.autotune`](https://nvidia.github.io/Model-Optimizer/reference/generated/modelopt.onnx.quantization.autotune.html) The API docs include detailed examples of: @@ -278,13 +278,13 @@ The API docs include detailed examples of: For comprehensive documentation on QDQ placement optimization, see: -- **User Guide**: [`docs/source/guides/9_qdq_placement.rst`](../../docs/source/guides/9_qdq_placement.rst) +- **User Guide**: [`9_autotune`](https://nvidia.github.io/Model-Optimizer/guides/9_autotune.html) - Detailed explanations of how the autotuner works - Advanced usage patterns and best practices - Configuration options and performance tuning - Troubleshooting common issues -- **API Reference**: [`docs/source/reference/2_qdq_placement.rst`](../../docs/source/reference/2_qdq_placement.rst) +- **API Reference**: [`modelopt.onnx.quantization.autotune`](https://nvidia.github.io/Model-Optimizer/reference/generated/modelopt.onnx.quantization.autotune.html) - Complete API documentation for all classes and functions - Low-level usage examples - State management and pattern cache details diff --git a/examples/pruning/cifar_resnet.ipynb b/examples/pruning/cifar_resnet.ipynb index c5a34cb1167..84a7bde256f 100644 --- a/examples/pruning/cifar_resnet.ipynb +++ b/examples/pruning/cifar_resnet.ipynb @@ -482,15 +482,15 @@ "source": [ "### Prune base model and store pruned net\n", "\n", - "Using [mtp.prune](../reference/generated/modelopt.torch.prune.pruning.rst#modelopt.torch.prune.pruning.prune) you can\n", + "Using [mtp.prune](https://nvidia.github.io/Model-Optimizer/reference/generated/modelopt.torch.prune.pruning.html#modelopt.torch.prune.pruning.prune) you can\n", "\n", "* generate a search space for pruning from your base model;\n", "* prune the model;\n", "* obtain a valid pytorch model that can be used for fine-tuning.\n", "\n", - "Let's say you have the ResNet20 model as our base model to prune from and we are looking for a model with at most 60M FLOPs. We can provide search constraints for `flops` and/or `params` by an upper bound. The values can either be absolute numbers (e.g. `60e6`) or a string percentage (e.g. `\"75%\"`). In addition, we should also provide our training data loader to [mtp.prune](../reference/generated/modelopt.torch.prune.pruning.rst#modelopt.torch.prune.pruning.prune). The training data loader will be used to calibrate the normalization layers in the model. Finally, we will also specify a custom config for configuring the pruning search space to get a more fine-grained selection of pruned nets.\n", + "Let's say you have the ResNet20 model as our base model to prune from and we are looking for a model with at most 60M FLOPs. We can provide search constraints for `flops` and/or `params` by an upper bound. The values can either be absolute numbers (e.g. `60e6`) or a string percentage (e.g. `\"75%\"`). In addition, we should also provide our training data loader to [mtp.prune](https://nvidia.github.io/Model-Optimizer/reference/generated/modelopt.torch.prune.pruning.html#modelopt.torch.prune.pruning.prune). The training data loader will be used to calibrate the normalization layers in the model. Finally, we will also specify a custom config for configuring the pruning search space to get a more fine-grained selection of pruned nets.\n", "\n", - "Finally, we can store the pruned architecture and weights using [mto.save](../reference/generated/modelopt.torch.opt.conversion.rst#modelopt.torch.opt.conversion.save).\n", + "Finally, we can store the pruned architecture and weights using [mto.save](https://nvidia.github.io/Model-Optimizer/reference/generated/modelopt.torch.opt.conversion.html#modelopt.torch.opt.conversion.save).\n", "\n", ".. note:: We are optimizing a relatively smaller model here. A finer-grained search could be more effective in such a case. This is why we are specifying custom configs. In general however, it is recommended to convert models with the default config itself." ] @@ -677,7 +677,7 @@ "source": [ "As we can see, the best subnet (59.3M FLOPs) has fitted our constraint of 60M FLOPs. We can also see a drop in validation accuracy for the searched model. This is very common after pruning and fine-tuning is necessary for this model.\n", "\n", - "#### Restore the pruned subnet using [mto.restore](../reference/generated/modelopt.torch.opt.conversion.rst#modelopt.torch.opt.conversion.restore)" + "#### Restore the pruned subnet using [mto.restore](https://nvidia.github.io/Model-Optimizer/reference/generated/modelopt.torch.opt.conversion.html#modelopt.torch.opt.conversion.restore)" ] }, { diff --git a/examples/windows/diffusers/qad_example/README.md b/examples/windows/diffusers/qad_example/README.md index af579325062..b102ad6db48 100644 --- a/examples/windows/diffusers/qad_example/README.md +++ b/examples/windows/diffusers/qad_example/README.md @@ -14,7 +14,7 @@ > (including quantized or distilled checkpoints) remain subject to the LTX Community License > Agreement and are **not** covered by Apache 2.0. -**Note:** This is a **sample script for illustrating the QAD pipeline**. It has been verified to run on a **Linux RTX 5090** system, but runs into **OOM (Out of Memory)** on that configuration. +**Note:** This is a **sample script for illustrating the QAD pipeline**. It has been tested on a **Linux RTX 5090** system and currently runs into **OOM (Out of Memory)** on that configuration. This example demonstrates **Quantization-Aware Distillation (QAD)** for [LTX-2](https://github.com/Lightricks/LTX-2) using the native LTX training loop and [NVIDIA ModelOpt](https://github.com/NVIDIA/Model-Optimizer). It combines: @@ -77,7 +77,7 @@ pip install torch accelerate safetensors pyyaml ### 1. Prepare your dataset -Run the LTX preprocessing script to extract latents and text embeddings from your videos. Use `preprocess_dataset.py` with the following arguments (matching the LTX training pipeline): +Run the LTX preprocessing script to extract latents and text embeddings from your videos. Use `process_dataset.py` with the following arguments (matching the LTX training pipeline): ```bash python scripts/process_dataset.py /path/to/dataset.json \ diff --git a/modelopt/torch/puzzletron/anymodel/README.md b/modelopt/torch/puzzletron/anymodel/README.md index 1c8c68b60bf..4c4ae11641a 100644 --- a/modelopt/torch/puzzletron/anymodel/README.md +++ b/modelopt/torch/puzzletron/anymodel/README.md @@ -84,7 +84,7 @@ Step 3: Configure YAML files Update the main model config YAML: - Set `descriptor` to match the name used in `@ModelDescriptorFactory.register_decorator("your_model_name")` -- See example: [llama_3_1_8b_instruct.yaml](../../../../tests/gpu/torch/puzzletron/resources/configs/llama_3_1_8b_instruct/llama_3_1_8b_instruct.yaml) +- See example: [Llama-3.1-8B-Instruct.yaml](../../../../tests/gpu/torch/puzzletron/resources/configs/meta-llama/Llama-3.1-8B-Instruct/Llama-3.1-8B-Instruct.yaml) Update pruning YAML files (`ffn_pruning.yaml`, `expert_pruning.yaml`, etc.): @@ -92,7 +92,7 @@ Update pruning YAML files (`ffn_pruning.yaml`, `expert_pruning.yaml`, etc.): - Set `layer_descriptor._target_` to your layer descriptor class - Set `hook_class` to the activation hook for scoring - Set `target_layer` in `activation_hooks_kwargs` to the layer name for hook attachment -- See examples in [configs/llama_3_1_8b_instruct/pruning/](../../../../tests/gpu/torch/puzzletron/resources/configs/llama_3_1_8b_instruct/pruning/) +- See examples in [configs/meta-llama/Llama-3.1-8B-Instruct/pruning/](../../../../tests/gpu/torch/puzzletron/resources/configs/meta-llama/Llama-3.1-8B-Instruct/pruning/) ## End-to-end example diff --git a/modelopt_recipes/configs/ptq/presets/README.md b/modelopt_recipes/configs/ptq/presets/README.md index 2aeec2d2b33..ed12bd66fa6 100644 --- a/modelopt_recipes/configs/ptq/presets/README.md +++ b/modelopt_recipes/configs/ptq/presets/README.md @@ -6,14 +6,9 @@ from `modelopt.torch.quantization.config` (e.g., `FP8_DEFAULT_CFG`, `FP8_KV_CFG`). Presets compose from the reusable snippets in `configs/numerics/` and -`configs/ptq/units/` via the `$import` system, and are split into two +`configs/ptq/units/` via the `$import` system, and are split into three kinds: -Preset files are also reusable snippets when imported by recipes or other -configs, so each preset must declare a `# modelopt-schema: ...` preamble. -Current preset files use `QuantizeConfig` so they are validated after their -own imports have been resolved. - - **`model/`** — *full* quantization presets. Each file is a complete, self-contained config (it sets `algorithm` and a full `quant_cfg` with a base-disable-all prefix + standard exclusions) and can be passed @@ -30,6 +25,11 @@ own imports have been resolved. attention and softmax quantizer choices that differ from the generic `model/` presets. +Preset files are also reusable snippets when imported by recipes or other +configs, so each preset must declare a `# modelopt-schema: ...` preamble. +Current preset files use `QuantizeConfig` so they are validated after their +own imports have been resolved. + **Note:** The main purpose of these presets is to support the existing `hf_ptq.py` script's `--qformat` / `--kv_cache_qformat` flags and other code paths that reference From 67aab935b7bb6d46a8b6fa130ba400b505f9300f Mon Sep 17 00:00:00 2001 From: noeyy-mino <174223378+noeyy-mino@users.noreply.github.com> Date: Thu, 22 Oct 2026 15:04:36 +0800 Subject: [PATCH 10/43] [NVBug 6967406] Fix misspellings in example READMEs and notebooks (#2928) 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: https://nvbugs/6967406 documentation Fixes 20 misspellings in 11 example READMEs and Jupyter notebooks (NVBug 6967406), e.g. demontrate, Qunatized, Dumpping, effciency, depployment, Dependancies, controled, contrainer, intellegant, cahce, Hugginface, LLama, ontop. Text-only change. No code, commands, links, or notebook outputs are modified. ### Usage N/A ### Testing N/A ### Before your PR is "*Ready for review*" - Is this change backward compatible?: N/A - 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?: ✅ ### Additional Information N/A ## Summary by CodeRabbit * **Documentation** * Corrected spelling, grammar, and capitalization across example guides and notebooks, including references to Hugging Face and model names. * Clarified wording around quantization, caching, and deployment. Documented commands, workflows, and guidance remain unchanged. --------- Signed-off-by: Noey Yang <174223378+noeyy-mino@users.noreply.github.com> Signed-off-by: Rohan Joshi --- examples/deepseek/README.md | 4 ++-- examples/diffusers/README.md | 4 ++-- examples/llm_distill/README.md | 2 +- examples/llm_qat/notebooks/QAT_QAD_Walkthrough.ipynb | 10 +++++----- examples/model_hub/README.md | 2 +- examples/pruning/cifar_resnet.ipynb | 2 +- examples/specdec_bench/README.md | 2 +- examples/speculative_decoding/README.md | 6 +++--- examples/windows/onnx_ptq/genai_llm/README.md | 4 ++-- examples/windows/onnx_ptq/sam2/README.md | 2 +- examples/windows/onnx_ptq/whisper/README.md | 2 +- modelopt/torch/export/trtllm/layer_utils.py | 2 +- 12 files changed, 21 insertions(+), 21 deletions(-) diff --git a/examples/deepseek/README.md b/examples/deepseek/README.md index 33848ed06aa..ccab5a09ef0 100644 --- a/examples/deepseek/README.md +++ b/examples/deepseek/README.md @@ -26,7 +26,7 @@ export HF_FP4_PATH={path_to_save_the_final_FP4_checkpoint} ### DeepSeek V3 R1 V3.1 ```bash -# download the FP8 checkpoint from Hugginface. This is an example of DeepSeek-R1 +# download the FP8 checkpoint from Hugging Face. This is an example of DeepSeek-R1 huggingface-cli download deepseek-ai/DeepSeek-R1 --local-dir $HF_FP8_CKPT # clone DeepSeek-V3 (base model of R1) Github repository for FP8 inference, @@ -36,7 +36,7 @@ git clone https://github.com/deepseek-ai/DeepSeek-V3.git && cd DeepSeek-V3 && gi ### [Experimental] DeepSeek V3.2 ```bash -# download the FP8 checkpoint from Hugginface. +# download the FP8 checkpoint from Hugging Face. huggingface-cli download deepseek-ai/DeepSeek-V3.2-Exp --local-dir $HF_FP8_CKPT # clone DeepSeek-V3.2 Github repository for FP8 inference, diff --git a/examples/diffusers/README.md b/examples/diffusers/README.md index 364aa0754ef..30402163092 100644 --- a/examples/diffusers/README.md +++ b/examples/diffusers/README.md @@ -16,7 +16,7 @@ Cache Diffusion is a technique that reuses cached outputs from previous diffusio | :------------: | :------------: | :------------: | :------------: | | Pre-Requisites | Required & optional packages to use this technique | \[[Link](#pre-requisites)\] | | | Getting Started | Learn how to optimize your models using quantization/cache diffusion to reduce precision and improve inference efficiency | \[[Link](#getting-started)\] | \[[docs](https://nvidia.github.io/Model-Optimizer/guides/1_quantization.html)\] | -| Support Matrix | View the support matrix to see quantization/cahce diffusion compatibility and feature availability across different models | \[[Link](#support-matrix)\] | \[[docs](https://nvidia.github.io/Model-Optimizer/guides/1_quantization.html)\] | +| Support Matrix | View the support matrix to see quantization/cache diffusion compatibility and feature availability across different models | \[[Link](#support-matrix)\] | \[[docs](https://nvidia.github.io/Model-Optimizer/guides/1_quantization.html)\] | | Sparse Attention (Skip-Softmax) | Skip-softmax sparse attention for diffusion models | \[[Link](#sparse-attention-skip-softmax)\] | | | Cache Diffusion | Caching technique to accelerate inference without compromising quality | \[[Link](#cache-diffusion)\] | | | Post Training Quantization (PTQ) | Example scripts on how to run PTQ on diffusion models | \[[Link](#post-training-quantization-ptq)\] | \[[docs](https://nvidia.github.io/Model-Optimizer/guides/1_quantization.html)\] | @@ -97,7 +97,7 @@ We support calibration for INT8, FP8 and FP4 precision and for both weights and We also provide instructions on deploying and running E2E diffusion pipelines with Model Optimizer quantized INT8 and FP8 Backbone to generate images and measure latency on target GPUs. Note, Jetson devices are not supported at this time due to the incompatibility of the software. > [!NOTE] -> Model calibration requires relatively more GPU computing power then deployment. It does not need to be on the same GPUs as the deployment target GPUs. Deprecated ONNX export and TensorRT engine instructions remain available in [`quantization/ONNX-TRT-Deployment.md`](./quantization/ONNX-TRT-Deployment.md). +> Model calibration requires relatively more GPU computing power than deployment. It does not need to be on the same GPUs as the deployment target GPUs. Deprecated ONNX export and TensorRT engine instructions remain available in [`quantization/ONNX-TRT-Deployment.md`](./quantization/ONNX-TRT-Deployment.md). ### Quantize scripts diff --git a/examples/llm_distill/README.md b/examples/llm_distill/README.md index 834ec15ada9..08f800a62bb 100644 --- a/examples/llm_distill/README.md +++ b/examples/llm_distill/README.md @@ -11,7 +11,7 @@ This section focuses on demonstrating how to apply Model Optimizer to perform kn | **Section** | **Description** | **Link** | **Docs** | | :------------: | :------------: | :------------: | :------------: | | Pre-Requisites | Required & optional packages to use this technique | \[[Link](#pre-requisites)\] | | -| Getting Started | Learn how to optimize your models using distillation to produce more intellegant smaller models | \[[Link](#getting-started)\] | \[[docs](https://nvidia.github.io/Model-Optimizer/guides/4_distillation.html)\] | +| Getting Started | Learn how to optimize your models using distillation to produce more intelligent smaller models | \[[Link](#getting-started)\] | \[[docs](https://nvidia.github.io/Model-Optimizer/guides/4_distillation.html)\] | | Support Matrix | View the support matrix to see compatibility and feature availability across different models | \[[Link](#support-matrix)\] | | | Distillation with Megatron-Bridge | Learn how to distill your models with Megatron-Bridge Framework | \[[Link](#knowledge-distillation-kd-in-nvidia-megatron-bridge-framework)\] | \[[docs](https://nvidia.github.io/Model-Optimizer/guides/4_distillation.html)\] | | Distillation with Megatron-LM | Learn how to distill your models with Megatron-LM Framework | \[[Link](#knowledge-distillation-kd-in-nvidia-megatron-lm-framework)\] | | diff --git a/examples/llm_qat/notebooks/QAT_QAD_Walkthrough.ipynb b/examples/llm_qat/notebooks/QAT_QAD_Walkthrough.ipynb index 6e65d2dcc9b..a9e486b39f6 100644 --- a/examples/llm_qat/notebooks/QAT_QAD_Walkthrough.ipynb +++ b/examples/llm_qat/notebooks/QAT_QAD_Walkthrough.ipynb @@ -23,7 +23,7 @@ "id": "c3f7f931-ac38-494e-aea8-ca2cd6d05794", "metadata": {}, "source": [ - "## Installing Prerequisites and Dependancies" + "## Installing Prerequisites and Dependencies" ] }, { @@ -31,14 +31,14 @@ "id": "d7d4f25f-e569-42cf-8022-bb7cc6f9ea6e", "metadata": {}, "source": [ - "If you haven't already, install the required dependencies for this notebook. Key dependancies include:\n", + "If you haven't already, install the required dependencies for this notebook. Key dependencies include:\n", "\n", "- nvidia-modelopt\n", "- torch\n", "- transformers\n", "- jupyterlab\n", "\n", - "This repo contains a `examples/llm_qat/notebooks/requirements.txt` file that can be used to install all required dependancies." + "This repo contains a `examples/llm_qat/notebooks/requirements.txt` file that can be used to install all required dependencies." ] }, { @@ -398,7 +398,7 @@ "id": "e471ef6c-1346-4e5e-8782-5e9f2bc38f8a", "metadata": {}, "source": [ - "Once you have quantized the model you can now start the post-training process for QAT. The training process will calculate validation loss at 50 step intervals and save the model. These can be controled by adjusting the `eval_steps` and `output_dir` above along with other `training_args`." + "Once you have quantized the model you can now start the post-training process for QAT. The training process will calculate validation loss at 50 step intervals and save the model. These can be controlled by adjusting the `eval_steps` and `output_dir` above along with other `training_args`." ] }, { @@ -779,7 +779,7 @@ "id": "eca645fb-d8d2-4c98-9cb2-afbae7d30d6c", "metadata": {}, "source": [ - "## Stop the TensorRT-LLM Docker contrainer" + "## Stop the TensorRT-LLM Docker container" ] }, { diff --git a/examples/model_hub/README.md b/examples/model_hub/README.md index 2653503c56f..0f70f0dc629 100644 --- a/examples/model_hub/README.md +++ b/examples/model_hub/README.md @@ -1,6 +1,6 @@ # Deploy quantized models from Nvidia's Hugging Face model hub with TensorRT-LLM, vLLM, and SGLang -The provided bash scripts show an example to deploy and run the [quantized LLama 3.1 8B Instruct FP8 model](https://huggingface.co/nvidia/Llama-3.1-8B-Instruct-FP8) from Nvidia's Hugging Face model hub on TensorRT-LLM, vLLM and SGLang respectively. +The provided bash scripts show an example to deploy and run the [quantized Llama 3.1 8B Instruct FP8 model](https://huggingface.co/nvidia/Llama-3.1-8B-Instruct-FP8) from Nvidia's Hugging Face model hub on TensorRT-LLM, vLLM and SGLang respectively. Before running the bash scripts, please make sure you have setup the environment properly: diff --git a/examples/pruning/cifar_resnet.ipynb b/examples/pruning/cifar_resnet.ipynb index 84a7bde256f..a3fd850cac8 100644 --- a/examples/pruning/cifar_resnet.ipynb +++ b/examples/pruning/cifar_resnet.ipynb @@ -450,7 +450,7 @@ "source": [ "We have now established a baseline model and accuracy that we will be comparing with using Model Optimizer.\n", "\n", - "So far, we have seen a regular PyTorch model trained without anything new. Now, lets optimize the model for our target constraints using Model Optimizer!\n", + "So far, we have seen a regular PyTorch model trained without anything new. Now, let's optimize the model for our target constraints using Model Optimizer!\n", "\n", "## FastNAS Pruning with Model Optimizer\n", "\n", diff --git a/examples/specdec_bench/README.md b/examples/specdec_bench/README.md index 9a2d2c464df..f8e69d47416 100644 --- a/examples/specdec_bench/README.md +++ b/examples/specdec_bench/README.md @@ -2,7 +2,7 @@ ## Installation -This benchmark is meant to be a lightweight layer ontop of an existing vLLM/SGLang/TRTLLM installation. For example, no install +This benchmark is meant to be a lightweight layer on top of an existing vLLM/SGLang/TRTLLM installation. For example, no install is required if one is running in the following dockers: `vllm/vllm-openai:v0.11.0` (vLLM), `lmsysorg/sglang:v0.5.4.post2` (SGLang), or `nvcr.io/nvidia/tensorrt-llm/release:1.2.0` (TRT-LLM). diff --git a/examples/speculative_decoding/README.md b/examples/speculative_decoding/README.md index 23ba03f6e39..d1811659e04 100644 --- a/examples/speculative_decoding/README.md +++ b/examples/speculative_decoding/README.md @@ -92,9 +92,9 @@ The saved modelopt checkpoint is similar in architecture to HF models. It can be For large models, you can export intermediate hidden states to disk and train only the draft model. This significantly reduces GPU memory requirements, but requires several to tens of terabytes of disk storage depending on dataset size. -### Dumpping Hidden States to Disk +### Dumping Hidden States to Disk -We support two backends for generating base model hidden states. For better effciency, it is recommended to use TRT-LLM: +We support two backends for generating base model hidden states. For better efficiency, it is recommended to use TRT-LLM: ```bash python collect_hidden_states/compute_hidden_states_trtllm.py \ @@ -187,7 +187,7 @@ Please refer to [TRT-LLM Doc: Speculative Decoding](https://nvidia.github.io/Ten Please refer to [VLLM Doc: Speculative Decoding](https://docs.vllm.ai/en/latest/features/spec_decode/) for detailed usage. -Optionally, you can convert the exported checkpoint to contain target model information, which is accepted by vLLM to simplify depployment: +Optionally, you can convert the exported checkpoint to contain target model information, which is accepted by vLLM to simplify deployment: ```bash python scripts/convert_to_vllm_ckpt.py --input --verifier --output diff --git a/examples/windows/onnx_ptq/genai_llm/README.md b/examples/windows/onnx_ptq/genai_llm/README.md index a03d4f94fb9..be64ad1f64b 100644 --- a/examples/windows/onnx_ptq/genai_llm/README.md +++ b/examples/windows/onnx_ptq/genai_llm/README.md @@ -128,7 +128,7 @@ The quantization strategy selects which layers to quantize to INT8 vs INT4: 1. **INT8 Layers** (Higher Precision): Important layers that significantly impact model quality. Quantized per-channel -2. **INT4 Layers** (Maximum Compression): All other layers. Qunatized blockwise. +2. **INT4 Layers** (Maximum Compression): All other layers. Quantized blockwise. This strategy preserves accuracy for the most sensitive layers while maintaining aggressive compression elsewhere. @@ -244,4 +244,4 @@ Refer to the following example scripts and tutorials for deployment: 1. **Error - Invalid Position-IDs input to the ONNX model** - The ONNX models produced using ONNX GenerativeAI (GenAI) have different IO bindings for models produced using different execution-providers (EPs). For instance, model built with DML EP has position-ids input in the ONNX model but models builts using CUDA EP or NvTensorRtRtx EP don't have position-ids inputs. So, if base model requires, use `--add_position_ids` command-line argument for enabling position_ids calibration input or set "add_position_ids" variable to `True` value (hard-code) in the quantize script if required. + The ONNX models produced using ONNX GenerativeAI (GenAI) have different IO bindings for models produced using different execution-providers (EPs). For instance, model built with DML EP has position-ids input in the ONNX model but models built using CUDA EP or NvTensorRtRtx EP don't have position-ids inputs. So, if base model requires, use `--add_position_ids` command-line argument for enabling position_ids calibration input or set "add_position_ids" variable to `True` value (hard-code) in the quantize script if required. diff --git a/examples/windows/onnx_ptq/sam2/README.md b/examples/windows/onnx_ptq/sam2/README.md index 8d9e5fed3eb..77ea8032072 100644 --- a/examples/windows/onnx_ptq/sam2/README.md +++ b/examples/windows/onnx_ptq/sam2/README.md @@ -1,6 +1,6 @@ # Scripts for ONNX SAM2 model -This repository contains an example to demontrate 8-bit quantization of SAM2 ONNX model. +This repository contains an example to demonstrate 8-bit quantization of SAM2 ONNX model. ## Table of Contents diff --git a/examples/windows/onnx_ptq/whisper/README.md b/examples/windows/onnx_ptq/whisper/README.md index 82ae7822001..4d257735adb 100644 --- a/examples/windows/onnx_ptq/whisper/README.md +++ b/examples/windows/onnx_ptq/whisper/README.md @@ -1,6 +1,6 @@ # Scripts for ONNX Whisper model -This repository contains an example to demontrate 8-bit quantization of Whisper ONNX model. +This repository contains an example to demonstrate 8-bit quantization of Whisper ONNX model. ## Table of Contents diff --git a/modelopt/torch/export/trtllm/layer_utils.py b/modelopt/torch/export/trtllm/layer_utils.py index 8f02ad7298f..006355f8e70 100755 --- a/modelopt/torch/export/trtllm/layer_utils.py +++ b/modelopt/torch/export/trtllm/layer_utils.py @@ -26,7 +26,7 @@ try: from transformers.activations import ACT2FN except Exception: - warn("Cannot find transformers package. Hugginface modules cannot be exported.") + warn("Cannot find transformers package. Hugging Face modules cannot be exported.") from modelopt.torch.models import is_moe from modelopt.torch.utils import distributed as dist From 740a77b784c5a3c0d7981c1cd2f810d77dc07f34 Mon Sep 17 00:00:00 2001 From: noeyy-mino <174223378+noeyy-mino@users.noreply.github.com> Date: Thu, 22 Oct 2026 15:05:22 +0800 Subject: [PATCH 11/43] [https://nvbugs/6820271] Fix skillevaluator schema findings in monitor and evaluation (#2885) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Type of change: Bug fix, https://nvbugs/6820271 Replace the placeholder in monitor's description, which the skillevaluator rejects as an XML tag, and give evaluation/SKILL.md a top-level H1 by promoting all headings one level. N/A N/A - Is this change backward compatible?: N/A - 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?: N/A N/A - **Documentation** - Improved the evaluation guide’s heading hierarchy for clearer navigation. - Updated monitoring guidance to use the “check on job” command with a SLURM job ID. Signed-off-by: Noey Yang <174223378+noeyy-mino@users.noreply.github.com> Signed-off-by: Rohan Joshi --- plugins/modelopt/skills/evaluation/SKILL.md | 32 ++++++++++----------- plugins/modelopt/skills/monitor/SKILL.md | 2 +- 2 files changed, 17 insertions(+), 17 deletions(-) diff --git a/plugins/modelopt/skills/evaluation/SKILL.md b/plugins/modelopt/skills/evaluation/SKILL.md index 342177ea29b..9f5aab58643 100644 --- a/plugins/modelopt/skills/evaluation/SKILL.md +++ b/plugins/modelopt/skills/evaluation/SKILL.md @@ -6,21 +6,21 @@ license: Apache-2.0 # https://github.com/NVIDIA-NeMo/Evaluator/tree/f1fa073/packages/nemo-evaluator-launcher/.claude/skills/nel-assistant --- -## NeMo Evaluator Launcher Assistant +# NeMo Evaluator Launcher Assistant Guide the user through creating NEL YAML configs, running evaluations, and monitoring progress. NEL manages serving for canaries and full runs; see [Step 3's serving-handoff check](references/launcher-workflow.md#step-3--configure-deployment) for external endpoints. -### Completed-run analysis +## Completed-run analysis For completed runs, skip launch/config; use launching-evals for artifacts/analysis and apply [run-validation.md](references/run-validation.md) before scores. Its policy overrides blanket uncapping and fixed token/context limits. -### Workspace integration +## Workspace integration If `MODELOPT_WORKSPACE_ROOT` is set, use the common skill's `workspace-management.md` and reuse existing workspaces (PTQ → checkpoint smoke test → NEL evaluation; carry compatibility requirements into NEL's deployment config). -### Workflow +## Workflow ```text - [ ] Step 0: Check workspace (if MODELOPT_WORKSPACE_ROOT set) @@ -38,7 +38,7 @@ If `MODELOPT_WORKSPACE_ROOT` is set, use the common skill's `workspace-managemen --- -### nel-next path (Terminal-Bench 2.x, SWE-bench, …) — branch here FIRST +## nel-next path (Terminal-Bench 2.x, SWE-bench, …) — branch here FIRST A few **agentic** AA benchmarks do **not** run on the currently validated `nemo-evaluator-launcher` 0.2.6 path (Steps 1–9 don't apply). They run on **nel-next** @@ -56,7 +56,7 @@ Steps 1–9 below are currently validated with 0.2.6 — use them for everything --- -### MRCR (NeMo Gym `simple_agent`) path — branch here too +## MRCR (NeMo Gym `simple_agent`) path — branch here too MRCR **does** run on the currently validated 0.2.6 `nel` launcher (as a `nemo_gym` task, not nel-next), so Steps 1–9 apply — but it is mechanically special and @@ -99,43 +99,43 @@ Read only the sections needed for the current stage; retain the dry-run → cana full-run gates. Existing configs must pass Step 3's serving-handoff check before Step 8. Paths in that reference are relative to this skill directory. -### Step 1 — Prerequisites +## Step 1 — Prerequisites Read [Step 1](references/launcher-workflow.md#step-1--prerequisites). -### Step 2 — Build base config (when not using shortcut) +## Step 2 — Build base config (when not using shortcut) Read [Step 2](references/launcher-workflow.md#step-2--build-base-config-when-not-using-shortcut). -### Step 3 — Configure deployment +## Step 3 — Configure deployment Read [Step 3](references/launcher-workflow.md#step-3--configure-deployment). -### Step 4 — Fill remaining ??? values +## Step 4 — Fill remaining ??? values Read [Step 4](references/launcher-workflow.md#step-4--fill-remaining--values). -### Step 5 — Confirm tasks (iterative) +## Step 5 — Confirm tasks (iterative) Read [Step 5](references/launcher-workflow.md#step-5--confirm-tasks-iterative). -### Step 6 — Multi-node +## Step 6 — Multi-node Read [Step 6](references/launcher-workflow.md#step-6--multi-node). -### Step 7 — Interceptors +## Step 7 — Interceptors Read [Step 7](references/launcher-workflow.md#step-7--interceptors). -### Step 7.5 — Container registry auth (SLURM private images only) +## Step 7.5 — Container registry auth (SLURM private images only) Read [Step 7.5](references/launcher-workflow.md#step-75--container-registry-auth-slurm-private-images-only). -### Step 8 — Run evaluation (gated dry-run → canary → full) +## Step 8 — Run evaluation (gated dry-run → canary → full) Read [Step 8](references/launcher-workflow.md#step-8--run-evaluation-gated-dry-run--canary--full). -### Step 9 — Verify completed run +## Step 9 — Verify completed run Apply [run-validation.md](references/run-validation.md)'s **Bounded Evaluation-Failure Policy** and **Timeout and Output-Limit Accounting** before reporting scores. For diff --git a/plugins/modelopt/skills/monitor/SKILL.md b/plugins/modelopt/skills/monitor/SKILL.md index 14ce4e14c32..01d619e11ad 100644 --- a/plugins/modelopt/skills/monitor/SKILL.md +++ b/plugins/modelopt/skills/monitor/SKILL.md @@ -1,6 +1,6 @@ --- name: monitor -description: Monitor submitted jobs (PTQ, evaluation, deployment) on SLURM clusters. Use when the user asks "check job status", "is my job done", "monitor my evaluation", "what's the status of the PTQ", "check on job ", or after any skill submits a long-running job. Also triggers on "nel status", "squeue", or any request to check progress of a previously submitted job. +description: Monitor submitted jobs (PTQ, evaluation, deployment) on SLURM clusters. Use when the user asks "check job status", "is my job done", "monitor my evaluation", "what's the status of the PTQ", "check on job" followed by a SLURM job ID, or after any skill submits a long-running job. Also triggers on "nel status", "squeue", or any request to check progress of a previously submitted job. --- # Job Monitor From 52999c1c4b9e0d78fae43b8dcc432c32a0659846 Mon Sep 17 00:00:00 2001 From: joe <222375587+joe-nv@users.noreply.github.com> Date: Thu, 22 Oct 2026 15:10:47 +0800 Subject: [PATCH 12/43] fix(deploy): unblock TensorRT 11 inference (#2908) ### What does this PR do? Type of change: Bug fix Remove the obsolete `IExecutionContext.all_shape_inputs_specified` assertion that aborts TensorRT inference initialization with `AttributeError` on TensorRT 11. TensorRT 10.16 documents that this legacy check does not work with name-based APIs such as `set_input_shape()` and recommends `infer_shapes()`. NVIDIA's [TensorRT 11 migration reference](https://docs.nvidia.com/deeplearning/tensorrt/11.4.0/api/migration/tensorrt-10x-to-11x-c-api-reference.html) lists the corresponding `allInputShapesSpecified()` API as "Removed (always true)". The change deletes only this three-line assertion. Existing name-based `infer_shapes()` validation, including its `io_shapes == {}` condition, is preserved. Input/output allocation, profile selection, inference execution, public APIs and dependency requirements are unchanged. ### Usage No new API or flag. Existing TensorRT deployment and evaluation calls continue to use the same entry points and optional `io_shapes` argument. ### Testing - A development-only CPU protocol probe executes the production initializer with simulated TensorRT binding interfaces and explicit I/O shapes. Before the change, the absent-property case reproduces the reported `AttributeError`, while the legacy-property control passes. After the change, both cases pass and retain the same binding/profile calls. This validates the Python interface regression, not native TensorRT/CUDA execution; the scaffolding is not included in this PR. - All applicable pre-commit hooks passed for the changed file, including Ruff, formatting, mypy, license checks and Bandit. `git diff --check` and Python syntax parsing passed. - Native TensorRT 10.16/11.x GPU inference and baseline/FP8/INT8 accuracy evaluation were not run on the development Mac, which has no TensorRT installation or CUDA device. Run the existing deployment runtime tests and affected ONNX evaluation cases on the final commit before merge. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: API-contract review supports removing the obsolete check for TensorRT 10.16/11.x; no public signature or existing shape-inference condition changes. Native regression validation is pending. - Copied third-party code or new PIP dependency: none. - Necessary tests: existing TensorRT compile/profile tests are retained. Development-only protocol regression checks passed; no mock/scaffolding tests are added to CI. - Changelog: not added; this addresses the TensorRT 11 compatibility issue reported during the current release-candidate testing. - Claude approval: not obtained. ### Additional Information Refs: bug/6915422 The commit includes DCO sign-off and an SSH signature verified locally. GitHub reports `unknown_key` because the signing public key is not registered on the account. Target: `main`, with `cherry-pick-0.48.0`. Keep as Draft pending native validation of the final commit. No merge or auto-merge requested. ## Summary by CodeRabbit * **Bug Fixes** * Updated TensorRT shape handling to allow execution when shape inputs are configured through tensor addresses and input shapes. * Retained validation for cases where no input/output shapes are provided. Signed-off-by: joe-nv <222375587+joe-nv@users.noreply.github.com> Signed-off-by: Rohan Joshi --- modelopt/torch/_deploy/_runtime/trt_client.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/modelopt/torch/_deploy/_runtime/trt_client.py b/modelopt/torch/_deploy/_runtime/trt_client.py index a9c300eca5f..1135dc1a205 100644 --- a/modelopt/torch/_deploy/_runtime/trt_client.py +++ b/modelopt/torch/_deploy/_runtime/trt_client.py @@ -180,9 +180,6 @@ def initialize_input_output_tensors(self, engine): tensor_name, output_tensors[idx - len(input_tensors)].data_ptr(), ) - assert self.execution_context.all_shape_inputs_specified, ( - "Not all shape inputs are specified." - ) # Set selected profile idx self.execution_context.set_optimization_profile_async(0, self.stream.cuda_stream) From 70ca1093c42cf9a92c9e3fedfe8080214de9e25b Mon Sep 17 00:00:00 2001 From: yingguo-trt <244492186+yingguo-trt@users.noreply.github.com> Date: Fri, 23 Oct 2026 09:43:31 +0800 Subject: [PATCH 13/43] [NVBug 6959002][fix] Supply image inputs for Wan I2V export fusion (#2929) ### What does this PR do? Type of change: Bug fix. Exporting a quantized Wan2.1 I2V transformer can report success even though the checkpoint fails to load because its QKV scales differ. The export probe supplies text context without the image features required by the model. Its forward pass fails on the empty image-attention path, and the exporter catches the exception and skips QKV scale fusion. Supply image features for image-conditioned `WanTransformer3DModel` configurations, retaining the 512-token text context expected by attention. Match the image projection's compute dtype using its bias, including when weights are packed, and account for first/last-frame positional embeddings. Reuse the existing strict fusion mode so a failed image-conditioned Wan probe raises instead of silently skipping fusion. This fixes the existing export adapter. It does not establish end-to-end support for every Wan variant. Configurations without `image_dim` retain their current probe inputs and error handling. Strict mode raises on a probe exception; it does not validate fusion-group coverage or generation quality. ### Usage No new API or flags. Use the existing `export_hf_checkpoint` call after quantization/calibration with the model's image-conditioned pipeline. ### Testing [Jenkins #57](http://dlswqa-nas:18880/job/modelopt_release_testing_v1/57/) completed successfully: **3 passed, 0 failed, 0 skipped** on one B100 GPU, using `nvcr.io/nvidia/tensorrt-llm/release:1.4.0rc0`, Diffusers 0.40.0, and Transformers 5.5.4. | Model | NVFP4 export and TensorRT-LLM deployment | QKV groups fused | | --- | --- | --- | | Wan2.1 I2V 14B 480P | Passed | 102 | | Wan2.1 I2V 14B 720P | Passed | 102 | | Wan2.2 I2V A14B | Passed for both transformer components | 68 per component | - Reused `TestDiffusersQuantization_hf_deployment::test_quantization_export_hf` from QA revision `6df5ab7d1c4c14da569cd914b764ff6d1061b336`. Each component was calibrated with 8 seeds, 20 denoising steps and 81 frames, then compressed and exported to a fresh checkpoint. TensorRT-LLM loaded the exported checkpoint, generated video, and passed shape/frame-count and saved-video decoding checks. - Audited archived static NVFP4 configuration and scale-tensor presence. The existing 720P/A14B cases assert scale completeness; the 480P archive was additionally checked because its case lacks that assertion: all 408 non-ignored packed Linear modules have `weight_scale`, `weight_scale_2`, and `input_scale`. A14B evidence confirms static quantization routing and execution of both noise-stage transformers. - The build checked out and installed ModelOpt `9d67af8087cebb1e3f945cc7db82c4c7e1d8bee3`. The current PR preserves the tested export implementation: follow-up edits restore copyright/comments and add CPU regression coverage, while the rebase adds an unrelated ONNX graph-surgery change. The fix commit is equivalent under `git range-diff`; the new test-only commit was validated locally, not rerun on the GPU. - All applicable pre-commit hooks passed, including Ruff, mypy, license checks, and Bandit. - Added one parameterized CPU regression, `test_wan_dummy_inputs_support_qkv_fusion[i2v/flf2v]`, using the existing tiny Wan helper. It checks image/text input shapes, runs real calibration and the strict default dummy forward, and verifies input/weight amax unification for self-attention QKV, text cross-attention KV, and image cross-attention KV. Per-tensor INT8 exercises real fusion on CPU; FP8 is excluded because its export format skips fusion. Distinct initial amax values prevent a skipped fusion from passing. - Local validation on Windows (Python 3.12.10, CPU PyTorch 2.12.0): **26 passed, 3 deselected** in `tests/unit/torch/export/test_export_diffusers.py`. The three deselected Transformers KV-metadata cases were previously confirmed to fail before export on both this branch and the unpatched baseline because the quantization plugin cannot import without Triton. Both new parameters pass in under one second each, and both fail when the dummy-input builder is reverted to the pre-fix implementation. - The CPU regression checks dummy inputs, forward execution, and amax fusion; it does not exercise compressed NVFP4 scale handling. Jenkins #57 provides the GPU NVFP4 export/deployment evidence above. These results do not assess visual quality, and FLF2V was not GPU-tested. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: API and checkpoint formats are unchanged. A failed image-conditioned Wan probe intentionally raises instead of warning and skipping fusion. Standard Diffusers Wan image embedders have the projection bias used for the compute dtype (checked in v0.33.0, v0.35.1, and v0.40.0). The existing input builder already dispatches Wan subclasses through `isinstance`; subclasses that replace the image embedder or remove that bias can now fail explicitly. This PR does not establish support for arbitrary custom subclasses, and the public exporter has no custom-probe override. - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`?: N/A; neither was added. - Did you write any new necessary tests?: Yes; one parameterized CPU regression covers I2V and FLF2V dummy inputs and actual QKV amax fusion. Reused existing GPU export/deployment cases for the three Wan I2V configurations. - Did you update Changelog?: N/A; reported against 0.48.0rc1, with no released-version regression established. - Did you get Claude approval on this PR?: No approval; the automated review comments from cjluo-nv are addressed with the CPU regression and compatibility clarification. ### Additional Information Tracks [NVBug 6959002](https://nvbugspro.nvidia.com/bug/6959002). ## Summary by CodeRabbit * **New Features** * Added support for exporting Wan image-to-video models with image conditioning, including configurations with custom positional embedding lengths. Image-conditioned exports now include the required image and text inputs. * **Bug Fixes** * Improved export reliability for image-conditioned Wan models by requiring attention fusion checks to succeed. Other model components retain their existing export behavior. --------- Signed-off-by: yingguo-trt <244492186+yingguo-trt@users.noreply.github.com> Signed-off-by: Rohan Joshi --- modelopt/torch/export/diffusers_utils.py | 21 ++++++++- modelopt/torch/export/unified_export_hf.py | 8 +++- .../torch/export/test_export_diffusers.py | 44 +++++++++++++++++++ 3 files changed, 70 insertions(+), 3 deletions(-) diff --git a/modelopt/torch/export/diffusers_utils.py b/modelopt/torch/export/diffusers_utils.py index 2892e16a02e..de4d6966a61 100644 --- a/modelopt/torch/export/diffusers_utils.py +++ b/modelopt/torch/export/diffusers_utils.py @@ -295,6 +295,7 @@ def _wan_inputs() -> dict[str, torch.Tensor]: # Requires: hidden_states, encoder_hidden_states, timestep in_channels = getattr(cfg, "in_channels", 16) text_dim = getattr(cfg, "text_dim", 4096) + image_dim = getattr(cfg, "image_dim", None) max_seq_len = getattr(cfg, "rope_max_seq_len", 512) patch_dtype = getattr(getattr(model, "patch_embedding", None), "weight", None) @@ -308,13 +309,14 @@ def _wan_inputs() -> dict[str, torch.Tensor]: # Wan expects num_frames = 4 * n + 1; keep n small for dummy forward num_frames = 5 - text_seq_len = min(max_seq_len, 512) + # I2V attention splits image context from the final 512 text tokens. + text_seq_len = 512 if image_dim is not None else min(max_seq_len, 512) # Keep spatial dims small and divisible by patch size (default 2x2) height = 8 width = 8 - return { + dummy_inputs = { "hidden_states": torch.randn( batch_size, in_channels, num_frames, height, width, device=device, dtype=patch_dtype ), @@ -324,6 +326,21 @@ def _wan_inputs() -> dict[str, torch.Tensor]: "timestep": torch.randint(0, 1000, (batch_size,), device=device), "return_dict": False, } + if image_dim is not None: + image_embedder = model.condition_embedder.image_embedder + # The projection bias retains its compute dtype when weights are packed. + image_dtype = image_embedder.ff.net[0].proj.bias.dtype + image_batch_size = batch_size + image_seq_len = 257 + pos_embed_seq_len = getattr(cfg, "pos_embed_seq_len", None) + if pos_embed_seq_len is not None: + # FLF2V concatenates first/last-frame features inside the image embedder. + image_batch_size *= 2 + image_seq_len = pos_embed_seq_len // 2 + dummy_inputs["encoder_hidden_states_image"] = torch.randn( + image_batch_size, image_seq_len, image_dim, device=device, dtype=image_dtype + ) + return dummy_inputs def _qwen_inputs() -> dict[str, Any]: # QwenImageTransformer2DModel does NOT take the standard diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 0ee1f1cd2eb..a6c4b38068b 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -1334,7 +1334,13 @@ def _export_diffusers_checkpoint( # dummy forward fails for it, fail loudly rather than silently skipping # fusion (which would export un-unified amax values). is_qwen_component = "qwen" in type(component).__name__.lower() - _fuse_qkv_linears_diffusion(component, strict=is_qwen_component) + # Image-conditioned Wan also requires a successful dummy forward for fusion. + is_wan_i2v_component = ( + HAS_DIFFUSERS + and isinstance(component, getattr(diffusers, "WanTransformer3DModel", ())) + and getattr(getattr(component, "config", None), "image_dim", None) is not None + ) + _fuse_qkv_linears_diffusion(component, strict=is_qwen_component or is_wan_i2v_component) # Process quantized modules (convert weights, register scales) _process_quantized_modules(component, component_dtype, is_modelopt_qlora=False) diff --git a/tests/unit/torch/export/test_export_diffusers.py b/tests/unit/torch/export/test_export_diffusers.py index eedeae677e1..ec7787687e5 100644 --- a/tests/unit/torch/export/test_export_diffusers.py +++ b/tests/unit/torch/export/test_export_diffusers.py @@ -286,6 +286,50 @@ def test_flux2_dummy_inputs_shape(): assert "guidance" in inputs +@pytest.mark.parametrize("pos_embed_seq_len", [None, 514], ids=["i2v", "flf2v"]) +@torch.no_grad() +def test_wan_dummy_inputs_support_qkv_fusion(pos_embed_seq_len): + torch.manual_seed(0) + model = get_tiny_wan22_transformer( + image_dim=32, added_kv_proj_dim=24, pos_embed_seq_len=pos_embed_seq_len + ).eval() + inputs = generate_diffusion_dummy_inputs(model, torch.device("cpu"), torch.float32) + assert inputs is not None + image_batch_size = 1 if pos_embed_seq_len is None else 2 + assert inputs["encoder_hidden_states_image"].shape == (image_batch_size, 257, 32) + assert inputs["encoder_hidden_states"].shape == (1, 512, 32) + + # Per-tensor INT8 exercises both input and weight amax fusion on CPU. + quant_cfg = copy.deepcopy(mtq.INT8_DEFAULT_CFG) + quant_cfg["quant_cfg"].append( + {"quantizer_name": "*weight_quantizer", "cfg": {"num_bits": 8, "axis": None}} + ) + mtq.quantize(model, quant_cfg, lambda m: m(**inputs)) + + groups = [] + for block in model.blocks: + groups.extend( + [ + (block.attn1.to_q, block.attn1.to_k, block.attn1.to_v), + (block.attn2.to_k, block.attn2.to_v), + (block.attn2.add_k_proj, block.attn2.add_v_proj), + ] + ) + # Distinct values ensure a skipped fusion cannot pass on already-equal scales. + for group in groups: + for index, module in enumerate(group, start=1): + module.input_quantizer.amax = torch.tensor(float(index)) + module.weight_quantizer.amax = torch.tensor(float(index)) + + unified_export_hf._fuse_qkv_linears_diffusion(model, strict=True) + + for group in groups: + expected_amax = torch.tensor(float(len(group))) + for module in group: + torch.testing.assert_close(module.input_quantizer.amax, expected_amax) + torch.testing.assert_close(module.weight_quantizer.amax, expected_amax) + + @pytest.mark.parametrize( ("model_factory", "export_attention_quantizers", "expect_attention_quantizers"), [ From f1b3d41c684cccd7e745a52835da0956366e7f74 Mon Sep 17 00:00:00 2001 From: Shengliang Xu <106840466+shengliangxu@users.noreply.github.com> Date: Thu, 22 Oct 2026 22:23:07 -0700 Subject: [PATCH 14/43] [OMNIML-6260] Pin the hosted Linux runners to ubuntu-24.04 (#2950) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Type of change: Bug fix (CI) Pins every `runs-on: ubuntu-latest` of the workflows to `ubuntu-24.04`. GitHub moved `ubuntu-latest` to Ubuntu 26.04 at about 00:45 UTC on 2026-10-23, and the hosted jobs that install torch into a nox/uv environment fail on it, so every pull request is red: | job (all attempts, from the run logs) | on ubuntu-24.04, until 10-23 00:xx UTC | on ubuntu-26.04, since 00:46 UTC | |---|---|---| | `linux` (`unit_tests.yml`) | 210 of 210 passed | 0 of 23 passed | | `build-docs` (`pages.yml`) | all passed | 0 of 25 passed | The failure, in "Run unit tests" and in the docs build alike, while uv installs torch 2.14.1: ``` warning: Failed to hardlink files; falling back to full copy. error: Failed to install: torch-2.14.1-cp312-cp312-manylinux_2_28_x86_64.whl (torch==2.14.1) cause: failed to copy file from /home/runner/.cache/uv/archive-v0/.../torch/lib/libtorch_cuda.so to /tmp/.nox//lib/python3.12/site-packages/torch/lib/.tmp...: Disk quota exceeded (os error 122) ``` The 24.04 logs have no hardlink warning, so there the uv cache and `/tmp/.nox` share a filesystem; on 26.04 they do not, uv copies the torch and CUDA libraries, and the copy exceeds the quota. The GPU, example and regression workflows wait for the unit check (`wait-checks`), so no pull request gets GPU CI while it fails. After this merges, a pull request needs a new `pull_request` run to pick it up (a push, or a rebase onto main): re-running a failed job keeps the old `runs-on` and lands on 26.04 again. Moving to 26.04 is a separate change that needs the install fixed first (untested ideas: a `UV_CACHE_DIR` on the filesystem of `/tmp/.nox`, or `NOX_ENVDIR` under the workspace). `tools/launcher/docs/testing.md` keeps `ubuntu-latest` because it is an example for users. This PR's own run is the test: its `linux` and `build-docs` jobs must run on `ubuntu-24.04` and pass. The changed YAML passes the pre-commit hooks (YAML format). Not run: a throwaway check on 26.04. 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?: N/A (CI) - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: no new dependency - Did you write any new necessary tests?: N/A (CI) - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A (CI) - Did you get Claude approval on this PR?: N/A (not run yet) Jira: [OMNIML-6260]. * **Chores** * Automated checks, tests, documentation builds, and release jobs now run on Ubuntu 24.04. This standardizes the operating system used across the project’s automation; no workflow conditions or other behavior changed. These updates affect project maintenance and do not change the app’s end-user features. Signed-off-by: Shengliang Xu Signed-off-by: Rohan Joshi --- .github/actions/ubuntu-setup/action.yml | 2 +- .github/workflows/_pr_gate.yml | 4 ++-- .github/workflows/bump_uv_lock.yml | 2 +- .github/workflows/check_model_hub_orgs.yml | 2 +- .github/workflows/claude.yml | 2 +- .github/workflows/claude_review.yml | 2 +- .github/workflows/close_inactive_issues_pr.yml | 2 +- .github/workflows/code_quality.yml | 2 +- .github/workflows/example_tests.yml | 4 ++-- .github/workflows/gpu_tests.yml | 2 +- .github/workflows/pages.yml | 8 ++++---- .github/workflows/regression_tests.yml | 2 +- .github/workflows/release.yml | 2 +- .github/workflows/unit_tests.yml | 16 ++++++++-------- 14 files changed, 26 insertions(+), 26 deletions(-) diff --git a/.github/actions/ubuntu-setup/action.yml b/.github/actions/ubuntu-setup/action.yml index 58fdc12b494..6a9c9fd81f5 100644 --- a/.github/actions/ubuntu-setup/action.yml +++ b/.github/actions/ubuntu-setup/action.yml @@ -1,5 +1,5 @@ name: Ubuntu runner setup -description: Standard setup for ubuntu-latest runners (setup Python, cleanup disk space) +description: Standard setup for ubuntu-24.04 runners (setup Python, cleanup disk space) inputs: python-version: diff --git a/.github/workflows/_pr_gate.yml b/.github/workflows/_pr_gate.yml index e3605202f7c..3720a3cd824 100644 --- a/.github/workflows/_pr_gate.yml +++ b/.github/workflows/_pr_gate.yml @@ -23,7 +23,7 @@ on: jobs: check-file-changes: - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 outputs: any_changed: ${{ steps.changed-tests.outputs.any_modified || steps.non-pr.outputs.any_changed }} steps: @@ -53,7 +53,7 @@ jobs: if: >- startsWith(github.ref, 'refs/heads/pull-request/') && needs.check-file-changes.outputs.any_changed == 'true' - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 permissions: checks: read steps: diff --git a/.github/workflows/bump_uv_lock.yml b/.github/workflows/bump_uv_lock.yml index ae624ddb211..01b8e50cc84 100644 --- a/.github/workflows/bump_uv_lock.yml +++ b/.github/workflows/bump_uv_lock.yml @@ -12,7 +12,7 @@ permissions: jobs: bump-uv-lock: - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 timeout-minutes: 15 steps: - uses: actions/checkout@v6 diff --git a/.github/workflows/check_model_hub_orgs.yml b/.github/workflows/check_model_hub_orgs.yml index fe4df1c5385..39d190b57a8 100644 --- a/.github/workflows/check_model_hub_orgs.yml +++ b/.github/workflows/check_model_hub_orgs.yml @@ -17,7 +17,7 @@ concurrency: jobs: check-model-hub-orgs: - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 timeout-minutes: 5 steps: - uses: actions/checkout@v6 diff --git a/.github/workflows/claude.yml b/.github/workflows/claude.yml index 46d0996515f..850383e66e3 100644 --- a/.github/workflows/claude.yml +++ b/.github/workflows/claude.yml @@ -47,7 +47,7 @@ jobs: contains(fromJson('["OWNER", "MEMBER", "COLLABORATOR"]'), github.event.issue.author_association) ) ) - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 timeout-minutes: 10 permissions: contents: write diff --git a/.github/workflows/claude_review.yml b/.github/workflows/claude_review.yml index 46f27063de8..b7abc076638 100644 --- a/.github/workflows/claude_review.yml +++ b/.github/workflows/claude_review.yml @@ -22,7 +22,7 @@ jobs: github.event.issue.pull_request && contains(github.event.comment.body, '/claude review') && contains(fromJson('["OWNER", "MEMBER", "COLLABORATOR"]'), github.event.comment.author_association) - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 timeout-minutes: 30 permissions: contents: read diff --git a/.github/workflows/close_inactive_issues_pr.yml b/.github/workflows/close_inactive_issues_pr.yml index 5c63c8316b9..53656b5d44a 100644 --- a/.github/workflows/close_inactive_issues_pr.yml +++ b/.github/workflows/close_inactive_issues_pr.yml @@ -7,7 +7,7 @@ on: jobs: stale: - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 permissions: issues: write pull-requests: write diff --git a/.github/workflows/code_quality.yml b/.github/workflows/code_quality.yml index 4743fef960b..80996cdc772 100644 --- a/.github/workflows/code_quality.yml +++ b/.github/workflows/code_quality.yml @@ -16,7 +16,7 @@ concurrency: jobs: code-quality: - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 timeout-minutes: 30 steps: - uses: actions/checkout@v6 diff --git a/.github/workflows/example_tests.yml b/.github/workflows/example_tests.yml index 8cc432a22ba..33cd33d3308 100644 --- a/.github/workflows/example_tests.yml +++ b/.github/workflows/example_tests.yml @@ -21,7 +21,7 @@ jobs: # One changed-files pass decides which lanes run. Lane granularity, not per-example: a job's # `if` cannot read `matrix`, so gating an individual example would need a job per example. pr-gate: - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 permissions: checks: read contents: read @@ -237,7 +237,7 @@ jobs: # Run even if example tests are skipped if: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && always() }} needs: [pr-gate, torch, vllm, trtllm, megatron, diffusers, onnx] - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 steps: - name: Required example tests did not succeed if: | diff --git a/.github/workflows/gpu_tests.yml b/.github/workflows/gpu_tests.yml index a0389da7b62..91577ca15b4 100644 --- a/.github/workflows/gpu_tests.yml +++ b/.github/workflows/gpu_tests.yml @@ -104,7 +104,7 @@ jobs: # Run even if gpu-tests is skipped if: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && always() }} needs: [pr-gate, gpu-tests] - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 steps: - name: Required GPU tests did not succeed if: ${{ needs.pr-gate.result != 'success' || (needs.pr-gate.outputs.any_changed == 'true' && needs.gpu-tests.result != 'success') }} diff --git a/.github/workflows/pages.yml b/.github/workflows/pages.yml index 2132cdd302e..7dd889f88ad 100644 --- a/.github/workflows/pages.yml +++ b/.github/workflows/pages.yml @@ -24,7 +24,7 @@ permissions: jobs: build-docs: if: github.event.action != 'closed' - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 timeout-minutes: 10 steps: - uses: actions/checkout@v6 @@ -40,7 +40,7 @@ jobs: changes: if: github.event_name == 'pull_request' && github.event.action != 'closed' - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 outputs: docs: ${{ steps.filter.outputs.docs }} steps: @@ -62,7 +62,7 @@ jobs: github.event.pull_request.head.repo.full_name == github.repository && (github.event.action == 'closed' || needs.changes.outputs.docs == 'true') needs: [build-docs, changes] - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 timeout-minutes: 10 # Per-PR concurrency without cancel-in-progress so 'closed' cleanup always runs concurrency: @@ -83,7 +83,7 @@ jobs: deploy-gh-pages: if: github.event_name == 'push' || github.event_name == 'schedule' || github.event_name == 'workflow_dispatch' needs: build-docs - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 steps: - uses: actions/checkout@v6 - name: Download docs artifact diff --git a/.github/workflows/regression_tests.yml b/.github/workflows/regression_tests.yml index 001c4e6af4a..a28ac16611c 100644 --- a/.github/workflows/regression_tests.yml +++ b/.github/workflows/regression_tests.yml @@ -74,7 +74,7 @@ jobs: # Run even if regression-tests is skipped if: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && always() }} needs: [pr-gate, regression-tests] - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 steps: - name: Required regression tests did not succeed if: | diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 002a08850aa..dbdb208f2ba 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -9,7 +9,7 @@ permissions: jobs: build-upload: - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 steps: - uses: actions/checkout@v6 - uses: ./.github/actions/ubuntu-setup diff --git a/.github/workflows/unit_tests.yml b/.github/workflows/unit_tests.yml index 6fbca185d9f..090c7065e12 100644 --- a/.github/workflows/unit_tests.yml +++ b/.github/workflows/unit_tests.yml @@ -28,7 +28,7 @@ concurrency: jobs: check-file-changes: - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 outputs: any_changed: ${{ steps.changed.outputs.any_modified || steps.non-pr.outputs.any_changed }} steps: @@ -58,7 +58,7 @@ jobs: .agents/skills/** plugins/modelopt/skills/** linux: - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 timeout-minutes: 20 steps: - uses: actions/checkout@v6 @@ -120,7 +120,7 @@ jobs: multi-version: if: needs.check-file-changes.outputs.any_changed == 'true' needs: [linux, check-file-changes] - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 timeout-minutes: 20 strategy: fail-fast: false @@ -150,7 +150,7 @@ jobs: partial-install: if: needs.check-file-changes.outputs.any_changed == 'true' needs: [linux, check-file-changes] - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 timeout-minutes: 15 strategy: fail-fast: false @@ -164,7 +164,7 @@ jobs: launcher: if: needs.check-file-changes.outputs.any_changed == 'true' needs: [linux, check-file-changes] - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 timeout-minutes: 15 steps: - uses: actions/checkout@v6 @@ -181,7 +181,7 @@ jobs: mcp: if: needs.check-file-changes.outputs.any_changed == 'true' needs: [linux, check-file-changes] - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 timeout-minutes: 15 steps: - uses: actions/checkout@v6 @@ -204,7 +204,7 @@ jobs: skills: if: needs.check-file-changes.outputs.any_changed == 'true' needs: [linux, check-file-changes] - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 timeout-minutes: 10 steps: - uses: actions/checkout@v6 @@ -222,7 +222,7 @@ jobs: # Run even if some jobs are skipped if: ${{ github.event_name == 'pull_request' && always() }} needs: [check-file-changes, linux, windows, multi-version, partial-install, launcher, mcp, skills] - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 steps: - name: Required unit tests did not succeed # windows is deliberately absent from this condition. It still runs and is still visible, From 674f97ee75f323ca4533357d7a675c4ff96e54a9 Mon Sep 17 00:00:00 2001 From: yingguo-trt <244492186+yingguo-trt@users.noreply.github.com> Date: Fri, 23 Oct 2026 17:00:08 +0800 Subject: [PATCH 15/43] [NVBug 6981670][docs] Fix Cache Diffusion prepare calls (#2972) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Type of change: documentation. Remove the obsolete `num_inference_steps` argument from the Cache Diffusion README's legacy TensorRT and Customize examples. Both calls now match the existing `prepare(pipe, config_list)` API and avoid `TypeError: prepare() takes 2 positional arguments but 3 were given`. ### Usage ```python cachify.prepare(pipe, SDXL_DEFAULT_CONFIG) ``` ### Testing - `pre-commit run --files examples/diffusers/README.md` and `git diff --check`. - B100 reproduction with official ModelOpt `0.48.0rc1`: the three-argument call fails; the corrected call succeeds on the same tiny `UNet2DConditionModel` on `cuda:0`, wrapping three cache modules. The cache implementation is unchanged between that release and this PR's base. No model forward or full SDXL/TensorRT inference was run. - Internal evidence: [Jenkins #65](http://dlswqa-nas:18880/job/modelopt_release_testing_v1/65/console). ### Before your PR is "Ready for review" - Is this change backward compatible?: ✅ Documentation only. - If you copied code from other sources or added a dependency, did you follow `CONTRIBUTING.md`?: N/A - Did you write any new necessary tests?: N/A Documentation correction; the actual API calls were verified above. - Did you update Changelog?: N/A - Did you get Claude approval on this PR?: N/A Small documentation correction. ### Additional Information Related: [NVBug 6981670](https://nvbugspro.nvidia.com/bug/6981670). Signed-off-by: yingguo-trt <244492186+yingguo-trt@users.noreply.github.com> Signed-off-by: Rohan Joshi --- examples/diffusers/README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/diffusers/README.md b/examples/diffusers/README.md index 30402163092..971b8045e3e 100644 --- a/examples/diffusers/README.md +++ b/examples/diffusers/README.md @@ -489,7 +489,7 @@ compile( engine_path=Path("./engine"), ) -cachify.prepare(pipe, num_inference_steps, SDXL_DEFAULT_CONFIG) +cachify.prepare(pipe, SDXL_DEFAULT_CONFIG) ``` Afterward, use it as a standard cache diffusion pipeline to generate the image. @@ -508,7 +508,7 @@ SDXL_DEFAULT_CONFIG = [ } ] -cachify.prepare(pipe, num_inference_steps, SDXL_DEFAULT_CONFIG) +cachify.prepare(pipe, SDXL_DEFAULT_CONFIG) ``` Two parameters are essential: `wildcard_or_filter_func` and `select_cache_step_func`. From 77e22d9d6abdf79c543248400837882d68afe034 Mon Sep 17 00:00:00 2001 From: Chenjie Luo <108829653+cjluo-nv@users.noreply.github.com> Date: Fri, 23 Oct 2026 09:11:26 -0700 Subject: [PATCH 16/43] Launch the fused NVFP4 GPTQ kernel on its tensors' GPU (#2877) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Type of change: Bug fix Standalone on `main`, split out of the GPTQ-for-MoE stack (#2840 merged ← #2864 ← #2945). It touches only the fused kernel, so it can merge in any order with that stack. GPTQ with `fused: true` runs the NVFP4 column update in a Triton kernel (`gptq_fused_block_scalar`, in modelopt since 0.45). Triton launches on the **current** CUDA device, and the wrapper launched without selecting the device of the tensors it was given. A model that `device_map="auto"` spreads over several GPUs keeps later decoder layers off the current GPU. Once layerwise GPTQ reached such a layer, the kernel ran on GPU 0 with pointers into GPU 1: - **NVLink systems (GB300):** `CUDA error: an illegal memory access was encountered`, reported later at an unrelated op. - **PCIe systems:** `ValueError: Pointer argument cannot be accessed from Triton`. The fix launches the kernel under `torch.cuda.device(w_block.device)`, as the other quantization Triton wrappers (`fp4_kernel.py`, `nvfp4_fp8_sweep.py`, `bmm2_qdq.py`) already do. Found while running layerwise GPTQ on Qwen3.8-Flash-Next with the experts' `down_proj` in static NVFP4 (`fused: true`). The model sits on 4 x GB300 with layers 0-22 on GPU 0. Three runs on three different nodes all died right after layer 22. No API change. - **New GPU test** `test_fused_nvfp4_gptq_runs_on_the_weights_gpu` (2 GPUs): the fused update with every tensor on `cuda:1` and `cuda:0` current matches the same update run on `cuda:0` bit for bit. It fails without the fix with the Triton pointer error on the PCIe GPUs I ran it on. - **Existing tests:** on `main`, `tests/gpu/torch/quantization/test_gptq.py` passes on 2x RTX PRO 6000 (11 tests, including the existing fused-vs-unfused NVFP4 ones). - **End to end:** the Qwen3.8-Flash-Next layerwise GPTQ that died after layer 22 completes with this change on oci-jhb; details in the comment below. - 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)?: ✅ (bug present since 0.45) - Did you get Claude approval on this PR?: ❌ 🤖 Generated with [Claude Code](https://claude.com/claude-code) * **Bug Fixes** * Fixed GPTQ fused quantization failing with CUDA or Triton errors when calibration runs across multiple GPUs. * Added a check that quantization results remain consistent when processing data on a different GPU than the current device. Signed-off-by: Chenjie Luo Co-authored-by: Claude Opus 5.5 (1M context) Signed-off-by: Rohan Joshi --- CHANGELOG.rst | 1 + .../quantization/gemm/gptq_fused_kernel.py | 29 ++++++++++--------- tests/gpu/torch/quantization/test_gptq.py | 14 +++++++++ 3 files changed, 31 insertions(+), 13 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index dd2d156d601..6969e9a2507 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -190,6 +190,7 @@ Changelog - 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. - The default PTQ recipes now leave multi-token-prediction (MTP) layers unquantized in Megatron vision-language models too (``language_model.mtp.*``), as they already did at the model root. - Fix ``examples/llm_sparsity/weight_sparsity/finetune.py`` saving a sharded state dict under FSDP2, the default FSDP version in Transformers 5, which made ``export_hf_ckpt.py`` fail to restore it. Re-run finetuning to produce a loadable ``finetuned_modelopt_state.pth``. +- Fix GPTQ with ``fused: true`` failing with a CUDA illegal memory access, or a Triton pointer error, on a model that ``device_map="auto"`` spreads over several GPUs once calibration reaches a layer outside the current GPU. 0.47.0 (2026-09-23) ^^^^^^^^^^^^^^^^^^^ diff --git a/modelopt/torch/kernels/quantization/gemm/gptq_fused_kernel.py b/modelopt/torch/kernels/quantization/gemm/gptq_fused_kernel.py index fb13386f851..3b410e97bae 100644 --- a/modelopt/torch/kernels/quantization/gemm/gptq_fused_kernel.py +++ b/modelopt/torch/kernels/quantization/gemm/gptq_fused_kernel.py @@ -119,18 +119,21 @@ def gptq_fused_block_scalar( qw_block = torch.empty_like(w_block) err_block = torch.empty_like(w_block) - _gptq_scalar_kernel[(num_rows,)]( - w_block.contiguous(), - qw_block, - err_block, - block_amax.contiguous(), - global_scale, - h_inv_cho_blk.contiguous(), - num_rows, - block_amax.shape[1], - quant_block_size, - block_start, - BLOCK_SIZE=block_size, - ) + # Triton launches on the current device; a model sharded with device_map="auto" keeps later + # layers on other GPUs, where launching from the current one reads another GPU's memory. + with torch.cuda.device(w_block.device): + _gptq_scalar_kernel[(num_rows,)]( + w_block.contiguous(), + qw_block, + err_block, + block_amax.contiguous(), + global_scale, + h_inv_cho_blk.contiguous(), + num_rows, + block_amax.shape[1], + quant_block_size, + block_start, + BLOCK_SIZE=block_size, + ) return qw_block, err_block diff --git a/tests/gpu/torch/quantization/test_gptq.py b/tests/gpu/torch/quantization/test_gptq.py index 496db08578a..b4b965b3154 100644 --- a/tests/gpu/torch/quantization/test_gptq.py +++ b/tests/gpu/torch/quantization/test_gptq.py @@ -208,6 +208,20 @@ def _run_fused_gptq_nvfp4( return w +@requires_triton +def test_fused_nvfp4_gptq_runs_on_the_weights_gpu(need_2_gpus): + """A layer device_map placed on another GPU is updated there, not from the current GPU.""" + torch.manual_seed(42) + weight, _, block_amax, global_scale, h_inv = _make_nvfp4_test_data(16, 64, 256) + expected = _run_fused_gptq_nvfp4(weight, block_amax, global_scale, h_inv, 128, 16) + + weight, block_amax, h_inv = (t.to("cuda:1") for t in (weight, block_amax, h_inv)) + got = _run_fused_gptq_nvfp4(weight, block_amax, global_scale, h_inv, 128, 16) + + assert torch.cuda.current_device() == 0 + torch.testing.assert_close(got.cpu(), expected.cpu(), rtol=0, atol=0) + + _NVFP4_QUANT_BLOCK_SIZES = [16, 128] _NVFP4_GPTQ_BLOCK_SIZES = [16, 128] From 3e8d41d08b9833bb510452f1a179785f9ee6286f Mon Sep 17 00:00:00 2001 From: Chenjie Luo <108829653+cjluo-nv@users.noreply.github.com> Date: Fri, 23 Oct 2026 09:23:39 -0700 Subject: [PATCH 17/43] [NVBug 6934537][fix] Re-tie embeddings in init_quantized_weights (#2951) 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 --low_memory_mode` crashes for any model with `tie_word_embeddings=True` (e.g. Qwen2.5-0.5B/1.5B, Llama-3.2-1B/3B) with `NotImplementedError: Cannot copy out of meta tensor` in `dispatch_model`. Root cause: `init_quantized_weights` builds the model with `from_config` inside accelerate's `init_empty_weights()`, which patches `register_parameter` to re-wrap every assigned parameter in a new meta `Parameter`. The `lm_head.weight = embed_tokens.weight` tie that `from_config` makes is therefore lost (`find_tied_parameters` returns `[]`). Tied checkpoints don't store `lm_head.weight`, and accelerate has no tie to restore after loading, so `lm_head.weight` stays on meta and `model.to(device)` fails. Fix: call `model.tie_weights()` after leaving `init_empty_weights()`, before `mtq.quantize`/`mtq.compress` — the pattern accelerate documents for big-model loading. The tied `Parameter` survives the `RealQuantLinear` conversion (only the `_parameters` dict is wrapped), and `load_checkpoint_in_model` re-ties it after loading `embed_tokens.weight`. Untied models are unaffected (`tie_weights()` is a no-op). ### Usage No new API or flags. ### Testing Parametrized the existing `test_real_quantize_linear` over `tie_word_embeddings` and asserted `lm_head.weight is embed_tokens.weight` iff tied. Ran on CPU (the test path needs no CUDA): | | transformers 4.57.6 / accelerate 1.5.1 | transformers 5.5.4 / accelerate 1.15.0 | |---|---|---| | without fix | tied case fails (`Cannot copy out of meta tensor`) | same | | with fix | 2 passed | 2 passed | Also loaded a tiny tied/untied Llama with `NVFP4_DEFAULT_CFG` and `FP8_DEFAULT_CFG`: with the fix, no parameter is left on meta and `lm_head.weight` equals the checkpoint's `embed_tokens.weight`; untied loads are unchanged. `tests/unit/torch/quantization/plugins/test_accelerate.py` passes. Not run: the end-to-end `hf_ptq.py --low_memory_mode` on a GPU (no GPU available to me); GPU CI covers the parametrized test. Note: if a config enables `lm_head` quantization on a tied model, `lm_head` now loads compressed from the embedding values and is untied from `embed_tokens` (previously it crashed). Default configs disable `lm_head` quantization, so `hf_ptq` doesn't hit this. ### 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?: ✅ (parametrized existing test) - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A - Did you get Claude approval on this PR?: ❌ ### Additional Information NVBug 6934537. Should land in 0.48.0 (`cherry-pick-0.48.0`). 🤖 Generated with [Claude Code](https://claude.com/claude-code) ## Summary by CodeRabbit * **Bug Fixes** * Quantized models now preserve the configured relationship between token embeddings and the output head during loading and checkpoint dispatch. When word embeddings are tied, both use the same weights after model loading and quantization; when they are untied, they remain independent. This keeps model behavior consistent with the selected embedding configuration across both setups. Signed-off-by: Chenjie Luo Co-authored-by: Claude Opus 5.5 (1M context) Signed-off-by: Rohan Joshi --- modelopt/torch/quantization/plugins/accelerate.py | 3 +++ tests/gpu/torch/quantization/test_real_quantize_cuda.py | 6 ++++-- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/quantization/plugins/accelerate.py b/modelopt/torch/quantization/plugins/accelerate.py index 35f6777fc8c..eb206db2c5c 100644 --- a/modelopt/torch/quantization/plugins/accelerate.py +++ b/modelopt/torch/quantization/plugins/accelerate.py @@ -233,6 +233,9 @@ def patched_from_pretrained(cls, /, pretrained_model_name_or_path, *args, **kwar "dtype", kwargs.get("torch_dtype", getattr(config, "torch_dtype", torch.float16)) ) model = cls.from_config(config, dtype=torch_dtype) + # init_empty_weights re-wraps every registered parameter, so the tie made inside + # from_config is lost; without it the absent tied checkpoint key (e.g. lm_head) stays meta. + model.tie_weights() mtq.quantize(model, quant_cfg) mtq.compress(model, config=mtq.CompressConfig(quant_gemm=quant_gemm)) diff --git a/tests/gpu/torch/quantization/test_real_quantize_cuda.py b/tests/gpu/torch/quantization/test_real_quantize_cuda.py index 74ff7e93048..5af7cee0c8f 100644 --- a/tests/gpu/torch/quantization/test_real_quantize_cuda.py +++ b/tests/gpu/torch/quantization/test_real_quantize_cuda.py @@ -197,10 +197,12 @@ def forward_loop(model, run_backward=False): @pytest.mark.parametrize("quant_config", [mtq.FP8_DEFAULT_CFG]) -def test_real_quantize_linear(quant_config, tmp_path): - tiny_llama_dir = create_tiny_llama_dir(tmp_path) +@pytest.mark.parametrize("tie_word_embeddings", [False, True]) +def test_real_quantize_linear(quant_config, tie_word_embeddings, tmp_path): + tiny_llama_dir = create_tiny_llama_dir(tmp_path, tie_word_embeddings=tie_word_embeddings) with init_quantized_weights(quant_config): model = AutoModelForCausalLM.from_pretrained(tiny_llama_dir) + assert (model.lm_head.weight is model.model.embed_tokens.weight) == tie_word_embeddings for _, module in model.named_modules(): if ( From a1ca420dfb75b22da4c4e301fa235551d8163079 Mon Sep 17 00:00:00 2001 From: anmaurya001 Date: Sat, 24 Oct 2026 22:44:01 +0530 Subject: [PATCH 18/43] Fix Windows ONNX external data corruption by upgrading ONNX (#2816) ### What does this PR do? Type of change: Bug fix Upgrade the ONNX extra from `~=1.21.0` to the locally validated `~=1.23.2` release line on Windows only. Non-Windows platforms retain `~=1.21.0`, as requested in review. On Windows, exporting multiple tensors into shared external data with ONNX 1.21.0 can truncate a trailing `0x1A` byte when reopening the file. Subsequent tensors overwrite earlier data, corrupting quantized weights and scales. This can produce incorrect output and runaway generation that eventually overflows the conversation context. ONNX already corrected Windows binary-mode file opening upstream in https://github.com/onnx/onnx/pull/7790. This change consumes the released fix and adds a regression test through ModelOpt's `save_onnx`, checking shared-file bytes and tensor round trips. Existing corrupted exports must be regenerated from clean weights. The lockfile is intentionally left to the weekly workflow, as required by CONTRIBUTING.md. ### Usage Reinstall the ONNX extra with the updated dependency and regenerate affected exports from clean weights. No API changes. ### Testing - Windows regression test on this branch: ONNX 1.23.2 passes all three cases. The unmodified ONNX 1.21.0 wheel fails the trailing-`0x1A` case (byte 1023 is overwritten), while the two control cases pass. - Changed-file pre-commit checks cover formatting, lint, typing, license headers, security lint, TOML and reStructuredText validation. - Focused command: `python -m pytest -o addopts='' --noconftest tests/unit/onnx/test_external_data.py -q`. The normal ONNX autouse C++ compilation fixture failed in this local environment; this run excludes repository conftest fixtures and optional coverage hooks. - Prior local end-to-end validation used ModelOpt 0.45.0rc0 with the unmodified ONNX 1.23.2 wheel: regenerated AWQ-lite INT4 model using 32 calibration samples; checked 341 external tensors with no overlaps or invalid scales; CUDA two-turn chat answered Paris and 4 and stopped normally. - Full local five-shot MMLU, all 57 subjects and 14,042 examples: regenerated INT4 8,099/14,042 (57.676969%); clean FP16 8,374/14,042 (59.635380%). These pipeline results are from the earlier ModelOpt configuration, not a full pipeline run on this PR's current main base. - Exact original CI reproduction and the complete current-main test suite have not been run. Release CI and reviewer validation remain required. ### Before your PR is "Ready for review" - Is this change backward compatible?: Yes; no API changes. The ONNX dependency release line changes on Windows only. - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md?: N/A; existing dependency upgraded, no copied upstream implementation. - Did you write any new necessary tests?: Yes. - Did you update Changelog?: Yes. - Did you get Claude approval on this PR?: Pending; keep as draft until required review is complete. ### Additional Information Addresses NVBugs [6309022](https://nvbugspro.nvidia.com/bug/6309022) and [6561663](https://nvbugspro.nvidia.com/bug/6561663) through the shared Windows external-data corruption fix. Local evidence supports this cause; original CI acceptance remains pending. ## Summary by CodeRabbit * **Bug Fixes** * Fixed an issue that could corrupt external tensor data when exporting ONNX models on Windows. Regenerate affected models from clean weights to ensure their data is intact. * **Updates** * Windows installations now use ONNX 1.23.2 for improved export compatibility. Non-Windows installations continue to use ONNX 1.21.0. --------- Signed-off-by: anmaurya001 Signed-off-by: Rohan Joshi --- CHANGELOG.rst | 1 + .../windows/_installation_windows_arm64.rst | 2 +- pyproject.toml | 3 +- tests/unit/onnx/test_external_data.py | 42 +++++++++++++++++++ 4 files changed, 46 insertions(+), 2 deletions(-) create mode 100644 tests/unit/onnx/test_external_data.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 6969e9a2507..079184bff09 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -133,6 +133,7 @@ Changelog **Bug Fixes** +- Fix corrupted ONNX external tensor data when exporting models on Windows by upgrading ONNX. Regenerate affected models from clean weights. - Fix ``import modelopt.torch`` failing with ``ModuleNotFoundError`` for ``requests`` or ``huggingface_hub`` when ``nvidia-modelopt`` is installed without the ``[hf]`` extra. - 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. diff --git a/docs/source/getting_started/windows/_installation_windows_arm64.rst b/docs/source/getting_started/windows/_installation_windows_arm64.rst index 2bafcf0be4e..e50f039a303 100644 --- a/docs/source/getting_started/windows/_installation_windows_arm64.rst +++ b/docs/source/getting_started/windows/_installation_windows_arm64.rst @@ -101,7 +101,7 @@ explicit dependency installation avoids selecting the Windows x64 ``onnxruntime- & $PythonExe -m pip install ` cppimport cupy-cuda13x lief ml_dtypes ` - "onnx~=1.21.0" "onnx-graphsurgeon>=0.6.1" ` + "onnx~=1.23.2" "onnx-graphsurgeon>=0.6.1" ` "onnxconverter-common~=1.16.0" onnxscript ` "onnxslim>=0.1.76" "polygraphy>=0.49.22" ` "onnxruntime>=1.24.2" onnxruntime-ep-nv-tensorrt-rtx-cu13 diff --git a/pyproject.toml b/pyproject.toml index 9a8d7cc4a76..648afc7b5c4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -59,7 +59,8 @@ onnx = [ "lief", "ml_dtypes", "onnx-graphsurgeon>=0.6.1", - "onnx~=1.21.0", + "onnx~=1.21.0; platform_system != 'Windows'", + "onnx~=1.23.2; platform_system == 'Windows'", "onnxconverter-common~=1.16.0", # Python 3.10 support for ModelOpt ONNX is deprecated and will be removed in a future release. # Native Windows ARM64 does not support ModelOpt ONNX on Python 3.10. diff --git a/tests/unit/onnx/test_external_data.py b/tests/unit/onnx/test_external_data.py new file mode 100644 index 00000000000..ddb77e997f8 --- /dev/null +++ b/tests/unit/onnx/test_external_data.py @@ -0,0 +1,42 @@ +# 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 numpy as np +import onnx +import pytest + +from modelopt.onnx.utils import save_onnx + + +@pytest.mark.parametrize("last_byte", [0x00, 0x0A, 0x1A]) +def test_save_external_data_preserves_tensor_boundaries(tmp_path, last_byte): + # Windows text-mode file opening can truncate a trailing Ctrl-Z before append. + first = np.full(1024, 0x41, dtype=np.uint8) + first[-1] = last_byte + second = np.full(1024, 0x42, dtype=np.uint8) + tensors = [ + onnx.numpy_helper.from_array(value, name) + for name, value in [("first", first), ("second", second)] + ] + graph = onnx.helper.make_graph([], "external_data", [], [], initializer=tensors) + model = onnx.helper.make_model(graph) + path = tmp_path / "model.onnx" + + save_onnx(model, str(path), save_as_external_data=True) + + assert (tmp_path / "model.onnx_data").read_bytes() == first.tobytes() + second.tobytes() + restored = onnx.load(path) + for actual, expected in zip(restored.graph.initializer, [first, second], strict=True): + np.testing.assert_array_equal(onnx.numpy_helper.to_array(actual), expected) From 74565b2ee5fd16533c70b54d8e5fd0929744951e Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Sun, 25 Oct 2026 14:25:47 +0530 Subject: [PATCH 19/43] Quantize FP8 weights in float32 in the functional export (#3005) 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 functional quantized-weight export (`export_quantized_weight_tensors`, used by Megatron-Bridge's `export_hf_weights_modelopt` / NeMo-RL refit) rounds some per-tensor FP8 weights to a different FP8 value than fake quantization uses. Root cause: the functional export state keeps the per-tensor FP8 `weight_scale` 0-d. In `to_quantized_weight`, `bf16_weight / scale` then stays in bf16, because a 0-d tensor does not promote a dimensioned tensor of the same category. The Megatron and HF exporters pass a 1-d scale, which promotes the division to float32. On a 256×256 bf16 linear quantized with `FP8_DEFAULT_CFG`, 1,584 of 65,536 exported weights differed from `weight_quantizer(weight)`, and a toy Qwen3 exported through Megatron-Bridge had FP8 weight tensors that differed from `export_mcore_gpt_to_hf`. Fix: divide in float32 explicitly. The Megatron and HF exporters already divided in float32 through promotion, so their outputs and peak memory are unchanged. Found while building the Megatron-Bridge quantized HF export (OMNIML-5976), whose parity check against `export_mcore_gpt_to_hf` this unblocks for FP8. ### Usage No API change. ### Testing - New unit test `test_functional_fp8_export_matches_the_fake_quantized_bf16_weight`: the dequantized export of a bf16 FP8 linear equals `weight_quantizer(weight)` bit for bit. Fails without the fix. - `tests/unit/torch/export/`: 513 passed; `test_streaming_export_two_cpu_ranks` fails the same way on main in my container. - `tests/gpu/torch/export/`: 188 passed. The 4 AWQ / SVDQuant cases of `test_unified_hf_export_and_check_safetensors.py` hit the 120 s timeout on my shared GPUs, and do the same on main. - Toy Qwen3 FP8 through Megatron-Bridge `export_hf_weights_modelopt` vs `export_mcore_gpt_to_hf`: all weight tensors bitwise equal with the fix (several FP8 weights differed before; per-tensor scales differ only in shape, `[]` vs `[1]`). ### 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 functional export API is new in 0.48, still in RC) - Did you get Claude approval on this PR?: ❌ ### Additional Information Part of OMNIML-5976 (Megatron-Bridge quantized export, epic OMNIML-5486). 🤖 Generated with [Claude Code](https://claude.com/claude-code) ## Summary by CodeRabbit * **Bug Fixes** * Fixed an issue with exporting FP8-quantized weights when their scale is a single value. Exported weights can now be reconstructed correctly using the exported scale, including for bfloat16 linear layers. This improves consistency between quantized weights and their exported representation. --------- Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 Signed-off-by: Rohan Joshi --- modelopt/torch/export/quant_utils.py | 4 ++++ .../torch/export/test_quantized_weight_state.py | 17 +++++++++++++++++ 2 files changed, 21 insertions(+) diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py index 91372bbd951..f19142cb650 100755 --- a/modelopt/torch/export/quant_utils.py +++ b/modelopt/torch/export/quant_utils.py @@ -897,6 +897,10 @@ def to_quantized_weight( # Clear GPU cache to avoid pontential GPU OOM issues for large models. clear_cuda_cache() return (weight / weights_scaling_factor.unsqueeze(-1)).to(torch.float8_e4m3fn) + # A 0-d scale would not promote a bf16 weight, so the division would round in bf16 instead of + # float32 like fake quantization. A 1-d scale promotes in-kernel, without a float32 copy. + if weights_scaling_factor.dim() == 0: + weights_scaling_factor = weights_scaling_factor.reshape(1) return (weight / weights_scaling_factor).to(torch.float8_e4m3fn) if quantization in [QUANTIZATION_INT8_SQ, QUANTIZATION_INT8_WO]: diff --git a/tests/unit/torch/export/test_quantized_weight_state.py b/tests/unit/torch/export/test_quantized_weight_state.py index 682bc0eb229..ef0734aff96 100644 --- a/tests/unit/torch/export/test_quantized_weight_state.py +++ b/tests/unit/torch/export/test_quantized_weight_state.py @@ -156,6 +156,23 @@ def test_functional_fp8_export_is_repeatable(): torch.testing.assert_close(base[name], repeated[name]) +def test_functional_fp8_export_matches_the_fake_quantized_bf16_weight(): + torch.manual_seed(0) + module = nn.Linear(64, 64, bias=False).to(torch.bfloat16) + mtq.quantize( + module, + copy.deepcopy(mtq.FP8_DEFAULT_CFG), + lambda m: m(torch.randn(2, 64, dtype=torch.bfloat16)), + ) + state = capture_quantized_weight_export_state(module) + + exported = export_quantized_weight_tensors(module.weight, state, torch.bfloat16) + + dequantized = exported["weight"].float() * exported["weight_scale"] + expected = module.weight_quantizer(module.weight) + torch.testing.assert_close(dequantized.to(torch.bfloat16), expected, rtol=0, atol=0) + + _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}}, From c955816c387c639891a6c323c62138f8b9dce486 Mon Sep 17 00:00:00 2001 From: yingguo-trt <244492186+yingguo-trt@users.noreply.github.com> Date: Mon, 26 Oct 2026 12:06:36 +0800 Subject: [PATCH 20/43] [NVBug 6982525] Honor save_modelopt_state in Diffusers export (#3012) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Type of change: Bug fix With HuggingFace checkpointing enabled, `export_hf_checkpoint(..., save_modelopt_state=False)` still writes `modelopt_state.pth` for Diffusers components. The Diffusers dispatch and component saves omit the flag, so the checkpointing hook uses its own default of `True`. Forward the public flag through the Diffusers exporter to both quantized and non-quantized model saves. This also makes the public default of `False` effective for Diffusers exports. Explicit `True` continues to use the existing checkpointing hook. The existing API now honors the option: ```python import modelopt.torch.opt as mto from modelopt.torch.export import export_hf_checkpoint mto.enable_huggingface_checkpointing() export_hf_checkpoint(model, export_dir="export", save_modelopt_state=False) ``` Reused `test_diffusers_export_respects_save_modelopt_state_false` from the existing QA reproduction at `0871d4fefdcb89929a37c7999d1e308e8f5cc3e3`. Invoked the unchanged test function with its CPU fixture directly, outside the Jenkins-wide pytest setup. Local, offline source-tree validation on Windows: Python 3.12.10, PyTorch 2.12.0+cpu, Diffusers 0.40.0, Transformers 5.5.4. Both runs calibrated a real tiny UNet with `FP8_DEFAULT_CFG` (45 quantizers, 30 enabled) and used fresh output directories. | Source | Direct save with False | Export with False | | --- | --- | --- | | Parent `378888145a` | No sidecar; control passes | Sidecar written; expected regression assertion fails | | Fix `5ff581d74c` | No sidecar; control passes | No sidecar; regression passes | Both exports produced config and safetensors files. This verifies the original file-saving bug on CPU; it is not a release-wheel or deployment-inference retest. - `git diff --check`: passed. - File-scoped pre-commit: Ruff, formatting, license and Bandit checks passed. Mypy reports the same existing Windows `os.sched_getaffinity` error in unchanged `prepare_megatron_data_blend.py:373` on both the parent and fixed revisions. The same Mypy version with the hook's arguments and `--platform linux` passed. Contributor and security guidance reviewed; the existing commit is signed and signed off. - Is this change backward compatible?: ✅ Preserves the API and makes the documented flag/default effective. - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A — neither. - Did you write any new necessary tests?: N/A — reused the existing QA regression with before/after evidence; no test files added or changed. - Did you update Changelog?: N/A — the issue is confirmed in 0.48.0rc1; impact on an earlier stable release has not been established. - Did you get Claude approval on this PR?: ❌ Not yet reviewed by Claude. Fixes [NVBug 6982525](https://nvbugspro.nvidia.com/bug/6982525). The change suppresses new sidecar writes; it retains the existing handling of files already present in the export directory. - **New Features** - Diffusers checkpoint exports can now optionally save ModelOpt state. This option is off by default and is available for both quantized and non-quantized components. 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 | 13 +++++++++++-- .../torch/export/test_export_diffusers.py | 19 +++++++++++++++++++ 3 files changed, 31 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 079184bff09..ed22b64a108 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -192,6 +192,7 @@ Changelog - The default PTQ recipes now leave multi-token-prediction (MTP) layers unquantized in Megatron vision-language models too (``language_model.mtp.*``), as they already did at the model root. - Fix ``examples/llm_sparsity/weight_sparsity/finetune.py`` saving a sharded state dict under FSDP2, the default FSDP version in Transformers 5, which made ``export_hf_ckpt.py`` fail to restore it. Re-run finetuning to produce a loadable ``finetuned_modelopt_state.pth``. - Fix GPTQ with ``fused: true`` failing with a CUDA illegal memory access, or a Triton pointer error, on a model that ``device_map="auto"`` spreads over several GPUs once calibration reaches a layer outside the current GPU. +- Diffusers ``export_hf_checkpoint`` now honors ``save_modelopt_state=False`` (the default) when HuggingFace checkpointing is enabled. Pass ``save_modelopt_state=True`` to retain the ModelOpt state file. 0.47.0 (2026-09-23) ^^^^^^^^^^^^^^^^^^^ diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index a6c4b38068b..9ac7e26eb2c 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -1250,6 +1250,7 @@ def _export_diffusers_checkpoint( components: list[str] | None, max_shard_size: int | str = "10GB", export_attention_quantizers: bool | None = None, + save_modelopt_state: bool = False, **kwargs, ) -> None: """Internal: Export diffusion(-like) model/pipeline checkpoint. @@ -1271,6 +1272,7 @@ def _export_diffusers_checkpoint( runtimes that consume legacy attention ``_amax`` keys, such as static FP8 FLUX. None preserves these keys for FLUX.1 and omits them for other diffusion components. Transformers components always retain them for KV-cache scales. + save_modelopt_state: Whether to save ModelOpt state with HF checkpointing enabled. **kwargs: Runtime-specific post-processing options forwarded to :func:`_postprocess_safetensors`. See its docstring for details. """ @@ -1378,7 +1380,9 @@ def _export_diffusers_checkpoint( if hasattr(component, "save_pretrained"): with hide_quantizers_from_state_dict(component, keep_attention_quantizers): component.save_pretrained( - component_export_dir, max_shard_size=max_shard_size + component_export_dir, + max_shard_size=max_shard_size, + save_modelopt_state=save_modelopt_state, ) else: with hide_quantizers_from_state_dict(component, keep_attention_quantizers): @@ -1407,7 +1411,11 @@ def _export_diffusers_checkpoint( _remove_promoted_quantizer_tensors(component) # Non-quantized component: just save as-is elif hasattr(component, "save_pretrained"): - component.save_pretrained(component_export_dir, max_shard_size=max_shard_size) + component.save_pretrained( + component_export_dir, + max_shard_size=max_shard_size, + save_modelopt_state=save_modelopt_state, + ) else: _save_component_state_dict_safetensors(component, component_export_dir) @@ -1994,6 +2002,7 @@ def export_hf_checkpoint( export_dir, components, max_shard_size, + save_modelopt_state=save_modelopt_state, **kwargs, ) return diff --git a/tests/unit/torch/export/test_export_diffusers.py b/tests/unit/torch/export/test_export_diffusers.py index ec7787687e5..001b4245b72 100644 --- a/tests/unit/torch/export/test_export_diffusers.py +++ b/tests/unit/torch/export/test_export_diffusers.py @@ -107,6 +107,25 @@ def test_export_diffusers_models_non_quantized(tmp_path, model_factory): assert "quantization_config" not in config_data +@pytest.mark.parametrize("quantized", [False, True]) +@pytest.mark.parametrize("save_modelopt_state", [False, True]) +@torch.no_grad() +def test_export_diffusers_save_modelopt_state( + tmp_path, enable_hf_checkpointing, quantized, save_modelopt_state +): + model = get_tiny_unet().eval() + inputs = generate_diffusion_dummy_inputs(model, torch.device("cpu"), torch.float32) + mtq.quantize(model, mtq.FP8_DEFAULT_CFG, lambda m: m(**inputs)) + if not quantized: + # Keep ModelOpt state while exercising the non-quantized save branch. + mtq.disable_quantizer(model, "*") + export_dir = tmp_path / "export_unet" + + export_hf_checkpoint(model, export_dir=export_dir, save_modelopt_state=save_modelopt_state) + + assert (export_dir / "modelopt_state.pth").is_file() == save_modelopt_state + + @pytest.fixture def tiny_pipeline(request, tmp_path, tiny_tokenizer): """Build and reload a tiny pipeline without accessing the Hugging Face Hub.""" From 8319ad82d20c6e6b08fea0c9e2e2ed70e913e8a6 Mon Sep 17 00:00:00 2001 From: yingguo-trt <244492186+yingguo-trt@users.noreply.github.com> Date: Mon, 26 Oct 2026 14:12:45 +0800 Subject: [PATCH 21/43] [NVBug 6980441] Fix Diffusers subfolder ModelOpt state restoration (#2971) 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 Loading a Diffusers component with `from_pretrained(parent, subfolder="unet")` finds its weights in the component directory, but ModelOpt still looks for `modelopt_state.pth` in the parent directory. The load can succeed while dropping the saved quantizers. Pass the component subfolder through ModelOpt's state lookup for local directories and existing Hub cache entries. Calls without `subfolder` and checkpoint saving retain their existing behavior; no download behavior is added. ### Usage ```python import modelopt.torch.opt as mto from diffusers import UNet2DConditionModel mto.enable_huggingface_checkpointing() # A quantized UNet was saved with save_pretrained("pipeline/unet"). restored = UNet2DConditionModel.from_pretrained( "pipeline", subfolder="unet", local_files_only=True ) ``` ### Testing Local CPU environment: Python 3.12.10, PyTorch 2.12.0+cpu, Diffusers 0.40.0, Transformers 5.5.4. - Existing `tests/unit/torch/opt/plugins/test_diffusers_save_load.py`: **1 passed**. - Existing `tests/unit/torch/opt/plugins/test_transformers_save_load.py::test_causal_lm_save_restore`: **2 passed** (`LlamaForCausalLM` and `AutoModelForCausalLM`). - Standalone tiny UNet INT8 reproducer, before and after the fix: direct-directory loading restores 30 enabled quantizers in both versions. Subfolder loading restores 0 before the fix and 30 afterward; its maximum absolute output difference from the original quantized model changes from approximately 0.020612 to 0. - `pre-commit run --files modelopt/torch/opt/plugins/diffusers.py modelopt/torch/opt/plugins/huggingface.py`: **passed**. No test files were added or modified. Fix validation uses a parent directory without component weights. Hub loading, GPU inference, and other Diffusers architectures were not validated. ### Before your PR is "*Ready for review*" Contributor and security guidance reviewed; the commit is signed and signed off. - Is this change backward compatible?: ✅ Existing calls without `subfolder` retain their behavior. - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A — neither. - Did you write any new necessary tests?: ❌ No new checked-in tests; existing tests and the standalone reproducer were run as described above. - Did you update Changelog?: N/A — no new feature or public API change; affected prior released versions have not been established. - Did you get Claude approval on this PR?: ❌ Not yet reviewed by Claude. ### Additional Information This change does not alter Diffusers' weight selection when both the parent and component directory contain same-named weights; that separate directory-layout edge case is discussed in huggingface/diffusers#7140. ## Summary by CodeRabbit * **Bug Fixes** * Model loading now correctly handles models stored in subfolders, including when retrieving associated optimization state from local directories or the Hugging Face Hub. Existing loading behavior remains unchanged when no subfolder is specified. --------- Signed-off-by: yingguo-trt <244492186+yingguo-trt@users.noreply.github.com> Signed-off-by: Rohan Joshi --- modelopt/torch/opt/plugins/diffusers.py | 4 +++- modelopt/torch/opt/plugins/huggingface.py | 14 +++++++---- .../opt/plugins/test_diffusers_save_load.py | 10 +++++--- .../torch/opt/plugins/test_hf_patching.py | 24 +++++++++++++++++++ 4 files changed, 43 insertions(+), 9 deletions(-) diff --git a/modelopt/torch/opt/plugins/diffusers.py b/modelopt/torch/opt/plugins/diffusers.py index 3942507522b..58ff9735f3d 100644 --- a/modelopt/torch/opt/plugins/diffusers.py +++ b/modelopt/torch/opt/plugins/diffusers.py @@ -26,7 +26,9 @@ def _new_from_pretrained(cls, /, pretrained_model_name_or_path, *args, **kwargs): """Patch for `cls.from_pretrained` method to restore ModelOpt state.""" - with _patch_model_init_for_modelopt(cls, pretrained_model_name_or_path): + with _patch_model_init_for_modelopt( + cls, pretrained_model_name_or_path, subfolder=kwargs.get("subfolder") or "" + ): model = types.MethodType(cls._modelopt_cache["from_pretrained"].__func__, cls)( pretrained_model_name_or_path, *args, **kwargs ) diff --git a/modelopt/torch/opt/plugins/huggingface.py b/modelopt/torch/opt/plugins/huggingface.py index afbe4a3fde9..67441298ed4 100644 --- a/modelopt/torch/opt/plugins/huggingface.py +++ b/modelopt/torch/opt/plugins/huggingface.py @@ -56,25 +56,29 @@ def register_for_patching(name: str, cls: type, patch_methods: list[tuple[str, A methods_list.append(patch_methods) -def _get_modelopt_state_path(model_name_or_path: str) -> str: +def _get_modelopt_state_path(model_name_or_path: str, subfolder: str = "") -> str: """Get the path to the ModelOpt state file or empty string if not found. Also handles HF model card as input path. However for hf hub models, we dont have modelopt_state at the moment. """ if os.path.isdir(model_name_or_path): - return os.path.join(model_name_or_path, _MODELOPT_STATE_SAVE_NAME) + return os.path.join(model_name_or_path, subfolder, _MODELOPT_STATE_SAVE_NAME) # Optional [hf] dependency: this module is imported even when huggingface_hub is not installed. from huggingface_hub import try_to_load_from_cache from huggingface_hub.errors import HFValidationError + filename = ( + f"{subfolder}/{_MODELOPT_STATE_SAVE_NAME}" if subfolder else _MODELOPT_STATE_SAVE_NAME + ) try: - return try_to_load_from_cache(model_name_or_path, _MODELOPT_STATE_SAVE_NAME) or "" + cache_path = try_to_load_from_cache(model_name_or_path, filename) + return cache_path if isinstance(cache_path, str) else "" except HFValidationError: return "" @contextmanager -def _patch_model_init_for_modelopt(cls, model_path): +def _patch_model_init_for_modelopt(cls, model_path, subfolder=""): """Patch for `cls.init` method to restore ModelOpt state after `init`.""" # Note: Keeping original config in local as the package will be shared among threads added_original_init = False @@ -89,7 +93,7 @@ def _patch_model_init_for_modelopt(cls, model_path): @functools.wraps(_original__init__) def new_init_fn(self, *args, **kwargs): - modelopt_state_path = _get_modelopt_state_path(model_path) + modelopt_state_path = _get_modelopt_state_path(model_path, subfolder=subfolder) _original__init__(self, *args, **kwargs) if os.path.isfile(modelopt_state_path): restore_from_modelopt_state(self, modelopt_state_path=modelopt_state_path) diff --git a/tests/unit/torch/opt/plugins/test_diffusers_save_load.py b/tests/unit/torch/opt/plugins/test_diffusers_save_load.py index cee5dbefbf8..c73e2a81707 100644 --- a/tests/unit/torch/opt/plugins/test_diffusers_save_load.py +++ b/tests/unit/torch/opt/plugins/test_diffusers_save_load.py @@ -23,14 +23,18 @@ @pytest.mark.parametrize("model_cls", [UNet2DConditionModel]) -def test_unet_save_restore(tmp_path, model_cls): +@pytest.mark.parametrize("use_subfolder", [False, True], ids=["direct", "subfolder"]) +def test_unet_save_restore(tmp_path, model_cls, use_subfolder): tiny_unet_dir = create_tiny_unet_dir(tmp_path) model_ref = model_cls.from_pretrained(tiny_unet_dir) # TODO: Add calibrate, compress mode to the test model_ref = apply_mode_with_sampling( model_ref, ["sparse_magnitude", "export_sparse", "quantize"] ) - model_ref.save_pretrained(tiny_unet_dir / "modelopt_model") + component_dir = tmp_path / "checkpoint" / "unet" + model_ref.save_pretrained(component_dir) - model_test = model_cls.from_pretrained(tiny_unet_dir / "modelopt_model") + load_path = component_dir.parent if use_subfolder else component_dir + load_kwargs = {"subfolder": "unet"} if use_subfolder else {} + model_test = model_cls.from_pretrained(load_path, local_files_only=True, **load_kwargs) df_modelopt_state_and_output_tester(model_ref, model_test) diff --git a/tests/unit/torch/opt/plugins/test_hf_patching.py b/tests/unit/torch/opt/plugins/test_hf_patching.py index 719550f7c78..ba568bb0790 100644 --- a/tests/unit/torch/opt/plugins/test_hf_patching.py +++ b/tests/unit/torch/opt/plugins/test_hf_patching.py @@ -13,6 +13,9 @@ # See the License for the specific language governing permissions and # limitations under the License. +from unittest.mock import Mock + +import huggingface_hub import pytest import torch import torch.nn as nn @@ -26,10 +29,31 @@ import modelopt.torch.distill as mtd import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq +from modelopt.torch.opt.plugins.huggingface import _get_modelopt_state_path from modelopt.torch.opt.plugins.transformers import _restore_qtensor_wrappers from modelopt.torch.quantization.qtensor import QTensorWrapper +@pytest.mark.parametrize( + ("load_kwargs", "cache_result", "filename", "expected"), + [ + ({}, "cached-state.pth", "modelopt_state.pth", "cached-state.pth"), + ({"subfolder": "unet"}, "cached-state.pth", "unet/modelopt_state.pth", "cached-state.pth"), + ({"subfolder": "unet"}, None, "unet/modelopt_state.pth", ""), + ({"subfolder": "unet"}, huggingface_hub._CACHED_NO_EXIST, "unet/modelopt_state.pth", ""), + ], + ids=["default-cached", "subfolder-cached", "subfolder-missing", "subfolder-known-missing"], +) +def test_get_modelopt_state_path_cached(monkeypatch, load_kwargs, cache_result, filename, expected): + cache_lookup = Mock(return_value=cache_result) + monkeypatch.setattr(huggingface_hub, "try_to_load_from_cache", cache_lookup) + + result = _get_modelopt_state_path("test-org/test-model", **load_kwargs) + + cache_lookup.assert_called_once_with("test-org/test-model", filename) + assert result == expected + + @pytest.mark.parametrize( ("model_cls", "teacher_model_type"), [ From 7484d0c334960ecd79c70849a3de461d48ac7721 Mon Sep 17 00:00:00 2001 From: yingguo-trt <244492186+yingguo-trt@users.noreply.github.com> Date: Mon, 26 Oct 2026 14:17:18 +0800 Subject: [PATCH 22/43] [NVBug 6982528] Save LTX full-mode checkpoints as safetensors (#3015) ### What does this PR do? Type of change: Bug fix. LTX-2 distillation with `training_mode: full` names its standalone weights `.safetensors`, but `Accelerator.save()` defaults to PyTorch serialization. Loading that file with safetensors fails with `Error while deserializing header: header too large`. Pass `safe_serialization=True` so the saved contents match the advertised format, and document the fix in the changelog. ### Usage No configuration changes are required. Full-mode training now writes valid safetensors weights. Existing incorrectly serialized files are not converted. ### Testing - Reused the existing external QA case `test_ltx_full_checkpoint_is_safetensors`; no new tests added. It executes the production save method extracted unchanged from the example, with a tiny CPU model and real Accelerate/safetensors APIs. The baseline fails with the header error; the fix passes with exact key and BF16 tensor equality. This is method-level validation, not a full LTX training run. - Environment: Python 3.12.10, PyTorch 2.12.0+cpu, Accelerate 1.15.0, safetensors 0.8.0, and ModelOpt 0.48.0rc1 for library imports; example source tested before and after this change. - Applicable pre-commit hooks and `git diff --check` pass. The mypy hook environment passes with `--platform linux`; the pre-existing Windows `os.sched_getaffinity` error also reproduces on the unmodified source and was excluded. - Full LTX GPU/FSDP training and inference were not run. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: Yes for the documented safetensors output contract; no API or configuration changes. - 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; reused the existing external QA regression described above. No in-repository regression test added. - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: Yes. - Did you get Claude approval on this PR?: Pending. ### Additional Information NVBug 6982528. ## Summary by CodeRabbit * **Bug Fixes** * Fixed the format used when saving full-mode LTX-2 distillation checkpoints. Checkpoint files now use the format indicated by their file extension, improving compatibility with tools that load them. LoRA checkpoint saving is unchanged. Signed-off-by: yingguo-trt <244492186+yingguo-trt@users.noreply.github.com> Signed-off-by: Rohan Joshi --- CHANGELOG.rst | 1 + examples/diffusers/distillation/distillation_trainer.py | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index ed22b64a108..5fefbb94741 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -133,6 +133,7 @@ Changelog **Bug Fixes** +- Fix LTX-2 distillation full-mode checkpoints to use the safetensors format expected by their file extension and loaders. - Fix corrupted ONNX external tensor data when exporting models on Windows by upgrading ONNX. Regenerate affected models from clean weights. - Fix ``import modelopt.torch`` failing with ``ModuleNotFoundError`` for ``requests`` or ``huggingface_hub`` when ``nvidia-modelopt`` is installed without the ``[hf]`` extra. - 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. diff --git a/examples/diffusers/distillation/distillation_trainer.py b/examples/diffusers/distillation/distillation_trainer.py index 38908f5c8f3..6c173ef0f79 100644 --- a/examples/diffusers/distillation/distillation_trainer.py +++ b/examples/diffusers/distillation/distillation_trainer.py @@ -1138,7 +1138,7 @@ def _save_checkpoint(self) -> Path | None: k: v.to(save_dtype) if isinstance(v, Tensor) else v for k, v in full_state_dict.items() } - self._accelerator.save(full_state_dict, saved_weights_path) + self._accelerator.save(full_state_dict, saved_weights_path, safe_serialization=True) rel_path = saved_weights_path.relative_to(self._config.output_dir) logger.info(f"Model weights for step {self._global_step} saved in {rel_path}") From 44950f1bd519cb1473324d82af69898d20e890c7 Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Mon, 26 Oct 2026 23:07:30 +0530 Subject: [PATCH 23/43] Require accelerate>=1.15.0 for llm_qat (#3006) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Type of change: Bug fix Raises the `llm_qat` example's accelerate requirement from `>=1.0.0,<1.14` to `>=1.15.0`. The `<1.14` cap (#2067) existed because accelerate 1.14.0 broke FSDP2 for models with tied input/output embeddings. 1.15.0 fixes that, and it also moves the example off accelerate versions affected by CVE-2026-69112 (path traversal in `load_checkpoint_in_model` / `load_checkpoint_and_dispatch` via unsanitized sharded-index `weight_map` entries). `uv.lock` already resolves accelerate 1.15.0. CI's tiny Qwen3 is untied, so it could not catch the 1.14 regression. The shared NVFP4 PTQ student for the QAT/QAD tests is now built from a tied-embedding tiny Qwen3 (`tied_qwen3_path`), so the existing tests cover it with no extra runs. No usage change. In `nvcr.io/nvidia/pytorch:26.07-py3` (the llm_qat CI image) on 2 GPUs, installed as CI does, running `pytest tests/examples/llm_qat -k fsdp2`: - accelerate 1.14.0: `test_qwen3_qat_nvfp4[fsdp2]` and `test_qwen3_qad_nvfp4[fsdp2]` **fail** (`ValueError: Parameter 'model.embed_tokens.weight' is shared with a parameter already managed by another FSDP group`). - accelerate 1.15.0: both pass. Not run: the `ddp` / `deepspeed` backends (nightly) and the rest of the llm_qat tests. - Is this change backward compatible?: ✅ (example-local requirement bump) - 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?: ✅ (existing QAT/QAD tests now use a tied-embedding model) - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A - Did you get Claude approval on this PR?: N/A 🤖 Generated with [Claude Code](https://claude.com/claude-code) * **Chores** * Updated the minimum required Accelerate version. * **Tests** * Expanded validation of quantization-aware workflows with models that use tied word embeddings. The updated checks also confirm that the model fixture remains unchanged during testing. These changes update a package requirement and internal test coverage; no new end-user functionality is noted. --------- Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Co-authored-by: Claude Sonnet 5.5 Signed-off-by: Rohan Joshi --- examples/llm_qat/requirements.txt | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/examples/llm_qat/requirements.txt b/examples/llm_qat/requirements.txt index 8b2ca580c07..1b9061b5596 100644 --- a/examples/llm_qat/requirements.txt +++ b/examples/llm_qat/requirements.txt @@ -1,5 +1,4 @@ -# Accelerate 1.14 regresses FSDP2 handling of shared parameters. -accelerate>=1.0.0,<1.14 +accelerate>=1.15.0 flash-attn>=2.6.0 liger-kernel>=0.5.0; platform_system != 'Darwin' and platform_system != 'Windows' # peft 0.21 stops writing adapter_model.safetensors where export.py expects it. From 58e6888203b26160b2d13d36a52fe22de5d3ed0c Mon Sep 17 00:00:00 2001 From: Connor Baker Date: Mon, 26 Oct 2026 11:49:59 -0700 Subject: [PATCH 24/43] Fall back to Python INT4 packing when the extension fails to build (#3026) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Type of change: Bug fix `modelopt/onnx/quantization/extensions.py` builds `modelopt_round_and_pack_ext` with cppimport when it is first imported. INT4 packing (`quant_utils.pack_float32_to_4bit_cpp_based`) imports it on first use. If the build fails, the module is meant to fall back to the Python implementation. But cppimport builds with `setuptools.setup()`, which raises `SystemExit` when the build fails, and `except Exception` doesn't catch that. So in a writable install without a working C++ compiler, the import ends the process: ``` INFO - Loading extension modelopt_round_and_pack_ext... error: command '/nonexistent/c++' failed: No such file or directory ``` (exit status 1). This catches `SystemExit` too, adds a unit test that fails a build the way setuptools does, and adds a changelog entry, since 0.47.0 and earlier have the bug. N/A On main at efda42cd: - `tests/unit/onnx/quantization/test_extensions.py` (new) fails on main (`SystemExit: error: command 'c++' failed: No such file or directory`) and passes with this PR. `test_quant_utils.py` passes with this PR. The test stubs `cppimport.imp` to raise the `SystemExit` that setuptools raises, so it doesn't depend on how each platform's compiler can be made to fail; the next item runs the real build. - Importing `modelopt.onnx.quantization.extensions` in a new process from a writable copy of the installed package, with `CC`/`CXX`/`LDSHARED` pointing at a compiler that doesn't exist: on main the process exits with status 1; with this PR it logs the warning and `round_and_pack_ext` is `None`. With a compiler, both build and load the extension, and INT4 packing with it matches the Python fallback. - `pre-commit run` with this repository's configuration (ruff 0.15.20, mypy 2.1.0, bandit, the rst hooks) passes on the changed files. - 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?: ❌ (not run yet) Found while packaging 0.47.0 for Nixpkgs (NixOS/nixpkgs#577198). `extensions.py` is the same in 0.47.0, `release/0.48.0` and main. `tests/unit/onnx/conftest.py`'s prebuild (#2461) suppresses only `Exception`, so it would not have caught this either. A failed build still leaves cppimport's temporary build directory in `$TMPDIR`, since cppimport removes it only when the build succeeds. Now that the process survives, that happens in every process that tries; **AI disclosure:** this change, the testing described above and this description were produced with Claude Code (Claude Opus 5.5). 🤖 Generated with [Claude Code](https://claude.com/claude-code) * **Bug Fixes** * ONNX INT4 quantization now falls back to the Python implementation if the optimized packing extension cannot be loaded, rather than terminating the process. * **Documentation** * Updated the 0.49.0 bug-fix notes to describe the fallback behavior. Signed-off-by: Connor Baker Co-authored-by: Claude Opus 5.5 Signed-off-by: Rohan Joshi --- CHANGELOG.rst | 1 + modelopt/onnx/quantization/extensions.py | 2 +- .../unit/onnx/quantization/test_extensions.py | 34 +++++++++++++++++++ 3 files changed, 36 insertions(+), 1 deletion(-) create mode 100644 tests/unit/onnx/quantization/test_extensions.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 5fefbb94741..974b0c2381b 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -194,6 +194,7 @@ Changelog - Fix ``examples/llm_sparsity/weight_sparsity/finetune.py`` saving a sharded state dict under FSDP2, the default FSDP version in Transformers 5, which made ``export_hf_ckpt.py`` fail to restore it. Re-run finetuning to produce a loadable ``finetuned_modelopt_state.pth``. - Fix GPTQ with ``fused: true`` failing with a CUDA illegal memory access, or a Triton pointer error, on a model that ``device_map="auto"`` spreads over several GPUs once calibration reaches a layer outside the current GPU. - Diffusers ``export_hf_checkpoint`` now honors ``save_modelopt_state=False`` (the default) when HuggingFace checkpointing is enabled. Pass ``save_modelopt_state=True`` to retain the ModelOpt state file. +- Fix ONNX INT4 quantization ending the Python process when the INT4 packing extension fails to build, for example without a C++ compiler. It now falls back to the Python implementation. 0.47.0 (2026-09-23) ^^^^^^^^^^^^^^^^^^^ diff --git a/modelopt/onnx/quantization/extensions.py b/modelopt/onnx/quantization/extensions.py index ba2c0e137dd..4bafe482e67 100644 --- a/modelopt/onnx/quantization/extensions.py +++ b/modelopt/onnx/quantization/extensions.py @@ -29,7 +29,7 @@ sys.path.append(path) round_and_pack_ext = cppimport.imp("modelopt_round_and_pack_ext") sys.path.remove(path) -except Exception as e: +except (Exception, SystemExit) as e: # setuptools exits on a failed build (e.g. no C++ compiler) logger.warning( f"{e}\nUnable to load `modelopt_round_and_pack_ext', falling back to python based optimized version. " ) diff --git a/tests/unit/onnx/quantization/test_extensions.py b/tests/unit/onnx/quantization/test_extensions.py new file mode 100644 index 00000000000..faab210ba11 --- /dev/null +++ b/tests/unit/onnx/quantization/test_extensions.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. + +import importlib + +import cppimport + +import modelopt.onnx.quantization.extensions as ext + + +def test_failed_build_falls_back_to_python(monkeypatch): + def build_fails(*args, **kwargs): + # What setuptools raises when, e.g., there is no C++ compiler. + raise SystemExit("error: command 'c++' failed: No such file or directory") + + monkeypatch.setattr(cppimport, "imp", build_fails) + try: + importlib.reload(ext) + assert ext.round_and_pack_ext is None + finally: + monkeypatch.undo() + importlib.reload(ext) From 2cce2ac3569c18389ae3f4c49570e7c73b9c1f8f Mon Sep 17 00:00:00 2001 From: Jiao Mingyang <105722645+wiederse@users.noreply.github.com> Date: Tue, 27 Oct 2026 03:12:28 +0800 Subject: [PATCH 25/43] Fix tensor-parallel gradient scaling in Top-K KL distillation (#2988) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Type of change: Bug fix With Top-K logits distillation and tensor parallelism enabled, each TP rank evaluates the same global KL loss. The autograd all-gather and full-vocabulary normalizer all-reduce sum gradients from these replicated losses, multiplying student gradients by the TP size even though the reported loss is correct. Scale only the Top-K loss's backward contribution by `1 / TP`. Forward loss values, loss-balancer weights, and dense KL behavior remain unchanged. Extend the existing ghost-token reference test to compare each vocabulary shard's gradients against a full-vocabulary reference, including Top-P truncation, non-default temperature, and reverse KL. No API or command-line changes. Existing `logit_kl_topk` / `--logit_kl_top_k` configurations automatically receive the corrected gradients. - A800/NCCL: **13 focused tests passed at each of TP=1, TP=2, and TP=4 (39 total)**, including real Megatron model forward/backward and the distributed gradient-reference tests. - Negative control: restoring the original loss implementation makes all **3 gradient-reference cases fail at TP=2**, at the gradient assertions. - CPU/Gloo: **14 focused numerical tests passed**, using a temporary TP=1/2/4 worker fixture; repeated successfully on this isolated PR branch. - `pre-commit run --files modelopt/torch/distill/plugins/megatron.py tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py CHANGELOG.rst`: passed. - A800 TP=2 forward/backward microbenchmark: BF16 logits, sequence length 4096, global vocabulary 32768, K=256, 3 warmups and 10 measured iterations. Peak allocated memory was unchanged at **1,086,358,528 bytes per rank**; measured time was **8.44 ms before / 7.81 ms after**. This short run checks overhead and does not establish a speedup. - Full training convergence was not evaluated. GPU test selection: ```bash python -m pytest tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py \ -k 'logits_kl_loss or logits_kl_losses_temperature_scaling or topk_logits_kl_loss' \ -q -o addopts='' ``` Run with 1, 2, or 4 visible GPUs. - Is this change backward compatible?: ✅ API and forward loss values are unchanged; incorrect TP gradient scaling 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. - Did you write any new necessary tests?: ✅ Extended the existing reference test with distributed gradient assertions and parameterized cases. - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅. - Did you get Claude approval on this PR?: ❌ Not run; this PR is a draft. No additional dependencies or communication collectives. * **Bug Fixes** * Corrected Megatron Top-K/Top-P distillation gradient scaling when tensor parallelism is enabled. Reported loss values remain unchanged. * **Tests** * Expanded coverage across tensor-parallel workers, including top-P truncation, reverse KL, and temperature scaling. --------- Signed-off-by: wiederse <1477618732@qq.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Co-authored-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com> Co-authored-by: Claude Sonnet 5.5 Signed-off-by: Rohan Joshi --- CHANGELOG.rst | 1 + modelopt/torch/distill/plugins/megatron.py | 5 ++ .../distill/plugins/test_distill_megatron.py | 65 +++++++++++++++---- 3 files changed, 58 insertions(+), 13 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 974b0c2381b..fa8bf6990a5 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -195,6 +195,7 @@ Changelog - Fix GPTQ with ``fused: true`` failing with a CUDA illegal memory access, or a Triton pointer error, on a model that ``device_map="auto"`` spreads over several GPUs once calibration reaches a layer outside the current GPU. - Diffusers ``export_hf_checkpoint`` now honors ``save_modelopt_state=False`` (the default) when HuggingFace checkpointing is enabled. Pass ``save_modelopt_state=True`` to retain the ModelOpt state file. - Fix ONNX INT4 quantization ending the Python process when the INT4 packing extension fails to build, for example without a C++ compiler. It now falls back to the Python implementation. +- Fix Megatron Top-K/Top-P distillation gradients being multiplied by the tensor-parallel size when tensor parallelism is enabled; reported loss values are unchanged. 0.47.0 (2026-09-23) ^^^^^^^^^^^^^^^^^^^ diff --git a/modelopt/torch/distill/plugins/megatron.py b/modelopt/torch/distill/plugins/megatron.py index 26b3c9cc5c8..63745ed48e0 100644 --- a/modelopt/torch/distill/plugins/megatron.py +++ b/modelopt/torch/distill/plugins/megatron.py @@ -549,6 +549,11 @@ def forward(self, predictions: Tensor, targets: Tensor) -> Tensor: kl = q.exp() * (q - p) loss = torch.sum(mask * kl, dim=-1) + if tp_size > 1: + # Autograd collectives sum gradients from identical losses on every TP rank. + # Average only backward, preserving reported loss and loss-balancer weights. + loss = loss.detach() + (loss - loss.detach()) / tp_size + # No need to reduce since all ranks compute same global Top-K return self.post_forward(loss, tp_reduce=False) 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 2bbcc108d69..18cc2807a8e 100644 --- a/tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py +++ b/tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py @@ -25,6 +25,7 @@ from _test_utils.torch.megatron.models import get_mcore_gpt_model from _test_utils.torch.megatron.utils import run_mcore_inference_with_dummy_input from _test_utils.torch.misc import set_seed +from megatron.core import parallel_state import modelopt.torch.distill as mtd from modelopt.torch.distill.plugins.megatron import ( @@ -465,26 +466,64 @@ def test_topk_logits_kl_loss_numerics_full_vocab_matches_dense(): assert torch.allclose(dense, topk, atol=1e-6) -def test_topk_logits_kl_loss_numerics_ghost_token_reference(): - """Top-K + ghost token matches a hand-written reference on the K+1 bucketed distributions.""" - cfg = SimpleNamespace(tensor_model_parallel_size=1) - student, teacher = _make_loss_inputs() +def _test_topk_logits_kl_loss_numerics_ghost_token_reference( + top_p, reverse, temperature, rank, size +): + parallel_state.initialize_model_parallel(tensor_model_parallel_size=size) + cfg = SimpleNamespace(tensor_model_parallel_size=size) + device = torch.device("cuda", rank) + # Equal vocab shards (torch.chunk is uneven when size does not divide the vocab) + student, teacher = _make_loss_inputs(vocab=16 * size) + student = student.detach().to(device).requires_grad_() + teacher = teacher.to(device) + student_shard = student.chunk(size, dim=-1)[rank].detach().contiguous().requires_grad_() + teacher_shard = teacher.chunk(size, dim=-1)[rank].contiguous() k = 4 - loss = TopLogitsKLLoss(cfg, top_k=k)(student, teacher)[0] + loss = TopLogitsKLLoss(cfg, top_k=k, top_p=top_p, reverse=reverse, temperature=temperature)( + student_shard, teacher_shard + )[0] - q_full = F.log_softmax(teacher, dim=-1) - p_full = F.log_softmax(student, dim=-1) + q_full = F.log_softmax(teacher / temperature, dim=-1) + p_full = F.log_softmax(student / temperature, dim=-1) _, idx = torch.topk(teacher, k, dim=-1) q_k, p_k = q_full.gather(-1, idx), p_full.gather(-1, idx) - q_rest = torch.log1p(-q_k.exp().sum(-1, keepdim=True)) - p_rest = torch.log1p(-p_k.exp().sum(-1, keepdim=True)) + keep = torch.ones_like(q_k, dtype=torch.bool) + if top_p is not None: + q_probs = q_k.exp() + keep = (q_probs.cumsum(-1) - q_probs) < top_p + assert not keep.all(), "test inputs should produce some truncation" + q_rest = torch.log1p(-(q_k.exp() * keep).sum(-1, keepdim=True)) + p_rest = torch.log1p(-(p_k.exp() * keep).sum(-1, keepdim=True)) q = torch.cat([q_k, q_rest], -1) p = torch.cat([p_k, p_rest], -1) - ref = (q.exp() * (q - p)).sum(-1).transpose(0, 1) - assert torch.allclose(loss, ref, atol=1e-6) + keep = torch.cat([keep, torch.ones_like(keep[..., :1])], -1) # Sanity: total mass within the K+1 buckets is 1 for both distributions. - assert torch.allclose(q.exp().sum(-1), torch.ones_like(q[..., 0]), atol=1e-5) - assert torch.allclose(p.exp().sum(-1), torch.ones_like(p[..., 0]), atol=1e-5) + assert torch.allclose((q.exp() * keep).sum(-1), torch.ones_like(q[..., 0]), atol=1e-5) + assert torch.allclose((p.exp() * keep).sum(-1), torch.ones_like(p[..., 0]), atol=1e-5) + if reverse: + p, q = q, p + ref = (keep * q.exp() * (q - p)).sum(-1).transpose(0, 1) + torch.testing.assert_close(loss, ref, atol=1e-6, rtol=1e-5) + loss.mean().backward() + ref.mean().backward() + torch.testing.assert_close( + student_shard.grad, student.grad.chunk(size, dim=-1)[rank], atol=1e-6, rtol=1e-5 + ) + + +@pytest.mark.parametrize( + ("top_p", "reverse", "temperature"), + [(None, False, 1.0), (0.5, False, 2.0), (0.5, True, 1.0)], +) +def test_topk_logits_kl_loss_numerics_ghost_token_reference( + dist_workers_size_2, top_p, reverse, temperature +): + """TP loss and gradients match a full-vocab reference on the K+1 bucketed distributions.""" + dist_workers_size_2.run( + partial( + _test_topk_logits_kl_loss_numerics_ghost_token_reference, top_p, reverse, temperature + ) + ) @pytest.mark.parametrize("temperature", [0.5, 2.0, 3.7]) From 5d8e81667f319a0de58eb1e9d43d911ffe2735a4 Mon Sep 17 00:00:00 2001 From: kaix-nv Date: Mon, 26 Oct 2026 12:24:30 -0700 Subject: [PATCH 26/43] Keep calibrated skip-softmax thresholds across save/restore (#3028) 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 A calibrated skip-softmax model loses its fitted thresholds after a ModelOpt checkpoint roundtrip. QA saw all 22 attention modules of TinyLlama-1.1B go from `dynamic_calibrated` to the static default after `mto.save` / `mto.restore` (max logit difference 2.09 on the same input), on both the PyTorch and Triton backends. **Root cause.** Calibration stores its results (`calibration_params`, the per-phase `a`/`b` fit, and `target_sparse_ratio`) as plain attributes on each module's `_sparse_method_instance`. They are not tensors, so the state dict doesn't carry them, and `update_sparse_attention_metadata` saved only `method` and `method_config`. On restore, `restore_sparse_attention_model` re-converts the model from the config, which builds fresh method instances with both attributes set to `None`. **Fix.** Follow the quantization design for non-tensor module state (`TensorQuantizer.get_modelopt_state` / `set_from_modelopt_state`, collected by `update_quantize_metadata` and applied by `restore_quantizer_state`): - `SparseAttentionModule.get_modelopt_state()` returns the module's state: `method`, `method_config`, the enabled flag, and the calibration results when present. `update_sparse_attention_metadata` stores it per module in the mode metadata. - `SparseAttentionModule.set_from_modelopt_state()` re-applies the enabled flag and calibration results after `convert` has rebuilt the method from the mode config. `restore_sparse_attention_state` calls it per module. - The old restore body is removed. It only set `_` attributes that nothing reads and called a `_setup()` that returns early, so removing it changes no behavior. The state stays in metadata rather than becoming buffers, so exported HF / vLLM weights gain no extra keys. The values are plain floats and dicts, so they load under `weights_only=True`. Checkpoints saved before this change carry none of the new keys and restore as before. Persisting the enabled flag also fixes a related loss: layers turned off with `disable_sparse_attention()` after `sparsify` came back enabled after restore (`TensorQuantizer` likewise saves its disabled flag). ### Usage No API change. Calibrated models now restore with their thresholds: ```python model = mtsa.sparsify(model, SKIP_SOFTMAX_TRITON_CALIB, forward_loop=calib_loop) mto.save(model, "ckpt.pth") restored = mto.restore(fresh_model, "ckpt.pth") # modules stay dynamic_calibrated ``` ### Testing - New unit test `test_restore_keeps_calibration_params`: injects calibration params, roundtrips through `mto.save` / `mto.restore` on a file and checks the params, target sparsity and `dynamic_calibrated` threshold type. It fails without the fix (`calibration_params` is `None` after restore) and passes with it. - New unit test `test_restore_keeps_disabled_modules`: a module disabled after `sparsify` stays disabled after restore. It fails without the fix. - New unit test `test_restore_legacy_state_without_calibration_keys`: metadata without the new keys still restores (enabled, uncalibrated). It passes both with and without the fix. - `pytest tests/unit/torch/sparsity/attention_sparsity`: 218 passed. - CPU repro on a 2-layer tiny Llama (`flash_skip_softmax`, PyTorch backend) through `mto.save` / `mto.restore`: before the fix, modules restored as `static` and a disabled layer came back enabled; with it, both are preserved. - Not re-run: QA's GPU TinyLlama-1.1B / `triton_skip_softmax` case. It goes through the same metadata path. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ (older checkpoints restore 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)?: ❌ - Did you get Claude approval on this PR?: ❌ (pending) ### Additional Information Needs cherry-pick to `release/0.48.0` (label `cherry-pick-0.48.0`). Known limitation, unchanged: attributes changed with `set_sparse_attention_attribute` after `sparsify` are not restored, since restore rebuilds the method from the original mode config. 🤖 Generated with [Claude Code](https://claude.com/claude-code) ## Summary by CodeRabbit * **Bug Fixes** * Sparse-attention models now retain calibrated thresholds and target sparsity when saved and restored, preserving their calibrated behavior. * Modules disabled before saving remain disabled after restoration. * Older saved states without enabled-status metadata continue to restore sparse attention as enabled, without calibration parameters. Signed-off-by: Kai Xu Signed-off-by: Rohan Joshi --- .../sparsity/attention_sparsity/conversion.py | 32 ++--------- .../attention_sparsity/sparse_attention.py | 35 ++++++++++++ .../test_sparse_attention_conversion.py | 54 +++++++++++++++++++ 3 files changed, 94 insertions(+), 27 deletions(-) diff --git a/modelopt/torch/sparsity/attention_sparsity/conversion.py b/modelopt/torch/sparsity/attention_sparsity/conversion.py index b2acd137930..aab3cbc74c7 100644 --- a/modelopt/torch/sparsity/attention_sparsity/conversion.py +++ b/modelopt/torch/sparsity/attention_sparsity/conversion.py @@ -332,18 +332,7 @@ def restore_sparse_attention_state(model: nn.Module, state_dict: dict[str, Any]) for name, module in get_named_sparse_attention_modules(model): module_name = get_unwrapped_name(name, model) if module_name in state_dict: - module_state = state_dict[module_name] - - # Restore method and config - if "method" in module_state: - module._method = module_state["method"] - if "method_config" in module_state: - # Restore config attributes - for key, val in module_state["method_config"].items(): - setattr(module, f"_{key}", val) - - # Re-setup with restored config - module._setup() + module.set_from_modelopt_state(state_dict[module_name]) def update_sparse_attention_metadata( @@ -356,21 +345,10 @@ def update_sparse_attention_metadata( config: Configuration used metadata: Metadata dict to update """ - sparse_state = {} - - for name, module in get_named_sparse_attention_modules(model): - module_name = get_unwrapped_name(name, model) - - # Save the method configuration that was used - # _method_config already contains the validated config dict - module_state = { - "method": module._sparse_method_instance.name, - "method_config": module._method_config.copy(), - } - - sparse_state[module_name] = module_state - - metadata["sparse_attention_state"] = sparse_state + metadata["sparse_attention_state"] = { + get_unwrapped_name(name, model): module.get_modelopt_state() + for name, module in get_named_sparse_attention_modules(model) + } metadata["sparse_attention_config"] = ( config.model_dump() if hasattr(config, "model_dump") else vars(config) ) diff --git a/modelopt/torch/sparsity/attention_sparsity/sparse_attention.py b/modelopt/torch/sparsity/attention_sparsity/sparse_attention.py index 4f29f382d2f..1979539e726 100644 --- a/modelopt/torch/sparsity/attention_sparsity/sparse_attention.py +++ b/modelopt/torch/sparsity/attention_sparsity/sparse_attention.py @@ -15,6 +15,7 @@ """Extensible sparse attention module.""" +import copy from typing import Any from modelopt.torch.opt.dynamic import DynamicModule, _DMRegistryCls @@ -23,6 +24,9 @@ from .methods import get_sparse_method from .stats_manager import SparseAttentionStatsManager +# Method-instance attributes set by calibration (see ``SparseAttentionMethod``). +_CALIBRATION_STATE_KEYS = ("calibration_params", "target_sparse_ratio") + class SparseAttentionModule(DynamicModule): """Generic sparse attention module wrapper for applying sparsity to attention layers. @@ -167,6 +171,37 @@ def get_threshold_info(self) -> dict[str, Any]: return self._sparse_method_instance.get_threshold_info() return {"type": "none", "value": None} + def get_modelopt_state(self) -> dict[str, Any]: + """Get the module state to save in the ModelOpt checkpoint metadata. + + Similar to ``TensorQuantizer.get_modelopt_state``. Calibration results live on the + method instance rather than in the state dict, so they are saved here. + """ + method = self._sparse_method_instance + modelopt_state = { + "method": method.name, + "method_config": self._method_config.copy(), + "enabled": self.is_enabled, + } + for key in _CALIBRATION_STATE_KEYS: + val = getattr(method, key, None) + if val is not None: + modelopt_state[key] = copy.deepcopy(val) + return modelopt_state + + def set_from_modelopt_state(self, modelopt_state: dict[str, Any]): + """Restore the state saved by :meth:`get_modelopt_state`. + + Similar to ``TensorQuantizer.set_from_modelopt_state``. The method and its config are + rebuilt from the mode config before this is called; this restores the runtime state on + top of them. + """ + if "enabled" in modelopt_state: + self._enabled = bool(modelopt_state["enabled"]) + for key in _CALIBRATION_STATE_KEYS: + if key in modelopt_state: + setattr(self._sparse_method_instance, key, copy.deepcopy(modelopt_state[key])) + def _setup(self): """Setup called by DynamicModule.""" # Apply default configuration if not yet configured diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attention_conversion.py b/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attention_conversion.py index 6cfdd2775b9..83f6f2c44ad 100644 --- a/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attention_conversion.py +++ b/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attention_conversion.py @@ -188,6 +188,60 @@ def test_restore_sparse_attention_model(self): assert hasattr(module, "_method") assert module._method == "flash_skip_softmax" + def test_restore_keeps_calibration_params(self, tmp_path): + """Calibrated thresholds survive an mto.save/mto.restore roundtrip.""" + calibration_params = {"prefill": {"a": 3.14, "b": 7.5}, "decode": {"a": 0.5, "b": 9.0}} + target_sparse_ratio = {"prefill": 0.4, "decode": 0.6} + + model = sparse_attn.sparsify(SimpleAttentionModel(), FLASH_SKIP_SOFTMAX_DEFAULT_CFG) + for module in model.modules(): + if isinstance(module, SparseAttentionModule): + module._sparse_method_instance.calibration_params = calibration_params + module._sparse_method_instance.target_sparse_ratio = target_sparse_ratio + + mto.save(model, tmp_path / "model.pth") + model_restored = mto.restore(SimpleAttentionModel(), tmp_path / "model.pth") + + sparse_modules = [ + m for m in model_restored.modules() if isinstance(m, SparseAttentionModule) + ] + assert sparse_modules + for module in sparse_modules: + method = module._sparse_method_instance + assert method.calibration_params == calibration_params + assert method.target_sparse_ratio == target_sparse_ratio + assert module.get_threshold_info()["type"] == "dynamic_calibrated" + + def test_restore_keeps_disabled_modules(self): + """Modules disabled after sparsify stay disabled after restore.""" + model = sparse_attn.sparsify(SimpleAttentionModel(), FLASH_SKIP_SOFTMAX_DEFAULT_CFG) + disable_sparse_attention(model, "*attention*") + + model_restored = mto.restore_from_modelopt_state( + SimpleAttentionModel(), mto.modelopt_state(model) + ) + + sparse_modules = [ + m for m in model_restored.modules() if isinstance(m, SparseAttentionModule) + ] + assert sparse_modules + assert not any(m.is_enabled for m in sparse_modules) + + def test_restore_legacy_state_without_calibration_keys(self): + """Metadata saved before calibration state was persisted still restores.""" + model = sparse_attn.sparsify(SimpleAttentionModel(), FLASH_SKIP_SOFTMAX_DEFAULT_CFG) + modelopt_state = mto.modelopt_state(model) + for _, mode_state in modelopt_state["modelopt_state_dict"]: + for module_state in mode_state["metadata"]["sparse_attention_state"].values(): + module_state.pop("enabled", None) + + model_restored = mto.restore_from_modelopt_state(SimpleAttentionModel(), modelopt_state) + + for module in model_restored.modules(): + if isinstance(module, SparseAttentionModule): + assert module.is_enabled + assert module._sparse_method_instance.calibration_params is None + class TestSparseAttentionModuleMethods: """Test SparseAttentionModule methods.""" From 8d52e2049ef2178ce7b1d438c71ee0ab7507779a Mon Sep 17 00:00:00 2001 From: Vu Dinh <8848239+dinhxuanvu@users.noreply.github.com> Date: Mon, 26 Oct 2026 15:50:19 -0400 Subject: [PATCH 27/43] fix: prevent UnboundLocalError masking real errors in fsdp2_aware_weight_update (#1860) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closes #1859 `fsdp2_aware_weight_update`'s `finally` block uses `fsdp_param_mapping`, `fsdp_param_group`, and `root_module`, which are only set after the `unshard()` inside the `try`. If `unshard()` fails first (e.g. a CUDA OOM on a large block), the `finally` raises `UnboundLocalError` — and since that happens inside `finally`, it replaces the real error. The user sees a confusing `UnboundLocalError` instead of the actual OOM. - Bind `fsdp_param_group` only after `create_fsdp_param_mapping` succeeds. - Skip the `finally` bookkeeping while it is unset, so an early failure (e.g. a CUDA OOM in `unshard()`) propagates and the FSDP param group is left untouched. The success path is unchanged. Added `test_fsdp2_weight_update_context_error_propagates`: forces the in-body `unshard()` and then `create_fsdp_param_mapping` to raise, and asserts each time that the original error surfaces and the FSDP param list is unchanged. * **Bug Fixes** * Fixed an issue where failures during FSDP setup, including unsharding errors, could be masked by a secondary cleanup error. * Original setup failures now propagate unchanged, resulting in more accurate error reporting. * Preserved FSDP parameter group state when setup fails. * **Tests** * Added regression coverage for unsharding and parameter-mapping failures, confirming that the original error is preserved and parameter groups remain unchanged. --------- Signed-off-by: Vu Dinh Signed-off-by: Rohan Joshi --- CHANGELOG.rst | 1 + .../torch/quantization/utils/core_utils.py | 10 ++++-- tests/gpu/torch/export/test_fsdp2_export.py | 35 +++++++++++++++++++ 3 files changed, 43 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index fa8bf6990a5..0f363e5d573 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -196,6 +196,7 @@ Changelog - Diffusers ``export_hf_checkpoint`` now honors ``save_modelopt_state=False`` (the default) when HuggingFace checkpointing is enabled. Pass ``save_modelopt_state=True`` to retain the ModelOpt state file. - Fix ONNX INT4 quantization ending the Python process when the INT4 packing extension fails to build, for example without a C++ compiler. It now falls back to the Python implementation. - Fix Megatron Top-K/Top-P distillation gradients being multiplied by the tensor-parallel size when tensor parallelism is enabled; reported loss values are unchanged. +- Fix ``fsdp2_aware_weight_update`` masking an error raised during its unshard (e.g. a CUDA OOM) with a spurious ``UnboundLocalError``. 0.47.0 (2026-09-23) ^^^^^^^^^^^^^^^^^^^ diff --git a/modelopt/torch/quantization/utils/core_utils.py b/modelopt/torch/quantization/utils/core_utils.py index 63f349dd10a..1c488c459e8 100644 --- a/modelopt/torch/quantization/utils/core_utils.py +++ b/modelopt/torch/quantization/utils/core_utils.py @@ -985,6 +985,9 @@ def fsdp2_aware_weight_update( Returns: None """ + # Bound only once setup succeeds; the finally skips FSDP bookkeeping until then so an + # early error (e.g. an OOM from the unshard below) propagates and FSDP state stays intact. + fsdp_param_group = None try: if isinstance(root_model, FSDPModule): # Get FSDP root module, if none is returned, then the update is not made to a submodule of an FSDPModule @@ -1006,8 +1009,9 @@ def fsdp2_aware_weight_update( root_module.unshard() # Get FSDPParam list - fsdp_param_group = fully_shard.state(root_module)._fsdp_param_group - fsdp_param_mapping = create_fsdp_param_mapping(fsdp_param_group.fsdp_params, root_model) + param_group = fully_shard.state(root_module)._fsdp_param_group + fsdp_param_mapping = create_fsdp_param_mapping(param_group.fsdp_params, root_model) + fsdp_param_group = param_group # Assert that all the modules in the module list are present in this fsdp_param_group if len(modules_to_update) > 1: @@ -1027,7 +1031,7 @@ def fsdp2_aware_weight_update( finally: from modelopt.torch.quantization.qtensor.base_qtensor import QFSDPParam, QTensorWrapper - if isinstance(root_model, FSDPModule): + if fsdp_param_group is not None: # Update FSDPParam list for module in modules_to_update: for param_name, param in module.named_parameters(): diff --git a/tests/gpu/torch/export/test_fsdp2_export.py b/tests/gpu/torch/export/test_fsdp2_export.py index 939fe3e581f..bbd1195d164 100644 --- a/tests/gpu/torch/export/test_fsdp2_export.py +++ b/tests/gpu/torch/export/test_fsdp2_export.py @@ -16,6 +16,7 @@ import json from functools import partial from pathlib import Path +from unittest.mock import patch import pytest import torch @@ -40,6 +41,7 @@ ) from modelopt.torch.export.unified_export_hf_streaming import _export_fsdp2_checkpoint_streaming from modelopt.torch.quantization.utils import ( + core_utils, enable_weight_access_and_writeback, fsdp2_aware_weight_update, module_name_maps, @@ -83,6 +85,35 @@ def _update_weight_test(rank, size): assert torch.allclose(torch.zeros(4, 4).to(output.device).to(output.dtype), output) +def _update_weight_error_propagates_test(rank, size): + """A setup error must propagate unchanged and leave the FSDP param group intact.""" + with patch_fsdp_mp_dtypes(): + model = ToyModel(dims=[4, 4], bias=False).to("cuda") + fully_shard(model.linears) + fully_shard(model) + + torch.distributed.barrier() + + class _SentinelError(RuntimeError): + pass + + def _raise(*args, **kwargs): + raise _SentinelError + + param_group = fully_shard.state(model.linears)._fsdp_param_group + fsdp_params = list(param_group.fsdp_params) + + # The unshard failure stands in for an OOM. + for target, attr in [(model.linears, "unshard"), (core_utils, "create_fsdp_param_mapping")]: + with ( + patch.object(target, attr, _raise), + pytest.raises(_SentinelError), + fsdp2_aware_weight_update(model, model.linears), + ): + pass + assert param_group.fsdp_params == fsdp_params + + def _compress_weight_test(rank, size): """Test fsdp2 weight update context for weight compression -> only value,shape and dtype changed""" with patch_fsdp_mp_dtypes(): @@ -262,6 +293,10 @@ def test_fsdp2_weight_update_context_for_export(dist_workers): dist_workers.run(_update_weight_test) +def test_fsdp2_weight_update_context_error_propagates(dist_workers): + dist_workers.run(_update_weight_error_propagates_test) + + @pytest.mark.parametrize( "quant_config", [ From 9ac931bcc24796c938a168aab42b0f990639284c Mon Sep 17 00:00:00 2001 From: Zhehao Hu <69692533+Zhehao-Hu@users.noreply.github.com> Date: Mon, 26 Oct 2026 13:08:05 -0700 Subject: [PATCH 28/43] [NVBug 6982529] Fix stale ModelOpt state when overwriting checkpoints with plain models (#3030) 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 Saving a plain model over a quantized checkpoint previously left modelopt_state.pth behind, causing reload to restore stale quantizers. Remove that file after a successful unconverted-model save, regardless of save_modelopt_state. Failed saves preserve the existing state file. ### Usage No API changes. Existing save_pretrained() calls now correctly replace quantized checkpoints with plain models. ### Scope and known limitation This PR addresses local checkpoint overwrites. With `push_to_hub=True`, Diffusers uploads the directory before the save hook removes an old `modelopt_state.pth`, so the Hub copy can still contain stale conversion state. Correcting upload ordering and remote checkpoint cleanup is deferred to a separate change, as noted in [CodeRabbit's review](https://github.com/NVIDIA/Model-Optimizer/pull/3030#discussion_r4381089257). ### Testing - Reproduced the stale-quantizer failure before applying the fix. - Added a tiny Diffusers UNet regression covering both save_modelopt_state settings, failed-save preservation, exact weight restoration, and output equality with a fresh-directory save/load control. - CPU checkpoint tests: 21 passed, 1 manual test skipped. - All applicable pre-commit checks passed. ### 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 - Did you get Claude approval on this PR?: ❌ - Not yet run ### Additional Information Fixes NVBug 6982529. ## Summary by CodeRabbit * **Bug Fixes** * Saving an unconverted model, or saving with ModelOpt state disabled, now removes any stale ModelOpt state file from the destination. * If a save fails, the existing ModelOpt state file is preserved. * Models saved without conversion restore as unquantized models with matching weights and outputs. * Exported checkpoints without ModelOpt state retain the expected quantization scales and visual-component weights. Signed-off-by: Zhehao Hu Signed-off-by: Rohan Joshi --- modelopt/torch/opt/plugins/huggingface.py | 4 +++ .../opt/plugins/test_diffusers_save_load.py | 34 +++++++++++++++++++ .../plugins/test_transformers_save_load.py | 30 ++++++++++++++++ 3 files changed, 68 insertions(+) diff --git a/modelopt/torch/opt/plugins/huggingface.py b/modelopt/torch/opt/plugins/huggingface.py index 67441298ed4..dc383c56566 100644 --- a/modelopt/torch/opt/plugins/huggingface.py +++ b/modelopt/torch/opt/plugins/huggingface.py @@ -20,6 +20,7 @@ import threading import warnings from contextlib import contextmanager +from pathlib import Path from typing import Any import torch @@ -117,6 +118,9 @@ def _new_save_pretrained(self, save_directory, *args, **kwargs): path = _get_modelopt_state_path(save_directory) torch.save(modelopt_state(self), path) print_rank_0(f"Saved ModelOpt state to {path}") + else: + # A checkpoint saved without ModelOpt state must not inherit an older sidecar. + Path(save_directory, _MODELOPT_STATE_SAVE_NAME).unlink(missing_ok=True) return outputs diff --git a/tests/unit/torch/opt/plugins/test_diffusers_save_load.py b/tests/unit/torch/opt/plugins/test_diffusers_save_load.py index c73e2a81707..5a6561fd28c 100644 --- a/tests/unit/torch/opt/plugins/test_diffusers_save_load.py +++ b/tests/unit/torch/opt/plugins/test_diffusers_save_load.py @@ -14,13 +14,19 @@ # limitations under the License. import pytest +import torch from _test_utils.torch.diffusers_models import ( create_tiny_unet_dir, df_modelopt_state_and_output_tester, + get_tiny_unet, + get_unet_dummy_inputs, ) from _test_utils.torch.opt.utils import apply_mode_with_sampling from diffusers import UNet2DConditionModel +import modelopt.torch.quantization as mtq +from modelopt.torch.quantization.nn import TensorQuantizer + @pytest.mark.parametrize("model_cls", [UNet2DConditionModel]) @pytest.mark.parametrize("use_subfolder", [False, True], ids=["direct", "subfolder"]) @@ -38,3 +44,31 @@ def test_unet_save_restore(tmp_path, model_cls, use_subfolder): load_kwargs = {"subfolder": "unet"} if use_subfolder else {} model_test = model_cls.from_pretrained(load_path, local_files_only=True, **load_kwargs) df_modelopt_state_and_output_tester(model_ref, model_test) + + +@pytest.mark.parametrize("save_modelopt_state", [True, False]) +@torch.no_grad() +def test_plain_save_replaces_quantized_checkpoint(tmp_path, save_modelopt_state): + quantized = get_tiny_unet().eval() + inputs = get_unet_dummy_inputs(quantized) + mtq.quantize(quantized, mtq.FP8_DEFAULT_CFG, lambda m: m(**inputs)) + quantized.save_pretrained(tmp_path) + assert (tmp_path / "modelopt_state.pth").is_file() + + plain = get_tiny_unet().eval() + original_state = (tmp_path / "modelopt_state.pth").read_bytes() + with pytest.raises(ValueError): + plain.save_pretrained(tmp_path, max_shard_size="invalid") + assert (tmp_path / "modelopt_state.pth").read_bytes() == original_state + + plain.save_pretrained(tmp_path / "fresh") + reference = UNet2DConditionModel.from_pretrained(tmp_path / "fresh", local_files_only=True) + plain.save_pretrained(tmp_path, save_modelopt_state=save_modelopt_state) + restored = UNet2DConditionModel.from_pretrained(tmp_path, local_files_only=True) + + assert not any(isinstance(module, TensorQuantizer) for module in restored.modules()) + assert not (tmp_path / "modelopt_state.pth").exists() + torch.testing.assert_close(restored.state_dict(), plain.state_dict(), rtol=0, atol=0) + torch.testing.assert_close( + restored(**inputs).sample, reference(**inputs).sample, rtol=0, atol=0 + ) diff --git a/tests/unit/torch/opt/plugins/test_transformers_save_load.py b/tests/unit/torch/opt/plugins/test_transformers_save_load.py index 1e526fc09c0..3faf3de79c1 100644 --- a/tests/unit/torch/opt/plugins/test_transformers_save_load.py +++ b/tests/unit/torch/opt/plugins/test_transformers_save_load.py @@ -22,11 +22,15 @@ from _test_utils.torch.opt.utils import apply_mode_with_sampling from _test_utils.torch.transformers_models import ( create_tiny_llama_dir, + get_tiny_qwen3vl, tf_modelopt_state_and_output_tester, ) from safetensors.torch import load_file from transformers import AutoConfig, AutoModelForCausalLM, LlamaForCausalLM +import modelopt.torch.opt as mto +import modelopt.torch.quantization as mtq +from modelopt.torch.export import export_hf_checkpoint from modelopt.torch.opt.plugins.transformers import _legacy_tied_weights_keys_as_dict @@ -44,6 +48,32 @@ def test_causal_lm_save_restore(tmp_path, model_cls): tf_modelopt_state_and_output_tester(model_ref, model_test) +@pytest.mark.parametrize( + "quant_config", [mtq.FP8_DEFAULT_CFG, mtq.NVFP4_DEFAULT_CFG], ids=["fp8", "nvfp4"] +) +@torch.no_grad() +def test_vlm_export_without_modelopt_state(tmp_path, quant_config): + """Export separately converted VLM components without requiring a single ModelOpt state.""" + model = get_tiny_qwen3vl().eval() + disabled_config = {"quant_cfg": [{"quantizer_name": "*", "enable": False}], "algorithm": "max"} + mtq.quantize(model.model.visual, disabled_config) + mtq.quantize(model.lm_head, disabled_config) + mtq.quantize( + model.model.language_model, + quant_config, + lambda _: model(input_ids=torch.ones(1, 4, dtype=torch.long)), + ) + state_path = tmp_path / "modelopt_state.pth" + torch.save(mto.modelopt_state(model.model.language_model), state_path) + + export_hf_checkpoint(model, export_dir=tmp_path, save_modelopt_state=False) + + assert not state_path.exists() + saved = load_file(tmp_path / "model.safetensors") + assert any(key.endswith(".weight_scale") for key in saved) + assert any("visual" in key for key in saved) + + @pytest.mark.parametrize("tie_word_embeddings", [False, True]) def test_save_pretrained_with_legacy_tied_weights_keys(tmp_path, tie_word_embeddings): """A model declaring 4.x list-style `_tied_weights_keys` must still save (nvbug 6518665). From 23f46c3aa93ee0e473b9a873ae9d74e63257024e Mon Sep 17 00:00:00 2001 From: Pratik Gandhi Date: Mon, 26 Oct 2026 13:17:33 -0700 Subject: [PATCH 29/43] docs: fix TensorRT-LLM and autotune links (#2612) ### What does this PR do? Type of change: documentation Fixes three broken links in example READMEs: - `examples/gpt-oss/README.md`: the TensorRT-LLM GPT-OSS blog is now `blog09_Deploying_GPT_OSS_on_TRTLLM.md` (was `blog9_...`). - `examples/hf_ptq/README.md`: the TensorRT-LLM Build API doc moved to `docs/source/legacy/architecture/workflow.md`; the `#build-apis` heading is still there. - `examples/onnx_ptq/autotune/README.md`: the User Guide link pointed to `docs/source/guides/9_qdq_placement.rst` with one `../` too few. The guide is now `docs/source/guides/9_autotune.rst` (added in 69c0d47), which is three levels up from this README. ### Usage N/A (docs only) ### Testing Docs only, no tests needed. Each new target was checked on the default branch of the relevant repo. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: N/A - 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?: N/A ### Additional Information The API Reference links to `docs/source/reference/2_qdq_placement.rst` in the same autotune README also look broken (that file does not exist under `docs/source/reference/`), but I was not sure what the right replacement is, so I left them alone. ## Summary by CodeRabbit * **Documentation** * Corrected the TensorRT-LLM setup guide link in the GPT-OSS example. * Updated the TensorRT-LLM Build API link in the HF PTQ example to point to the appropriate architecture-specific documentation. --------- Signed-off-by: Pratik Gandhi Signed-off-by: Rohan Joshi --- examples/gpt-oss/README.md | 2 +- examples/hf_ptq/README.md | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/gpt-oss/README.md b/examples/gpt-oss/README.md index be3dea6b558..6779d7e29ba 100644 --- a/examples/gpt-oss/README.md +++ b/examples/gpt-oss/README.md @@ -115,7 +115,7 @@ Note: Model Optimizer currently exports quantized checkpoints in formats other t
Deployment on TensorRT-LLM -To setup TensorRT-LLM, follow the official guide: [Deploying GPT-OSS on TensorRT-LLM](https://github.com/NVIDIA/TensorRT-LLM/blob/main/docs/source/blogs/tech_blog/blog9_Deploying_GPT_OSS_on_TRTLLM.md) +To setup TensorRT-LLM, follow the official guide: [Deploying GPT-OSS on TensorRT-LLM](https://github.com/NVIDIA/TensorRT-LLM/blob/main/docs/source/blogs/tech_blog/blog09_Deploying_GPT_OSS_on_TRTLLM.md) Once installed, launch an OpenAI-compatible endpoint using: ```bash diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index 107ce82c6a7..a0877b52e39 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -732,7 +732,7 @@ with torch.inference_mode(): ### Build the TensorRT-LLM engines -After the TensorRT-LLM checkpoint export, you can use the `trtllm-build` build command to build the engines from the exported checkpoints. Please check the [TensorRT-LLM Build API](https://github.com/NVIDIA/TensorRT-LLM/blob/main/docs/source/architecture/workflow.md#build-apis) documentation for reference. +After the TensorRT-LLM checkpoint export, you can use the `trtllm-build` build command to build the engines from the exported checkpoints. Please check the [TensorRT-LLM Build API](https://github.com/NVIDIA/TensorRT-LLM/blob/main/docs/source/legacy/architecture/workflow.md#build-apis) documentation for reference. ## Pre-Quantized Checkpoints From 4a5af5523d3c1fffbe2d1e0377112665ba75fff9 Mon Sep 17 00:00:00 2001 From: Zhehao Hu <69692533+Zhehao-Hu@users.noreply.github.com> Date: Mon, 26 Oct 2026 13:20:15 -0700 Subject: [PATCH 30/43] [NVBug 6980192] Fix NVFP4 Conv3d activation quantization disable state (#3033) 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 Calling `input_quantizer.disable_quant()` previously left activation quantization active in the fused NVFP4 Conv3d kernel because dispatch checked only `is_enabled`. Check both `is_enabled` and `_if_quant` before enabling activation quantization in the kernel. This fixes regular Conv3d and Wan causal Conv3d while preserving independent weight quantization. ### Usage Existing calls now behave as expected: ```python model.input_quantizer.disable_quant() output = model(x) # Fused Conv3d skips activation quantization. ``` ### Testing Validated using the real CUDA kernel on RTX 6000 Ada with PyTorch 26.07 and Diffusers 0.39.0: - 80 focused checks passed, including existing convolution tests. - Six new GPU cases cover regular Conv3d and Wan Conv3d with and without cached frames, with weight quantization enabled or disabled. - Each case verifies equivalence to bypassing the activation quantizer and restoration of the original output after re-enabling. - All six cases fail with the fix reverted and pass after restoration. - All pre-commit checks passed. B300 and full-model video generation were not tested. ### 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 - Did you get Claude approval on this PR?: N/A - Not yet run ### Additional Information Fixes NVBug 6980192. ## Summary by CodeRabbit * **Bug Fixes** * Conv3D now correctly respects the activation quantizer’s active state when applying quantization, making disabled activation quantization behave as expected. * **Tests** * Added GPU regression coverage for NVFP4 Conv3D behavior, including activation quantization toggling and optional Wan convolution variants. Signed-off-by: Zhehao Hu Signed-off-by: Rohan Joshi --- .../quantization/nn/modules/quant_conv.py | 2 +- .../torch/quantization/test_quant_conv_gpu.py | 65 +++++++++++++++++++ 2 files changed, 66 insertions(+), 1 deletion(-) create mode 100644 tests/gpu/torch/quantization/test_quant_conv_gpu.py diff --git a/modelopt/torch/quantization/nn/modules/quant_conv.py b/modelopt/torch/quantization/nn/modules/quant_conv.py index 375fdc96e0c..8842540c87e 100644 --- a/modelopt/torch/quantization/nn/modules/quant_conv.py +++ b/modelopt/torch/quantization/nn/modules/quant_conv.py @@ -118,7 +118,7 @@ def _implicit_gemm_forward(self, input): padding=self.padding, dilation=self.dilation, act_amax=act_amax, - quant_act=self.input_quantizer.is_enabled, + quant_act=self.input_quantizer.is_enabled and self.input_quantizer._if_quant, fp4_block_size=fp4_block_size, ) return self.output_quantizer(output) diff --git a/tests/gpu/torch/quantization/test_quant_conv_gpu.py b/tests/gpu/torch/quantization/test_quant_conv_gpu.py new file mode 100644 index 00000000000..845db583c77 --- /dev/null +++ b/tests/gpu/torch/quantization/test_quant_conv_gpu.py @@ -0,0 +1,65 @@ +# 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. + +"""GPU regression tests for quantized convolutions.""" + +import pytest +import torch + +import modelopt.torch.quantization as mtq + + +@pytest.mark.parametrize("conv_kind", ["conv3d", "wan", "wan_cached"]) +@pytest.mark.parametrize("quantize_weight", [False, True]) +@torch.no_grad() +def test_nvfp4_conv3d_disable_quant(cuda_capability, conv_kind, quantize_weight): + """Disabling activation QDQ must match bypassing its quantizer, even in the fused kernel.""" + if cuda_capability < (8, 9): + pytest.skip("NVFP4 quantization requires compute capability >= 8.9") + + torch.manual_seed(1234) + conv_cls = torch.nn.Conv3d + if conv_kind != "conv3d": + # Diffusers and its quantization plugin are optional for the regular Conv3d case. + pytest.importorskip("onnx") + pytest.importorskip("diffusers") + from diffusers.models.autoencoders.autoencoder_kl_wan import WanCausalConv3d + + import modelopt.torch.quantization.plugins.diffusion.diffusers # noqa: F401 + + conv_cls = WanCausalConv3d + + model = conv_cls(16, 16, kernel_size=3, padding=1).to("cuda", torch.bfloat16).eval() + x = torch.randn(1, 16, 3, 6, 6, device="cuda", dtype=torch.bfloat16) + kwargs = {} + if conv_kind == "wan_cached": + kwargs["cache_x"] = torch.randn(1, 16, 1, 6, 6, device="cuda", dtype=torch.bfloat16) + model = mtq.quantize(model, mtq.NVFP4_DEFAULT_CFG, lambda m: m(x, **kwargs)) + if not quantize_weight: + model.weight_quantizer.disable_quant() + + enabled = model(x, **kwargs) + model.input_quantizer.disable_quant() + disabled = model(x, **kwargs) + + # Bypass the activation quantizer entirely as a control, keeping the weight state fixed. + model.input_quantizer.disable() + bypassed = model(x, **kwargs) + torch.testing.assert_close(disabled, bypassed, rtol=0, atol=0) + assert not torch.equal(enabled, disabled) + + model.input_quantizer.enable() + model.input_quantizer.enable_quant() + torch.testing.assert_close(model(x, **kwargs), enabled, rtol=0, atol=0) From 92f78b55aaa94681a55c3fbdc387a631e6479716 Mon Sep 17 00:00:00 2001 From: Michael R <89427972+Michael-RDev@users.noreply.github.com> Date: Mon, 26 Oct 2026 16:29:50 -0400 Subject: [PATCH 31/43] Fix nested distillation forward contexts (#2653) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Type of change: Bug fix I found that nesting `only_student_forward()` or `only_teacher_forward()` causes the outer context to stop working when the inner one exits. This also happens when the inner context uses `enable=False`. Both context managers reset their flags to `False` on exit. This change saves the previous value and restores it in `finally`, so the enclosing context keeps its execution mode, including after an exception. Public APIs, checkpoint handling, and loss calculations are unchanged. Added eight regression cases using small local models and real forward passes. They cover: - Both teacher-only and student-only contexts. - Enabled and disabled inner contexts. - Normal exits and exceptions. - Normal execution after leaving a context. - Overlapping teacher/student contexts. All eight cases fail before the fix and pass afterward. Verified with CPU PyTorch: - Distillation and layerwise test modules — 39 passed. - All applicable pre-commit checks — passed, including Ruff, mypy, Bandit, and license checks. GPU and distributed integration tests were not run. - Backward compatible: Yes. - Copied code or new dependencies: None. - Necessary tests added: Yes. - Changelog updated: Yes. - Claude approval: Not obtained, as I don't use Claude. * **Bug Fixes** * Nested student-only and teacher-only forwarding contexts now preserve the surrounding execution mode when they exit, including when an inner context is disabled or raises an exception. * After nested contexts finish, subsequent model calls use the expected forwarding behavior. This applies when contexts are nested in either order and when exceptions occur inside either scope. Signed-off-by: Michael Rusu Signed-off-by: Michael-Rusu <89427972+Michael-Rusu@users.noreply.github.com> Co-authored-by: Michael-Rusu <89427972+Michael-Rusu@users.noreply.github.com> Signed-off-by: Rohan Joshi --- CHANGELOG.rst | 1 + modelopt/torch/distill/distillation_model.py | 7 +- tests/unit/torch/distill/test_distill.py | 69 ++++++++++++++++++++ 3 files changed, 75 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 0f363e5d573..a25d02ba0c2 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -197,6 +197,7 @@ Changelog - Fix ONNX INT4 quantization ending the Python process when the INT4 packing extension fails to build, for example without a C++ compiler. It now falls back to the Python implementation. - Fix Megatron Top-K/Top-P distillation gradients being multiplied by the tensor-parallel size when tensor parallelism is enabled; reported loss values are unchanged. - Fix ``fsdp2_aware_weight_update`` masking an error raised during its unshard (e.g. a CUDA OOM) with a spurious ``UnboundLocalError``. +- Nested ``only_student_forward()`` and ``only_teacher_forward()`` contexts now preserve the enclosing execution mode, including when an inner context is disabled or raises an exception. 0.47.0 (2026-09-23) ^^^^^^^^^^^^^^^^^^^ diff --git a/modelopt/torch/distill/distillation_model.py b/modelopt/torch/distill/distillation_model.py index 3d4fc1cb6a3..2bd2a2d4094 100644 --- a/modelopt/torch/distill/distillation_model.py +++ b/modelopt/torch/distill/distillation_model.py @@ -166,22 +166,25 @@ def hide_loss_modules(self, enable=True): @contextmanager def only_teacher_forward(self, enable=True): """Context manager to temporarily disable forward passes on the student model.""" + # restore the previous mode so nested contexts keep working + only_teacher_fwd = self._only_teacher_fwd if enable: self._only_teacher_fwd = True try: yield finally: - self._only_teacher_fwd = False + self._only_teacher_fwd = only_teacher_fwd @contextmanager def only_student_forward(self, enable=True): """Context manager to temporarily run forward passes only on the student model.""" + only_student_fwd = self._only_student_fwd if enable: self._only_student_fwd = True try: yield finally: - self._only_student_fwd = False + self._only_student_fwd = only_student_fwd def train(self, mode: bool = True): """Override to prevent warnings of stored intermediate outputs in future forwards.""" diff --git a/tests/unit/torch/distill/test_distill.py b/tests/unit/torch/distill/test_distill.py index 7c938aee90d..721ff4ef9a6 100644 --- a/tests/unit/torch/distill/test_distill.py +++ b/tests/unit/torch/distill/test_distill.py @@ -15,6 +15,7 @@ import inspect import warnings +from contextlib import nullcontext import pytest import torch @@ -26,6 +27,17 @@ import modelopt.torch.opt as mto +class LinearModel(nn.Module): + def __init__(self, scale): + super().__init__() + self.linear = nn.Linear(2, 2, bias=False) # Wx + with torch.no_grad(): + self.linear.weight.copy_(scale * torch.eye(2)) + + def forward(self, inputs): + return self.linear(inputs) + + def get_input_tensor(): """Dummy input tensor.""" return torch.rand(2, 3, 112, 112) @@ -310,6 +322,63 @@ def test_student_fwd_only(distillation_model): assert distillation_model._intermediate_output is not None +@pytest.mark.parametrize("context_name", ["only_student_forward", "only_teacher_forward"]) +@pytest.mark.parametrize("inner_enabled", [False, True]) +@pytest.mark.parametrize("raise_in_inner", [False, True]) +def test_nested_forward_contexts(context_name, inner_enabled, raise_in_inner): + student, teacher = LinearModel(1), LinearModel(2) + model = mtd.convert( + student, mode=[("kd_loss", {"teacher_model": teacher, "criterion": nn.MSELoss()})] + ) + context = getattr(model, context_name) + inputs = torch.ones(1, 2) + calls = {"student": 0, "teacher": 0} + + def count_student(module, inputs): + calls["student"] += 1 + + def count_teacher(module, inputs): + calls["teacher"] += 1 + + model.linear.register_forward_pre_hook(count_student) + teacher.linear.register_forward_pre_hook(count_teacher) + expected = inputs if context_name == "only_student_forward" else 2 * inputs + + with context(): + torch.testing.assert_close(model(inputs), expected) + with ( + pytest.raises(RuntimeError, match="inner failure") if raise_in_inner else nullcontext(), + context(enable=inner_enabled), + ): + torch.testing.assert_close(model(inputs), expected) + if raise_in_inner: + raise RuntimeError("inner failure") + torch.testing.assert_close(model(inputs), expected) + + assert calls["teacher"] == (0 if context_name == "only_student_forward" else 3) + assert calls["student"] == (3 if context_name == "only_student_forward" else 0) + torch.testing.assert_close(model(inputs), inputs) + assert calls["teacher"] == (1 if context_name == "only_student_forward" else 4) + assert calls["student"] == (4 if context_name == "only_student_forward" else 1) + + with pytest.raises(RuntimeError, match="outer failure"), context(): + raise RuntimeError("outer failure") + with context(enable=False): + torch.testing.assert_close(model(inputs), inputs) + assert calls["teacher"] == (2 if context_name == "only_student_forward" else 5) + assert calls["student"] == (5 if context_name == "only_student_forward" else 2) + + other_context = ( + model.only_teacher_forward + if context_name == "only_student_forward" + else model.only_student_forward + ) + with context(): + with other_context(): + torch.testing.assert_close(model(inputs), inputs) + torch.testing.assert_close(model(inputs), expected) + + def test_train_eval_mode_switch(distillation_model): distillation_model.train() assert distillation_model._intermediate_output is None From b4c4f3595a6a158e0baa7e948169106912343471 Mon Sep 17 00:00:00 2001 From: namgyu-youn Date: Tue, 27 Oct 2026 05:31:20 +0900 Subject: [PATCH 32/43] Tile static NVFP4 fake quant and use int64 offsets (#2599) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Type of change: Bug fix `static_blockwise_fp4_fake_quant` launched one Triton program per 16-element block and used int32 offsets. It ran at about 4% of HBM bandwidth and hit an illegal memory access once a tensor had more than 2^31 elements, as stacked MoE expert weights do. Each program now handles 64 blocks with int64 offsets; output is bit-identical. N/A (no API change) bf16, H100, `triton.testing.do_bench`: | shape | before (ms) | after (ms) | |---|---|---| | 4096×12288 | 1.947 | 0.139 | | 18432×7168 | 5.063 | 0.315 | | 64×7168×2048 | 35.78 | 2.029 | ``` pytest tests/gpu/torch/quantization/test_tensor_quant_cuda.py → 135 passed, 2 skipped pytest tests/gpu/torch/quantization/test_nvfp4_fp8_sweep_kernel.py → 45 passed pytest tests/gpu/torch/quantization/test_gptq.py → 10 passed ``` 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?: ❌ * **Bug Fixes** * Fixed an illegal memory access when static NVFP4 fake quantization processes tensors larger than \(2^{31}\) elements, including very large stacked expert weights. * Improved the performance of this quantization operation by approximately 14×. Results for the final portion of very large tensors are now handled correctly, avoiding errors that could occur when processing data beyond the \(2^{31}\)-element boundary. Signed-off-by: namgyu-youn Signed-off-by: Rohan Joshi --- CHANGELOG.rst | 1 + .../kernels/quantization/gemm/fp4_kernel.py | 26 +++++++++++-------- .../quantization/test_tensor_quant_cuda.py | 9 +++++++ 3 files changed, 25 insertions(+), 11 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index a25d02ba0c2..7f03618f750 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -198,6 +198,7 @@ Changelog - Fix Megatron Top-K/Top-P distillation gradients being multiplied by the tensor-parallel size when tensor parallelism is enabled; reported loss values are unchanged. - Fix ``fsdp2_aware_weight_update`` masking an error raised during its unshard (e.g. a CUDA OOM) with a spurious ``UnboundLocalError``. - Nested ``only_student_forward()`` and ``only_teacher_forward()`` contexts now preserve the enclosing execution mode, including when an inner context is disabled or raises an exception. +- Fix an illegal memory access in static NVFP4 fake quantization for tensors with more than 2^31 elements, such as stacked MoE expert weights. The kernel is also about 14x faster. 0.47.0 (2026-09-23) ^^^^^^^^^^^^^^^^^^^ diff --git a/modelopt/torch/kernels/quantization/gemm/fp4_kernel.py b/modelopt/torch/kernels/quantization/gemm/fp4_kernel.py index c9499b612f6..7f084ed089a 100644 --- a/modelopt/torch/kernels/quantization/gemm/fp4_kernel.py +++ b/modelopt/torch/kernels/quantization/gemm/fp4_kernel.py @@ -190,6 +190,9 @@ def fp4_dequantize( return output +_STATIC_FP4_BLOCKS_PER_PROGRAM = 64 + + @triton.jit def static_blockwise_fp4_fake_quant_kernel( x_ptr, # [NUM_FP4_BLOCKS * BLOCK_SIZE] @@ -197,21 +200,21 @@ def static_blockwise_fp4_fake_quant_kernel( scale_ptr, # [NUM_FP4_BLOCKS] NUM_FP4_BLOCKS, BLOCK_SIZE: tl.constexpr, + BLOCKS_PER_PROGRAM: tl.constexpr, OUT_DTYPE: tl.constexpr, ): - pid = tl.program_id(axis=0) - if pid >= NUM_FP4_BLOCKS: - return - - block_offset = pid * BLOCK_SIZE - idx = block_offset + tl.arange(0, BLOCK_SIZE) + # int64 offsets: stacked MoE weights can exceed 2**31 elements. + pid = tl.program_id(axis=0).to(tl.int64) + block_idx = pid * BLOCKS_PER_PROGRAM + tl.arange(0, BLOCKS_PER_PROGRAM) + block_mask = block_idx < NUM_FP4_BLOCKS + idx = block_idx[:, None] * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)[None, :] - scale = tl.load(scale_ptr + pid).to(tl.float32) - x = tl.load(x_ptr + idx).to(tl.float32) + scale = tl.load(scale_ptr + block_idx, mask=block_mask, other=0.0).to(tl.float32) + x = tl.load(x_ptr + idx, mask=block_mask[:, None], other=0.0).to(tl.float32) - x_quant = nvfp4_scalar_quant(x, scale, BLOCK_SIZE) + x_quant = nvfp4_scalar_quant(x, scale[:, None], BLOCK_SIZE) - tl.store(y_ptr + idx, x_quant.to(OUT_DTYPE)) + tl.store(y_ptr + idx, x_quant.to(OUT_DTYPE), mask=block_mask[:, None]) def compute_fp4_scales( @@ -301,7 +304,7 @@ def static_blockwise_fp4_fake_quant( tl_out_dtype = _torch_dtype_to_tl(out_dtype) - grid = (NUM_FP4_BLOCKS,) + grid = (triton.cdiv(NUM_FP4_BLOCKS, _STATIC_FP4_BLOCKS_PER_PROGRAM),) with torch.cuda.device(x.device): static_blockwise_fp4_fake_quant_kernel[grid]( @@ -310,6 +313,7 @@ def static_blockwise_fp4_fake_quant( scale_flat, NUM_FP4_BLOCKS, BLOCK_SIZE, + BLOCKS_PER_PROGRAM=_STATIC_FP4_BLOCKS_PER_PROGRAM, OUT_DTYPE=tl_out_dtype, ) diff --git a/tests/gpu/torch/quantization/test_tensor_quant_cuda.py b/tests/gpu/torch/quantization/test_tensor_quant_cuda.py index c1cbd6a09fa..3ab5edd0db4 100644 --- a/tests/gpu/torch/quantization/test_tensor_quant_cuda.py +++ b/tests/gpu/torch/quantization/test_tensor_quant_cuda.py @@ -219,6 +219,15 @@ def test_zero_block_scale_zeroes_block(self): assert torch.equal(output, torch.zeros_like(output)) + @pytest.mark.skipif(not triton_kernel.IS_AVAILABLE, reason="triton kernel is not available") + def test_static_blockwise_fp4_over_int32_numel(self): + x = torch.randn(2**31 + 2**20, device="cuda", dtype=torch.bfloat16) + amax = x.view(-1, 16).abs().amax(-1).float() + output = triton_kernel.static_blockwise_fp4_fake_quant(x, amax, amax.max()) + tail = triton_kernel.static_blockwise_fp4_fake_quant(x[-4096:], amax[-256:], amax.max()) + + assert torch.equal(output[-4096:], tail) + @pytest.mark.skipif(get_cuda_ext_mx() is None, reason="cuda_ext_mx is not available") @pytest.mark.parametrize( "set_torch_dtype", [torch.float, torch.float16, torch.bfloat16], indirect=True From 8ef5d2f302712f21e9d44b7e086225a2d7683cbe Mon Sep 17 00:00:00 2001 From: Sylvester Kaczmarek <16242628+sylvesterkaczmarek@users.noreply.github.com> Date: Mon, 26 Oct 2026 20:48:59 +0000 Subject: [PATCH 33/43] Preserve histogram counts during entropy calibration (#2655) Fixes #2654. Make the entropy calibration helper copy its input histogram before applying its existing first-bin adjustment. A NumPy slice is a view, so the previous implementation modified the collector's persistent counts. Subsequent percentile/MSE queries and continued collection could then use the altered distribution. The regression uses the public `HistogramCalibrator` with both histogram backends. It checks unchanged counts and edges, repeated entropy queries, unchanged percentile thresholds, and equivalence with an untouched reference collector after another batch. ```sh python -m pytest -q tests/unit/torch/quantization/test_calibrator.py tests/unit/torch/quantization/test_tensor_quantizer_cpu.py tests/unit/torch/quantization/test_print.py ``` 79 passed, 1 existing skip on CPU. The focused regression on unchanged upstream gives 1 failure for NumPy histograms and 1 pass for PyTorch histograms. Repository pre-commit checks pass for the changed files. No new dependencies or public API changes. CUDA execution was not run. This is independent of #2369: its zero-prefix changes affect collection, while this patch changes only the entropy calculation. The changelog is updated. * **Bug Fixes** * Fixed entropy-based calibration so repeated calculations preserve collected histogram counts and bin edges. Percentile and entropy results remain consistent, and collecting additional data after a calculation now produces calibration results consistent with uninterrupted collection. The fix applies to both NumPy- and Torch-based histograms, ensuring calibration state remains stable across calculation and collection steps. Signed-off-by: Sylvester Kaczmarek <16242628+sylvesterkaczmarek@users.noreply.github.com> Signed-off-by: Rohan Joshi --- CHANGELOG.rst | 1 + .../torch/quantization/calib/histogram.py | 2 +- .../torch/quantization/test_calibrator.py | 26 +++++++++++++++++++ 3 files changed, 28 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 7f03618f750..0f91e5e450a 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -199,6 +199,7 @@ Changelog - Fix ``fsdp2_aware_weight_update`` masking an error raised during its unshard (e.g. a CUDA OOM) with a spurious ``UnboundLocalError``. - Nested ``only_student_forward()`` and ``only_teacher_forward()`` contexts now preserve the enclosing execution mode, including when an inner context is disabled or raises an exception. - Fix an illegal memory access in static NVFP4 fake quantization for tensors with more than 2^31 elements, such as stacked MoE expert weights. The kernel is also about 14x faster. +- Preserve accumulated histogram counts when computing an entropy calibration threshold with NumPy histograms. 0.47.0 (2026-09-23) ^^^^^^^^^^^^^^^^^^^ diff --git a/modelopt/torch/quantization/calib/histogram.py b/modelopt/torch/quantization/calib/histogram.py index e27a5471e30..73ebc327892 100644 --- a/modelopt/torch/quantization/calib/histogram.py +++ b/modelopt/torch/quantization/calib/histogram.py @@ -218,7 +218,7 @@ def _normalize_distr(distr): if summ != 0: distr = distr / summ - bins = calib_hist[:] + bins = calib_hist.copy() bins[0] = bins[1] total_data = np.sum(bins) diff --git a/tests/unit/torch/quantization/test_calibrator.py b/tests/unit/torch/quantization/test_calibrator.py index 9f7d77ce6f7..2d489b14543 100644 --- a/tests/unit/torch/quantization/test_calibrator.py +++ b/tests/unit/torch/quantization/test_calibrator.py @@ -89,6 +89,32 @@ def test_track_amax_raises(self): class TestHistogramCalibrator: + @pytest.mark.parametrize("torch_hist", [False, True]) + def test_entropy_preserves_collected_histogram(self, torch_hist): + calibrator = calib.HistogramCalibrator(4, None, False, num_bins=32, torch_hist=torch_hist) + reference = calib.HistogramCalibrator(4, None, False, num_bins=32, torch_hist=torch_hist) + first_batch = torch.cat((torch.zeros(100), torch.arange(1, 33).float())) + calibrator.collect(first_batch) + reference.collect(first_batch) + expected_hist = np.asarray(calibrator._calib_hist).copy() + expected_edges = np.asarray(calibrator._calib_bin_edges).copy() + percentile = calibrator.compute_amax("percentile", percentile=50) + + first_entropy = calibrator.compute_amax("entropy", start_bin=16) + second_entropy = calibrator.compute_amax("entropy", start_bin=16) + + np.testing.assert_array_equal(calibrator._calib_hist, expected_hist) + np.testing.assert_array_equal(calibrator._calib_bin_edges, expected_edges) + assert calibrator.compute_amax("percentile", percentile=50) == percentile + assert first_entropy == second_entropy + next_batch = torch.tensor([0.0, 2.0, 4.0, 64.0]) + calibrator.collect(next_batch) + reference.collect(next_batch) + np.testing.assert_array_equal(calibrator._calib_hist, reference._calib_hist) + assert calibrator.compute_amax("percentile", percentile=50) == reference.compute_amax( + "percentile", percentile=50 + ) + @pytest.mark.skip(reason="TODO: Fix assertions in test_grow") def test_grow(self, verbose): x_1 = torch.tensor([0, 255, 255, 255, 255, 255]) From f5e8d0bf28f246ebf17dc27f738aa9d1096770f2 Mon Sep 17 00:00:00 2001 From: Javier De Jesus Date: Mon, 26 Oct 2026 22:20:12 +0100 Subject: [PATCH 34/43] Fix q/k/v quantizer wiring in register_attention_for_kv_quant (#2618) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Type of change: Bug fix `register_attention_for_kv_quant` collected the two matmul nodes with breadth-first `ast.walk` and patched them by list index, which assumes the `attn @ v` node comes first. Within one method that only holds when the score matmul is nested deeper in the AST than the output matmul. In a flat layout (`scores = q @ k.T`, `probs = softmax(scores)`, `probs @ v`) the score matmul comes first, so `v_bmm_quantizer` was applied to k, `k_bmm_quantizer` to v and `q_bmm_quantizer` to the softmax output. The cause is the node order, not the transpose wrapper the issue's triage points at. The nodes are now collected in source order (children before parents), which puts the score matmul first when both matmuls are in one method, and the two index pairs are swapped. The `transpose` wrapper on k and the quantizer-to-operand design are unchanged. `baddbmm` is no longer matched. `patch` wires `args[0]` and `args[1]`, but `torch.baddbmm(input, batch1, batch2)` takes the additive input first, so the q quantizer landed on the bias and k was never quantized (`GPTNeoXJapaneseAttention`), and keyword calls raised `IndexError` (`BloomAttention`). Those classes now return `False` and get the existing warning to register a custom quantized attention class. Flat-layout attention classes will calibrate differently after this change. For them the k and v scales that export reads (`get_kv_cache_scaling_factor`) were swapped, so models calibrated or exported with the old wiring need to be recalibrated and re-exported. In the `modeling_.py` modules of transformers 5.14.1 that have no `ALL_ATTENTION_FUNCTIONS`, 72 of the 178 `*Attention` classes other than `T5Attention` are accepted by this function; the generated code changes for 65 and is byte-identical for 7. No API change. - New `test_kv_quant_operand_wiring` in `tests/unit/torch/quantization/plugins/test_attention_quant.py`: flat and nested layouts for `torch.matmul`, `torch.bmm` and `@`, with q, k and v of different scales, asserting that each quantizer's `amax` equals the abs-max of its own operand. On the base commit the three flat cases fail and the three nested cases pass; with the fix all six pass. - New `test_kv_quant_rejects_baddbmm`: positional and keyword `baddbmm` scores. Before the change the first registers with the wrong wiring and the second raises `IndexError`; with it both return `False`. On transformers 5.18, `BloomAttention` goes from `IndexError` to `False` and `GPTNeoXJapaneseAttention` from `True` to `False`; no other class that reaches this fallback calls `baddbmm`. - 17 real classes (15 changed, 2 unchanged): the operands of the two matmul calls recorded from the original forward and compared with the calibrated `amax`. On the base the 15 changed ones are miswired (`BarkSelfAttention` and `MptAttention`, both nested, are correct); with the fix all 17 are correct. For example `MegatronBertSelfAttention`: | | q | k | v | |---|---|---|---| | abs-max of the operand | 2.3635 | 2.0925 | 2.2565 | | amax on base | 0.3375 | 2.2565 | 2.0925 | | amax with the fix | 2.3635 | 2.0925 | 2.2565 | - `pytest tests/unit`: 4375 passed, 28 skipped, 13 xfailed. - `pre-commit run --all-files`: every hook passes except `generate-arguments-md`, which cannot import the Hugging Face plugin on torch 2.14 here (CI pins `torch<2.14` for that job); running its generator directly leaves `examples/llm_qat/ARGUMENTS.md` unchanged. - Not run: `tests/gpu` (no GPU). `tests/gpu/torch/quantization/test_hadamard.py` is the only GPU test that calls this function directly and it uses an SDPA-only class, which this change does not touch. 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?: ❌ Flat-layout attention classes calibrate differently, as described above. A class that defines the score matmul in a helper placed after `forward` can change from correct to wrong (see Additional Information). `BloomAttention` and `GPTNeoXJapaneseAttention` no longer get KV-cache quantizers from this fallback. - 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)?: ✅ Fixes #1064. Not covered: real `baddbmm` support (matrix operands at `args[1]` and `args[2]`, keyword arguments), and classes whose two matmuls sit in different methods, which are ordered by source position instead of call order. None of the 71 classes with two matmuls in the `modeling_.py` modules of transformers 5.14.1 does this, but a helper defined after `forward` that was wired correctly by accident before would now be miswired. * **Bug Fixes** * Corrected query, key, and value quantizer assignments for supported attention layouts using fallback quantization, including separated and nested matrix multiplication, batch matrix multiplication, and `@` operations. Quantizers now track their corresponding inputs. * Attention implementations using `torch.baddbmm` are no longer patched by the fallback and receive a warning to register a custom quantized attention class. Models quantized with the affected fallback must be recalibrated and re-exported. The SDPA quantization path is unchanged. --------- Signed-off-by: Javier de Jesus Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Signed-off-by: javierdejesusda Co-authored-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Signed-off-by: Rohan Joshi --- CHANGELOG.rst | 1 + .../torch/quantization/plugins/attention.py | 20 +++-- .../plugins/test_attention_quant.py | 79 +++++++++++++++++++ 3 files changed, 94 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 0f91e5e450a..3159f11b585 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -200,6 +200,7 @@ Changelog - Nested ``only_student_forward()`` and ``only_teacher_forward()`` contexts now preserve the enclosing execution mode, including when an inner context is disabled or raises an exception. - Fix an illegal memory access in static NVFP4 fake quantization for tensors with more than 2^31 elements, such as stacked MoE expert weights. The kernel is also about 14x faster. - Preserve accumulated histogram counts when computing an entropy calibration threshold with NumPy histograms. +- Fix KV-cache quantization through ``register_attention_for_kv_quant``, the fallback for Hugging Face attention classes that do not use the attention interface, attaching the ``q``, ``k`` and ``v`` bmm quantizers to the wrong tensors for most attention layouts. Recalibrate and re-export models quantized through this fallback. Attention classes that compute their scores with ``torch.baddbmm`` (e.g. ``BloomAttention``, ``GPTNeoXJapaneseAttention``) are no longer patched, since their operands were always wired wrong (or raised ``IndexError``), and now get the existing warning to register a custom quantized attention class. 0.47.0 (2026-09-23) ^^^^^^^^^^^^^^^^^^^ diff --git a/modelopt/torch/quantization/plugins/attention.py b/modelopt/torch/quantization/plugins/attention.py index 2113edea8a7..e4843a3a2ec 100644 --- a/modelopt/torch/quantization/plugins/attention.py +++ b/modelopt/torch/quantization/plugins/attention.py @@ -52,7 +52,9 @@ def register_attention_for_kv_quant(attention_cls: type) -> bool: model_module = inspect.getmodule(attention_cls) head = ast.parse(source_code) - bmm_ops = ("matmul", "bmm", "baddbmm") + # No baddbmm: patch() wires args[0] and args[1], which miss its batch1 and batch2 operands + # (torch.baddbmm passes the additive input first, and Bloom passes them as keywords) + bmm_ops = ("matmul", "bmm") sdpa_ops = ("scaled_dot_product_attention",) def is_bmm(node): @@ -72,6 +74,12 @@ def is_sdpa(node): def is_bin_matmul(node): return isinstance(node, ast.BinOp) and isinstance(node.op, ast.MatMult) + def walk_post_order(node): + # Source order, children first; expects the q/k score matmul before the matmul consuming it + for child in ast.iter_child_nodes(node): + yield from walk_post_order(child) + yield node + def patch(node, quantizer_names, transpose=False): for index, quantizer_name in enumerate(quantizer_names): if quantizer_name is None: @@ -165,7 +173,7 @@ def patch_binop(node, quantizer_names, transpose=False): bmm_nodes = [] sdpa_nodes = [] bin_matmul_nodes = [] - for node in ast.walk(head): + for node in walk_post_order(head): if is_bmm(node): bmm_nodes.append(node) if is_sdpa(node): @@ -186,16 +194,16 @@ def patch_binop(node, quantizer_names, transpose=False): # after transpose, the quantization will be per-token, i.e., # self.k_bmm_quantizer(key_states.transpose(-1, -2).transpose(-1, -2)).transpose(-1, -2) # removing the additional transpose is doable but not trivial - patch(bmm_nodes[0], quantizer_names=(None, "v_bmm_quantizer")) - patch(bmm_nodes[1], quantizer_names=("q_bmm_quantizer", "k_bmm_quantizer"), transpose=True) + patch(bmm_nodes[0], quantizer_names=("q_bmm_quantizer", "k_bmm_quantizer"), transpose=True) + patch(bmm_nodes[1], quantizer_names=(None, "v_bmm_quantizer")) print("Patching 2 BMM/Matmul operators with quantizers") if len(bin_matmul_nodes) == 2: patch_binop( - bin_matmul_nodes[1], + bin_matmul_nodes[0], quantizer_names=("q_bmm_quantizer", "k_bmm_quantizer"), transpose=True, ) - patch_binop(bin_matmul_nodes[0], quantizer_names=(None, "v_bmm_quantizer")) + patch_binop(bin_matmul_nodes[1], quantizer_names=(None, "v_bmm_quantizer")) print("Patching 2 @ operators with quantizers") if len(sdpa_nodes) == 1: diff --git a/tests/unit/torch/quantization/plugins/test_attention_quant.py b/tests/unit/torch/quantization/plugins/test_attention_quant.py index 22a86971a6d..ff9e9404ea4 100644 --- a/tests/unit/torch/quantization/plugins/test_attention_quant.py +++ b/tests/unit/torch/quantization/plugins/test_attention_quant.py @@ -52,6 +52,54 @@ def forward(self, hidden_states, **kwargs): return F.scaled_dot_product_attention(q, k, v), None +class FlatMatmulAttention(nn.Module): + def forward(self, q, k, v): + scores = torch.matmul(q, k.transpose(-2, -1)) + probs = torch.softmax(scores, dim=-1) + return torch.matmul(probs, v) + + +class FlatBMMAttention(nn.Module): + def forward(self, q, k, v): + scores = torch.bmm(q, k.transpose(-2, -1)) + probs = torch.softmax(scores, dim=-1) + return torch.bmm(probs, v) + + +class FlatBinMatmulAttention(nn.Module): + def forward(self, q, k, v): + scores = q @ k.transpose(-2, -1) + probs = torch.softmax(scores, dim=-1) + return probs @ v + + +class NestedMatmulAttention(nn.Module): + def forward(self, q, k, v): + return torch.matmul(torch.softmax(torch.matmul(q, k.transpose(-2, -1)), dim=-1), v) + + +class NestedBMMAttention(nn.Module): + def forward(self, q, k, v): + return torch.bmm(torch.softmax(torch.bmm(q, k.transpose(-2, -1)), dim=-1), v) + + +class NestedBinMatmulAttention(nn.Module): + def forward(self, q, k, v): + return torch.softmax(q @ k.transpose(-2, -1), dim=-1) @ v + + +class BAddBMMAttention(nn.Module): + def forward(self, q, k, v, bias): + scores = torch.baddbmm(bias, q, k.transpose(-2, -1)) + return torch.bmm(torch.softmax(scores, dim=-1), v) + + +class BAddBMMKeywordAttention(nn.Module): + def forward(self, q, k, v, bias): + scores = torch.baddbmm(input=bias, batch1=q, batch2=k.transpose(-2, -1)) + return torch.bmm(torch.softmax(scores, dim=-1), v) + + kv_cache_config = { "quant_cfg": [ {"quantizer_name": "*[kv]_bmm_quantizer", "cfg": {"num_bits": 4}, "enable": True}, @@ -115,6 +163,37 @@ def test_kv_quant_hf(model_getter, attn_cls): mtq.unregister(attn_cls) +@pytest.mark.parametrize( + "attn_cls", + [ + FlatMatmulAttention, + FlatBMMAttention, + FlatBinMatmulAttention, + NestedMatmulAttention, + NestedBMMAttention, + NestedBinMatmulAttention, + ], +) +def test_kv_quant_operand_wiring(attn_cls): + """Each q/k/v_bmm_quantizer must calibrate on its own operand, whatever the layout.""" + q, k, v = torch.randn(1, 4, 8), 10 * torch.randn(1, 4, 8), 100 * torch.randn(1, 4, 8) + assert mtq.plugins.register_attention_for_kv_quant(attn_cls) + try: + model = attn_cls() + mtq.quantize(model, kv_cache_config, lambda model: model(q, k, v)) + for name, operand in zip("qkv", (q, k, v)): + quantizer = getattr(model, f"{name}_bmm_quantizer") + assert quantizer.amax == operand.abs().max(), f"{name}_bmm_quantizer saw another tensor" + finally: + mtq.unregister(attn_cls) + + +@pytest.mark.parametrize("attn_cls", [BAddBMMAttention, BAddBMMKeywordAttention]) +def test_kv_quant_rejects_baddbmm(attn_cls): + """baddbmm takes the additive input first, so its operands cannot be wired by position.""" + assert not mtq.plugins.register_attention_for_kv_quant(attn_cls) + + def test_kv_quant_bert(): """Test KV cache quantization on BERT model with decorated attention.""" model_test = get_tiny_bert() From bad3cb329a3c2f12abed7f5fe2930cb6a6527460 Mon Sep 17 00:00:00 2001 From: Zhiyu Date: Mon, 26 Oct 2026 16:16:23 -0700 Subject: [PATCH 35/43] Emit VisualGen-compatible sparse_attention_config for diffusers skip-softmax export (#2868) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Type of change: Bug fix The diffusers skip-softmax HF export wrote the LLM phase-keyed schema (`threshold_scale_factor.prefill` and a `{phase: ratio}` `target_sparsity`). TRT-LLM VisualGen (PR #14687) expects `threshold_scale_factor.coefficients` and a scalar `target_sparsity`, so the exported Wan2.2 checkpoint fails to load (`float() ... not 'dict'`). NVBug 6927804. - `export_sparse_attention_config(model, *, for_diffusion=False)`: with `for_diffusion=True`, emit one `coefficients` block and a scalar `target_sparsity`. The LLM path and the vLLM calibration plugin are unchanged. `_export_diffusers_checkpoint` passes `for_diffusion=True`. - Fix the READMEs, which referenced the removed `resolve_for_target_sparsity`. Based on #1816 (closed, approved). Deliberately dropped its `initial_disabled_steps` -> `disabled_until_timestep` rename: `initial_disabled_steps` shipped in 0.46 and the vLLM calibration path now depends on it, so the rename is a separate, breaking change. ```bash python examples/diffusers/sparsity/wan22_skip_softmax.py --calibrate --target-sparsity 0.5 --export-dir ckpt ... ``` - Extended `test_sparse_attention_conversion.py` (diffusion schema) and the Wan calibrated-export example test. - **Not run:** no torch environment was available where this was written, so pytest has not been executed; only pre-commit passed and the conversion logic was checked standalone. Relying on CI. No TRT-LLM VisualGen load was tested. - Is this change backward compatible?: ❌ Diffusers skip-softmax checkpoints are now emitted in the VisualGen schema (previous ones did not load in VisualGen); LLM export is 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?: ✅ - Did you get Claude approval on this PR?: ❌ NVBug 6927804; supersedes #1816. 🤖 Generated with [Claude Code](https://claude.com/claude-code) * **Bug Fixes** * Fixed skip-softmax Hugging Face exports for diffusion models to use the TRT-LLM VisualGen format, including grouped threshold coefficients and scalar target sparsity. * Existing affected checkpoints must be re-exported to apply the corrected format. * **Documentation** * Updated diffusion sparsity documentation to describe the exported checkpoint schema and its calibration and runtime-sparsity settings. --------- Signed-off-by: Zhiyu Cheng Co-authored-by: Claude Sonnet 5.5 Signed-off-by: Rohan Joshi --- CHANGELOG.rst | 1 + examples/diffusers/README.md | 2 +- examples/diffusers/sparsity/README.md | 7 +++--- modelopt/torch/export/unified_export_hf.py | 2 +- .../sparsity/attention_sparsity/conversion.py | 23 ++++++++++++++++++- .../diffusers/sparsity/test_sparsity.py | 5 ++-- .../test_sparse_attention_conversion.py | 9 ++++++++ 7 files changed, 40 insertions(+), 9 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 3159f11b585..4b3219096ca 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -201,6 +201,7 @@ Changelog - Fix an illegal memory access in static NVFP4 fake quantization for tensors with more than 2^31 elements, such as stacked MoE expert weights. The kernel is also about 14x faster. - Preserve accumulated histogram counts when computing an entropy calibration threshold with NumPy histograms. - Fix KV-cache quantization through ``register_attention_for_kv_quant``, the fallback for Hugging Face attention classes that do not use the attention interface, attaching the ``q``, ``k`` and ``v`` bmm quantizers to the wrong tensors for most attention layouts. Recalibrate and re-export models quantized through this fallback. Attention classes that compute their scores with ``torch.baddbmm`` (e.g. ``BloomAttention``, ``GPTNeoXJapaneseAttention``) are no longer patched, since their operands were always wired wrong (or raised ``IndexError``), and now get the existing warning to register a custom quantized attention class. +- Fix the diffusers skip-softmax HF export to emit the schema TensorRT-LLM VisualGen loads (``threshold_scale_factor.coefficients`` and a scalar ``target_sparsity``). Re-export checkpoints that failed to load. 0.47.0 (2026-09-23) ^^^^^^^^^^^^^^^^^^^ diff --git a/examples/diffusers/README.md b/examples/diffusers/README.md index 971b8045e3e..62f8e5f1df0 100644 --- a/examples/diffusers/README.md +++ b/examples/diffusers/README.md @@ -372,7 +372,7 @@ By following these steps, your PEFT LoRA model should be efficiently quantized u ## Sparse Attention (Skip-Softmax) -Skip-softmax sparse attention skips KV tiles whose attention scores are negligible during the softmax computation, reducing FLOPs without retraining. An exponential model (`scale_factor = a * exp(b * target_sparsity)`) is calibrated once, then the target sparsity can be adjusted at runtime without recalibration. Calibrated coefficients can be exported as a Hugging Face checkpoint (embedded in each component's `config.json` under `sparse_attention_config`) consumed directly by TRT-LLM's `SkipSoftmaxAttentionConfig.resolve_for_target_sparsity` — no extra conversion needed downstream. +Skip-softmax sparse attention skips KV tiles whose attention scores are negligible during the softmax computation, reducing FLOPs without retraining. An exponential model (`scale_factor = a * exp(b * target_sparsity)`) is calibrated once, then the target sparsity can be adjusted at runtime without recalibration. Calibrated coefficients can be exported as a Hugging Face checkpoint (embedded in each component's `config.json` under `sparse_attention_config`) in the schema TRT-LLM VisualGen loads (`threshold_scale_factor.coefficients` plus a scalar `target_sparsity`) — no extra conversion needed downstream. ### Getting Started diff --git a/examples/diffusers/sparsity/README.md b/examples/diffusers/sparsity/README.md index e102f1b6d5d..bc1188d2ea0 100644 --- a/examples/diffusers/sparsity/README.md +++ b/examples/diffusers/sparsity/README.md @@ -65,10 +65,9 @@ python wan22_skip_softmax.py \ `--export-dir` writes a Hugging Face checkpoint with the calibrated `threshold_scale_factor` block embedded in each component's `config.json` -(under the `sparse_attention_config` key). TRT-LLM's -`SkipSoftmaxAttentionConfig.resolve_for_target_sparsity` reads the -`(a, b)` directly via `coeffs['a'] * math.exp(coeffs['b'] * sparsity)` — -no extra conversion needed downstream. +(under the `sparse_attention_config` key), in the schema TRT-LLM VisualGen +loads: `(a, b)` under `threshold_scale_factor.coefficients` and a scalar +`target_sparsity`. No extra conversion is needed downstream. ## Threshold Modes diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 9ac7e26eb2c..842ce4dacce 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -1421,7 +1421,7 @@ def _export_diffusers_checkpoint( # Update config.json with sparse attention info (both quantized and non-quantized) if export_sparse_attention_config is not None: - sparse_attn_config = export_sparse_attention_config(component) + sparse_attn_config = export_sparse_attention_config(component, for_diffusion=True) if sparse_attn_config is not None: config_path = component_export_dir / "config.json" if config_path.exists(): diff --git a/modelopt/torch/sparsity/attention_sparsity/conversion.py b/modelopt/torch/sparsity/attention_sparsity/conversion.py index aab3cbc74c7..38944517740 100644 --- a/modelopt/torch/sparsity/attention_sparsity/conversion.py +++ b/modelopt/torch/sparsity/attention_sparsity/conversion.py @@ -372,7 +372,9 @@ def _get_sparse_softmax_export_config(module: SparseAttentionModule) -> dict[str return sparse_softmax_config -def export_sparse_attention_config(model: nn.Module) -> dict[str, Any] | None: +def export_sparse_attention_config( + model: nn.Module, *, for_diffusion: bool = False +) -> dict[str, Any] | None: """Extract sparse attention config for export to config.json. Extracts calibrated skip-softmax parameters and N:M sparse-softmax metadata @@ -384,6 +386,10 @@ def export_sparse_attention_config(model: nn.Module) -> dict[str, Any] | None: Args: model: Model with sparse attention applied + for_diffusion: Emit the single-phase schema read by the TRT-LLM VisualGen + runtime: one ``threshold_scale_factor.coefficients`` block and a scalar + ``target_sparsity``. The default is the LLM schema, with per-phase + (``prefill``/``decode``) blocks and a ``{phase: ratio}`` dict. Returns: Dictionary with sparse attention config for HuggingFace config.json export. @@ -473,6 +479,21 @@ def export_sparse_attention_config(model: nn.Module) -> dict[str, Any] | None: skip_group["threshold_scale_factor"] = export_threshold_scale_factor(calibration_params) if target_sparse_ratio is not None: skip_group["target_sparsity"] = target_sparse_ratio + if for_diffusion: + # Diffusion calibrates a single (prefill) phase; VisualGen rejects phase-keyed dicts. + tsf = skip_group["threshold_scale_factor"] + phases = {p: tsf.pop(p) for p in ("prefill", "decode") if p in tsf} + if phases: + tsf["coefficients"] = phases.get("prefill") or next(iter(phases.values())) + ratio = skip_group.get("target_sparsity") + if isinstance(ratio, dict): + scalar = ( + ratio["prefill"] if "prefill" in ratio else next(iter(ratio.values()), None) + ) + if scalar is None: + del skip_group["target_sparsity"] # VisualGen needs a scalar + else: + skip_group["target_sparsity"] = scalar config_groups[f"group_{group_idx}"] = skip_group group_idx += 1 diff --git a/tests/examples/diffusers/sparsity/test_sparsity.py b/tests/examples/diffusers/sparsity/test_sparsity.py index bca94e3dafb..4067251cc01 100644 --- a/tests/examples/diffusers/sparsity/test_sparsity.py +++ b/tests/examples/diffusers/sparsity/test_sparsity.py @@ -229,8 +229,9 @@ def test_wan22_calibrated_export(tiny_wan22_path, tmp_path): # threshold_scale_factor lives inside the skip_softmax group. tsf = group_0["threshold_scale_factor"] assert tsf["formula"] == "a * exp(b * target_sparsity)" - assert tsf["prefill"]["a"] == pytest.approx(math.exp(test_log_a)) - assert tsf["prefill"]["b"] == pytest.approx(test_b) + assert tsf["coefficients"]["a"] == pytest.approx(math.exp(test_log_a)) + assert tsf["coefficients"]["b"] == pytest.approx(test_b) + assert isinstance(group_0["target_sparsity"], float) # scalar, as VisualGen expects # Calibrated mode — no raw_threshold. assert "raw_threshold" not in group_0 diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attention_conversion.py b/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attention_conversion.py index 83f6f2c44ad..e70f2b7de65 100644 --- a/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attention_conversion.py +++ b/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attention_conversion.py @@ -441,6 +441,15 @@ def test_exports_when_calibration_present(self): assert group_0["target_sparsity"] == {"prefill": 0.4, "decode": 0.6} assert out["producer"]["name"] == "modelopt" + # Diffusion (VisualGen) schema: one coefficient block, scalar target_sparsity. + group_0 = export_sparse_attention_config(model, for_diffusion=True)["config_groups"][ + "group_0" + ] + tsf = group_0["threshold_scale_factor"] + assert tsf["coefficients"] == {"a": 3.14, "b": 7.5} + assert "prefill" not in tsf and "decode" not in tsf + assert group_0["target_sparsity"] == 0.4 + def test_exports_sparse_softmax_metadata(self): """N:M sparse-softmax params are exported without calibration metadata.""" model = SimpleAttentionModel() From c342882fb3b2b86417f9b3430e5f67b678c0e4a2 Mon Sep 17 00:00:00 2001 From: Zhehao Hu <69692533+Zhehao-Hu@users.noreply.github.com> Date: Mon, 26 Oct 2026 17:18:16 -0700 Subject: [PATCH 36/43] [NVBug 6980436] Fix per-step averaging in diffusion min-mean calibration (#3032) 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 Diffusers `min-mean` calibrator to divide each denoising step’s accumulated activation maxima by its observation count instead of the number of denoising steps. Previously, repeating identical calibration data changed the resulting quantization scale: four-step data `[2, 4, 6, 8]` produced `amax=0.5` for one cycle and `amax=1.5` for three identical cycles, instead of `2.0` in both cases. Track observation counts per step so averaging also handles partial final cycles. Other collection methods retain their existing calculations. ### Usage No API or configuration changes. Applies when using `--collect-method min-mean`. ### Testing - Added a parameterized CPU regression test covering one cycle, repeated identical cycles, varying observations, and a partial final cycle. - All four cases fail on the original implementation and pass with the fix. - All applicable pre-commit checks pass. - Full-model image/video quality testing was not performed. ### 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 - Did you get Claude approval on this PR?: N/A - Not yet run ### Additional Information Fixes NVBug 6980436. ## Summary by CodeRabbit * **Bug Fixes** * Improved calibration results for min-mean mode when steps contain different numbers of samples, including partial cycles. * **Tests** * Added coverage for calibration results across single, repeated, varying, and partial cycles. Signed-off-by: Zhehao Hu Signed-off-by: Rohan Joshi --- .../quantization/calib/plugin_calib.py | 4 +- tests/unit/examples/test_plugin_calib.py | 38 +++++++++++++++++++ 2 files changed, 41 insertions(+), 1 deletion(-) create mode 100644 tests/unit/examples/test_plugin_calib.py diff --git a/examples/diffusers/quantization/calib/plugin_calib.py b/examples/diffusers/quantization/calib/plugin_calib.py index b8d82734f21..7490ad15ad6 100644 --- a/examples/diffusers/quantization/calib/plugin_calib.py +++ b/examples/diffusers/quantization/calib/plugin_calib.py @@ -26,6 +26,7 @@ def __init__(self, num_bits=8, axis=None, unsigned=False, track_amax=False, **kw self.total_step = kwargs["total_step"] self.collect_method = kwargs["collect_method"] self.data = {} + self._counts = {} self.i = 0 def collect(self, x): @@ -49,6 +50,7 @@ def collect(self, x): self.data[_cur_step] = torch.max(self.data[_cur_step], local_amax) else: self.data[_cur_step] += local_amax + self._counts[_cur_step] = self._counts.get(_cur_step, 0) + 1 if self._track_amax: raise NotImplementedError self.i += 1 @@ -57,7 +59,7 @@ def compute_amax(self): """Return the absolute max of all tensors collected.""" up_lim = int(self.total_step * self.percentile) if self.collect_method == "min-mean": - amaxs_values = [self.data[i] / self.total_step for i in range(up_lim)] + amaxs_values = [self.data[i] / self._counts[i] for i in range(up_lim)] else: amaxs_values = [self.data[i] for i in range(up_lim)] if self.collect_method == "mean-max": diff --git a/tests/unit/examples/test_plugin_calib.py b/tests/unit/examples/test_plugin_calib.py new file mode 100644 index 00000000000..01eee88e4ec --- /dev/null +++ b/tests/unit/examples/test_plugin_calib.py @@ -0,0 +1,38 @@ +# 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 pytest +import torch + +from examples.diffusers.quantization.calib.plugin_calib import PercentileCalibrator + + +@pytest.mark.parametrize( + ("observations", "expected_amax"), + [ + pytest.param([2.0, 4.0, 6.0, 8.0], 2.0, id="one-cycle"), + pytest.param([2.0, 4.0, 6.0, 8.0] * 3, 2.0, id="repeated-cycles"), + pytest.param([2.0, 4.0, 6.0, 8.0, 6.0, 8.0, 10.0, 12.0], 4.0, id="varying-cycles"), + pytest.param([2.0, 8.0, 6.0, 10.0, 6.0], 4.0, id="partial-final-cycle"), + ], +) +def test_min_mean_averages_observations_per_step(observations, expected_amax): + calibrator = PercentileCalibrator( + num_bits=8, axis=None, percentile=1.0, total_step=4, collect_method="min-mean" + ) + for value in observations: + calibrator.collect(torch.tensor([value])) + + torch.testing.assert_close(calibrator.compute_amax(), torch.tensor([expected_amax])) From d9069e5f98b8bc268acdadd786417e45f0667837 Mon Sep 17 00:00:00 2001 From: Zhehao Hu <69692533+Zhehao-Hu@users.noreply.github.com> Date: Mon, 26 Oct 2026 17:18:37 -0700 Subject: [PATCH 37/43] [NVBug 6982527] Fix duplicate final checkpoint save in LTX distillation (#3027) 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 When the final LTX distillation step coincides with a checkpoint interval, the trainer saves the same step twice. The second save fails while renaming its temporary directory onto the existing nonempty checkpoint directory, preventing final inference export. Time-limit exits can encounter the same collision. Reuse the existing `saved_this_step` flag to skip the final training-state save when the last step was already checkpointed. Initialize the flag before the loop and set it after a time-limit save. Final inference export still runs, and steps not already checkpointed still receive a final save. ### Usage No configuration changes required. For example, `optimization.steps=1` with `checkpoints.interval=1` now saves the training state once and proceeds to inference export. ### Testing ```bash pytest tests/examples/diffusers/test_distillation_trainer.py -q ``` - The checked-in parametrized regression test covers interval-aligned completion, a nonaligned final step, and time-limit exits with and without an interval save, each with gradient accumulation of one and two. - All **8 cases passed locally on CPU** (pytest-reported duration: **1.87 seconds**). The test imports the complete example and exercises its real training loop, Accelerate training-state save, checkpoint retention, and safetensors inference export using a tiny model. It verifies exactly one training-state save per expected step, one inference export, and exported weights matching the trained model. - With the production fix reverted, **6 cases fail with the original directory collision**; the two nonaligned controls still pass. - All applicable pre-commit checks passed. The test substitutes the optional LTX dependency/setup and UI; it requires no LTX model downloads. Full LTX model execution and GPU/FSDP validation were not run. ### 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?: ✅ — eight parametrized checkpoint-finalization cases in `tests/examples/diffusers/test_distillation_trainer.py`. - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A - Did you get Claude approval on this PR?: N/A - Not yet run ### Additional Information Fixes NVBug 6982527. ## Summary by CodeRabbit - **Bug Fixes** - Training now avoids saving duplicate checkpoints when an interval or time-limit save already covers the final step. - A final training-state checkpoint is still saved when no earlier save covers the last step, helping ensure the latest training progress is retained. Signed-off-by: Zhehao Hu Signed-off-by: Rohan Joshi --- .../distillation/distillation_trainer.py | 7 +- .../diffusers/test_distillation_trainer.py | 161 ++++++++++++++++++ 2 files changed, 166 insertions(+), 2 deletions(-) create mode 100644 tests/examples/diffusers/test_distillation_trainer.py diff --git a/examples/diffusers/distillation/distillation_trainer.py b/examples/diffusers/distillation/distillation_trainer.py index 6c173ef0f79..c1b4cc60e9a 100644 --- a/examples/diffusers/distillation/distillation_trainer.py +++ b/examples/diffusers/distillation/distillation_trainer.py @@ -1504,6 +1504,7 @@ def train( # Calculate how many raw steps to skip and how many to run total_raw_steps = cfg.optimization.steps * cfg.optimization.gradient_accumulation_steps skip_raw_steps = resume_step * cfg.optimization.gradient_accumulation_steps + saved_this_step = False with progress: # Initial validation before training starts (skip if resuming) @@ -1616,6 +1617,7 @@ def train( ) if not saved_this_step: self._save_training_state() + saved_this_step = True # Break out of the training loop; post-loop code # will collect stats and return. break @@ -1700,8 +1702,9 @@ def train( global_batch_size=cfg.optimization.batch_size * self._accelerator.num_processes, ) - # Save final training state (for potential resume) - self._save_training_state() + # Interval and time-limit saves may already have checkpointed the final step. + if not saved_this_step: + self._save_training_state() # Save inference-ready model weights (standalone safetensors file) saved_path = self._save_checkpoint() diff --git a/tests/examples/diffusers/test_distillation_trainer.py b/tests/examples/diffusers/test_distillation_trainer.py new file mode 100644 index 00000000000..e9753469799 --- /dev/null +++ b/tests/examples/diffusers/test_distillation_trainer.py @@ -0,0 +1,161 @@ +# 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. + +"""Exercise LTX checkpoint finalization without LTX model assets or a GPU.""" + +import importlib.util +import itertools +import json +import logging +import sys +from contextlib import nullcontext +from pathlib import Path +from types import ModuleType, SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch +from accelerate import Accelerator +from pydantic import BaseModel +from safetensors.torch import load_file + + +@pytest.fixture +def trainer_module(monkeypatch): + # LTX is an optional, separately licensed dependency. Substitute its model setup + # and UI, but import the complete example and execute its train/save/export methods. + progress = nullcontext() + progress.update_training = lambda **kwargs: None + for name, attributes in { + "ltx_trainer": {"logger": logging.getLogger(__name__)}, + "ltx_trainer.config": {"ConfigBaseModel": BaseModel, "LtxTrainerConfig": BaseModel}, + "ltx_trainer.model_loader": {"load_transformer": None}, + "ltx_trainer.trainer": { + "IS_MAIN_PROCESS": True, + "LtxvTrainer": object, + "TrainingStats": SimpleNamespace, + }, + "ltx_trainer.gpu_utils": {"get_gpu_memory_gb": lambda device: 0}, + "ltx_trainer.hf_hub_utils": {"push_to_hub": None}, + "ltx_trainer.progress": {"TrainingProgress": lambda **kwargs: progress}, + }.items(): + module = ModuleType(name) + module.__dict__.update(attributes) + monkeypatch.setitem(sys.modules, name, module) + + path = ( + Path(__file__).resolve().parents[3] + / "examples/diffusers/distillation/distillation_trainer.py" + ) + spec = importlib.util.spec_from_file_location("ltx_checkpoint_example", path) + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, spec.name, module) + spec.loader.exec_module(module) + monkeypatch.setattr(module, "is_global_rank0", lambda: True) + # Deterministically cross the one-minute deadline during the first optimizer step. + clock = itertools.count(step=60) + monkeypatch.setattr(module, "time", SimpleNamespace(time=lambda: next(clock))) + return module + + +@pytest.mark.parametrize("grad_accum", [1, 2]) +@pytest.mark.parametrize( + ("steps", "interval", "deadline", "saved_steps"), + [ + pytest.param(2, 1, None, [1, 2], id="interval-aligned"), + pytest.param(3, 2, None, [2, 3], id="nonaligned-final"), + pytest.param(3, 1, 1, [1], id="time-limit-with-interval"), + pytest.param(3, 2, 1, [1], id="time-limit-without-interval"), + ], +) +def test_train_final_checkpoint( + trainer_module, tmp_path, grad_accum, steps, interval, deadline, saved_steps +): + class TinyTrainer(trainer_module.DistillationTrainer): + def __init__(self): + self._accelerator = Accelerator(cpu=True, gradient_accumulation_steps=grad_accum) + self._config = SimpleNamespace( + seed=42, + output_dir=str(tmp_path), + model=SimpleNamespace(training_mode="full"), + optimization=SimpleNamespace( + steps=steps, + gradient_accumulation_steps=grad_accum, + max_grad_norm=1.0, + batch_size=1, + ), + checkpoints=SimpleNamespace(interval=interval, keep_last_n=1, precision="float32"), + validation=SimpleNamespace(interval=None, skip_initial_validation=True), + wandb=SimpleNamespace(log_validation_videos=False), + hub=SimpleNamespace(push_to_hub=False), + ) + self._distillation_config = trainer_module.DistillationConfig(must_save_by=deadline) + self._transformer = torch.nn.Linear(2, 1) + self._lr_scheduler = None + self._wandb_run = None + self._checkpoint_paths = [] + + def _init_optimizer(self): + optimizer = torch.optim.AdamW(self._transformer.parameters(), lr=0.01) + self._transformer, self._optimizer = self._accelerator.prepare( + self._transformer, optimizer + ) + self._trainable_params = list(self._transformer.parameters()) + + def _init_dataloader(self): + self._dataloader = [torch.ones(1, 2)] * (steps * grad_accum) + + def _training_step(self, batch): + return self._transformer(batch).square().mean() + + def _init_timestep_sampler(self): + pass + + def _save_config(self): + pass + + def _log_metrics(self, *args): + pass + + def _log_training_stats(self, *args): + pass + + trainer = TinyTrainer() + save_state = trainer._save_training_state + observed_steps = [] + + def record_save(): + observed_steps.append(trainer._global_step) + return save_state() + + trainer._save_training_state = record_save + trainer._save_checkpoint = Mock(wraps=trainer._save_checkpoint) + try: + exported_path, _ = trainer.train(disable_progress_bars=True) + assert observed_steps == saved_steps + trainer._save_checkpoint.assert_called_once() + assert trainer._checkpoint_paths == [exported_path] + exported = load_file(exported_path) + assert all( + torch.equal(exported[k], v) for k, v in trainer._transformer.state_dict().items() + ) + checkpoint_dir = tmp_path / "checkpoints" + assert sorted(p.name for p in checkpoint_dir.glob("step_*")) == [ + f"step_{saved_steps[-1]:06d}" + ] + metadata = checkpoint_dir / f"step_{saved_steps[-1]:06d}" / "distillation_metadata.json" + assert json.loads(metadata.read_text())["global_step"] == saved_steps[-1] + finally: + trainer._accelerator.free_memory() From 23ec00299ea40d5b2b0c235b4efda97567fddba4 Mon Sep 17 00:00:00 2001 From: Zhiyu Date: Mon, 26 Oct 2026 17:19:27 -0700 Subject: [PATCH 38/43] LTX-2 quantization example: filter calibration extra-params, document supported LTX commit (#2869) 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 (partial; the port to the new LTX API is a follow-up). NVBug 6932665. - `_create_ltx2_pipeline` no longer forwards calibration-only `--extra-param` keys (`height`, `width`, `num_frames`, `frame_rate`, `seed`, `images`, `tiling_config`) to the `TI2VidTwoStagesPipeline` constructor, where they raised `TypeError`. `calibration.py` still reads them. - The LTX-2 README section now states the LTX-2 commit the example targets ([`28c3c73`](https://github.com/Lightricks/LTX-2/commit/28c3c73fe557666c3de176e1e50a5220152ccfca), which has `QuantizationPolicy.fp8_cast()` and `checkpoint_path`/`gemma_root`), with a pip install line, and says newer releases are unsupported. **Not fixed here (follow-up):** the example still fails on current LTX-2 releases (`QuantizationPolicy.fp8_cast` removed; `TI2VidTwoStagesPipeline` now takes `model_paths: ModelPaths`). Porting `pipeline_manager.py` to the new API needs real LTX packages and model assets to validate. The pinned commit predates LTX's Transformers-v5 adaptation, so it may not install alongside ModelOpt's `transformers>=5.5`; the pin documents the API the example was written against, not a verified combination. ### Usage ```sh python quantize.py --model ltx-2 --format fp4 ... --extra-param height=512 ``` ### Testing Pre-commit passed. **Not run:** no LTX packages or assets were available, so neither the filter nor the pin has been exercised. No test added (it needs the LTX packages). ### 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?: ❌ needs LTX packages/assets - Did you update Changelog?: N/A - Did you get Claude approval on this PR?: ❌ ### Additional Information NVBug 6932665. 🤖 Generated with [Claude Code](https://claude.com/claude-code) ## Summary by CodeRabbit * **Bug Fixes** * LTX-2 FP4 pipeline setup now excludes calibration-only options from initialization and correctly applies the available FP8 transformer setting. * **Documentation** * Updated the LTX-2 FP4 example with compatible package versions and required configuration details. The guide also warns that newer package releases are incompatible with the example. --------- Signed-off-by: Zhiyu Cheng Co-authored-by: Claude Sonnet 5.5 Signed-off-by: Rohan Joshi --- examples/diffusers/README.md | 8 ++++++++ .../diffusers/quantization/pipeline_manager.py | 17 ++++++++++++++--- 2 files changed, 22 insertions(+), 3 deletions(-) diff --git a/examples/diffusers/README.md b/examples/diffusers/README.md index 62f8e5f1df0..b67d9630b8d 100644 --- a/examples/diffusers/README.md +++ b/examples/diffusers/README.md @@ -165,6 +165,14 @@ python quantize.py \ > (including quantized or distilled checkpoints) remain subject to the LTX Community License > Agreement and are **not** covered by Apache 2.0. +> [!NOTE] +> This example targets the LTX-2 packages at commit +> [`28c3c73`](https://github.com/Lightricks/LTX-2/commit/28c3c73fe557666c3de176e1e50a5220152ccfca) +> (`QuantizationPolicy.fp8_cast()` and `checkpoint_path`/`gemma_root` pipeline arguments), e.g. +> `pip install "ltx-core @ git+https://github.com/Lightricks/LTX-2.git@28c3c73#subdirectory=packages/ltx-core" "ltx-pipelines @ git+https://github.com/Lightricks/LTX-2.git@28c3c73#subdirectory=packages/ltx-pipelines"`. +> Newer LTX-2 releases removed `fp8_cast()` and changed the pipeline constructor, so this example does +> not run on them yet. + This example produces three outputs: a PyTorch checkpoint (`--quantized-torch-ckpt-save-path`), a Hugging Face checkpoint (`--hf-ckpt-dir`), and a ComfyUI-compatible merged safetensor (`--extra-param merged_base_safetensor_path`). ```sh diff --git a/examples/diffusers/quantization/pipeline_manager.py b/examples/diffusers/quantization/pipeline_manager.py index af89ed568ff..c03af7b3ce9 100644 --- a/examples/diffusers/quantization/pipeline_manager.py +++ b/examples/diffusers/quantization/pipeline_manager.py @@ -217,13 +217,24 @@ def _create_ltx2_pipeline(self) -> Any: distilled_lora_strength = params.pop("distilled_lora_strength", 0.8) spatial_upsampler_path = params.pop("spatial_upsampler_path", None) gemma_root = params.pop("gemma_root", None) - fp8_quantization = params.pop("fp8_quantization", None) or params.pop( - "fp8transformer", False - ) + # Pop both aliases: `or` would short-circuit and leave `fp8transformer` in the kwargs. + fp8_quantization = params.pop("fp8_quantization", None) + fp8_quantization = params.pop("fp8transformer", False) or fp8_quantization params.pop("merged_base_safetensor_path", None) params.pop("enable_swizzle_layout", None) params.pop("padding_strategy", None) params.pop("enable_layerwise_quant_metadata", None) + # Calibration-only knobs read in calibration.py; the pipeline constructor rejects them. + for key in ( + "height", + "width", + "num_frames", + "frame_rate", + "seed", + "images", + "tiling_config", + ): + params.pop(key, None) if not checkpoint_path: raise ValueError("Missing required extra_param: checkpoint_path.") From c38821ade2fc07e3dc142f1877a5defa543fb9e9 Mon Sep 17 00:00:00 2001 From: Shengliang Xu <106840466+shengliangxu@users.noreply.github.com> Date: Mon, 26 Oct 2026 17:27:24 -0700 Subject: [PATCH 39/43] [NVBug 6938136] Use private default locations for TensorRT build artifacts and the AutoTune timing cache (#2984) 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 TensorRT deployment runtime and the ONNX AutoTune benchmark kept their working files at fixed, shared locations under the system temp directory: `/modelopt_build/trt_artifacts` for engine builds and profiling, and `/tmp/trtexec_timing.cache` for the AutoTune timing cache. On a host with several users, another local user could create that location first and plant symlinks in it, which makes a build overwrite files of the user who runs it, or swap the engine file between its write and its read-back, so that bytes the builder did not produce reach TensorRT's engine deserializer. - `build_engine` and `profile_engine` default to one private `mkdtemp` directory (mode 0700) per process. An explicit `output_dir` behaves as before, and `DEFAULT_ARTIFACT_DIR` is removed. - The AutoTune timing cache defaults to `trtexec_timing.cache` in one private `mkdtemp` directory (mode 0700) per process, for `Benchmark` and for the CLI. The CLI resolves that default at run time instead of at import. An explicit `--timing_cache` or `timing_cache_file` behaves as before. - 0.48.0 changelog entry, near the top of the Bug Fixes list so it cherry-picks cleanly to `release/0.48.0`. Two items of the original report are deliberately not taken: opening files with `O_NOFOLLOW`, and hashing the engine from in-memory bytes. With a private directory no other user can touch the files, and `trtexec` writes the engine to disk in a subprocess, so there are no in-memory build bytes to hash. Behavior change: default builds no longer reuse a timing cache across processes, so a repeated run searches tactics again and its engine-build or AutoTune time is not comparable with earlier measurements, and each process that uses a default leaves its own temp directory. Pass `timing_cache_path` (runtime) or `--timing_cache` (AutoTune) to keep a cache across runs. ### Usage No API change; only the default locations move. ### Testing | Check | Result | |---|---| | New tests `test_engine_builder.py::test_default_artifact_dir_is_private_and_reused_within_a_process` and `autotune/test_benchmark.py::test_default_timing_cache_is_private_and_shared_within_a_process` with the old fixed locations restored | both fail, as intended | | `tests/unit/onnx/quantization/autotune` and `tests/unit/torch/deploy` | pass (278) | | The same unit tests on `release/0.48.0` plus this commit | pass (278); the changelog hunk needs one small resolution | | pre-commit (ruff, mypy, bandit, rst, license) over the branch | pass | | `tests/gpu/torch/deploy`, `tests/gpu/onnx/quantization/autotune` | not run locally; covered by the GPU lanes of this PR's CI | ### 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?: ✅ Explicit paths are unchanged; only the default locations move (see the behavior change above). - 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 NVBug 6938136. Targets the 0.48.0 release: cherry-pick this one commit to `release/0.48.0`. A shared per-user cache directory for the AutoTune timing cache was considered separately (#3031, draft) and is not part of this fix. ## Summary by CodeRabbit * **Bug Fixes** * TensorRT builds and profiling now use a private, process-specific temporary directory by default instead of a shared artifact path. * AutoTune timing caches now use a private, process-specific temporary directory by default instead of a shared system temporary path. * **Usage** * Set `timing_cache_path` or pass `--timing_cache` to retain and reuse an AutoTune timing cache across processes. --------- Signed-off-by: Shengliang Xu Signed-off-by: Rohan Joshi --- CHANGELOG.rst | 1 + .../onnx/quantization/autotune/__main__.py | 8 ++-- .../onnx/quantization/autotune/benchmark.py | 19 ++++++-- .../onnx/quantization/autotune/workflows.py | 12 ++++-- .../_deploy/_runtime/tensorrt/constants.py | 3 -- .../_runtime/tensorrt/engine_builder.py | 37 ++++++++++++---- .../quantization/autotune/test_benchmark.py | 43 +++++++++++++++++++ .../_runtime/tensorrt/test_engine_builder.py | 39 +++++++++++++++-- 8 files changed, 137 insertions(+), 25 deletions(-) create mode 100644 tests/unit/onnx/quantization/autotune/test_benchmark.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 4b3219096ca..0a686618511 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -135,6 +135,7 @@ Changelog - Fix LTX-2 distillation full-mode checkpoints to use the safetensors format expected by their file extension and loaders. - Fix corrupted ONNX external tensor data when exporting models on Windows by upgrading ONNX. Regenerate affected models from clean weights. +- Harden the default locations of the TensorRT deployment runtime's build artifacts and the ONNX AutoTune timing cache: they no longer sit at fixed paths under the shared system temp directory, where another local user could pre-create them to plant symlinks or swap a built engine. Both now live in a private temporary directory that is created once per process; pass ``timing_cache_path`` (runtime) or ``--timing_cache`` (AutoTune) to keep a timing cache across processes. - Fix ``import modelopt.torch`` failing with ``ModuleNotFoundError`` for ``requests`` or ``huggingface_hub`` when ``nvidia-modelopt`` is installed without the ``[hf]`` extra. - 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. diff --git a/modelopt/onnx/quantization/autotune/__main__.py b/modelopt/onnx/quantization/autotune/__main__.py index ccb826c291b..7c0357ca93e 100644 --- a/modelopt/onnx/quantization/autotune/__main__.py +++ b/modelopt/onnx/quantization/autotune/__main__.py @@ -17,10 +17,10 @@ import argparse import sys -import tempfile from pathlib import Path from modelopt.onnx.logging_config import logger +from modelopt.onnx.quantization.autotune.benchmark import default_timing_cache_file from modelopt.onnx.quantization.autotune.utils import ( DEFAULT_NUM_SCHEMES, DEFAULT_TIMING_RUNS, @@ -38,7 +38,6 @@ DEFAULT_OUTPUT_DIR = "./autotuner_output" DEFAULT_QUANT_TYPE = "int8" DEFAULT_DQ_DTYPE = "float32" -DEFAULT_TIMING_CACHE = str(Path(tempfile.gettempdir()) / "trtexec_timing.cache") def apply_mode_presets(args) -> None: @@ -94,6 +93,7 @@ def run_autotune() -> int: """ args = get_parser().parse_args() apply_mode_presets(args) + args.timing_cache = args.timing_cache or default_timing_cache_file() model_path = validate_file_path(args.onnx_path, "Model file") validate_file_path(args.qdq_baseline, "QDQ baseline model") output_dir = Path(args.output_dir) @@ -279,8 +279,8 @@ def get_parser() -> argparse.ArgumentParser: trt_group.add_argument( "--timing_cache", type=str, - default=DEFAULT_TIMING_CACHE, - help=f"TensorRT timing cache file (default: {DEFAULT_TIMING_CACHE})", + default=None, + help="TensorRT timing cache file (default: a file in a private temporary directory)", ) trt_group.add_argument( "--warmup_runs", diff --git a/modelopt/onnx/quantization/autotune/benchmark.py b/modelopt/onnx/quantization/autotune/benchmark.py index 6c62ee4e628..d16220b8094 100644 --- a/modelopt/onnx/quantization/autotune/benchmark.py +++ b/modelopt/onnx/quantization/autotune/benchmark.py @@ -27,6 +27,7 @@ """ import ctypes +import functools import importlib.util import os import re @@ -50,6 +51,16 @@ TORCH_CUDA_AVAILABLE = torch.cuda.is_available() +@functools.cache +def default_timing_cache_file() -> str: + """Return the default timing-cache path: a file in one private temporary directory per process. + + A fixed path under the shared temp dir would let another local user pre-create the file or plant + a symlink there. Pass ``timing_cache_file`` to keep a cache across runs. + """ + return str(Path(tempfile.mkdtemp(prefix="modelopt_autotune_")) / "trtexec_timing.cache") + + def _validate_shape_range(min_shape: list, opt_shape: list, max_shape: list) -> None: """Raise ValueError if shape lengths differ or if min <= opt <= max fails at any dimension.""" if len(min_shape) != len(opt_shape) or len(opt_shape) != len(max_shape): @@ -92,7 +103,8 @@ def __init__( Args: timing_cache_file: Path to timing cache file to accelerate engine builds. - If None, uses '/tmp/trtexec_timing.cache' as default. + If None, uses 'trtexec_timing.cache' in a private temporary directory + that is created once per process. warmup_runs: Number of warmup iterations before timing measurements. timing_runs: Number of iterations for latency measurement. Results are averaged across these runs. @@ -100,7 +112,7 @@ def __init__( These plugins will be loaded during engine building. If None, no custom plugins are loaded. """ - self.timing_cache_file = timing_cache_file or "/tmp/trtexec_timing.cache" # nosec B108 + self.timing_cache_file = timing_cache_file or default_timing_cache_file() self.warmup_runs = warmup_runs self.timing_runs = timing_runs self.plugin_libraries = plugin_libraries or [] @@ -328,7 +340,8 @@ def __init__( Args: timing_cache_file: Path to TensorRT timing cache file. If None, - defaults to '/tmp/trtexec_timing.cache'. + defaults to 'trtexec_timing.cache' in a private temporary directory + that is created once per process. warmup_runs: Number of warmup iterations before timing measurements. timing_runs: Number of iterations for latency measurement. plugin_libraries: List of paths to TensorRT plugin shared libraries (.so files). diff --git a/modelopt/onnx/quantization/autotune/workflows.py b/modelopt/onnx/quantization/autotune/workflows.py index a863d323527..21ee87313d9 100644 --- a/modelopt/onnx/quantization/autotune/workflows.py +++ b/modelopt/onnx/quantization/autotune/workflows.py @@ -29,7 +29,11 @@ from modelopt.onnx.logging_config import logger from modelopt.onnx.quantization.autotune.autotuner import QDQAutotuner -from modelopt.onnx.quantization.autotune.benchmark import TensorRTPyBenchmark, TrtExecBenchmark +from modelopt.onnx.quantization.autotune.benchmark import ( + TensorRTPyBenchmark, + TrtExecBenchmark, + default_timing_cache_file, +) from modelopt.onnx.quantization.autotune.common import Config, PatternCache from modelopt.onnx.quantization.qdq_utils import get_quantized_tensors @@ -97,7 +101,8 @@ def init_benchmark_instance( These plugins will be loaded by trtexec or TensorRT Python API during engine building. If None, no custom plugins are loaded. timing_cache_file: Path to TensorRT timing cache file for faster engine builds. - If None, uses default "trtexec_timing.cache" (default: None) + If None, uses "trtexec_timing.cache" in a private temporary directory + (default: None) warmup_runs: Number of warmup inference iterations before measurement. Allows GPU to reach stable performance state (default: 5) timing_runs: Number of timed inference iterations for latency measurement. @@ -126,7 +131,8 @@ def init_benchmark_instance( logger.info("TensorRT Python API benchmark initialized") logger.debug( f"Settings: warmup={warmup_runs}, timing={timing_runs}, " - f"cache={timing_cache_file or 'trtexec_timing.cache'}, plugin_libraries={plugin_libraries}" + f"cache={timing_cache_file or default_timing_cache_file()}, " + f"plugin_libraries={plugin_libraries}" ) return _benchmark_instance except Exception as e: diff --git a/modelopt/torch/_deploy/_runtime/tensorrt/constants.py b/modelopt/torch/_deploy/_runtime/tensorrt/constants.py index d9ace1645a3..eecc7073e2c 100644 --- a/modelopt/torch/_deploy/_runtime/tensorrt/constants.py +++ b/modelopt/torch/_deploy/_runtime/tensorrt/constants.py @@ -31,9 +31,6 @@ ONE_MEBI_IN_BYTES = 1 << 20 ONE_GIBI_IN_BYTES = 1 << 30 -# TensorRT conversion tool names -DEFAULT_ARTIFACT_DIR = "modelopt_build/trt_artifacts" - # Default conversion params DEFAULT_VALIDATION_THRESHOLD = 1e-4 DEFAULT_MAX_BATCH_SIZE = 1 diff --git a/modelopt/torch/_deploy/_runtime/tensorrt/engine_builder.py b/modelopt/torch/_deploy/_runtime/tensorrt/engine_builder.py index bb8bbd292b8..d00272ce158 100644 --- a/modelopt/torch/_deploy/_runtime/tensorrt/engine_builder.py +++ b/modelopt/torch/_deploy/_runtime/tensorrt/engine_builder.py @@ -13,17 +13,17 @@ # See the License for the specific language governing permissions and # limitations under the License. +import functools import logging import shutil import sys from pathlib import Path -from tempfile import TemporaryDirectory, gettempdir +from tempfile import TemporaryDirectory, mkdtemp from ..._runtime.common import read_bytes, timeit, write_bytes, write_string from ..._runtime.tensorrt.layerwise_profiling import process_layerwise_result from ...utils import OnnxBytes from .constants import ( - DEFAULT_ARTIFACT_DIR, DEFAULT_NUM_INFERENCE_PER_RUN, SHA_256_HASH_LENGTH, TRT_MODE_FLAGS, @@ -135,6 +135,27 @@ def _update_dynamic_shapes(dynamic_shapes: dict, cmd: list[str]) -> None: cmd.extend([f"--{shape}={convert_shape_to_string(dynamic_shapes[shape])}"]) +@functools.cache +def _default_artifact_dir() -> Path: + """Return this process's private temporary directory for builds without an ``output_dir``. + + The default must not be a fixed path under the shared temp dir: another local user could + pre-create it, plant symlinks in it, or swap the engine between its write and read-back. One + unpredictable directory per process keeps timing-cache reuse across calls and bounds what is + left behind to one directory, however many engines the process builds. + """ + return Path(mkdtemp(prefix="modelopt_trt_")) + + +def _artifact_dir(output_dir: Path | None) -> Path: + """Return ``output_dir`` (created if missing), or this process's private temporary directory.""" + if not output_dir: + return _default_artifact_dir() + path = Path(output_dir) + path.mkdir(parents=True, exist_ok=True) + return path + + @timeit def build_engine( onnx_bytes: OnnxBytes, @@ -183,7 +204,8 @@ def build_engine( For more details, please refer to: https://docs.nvidia.com/deeplearning/tensorrt/latest/_static/python-api/infer/Core/BuilderConfig.html output_dir: Directory to save the engine file and trtexec artifacts. - If not provided, the artifacts will be saved in a temporary directory. + If not provided, this process's private temporary directory is used (created on first + use, not removed at exit). verbose: Set trtexec command to verbose or not. Returns: @@ -230,8 +252,7 @@ def _setup_files_and_paths( onnx_bytes.write_to_disk(str(tmp_onnx_dir)) onnx_path = tmp_onnx_dir / f"{onnx_bytes.model_name}.onnx" - final_output_dir = Path(output_dir or Path(gettempdir()) / DEFAULT_ARTIFACT_DIR) - final_output_dir.mkdir(parents=True, exist_ok=True) + final_output_dir = _artifact_dir(output_dir) engine_path = ( Path(engine_path) if engine_path @@ -294,7 +315,8 @@ def profile_engine( "maxShapes": {"input": [32,3,244,244]} } output_dir: Directory to save the engine file and trtexec artifacts. - If not provided, the artifacts will be saved in a temporary directory. + If not provided, this process's private temporary directory is used (created on first + use, not removed at exit). Returns: Layerwise profiling output as a json string. Stdout log produced by trtexec tool. @@ -322,8 +344,7 @@ def _setup_files_and_paths(tmp_dir_path: Path, engine_hash: str) -> tuple[Path, engine_path = tmp_dir_path / f"{engine_hash}-model.engine" write_bytes(engine_bytes[SHA_256_HASH_LENGTH:], engine_path) - final_output_dir = Path(output_dir or Path(gettempdir()) / DEFAULT_ARTIFACT_DIR) - final_output_dir.mkdir(parents=True, exist_ok=True) + final_output_dir = _artifact_dir(output_dir) profile_path = final_output_dir / f"{engine_hash}-profile.json" layer_info_path = final_output_dir / f"{engine_hash}-layerInfo.json" diff --git a/tests/unit/onnx/quantization/autotune/test_benchmark.py b/tests/unit/onnx/quantization/autotune/test_benchmark.py new file mode 100644 index 00000000000..c6d98a164c4 --- /dev/null +++ b/tests/unit/onnx/quantization/autotune/test_benchmark.py @@ -0,0 +1,43 @@ +# 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 the default timing-cache file in autotune/benchmark.py.""" + +import stat +import sys +import tempfile +from pathlib import Path + +from modelopt.onnx.quantization.autotune.benchmark import Benchmark, default_timing_cache_file + + +class _NoOpBenchmark(Benchmark): + def run(self, path_or_bytes, log_file=None): + return 0.0 + + +def test_default_timing_cache_is_private_and_shared_within_a_process(tmp_path, monkeypatch): + monkeypatch.setattr(tempfile, "tempdir", str(tmp_path)) + default_timing_cache_file.cache_clear() + first = Path(_NoOpBenchmark().timing_cache_file) + second = Path(_NoOpBenchmark().timing_cache_file) + default_timing_cache_file.cache_clear() + + # A fixed shared path would let another local user pre-create the file or plant a symlink. + assert first == second + assert first.name == "trtexec_timing.cache" + assert first.parent.parent == tmp_path + if sys.platform != "win32": # POSIX file modes + assert stat.S_IMODE(first.parent.stat().st_mode) == 0o700 diff --git a/tests/unit/torch/deploy/_runtime/tensorrt/test_engine_builder.py b/tests/unit/torch/deploy/_runtime/tensorrt/test_engine_builder.py index ff7f77cf617..4385ef133bf 100755 --- a/tests/unit/torch/deploy/_runtime/tensorrt/test_engine_builder.py +++ b/tests/unit/torch/deploy/_runtime/tensorrt/test_engine_builder.py @@ -15,9 +15,10 @@ import json import os +import stat import sys +import tempfile from pathlib import Path -from tempfile import gettempdir from types import ModuleType from unittest import mock @@ -30,12 +31,15 @@ tensorrt.__version__ = "10.10" sys.modules["tensorrt"] = tensorrt -from modelopt.torch._deploy._runtime.tensorrt.engine_builder import build_engine, profile_engine +from modelopt.torch._deploy._runtime.tensorrt.engine_builder import ( + _default_artifact_dir, + build_engine, + profile_engine, +) @pytest.fixture -def setup_mocks(): - tmp_path = Path(gettempdir()) / "modelopt_build/trt_artifacts" +def setup_mocks(tmp_path): engine_bytes = b"engine_bytes" dummy_hash = "997930f2" expected_engine_bytes = b"\x99y0\xf2" + b"A" * 28 + b",engine_bytes" @@ -125,3 +129,30 @@ def test_profile_engine_variants(setup_mocks, enable_layerwise, profiling_runs): assert profiled_data == {} assert out == "success" + + +def _write_engine_like_trtexec(cmd): + for arg in cmd: + if arg.startswith("--saveEngine="): + Path(arg.split("=", 1)[1]).write_bytes(b"engine_bytes") + return 0, "success" + + +def test_default_artifact_dir_is_private_and_reused_within_a_process(setup_mocks, monkeypatch): + tmp_path, mock_run, mock_onnx, _ = setup_mocks + shared_tmp = tmp_path / "shared_tmp" + shared_tmp.mkdir() + monkeypatch.setattr(tempfile, "tempdir", str(shared_tmp)) + _default_artifact_dir.cache_clear() + mock_run.side_effect = _write_engine_like_trtexec + + build_engine(onnx_bytes=mock_onnx) + build_engine(onnx_bytes=mock_onnx) + profile_engine(engine_bytes=b"\x99y0\xf2" + b"A" * 28 + b",engine_bytes") + _default_artifact_dir.cache_clear() + + # A fixed shared path would let another local user pre-create it or plant symlinks in it. One + # unpredictable private directory per process keeps the timing cache shared across calls. + (artifact_dir,) = shared_tmp.iterdir() + if sys.platform != "win32": # POSIX file modes + assert stat.S_IMODE(artifact_dir.stat().st_mode) == 0o700 From 0174b5d7c82d4280ad3fb353cb205d67aabf65a1 Mon Sep 17 00:00:00 2001 From: Juhi Mittal <39641197+juhi10071998@users.noreply.github.com> Date: Mon, 26 Oct 2026 19:19:18 -0700 Subject: [PATCH 40/43] [NVBug 6901733] fix(hf_ptq): use bundled code for dense Nemotron-H on Transformers 5.5 (#2975) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Type of change: Bug fix `get_model()` picked the built-in `NemotronHForCausalLM` whenever Transformers ships one, ignoring `--trust_remote_code`. On Transformers 5.5 the built-in class implements only the MoE variant, so dense Nemotron-H checkpoints (Nano-9B/12B-v2, H-8B/47B/56B) failed with `KeyError: 'mlp'`. It now uses the checkpoint's bundled code when `trust_remote_code` is set and its config can't be re-derived as a built-in one. From 5.6 the built-in class handles dense models and nothing changes. - New `test_get_model_selects_bundled_or_builtin_class` cases pass; the dense case fails without the fix. - Transformers 5.5.4 (TRT-LLM 1.3.0rc29, aarch64): `hf_ptq.py` FP8 PTQ of Nano-9B-v2 exports a valid checkpoint (9.4 GB, 120 `F8_E4M3` tensors), single GPU. - Not run: multi-GPU (the default 4-GPU split crashes in a `mamba-ssm` triton kernel). - Known, separate: Transformers 5.14.x fails upstream (`DynamicCache` has no `mlp` layer type). The bundled path needs `mamba-ssm`, whose current releases don't import on Transformers 5.x; I used a local shim to test. - 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?: ❌ Fixes NVBug 6901733. 🤖 Generated with [Claude Code](https://claude.com/claude-code) * **Bug Fixes** * Fixed loading dense Nemotron-H checkpoints with trusted checkpoint code enabled when the built-in configuration cannot be reconstructed. * Preserved model selection for DeepSeek checkpoints and Nemotron-H checkpoints whose built-in configuration loads successfully. * Checkpoints with layer types unsupported by the built-in Nemotron-H model now use the trusted checkpoint code. * Unrelated configuration-loading errors continue to be reported rather than triggering a fallback. This resolves a loading failure observed with Transformers 5.5. --------- Signed-off-by: Juhi Mittal Co-authored-by: Claude Sonnet 5.5 Signed-off-by: Rohan Joshi --- CHANGELOG.rst | 1 + examples/hf_ptq/example_utils.py | 16 +++- modelopt/torch/models/hf.py | 9 +++ modelopt/torch/models/nemotron_h/mtp.py | 25 ++++++ tests/examples/hf_ptq/test_example_utils.py | 85 +++++++++++++++++---- 5 files changed, 120 insertions(+), 16 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 0a686618511..e226edae3bd 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -143,6 +143,7 @@ Changelog - 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 ``examples/hf_ptq`` failing with ``KeyError: 'mlp'`` on dense Nemotron-H checkpoints (Nano-9B/12B-v2, H-8B/47B/56B) with Transformers 5.5, whose built-in ``NemotronHForCausalLM`` supports only the MoE variant. ``get_model`` now uses the checkpoint's bundled code when ``--trust_remote_code`` is set and the built-in config can't be built. - 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. diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index 52241079d39..dc774b12da1 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -47,7 +47,11 @@ from modelopt.torch.export import has_spec_opt from modelopt.torch.export.model_utils import is_multimodal_model from modelopt.torch.models import hf_model_type -from modelopt.torch.models.hf import checkpoint_has_mtp, prepare_model_for_loading +from modelopt.torch.models.hf import ( + checkpoint_has_mtp, + prefers_bundled_code, + prepare_model_for_loading, +) from modelopt.torch.utils import distributed as dist_utils from modelopt.torch.utils.mlflow import Tool, resolved_recipe_texts, tracked_run from modelopt.torch.utils.plugins.model_load_utils import record_unplaced_source_keys @@ -722,7 +726,15 @@ def has_pack_quantized_config(config): # DeepSeek ships bundled modeling code, but the built-in class is what the # disk-offload and streaming-export paths are validated against. - use_bundled_code = trust_remote_code and "Deepseek" in architecture + # Also use it when the model plugin says the built-in class can't represent the checkpoint + # (dense Nemotron-H on Transformers 5.5, whose built-in class is MoE-only). + use_bundled_code = trust_remote_code and ( + "Deepseek" in architecture + or ( + hasattr(transformers, architecture) + and prefers_bundled_code(hf_config.model_type, hf_config, ckpt_path, config_kwargs) + ) + ) if not hasattr(transformers, architecture) or use_bundled_code: if not hasattr(transformers, architecture): diff --git a/modelopt/torch/models/hf.py b/modelopt/torch/models/hf.py index f9bdd49a931..3b765ddd2ba 100644 --- a/modelopt/torch/models/hf.py +++ b/modelopt/torch/models/hf.py @@ -20,6 +20,7 @@ __all__ = [ "checkpoint_has_mtp", + "prefers_bundled_code", "prepare_model_for_calibration", "prepare_model_for_loading", ] @@ -53,6 +54,14 @@ def checkpoint_has_mtp(model_type, checkpoint_path: str) -> bool: return plugin is not None and plugin.has_mtp_weights(checkpoint_path) +def prefers_bundled_code(model_type, hf_config, checkpoint_path: str, config_kwargs: dict) -> bool: + """Use the model plugin to tell whether the checkpoint's bundled code must replace the built-in class.""" + plugin = _get_plugin(model_type) + return plugin is not None and plugin.prefers_bundled_code( + hf_config, checkpoint_path, config_kwargs + ) + + def prepare_model_for_calibration(model) -> None: """Run the optional model-specific hook that augments the default calibration forward.""" plugin = _get_plugin(getattr(getattr(model, "config", None), "model_type", None)) diff --git a/modelopt/torch/models/nemotron_h/mtp.py b/modelopt/torch/models/nemotron_h/mtp.py index b3a234ac21f..736ee4a5322 100644 --- a/modelopt/torch/models/nemotron_h/mtp.py +++ b/modelopt/torch/models/nemotron_h/mtp.py @@ -29,6 +29,7 @@ import torch from huggingface_hub import snapshot_download from safetensors import SafetensorError +from transformers import AutoConfig from transformers.dynamic_module_utils import get_class_from_dynamic_module from transformers.masking_utils import create_causal_mask @@ -42,6 +43,7 @@ __all__ = [ "has_mtp_weights", + "prefers_bundled_code", "prepare_for_calibration", "prepare_for_loading", ] @@ -67,6 +69,29 @@ def has_mtp_weights(checkpoint_path: str) -> bool: return _get_mtp_layout(_resolve_checkpoint_path(checkpoint_path)) is not None +def prefers_bundled_code(hf_config, checkpoint_path: str, config_kwargs: dict) -> bool: + """True for a remote-code config the built-in Nemotron-H class cannot represent. + + Built-in Nemotron-H implements only the MoE variant on some Transformers releases, so a + dense checkpoint's config fails to parse as a built-in one and needs its bundled code. + """ + if getattr(hf_config, "model_type", None) != "nemotron_h": + return False + if not type(hf_config).__module__.startswith("transformers_modules"): + return False + builtin_kwargs = {k: v for k, v in config_kwargs.items() if k != "trust_remote_code"} + try: + builtin_config = AutoConfig.from_pretrained(checkpoint_path, **builtin_kwargs) + except (KeyError, ValueError, TypeError, AttributeError): + # Config-parsing incompatibilities only; e.g. an OSError from the Hub must propagate. + return True + # Parsing can succeed while model construction fails on a block type the built-in lacks. + # Native Nemotron-H is optional: older Transformers can still load unrelated models. + from transformers.models.nemotron_h.modeling_nemotron_h import MIXER_TYPES + + return any(t not in MIXER_TYPES for t in getattr(builtin_config, "layers_block_type", ())) + + def _normalize_mtp_block_types(block_types, mixer_types): """Normalize checkpoint attention names to the installed Transformers registry.""" aliases = {"attention": "full_attention", "full_attention": "attention"} diff --git a/tests/examples/hf_ptq/test_example_utils.py b/tests/examples/hf_ptq/test_example_utils.py index a6e44a0c363..0ed9f41c730 100644 --- a/tests/examples/hf_ptq/test_example_utils.py +++ b/tests/examples/hf_ptq/test_example_utils.py @@ -371,21 +371,55 @@ def test_is_diffusion_gemma(hf_config, expected): @pytest.mark.parametrize( - ("trust_remote_code", "expect_bundled_code"), - [(True, True), (False, False)], + ( + "architecture", + "remote_config", + "builtin_config_error", + "trust_remote_code", + "expect_bundled", + ), + [ + ("DeepseekV3ForCausalLM", False, None, True, True), + ("DeepseekV3ForCausalLM", False, None, False, False), + # Dense Nemotron-H: the built-in class cannot parse the remote config -> bundled code. + ("NemotronHForCausalLM", True, KeyError("-"), True, True), + # MoE Nemotron-H (or a newer Transformers): the built-in config parses -> built-in class. + ("NemotronHForCausalLM", True, None, True, False), + # An unrelated failure (e.g. Hub retrieval) is not an incompatibility: it propagates. + ("NemotronHForCausalLM", True, OSError("hub down"), True, None), + ], ) -def test_get_model_deepseek_honors_trust_remote_code( - monkeypatch, trust_remote_code, expect_bundled_code +def test_get_model_selects_bundled_or_builtin_class( + monkeypatch, + tmp_path, + *, + architecture, + remote_config, + builtin_config_error, + trust_remote_code, + expect_bundled, ): - """DeepSeek ships bundled modeling code; --trust_remote_code selects it, else built-in.""" + """--trust_remote_code selects bundled code for DeepSeek and for a config the built-in + class cannot represent; otherwise the built-in class is used.""" used = {} - hf_config = SimpleNamespace( - architectures=["DeepseekV3ForCausalLM"], + config_cls = type( + "_Config", + (SimpleNamespace,), + {"__module__": "transformers_modules.ckpt.config" if remote_config else __name__}, + ) + hf_config = config_cls( + architectures=[architecture], dtype=torch.bfloat16, - model_type="deepseek_v3", + model_type="nemotron_h" if "NemotronH" in architecture else "deepseek_v3", torch_dtype=torch.bfloat16, ) + def _load_config(*args, **kwargs): + # The re-derivation of a built-in config is the call made without trust_remote_code. + if builtin_config_error and not kwargs.get("trust_remote_code"): + raise builtin_config_error + return hf_config + class FakeModel: def eval(self): return None @@ -407,20 +441,43 @@ def from_pretrained(*args, output_loading_info=False, **kwargs): return Fake - monkeypatch.setattr(example_utils.AutoConfig, "from_pretrained", lambda *a, **k: hf_config) + monkeypatch.setattr(example_utils.AutoConfig, "from_pretrained", _load_config) monkeypatch.setattr(example_utils, "AutoModelForCausalLM", _record("bundled")) - monkeypatch.setattr( - example_utils.transformers, "DeepseekV3ForCausalLM", _record("builtin"), raising=False - ) + monkeypatch.setattr(example_utils.transformers, architecture, _record("builtin"), raising=False) + monkeypatch.setattr(example_utils, "prepare_model_for_loading", lambda *a, **k: nullcontext()) monkeypatch.setattr(example_utils, "is_nemotron_vl", lambda config: False) monkeypatch.setattr(example_utils, "is_speculative", lambda config: False) monkeypatch.setattr(example_utils, "init_empty_weights", lambda include_buffers: nullcontext()) monkeypatch.setattr(example_utils, "get_max_memory", lambda: {0: 1024}) monkeypatch.setattr(example_utils, "infer_auto_device_map", lambda model, max_memory: {"": 0}) - example_utils.get_model("checkpoint", device="cpu", trust_remote_code=trust_remote_code) + if expect_bundled is None: + with pytest.raises(OSError): + example_utils.get_model( + str(tmp_path), device="cpu", trust_remote_code=trust_remote_code + ) + return + example_utils.get_model(str(tmp_path), device="cpu", trust_remote_code=trust_remote_code) + + assert used["path"] == ("bundled" if expect_bundled else "builtin") + + +def test_nemotron_h_prefers_bundled_code_for_unsupported_block_type(monkeypatch): + """A built-in config that parses but names a block type the built-in model lacks needs the + checkpoint's bundled code.""" + modeling = pytest.importorskip("transformers.models.nemotron_h.modeling_nemotron_h") + remote_cls = type("_Config", (SimpleNamespace,), {"__module__": "transformers_modules.c.cfg"}) + hf_config = remote_cls(model_type="nemotron_h") + monkeypatch.setattr( + example_utils.AutoConfig, + "from_pretrained", + lambda *a, **k: SimpleNamespace(layers_block_type=["mamba", "mlp"]), + ) - assert used["path"] == ("bundled" if expect_bundled_code else "builtin") + monkeypatch.setattr(modeling, "MIXER_TYPES", {"mamba": object}) + assert example_utils.prefers_bundled_code("nemotron_h", hf_config, "ckpt", {}) + monkeypatch.setattr(modeling, "MIXER_TYPES", {"mamba": object, "mlp": object}) + assert not example_utils.prefers_bundled_code("nemotron_h", hf_config, "ckpt", {}) # ---------- get_tokenizer ---------------------------------------------------- From 7a6183f5c1075a7e4da4f37f1086d808a65e8422 Mon Sep 17 00:00:00 2001 From: Dmitry Rantovov Date: Tue, 27 Oct 2026 00:03:56 -0700 Subject: [PATCH 41/43] Fix selected layers being skipped by post-init sparsity (#2345) 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 `do_sparsity()` silently skips selected layers such as `dense` and `gate`: `rstrip(".weight")` treats its argument as a set of characters, turning `dense.weight` into `dens` and `gate.weight` into `ga`. Those names no longer match the selected modules. Remove the literal suffix instead. ### Usage No API changes. The regression test builds a small model and applies `SparsityMethod2o4().do_sparsity(model)` to each selected layer. ### Testing On main, the regression test fails for `dense` and `gate` because no pruning mask is attached; the `proj` control passes. With the fix, all three cases pass and verify the mask and resulting weights. All 13 tests in `test_post_init_sparse.py` and `test_common.py` pass on CPU. Pre-commit checks pass for both changed files. ### 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 - Did you get Claude approval on this PR?: N/A ### Additional Information Found while inspecting the layer-name filtering in the post-initialization sparsity helper. ## Summary by CodeRabbit * **Bug Fixes** * Fixed sparsity processing for layers whose names include `.weight`, ensuring the correct model parameters are matched and updated. * **Tests** * Added coverage for structured 2:4 sparsity across `dense`, `gate`, and `proj` MLP layers, including mask placement and resulting weight values. Signed-off-by: Dmitry Rantovov Co-authored-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Signed-off-by: Rohan Joshi --- .../puzzletron/tools/post_init_sparse.py | 4 +- .../torch/puzzletron/test_post_init_sparse.py | 52 +++++++++++++++++++ 2 files changed, 54 insertions(+), 2 deletions(-) create mode 100644 tests/unit/torch/puzzletron/test_post_init_sparse.py diff --git a/modelopt/torch/puzzletron/tools/post_init_sparse.py b/modelopt/torch/puzzletron/tools/post_init_sparse.py index ab939e3bb4b..0ff8f841dab 100644 --- a/modelopt/torch/puzzletron/tools/post_init_sparse.py +++ b/modelopt/torch/puzzletron/tools/post_init_sparse.py @@ -82,9 +82,9 @@ def do_sparsity(self, model: nn.Module, mask_dict=None): if mask_dict is None: state_dict_for_sparsifying = { - k.rstrip(".weight"): v + k.removesuffix(".weight"): v for k, v in model.state_dict().items() - if k.rstrip(".weight") in full_name_layers + if k.removesuffix(".weight") in full_name_layers } mask_dict = self.calculate_masks(state_dict_for_sparsifying) # print('Apply sparsity') diff --git a/tests/unit/torch/puzzletron/test_post_init_sparse.py b/tests/unit/torch/puzzletron/test_post_init_sparse.py new file mode 100644 index 00000000000..5cab9b5869c --- /dev/null +++ b/tests/unit/torch/puzzletron/test_post_init_sparse.py @@ -0,0 +1,52 @@ +# 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 post-initialization structured sparsity.""" + +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from modelopt.torch.puzzletron.tools.post_init_sparse import SparsityMethod2o4 + + +@pytest.mark.parametrize("layer_name", ["dense", "gate", "proj"]) +def test_do_sparsity_preserves_layer_name(layer_name): + block = nn.Module() + block.mlp = nn.Module() + linear = nn.Linear(4, 2, bias=False) + block.mlp.add_module(layer_name, linear) + + model = nn.Module() + model.model = nn.Module() + model.model.layers = nn.ModuleList([block]) + model.config = SimpleNamespace( + block_configs=[ + SimpleNamespace( + ffn=SimpleNamespace(sparsify=[layer_name]), + attention=SimpleNamespace(sparsify=[]), + ) + ] + ) + with torch.no_grad(): + linear.weight.copy_(torch.arange(1, 9, dtype=torch.float32).reshape(2, 4)) + + SparsityMethod2o4().do_sparsity(model) + + expected_mask = torch.tensor([[0, 0, 1, 1], [0, 0, 1, 1]], dtype=torch.float32) + torch.testing.assert_close(linear.weight_mask, expected_mask) + torch.testing.assert_close(linear.weight, linear.weight_orig * expected_mask) From 7f48c9db23ffdb7069469d5d19667c48a986933b Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Tue, 27 Oct 2026 18:39:20 +0800 Subject: [PATCH 42/43] fix(speculative): pin and hash-verify the Kimi-K2 remote modeling code (#3020) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Type of change: Bug fix `eagle_decoder_type="kimik2"` builds the EAGLE decoder layer from `modeling_deepseek.py` and `configuration_deepseek.py` in `moonshotai/Kimi-K2-Thinking`, which ModelOpt downloads and imports at runtime (`_setup_kimi_k2_decoder` in `modelopt/torch/speculative/utils.py`). The download tracked the repo's current HEAD with no integrity check, so whatever that third-party repo published next would be imported unreviewed. This PR: - Pins the download to commit `a51ccc050d73dab088bf7b0e2dd9b30ae85a4e55` and checks both files' SHA-256 against values stored in ModelOpt before executing either of them. A mismatch in either file raises `RuntimeError` before any downloaded code runs. - Executes exactly the bytes that were hashed, through a small `importlib.abc.InspectLoader`. The default `SourceFileLoader` would re-read the file and could run a `__pycache__/*.pyc` whose recorded mtime/size match the source, bypassing the hash. This also avoids a bare `exec` (Bandit B102). - Downloads the two files with `hf_hub_download` instead of `snapshot_download`. With a commit-hash revision, `snapshot_download` skips its local-cache fallback and calls the Hub tree API even when `HF_HUB_OFFLINE=1` (seen on huggingface_hub 1.28.0 and 1.33.0), so cached offline runs would have started failing; `hf_hub_download` serves a pinned commit straight from the cache. `tokenization_kimi.py`, which was never imported, is no longer downloaded. - Applies the `DynamicCache` compatibility patch only after both files are verified and imported, so a refusal leaves `sys.modules` and global state untouched. A module that raises while importing is removed from `sys.modules`, as `importlib` does. **No change in the executed code.** Both files have the same git blob ids (`79b86a08…`, `1192a006…`) in every upstream commit from the first one (2025-11-04) through `a51ccc05` (current HEAD, 2026-01-30), and a fresh download at the pinned commit matches the recorded SHA-256 values. Moving the pin later means bumping `KIMI_K2_REVISION` and both hashes in `KIMI_K2_FILE_SHA256` together after reviewing the upstream diff. Offline environments need the pinned commit in the Hugging Face cache. Caches populated after 2026-01-30 already have it; older ones need a one-time fetch of the two files. **Known pre-existing limitation, unchanged by this PR:** the upstream `modeling_deepseek.py` imports `transformers.utils.import_utils.is_torch_fx_available`, which is gone from every supported transformers version (checked 5.5.4, 5.12.1, 5.18.0). The HF `kimik2` decoder path therefore fails at import on `main`, before and after this change. A possible follow-up is porting this decoder to transformers' native `DeepseekV3` implementation, which would drop the remote code entirely and restore transformers 5 support; that changes numerics and needs its own validation, so it is out of scope here. N/A: no API change. `eagle_decoder_type="kimik2"` keeps working the same way and now fetches the pinned, hash-verified files. - `pre-commit run` on the changed files (ruff, ruff-format, mypy, bandit, rst checks): pass. - `tests/unit/torch/speculative/test_speculative_utils.py` on transformers 5.5.4, 5.12.1 and 5.18.0 (CPU): 9 passed. - Three new tests, all offline (`hf_hub_download` is monkeypatched) and about 0.01 s each: - `test_kimi_k2_decoder_refuses_unverified_code` serves stand-in files whose top-level code would create a marker file. The config file matches its pinned hash and only the modeling file does not, so the test asserts a `RuntimeError` and that the marker never appears. It fails with the hash comparison disabled, and with the first revision's verify-then-import-per-file order, where the valid config ran before the modeling check. - `test_kimi_k2_decoder_imports_verified_code` loads verified stand-ins and checks that the modeling file's relative `from .configuration_deepseek import ...` resolves. CI never imports the real files, so this is the only coverage of the loader's success path. - `test_kimi_k2_decoder_drops_failed_module` checks that a module raising while importing is not left half-initialized in `sys.modules`; it fails without the cleanup. - Pinned download with huggingface_hub 1.28.0 and 1.33.0, offline (`HF_HUB_OFFLINE=1`) and online: both files resolve to the `a51ccc05` snapshot. The same offline request through `snapshot_download(revision=)` fails on both versions. - End-to-end smoke on CPU, offline, converting a tiny Llama base with `eagle_decoder_type="kimik2"`: download and hash verification pass and the upstream module runs through the new loader until the pre-existing `is_torch_fx_available` ImportError above, the same failure as unmodified `main`. With that symbol shimmed in the smoke script only, `_setup_kimi_k2_decoder` completes and returns `DeepseekV3DecoderLayer`; building the layer then hits a further transformers 5 config incompatibility (`rope_theta`), also pre-existing. - Not run: GPU or multi-node training with `kimik2`, since the path does not import on supported transformers versions. 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?: ✅ (same executed code; offline caches older than 2026-01-30 need a one-time fetch of the pinned commit) - 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?: ✅ (`/claude review` approved the first revision) The pinned revision and hashes live in `KIMI_K2_REVISION` and `KIMI_K2_FILE_SHA256` in `modelopt/torch/speculative/utils.py`. 🤖 Generated with [Claude Code](https://claude.com/claude-code) - **Bug Fixes** - Kimi-K2 modeling files used for EAGLE training are now downloaded at a pinned revision and verified before execution. Offline runs require that revision to be available in the Hugging Face cache. - Fixed issues across quantization, export, calibration, and serving. - **Documentation** - Documented updated ONNX Runtime defaults for Linux x86-64 with Python 3.11+, including CUDA 13 packages, and noted breaking changes to Megatron HF export and diffusion attention-quantizer export. --------- Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 Co-authored-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Signed-off-by: Rohan Joshi --- CHANGELOG.rst | 1 + modelopt/torch/speculative/utils.py | 88 +++++++++++++++---- .../speculative/test_speculative_utils.py | 74 +++++++++++++++- 3 files changed, 143 insertions(+), 20 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index e226edae3bd..b492b77c6e8 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -204,6 +204,7 @@ Changelog - Preserve accumulated histogram counts when computing an entropy calibration threshold with NumPy histograms. - Fix KV-cache quantization through ``register_attention_for_kv_quant``, the fallback for Hugging Face attention classes that do not use the attention interface, attaching the ``q``, ``k`` and ``v`` bmm quantizers to the wrong tensors for most attention layouts. Recalibrate and re-export models quantized through this fallback. Attention classes that compute their scores with ``torch.baddbmm`` (e.g. ``BloomAttention``, ``GPTNeoXJapaneseAttention``) are no longer patched, since their operands were always wired wrong (or raised ``IndexError``), and now get the existing warning to register a custom quantized attention class. - Fix the diffusers skip-softmax HF export to emit the schema TensorRT-LLM VisualGen loads (``threshold_scale_factor.coefficients`` and a scalar ``target_sparsity``). Re-export checkpoints that failed to load. +- EAGLE training with ``eagle_decoder_type="kimik2"`` now downloads the Kimi-K2 modeling code from a pinned ``moonshotai/Kimi-K2-Thinking`` revision and verifies its SHA-256 hashes before importing it. Offline runs need that revision in the Hugging Face cache. 0.47.0 (2026-09-23) ^^^^^^^^^^^^^^^^^^^ diff --git a/modelopt/torch/speculative/utils.py b/modelopt/torch/speculative/utils.py index bf33761eb74..910536fdb10 100644 --- a/modelopt/torch/speculative/utils.py +++ b/modelopt/torch/speculative/utils.py @@ -17,9 +17,10 @@ import contextlib import copy +import hashlib +import importlib.abc import importlib.util import json -import os import sys import warnings from collections import Counter, defaultdict, deque @@ -28,12 +29,19 @@ import torch import torch.distributed import transformers -from huggingface_hub import snapshot_download +from huggingface_hub import hf_hub_download from torch import nn from torch.nn.attention import SDPBackend, sdpa_kernel from transformers.cache_utils import DynamicCache KIMI_K2_REPO_ID = "moonshotai/Kimi-K2-Thinking" +# These third-party files are executed, so pin them; bump the revision and hashes together after +# reviewing the upstream diff. +KIMI_K2_REVISION = "a51ccc050d73dab088bf7b0e2dd9b30ae85a4e55" +KIMI_K2_FILE_SHA256 = { + "configuration_deepseek.py": "54f29dd0d56e9a13f639c703e524daa71794e360050012aea3788ad7cfc38b17", + "modeling_deepseek.py": "ce753fc9076ce13ce9557e680e2f5bd68ae8929a8225772968a9a22a1d17b7d4", +} KIMI_K2_PACKAGE_NAME = "kimi_k2_temp" @@ -499,39 +507,81 @@ def _patch_dynamic_cache_compatibility() -> None: ) -def _import_module_from_path(module_path: str, module_name: str, package_name: str): - """Dynamically import a module from file path.""" - spec = importlib.util.spec_from_file_location(module_name, module_path) - if spec is None or spec.loader is None: +class _VerifiedSourceLoader(importlib.abc.InspectLoader): + """Execute already-verified source bytes, never re-reading the file or its ``__pycache__``.""" + + def __init__(self, path: str, source: bytes): + self.path = path + self.source = source + + def get_source(self, fullname): + return importlib.util.decode_source(self.source) + + def get_code(self, fullname): + return self.source_to_code(self.source, self.path) + + +def _download_verified_kimi_k2_file(filename: str) -> tuple[str, bytes]: + """Download ``filename`` at the pinned revision and return its path and hash-checked content.""" + # Not snapshot_download: for a commit revision it calls the Hub even with HF_HUB_OFFLINE set, + # so cached runs would fail offline. + path = hf_hub_download(KIMI_K2_REPO_ID, filename, revision=KIMI_K2_REVISION) + with open(path, "rb") as f: + source = f.read() + digest = hashlib.sha256(source).hexdigest() + expected = KIMI_K2_FILE_SHA256[filename] + if digest != expected: + raise RuntimeError( + f"sha256 of {path} is {digest}, expected {expected}; refusing to execute it." + ) + return path, source + + +def _import_module_from_source( + module_path: str, source: bytes, module_name: str, package_name: str +): + """Import already-verified ``source`` as ``module_name``, with ``module_path`` as its file.""" + loader = _VerifiedSourceLoader(module_path, source) + spec = importlib.util.spec_from_file_location(module_name, module_path, loader=loader) + if spec is None: raise ImportError(f"Failed to load module {module_name} from {module_path}") module = importlib.util.module_from_spec(spec) module.__package__ = package_name sys.modules[module_name] = module - spec.loader.exec_module(module) + try: + loader.exec_module(module) + except BaseException: + # Like importlib, don't leave a half-initialized module registered + sys.modules.pop(module_name, None) + raise return module def _setup_kimi_k2_decoder(): """Setup Kimi-K2 decoder with proper imports and patches.""" - # Download repository files - kimi_k2_path = snapshot_download(repo_id=KIMI_K2_REPO_ID, allow_patterns="*.py") - - # Patch deprecated usage in Kimi implementation - _patch_dynamic_cache_compatibility() + # Verify both files before executing either, so a refusal runs none of the downloaded code + config_module_path, config_source = _download_verified_kimi_k2_file("configuration_deepseek.py") + model_module_path, model_source = _download_verified_kimi_k2_file("modeling_deepseek.py") # Import required modules - config_module_path = os.path.join(kimi_k2_path, "configuration_deepseek.py") - model_module_path = os.path.join(kimi_k2_path, "modeling_deepseek.py") - - _import_module_from_path( - config_module_path, f"{KIMI_K2_PACKAGE_NAME}.configuration_deepseek", KIMI_K2_PACKAGE_NAME + _import_module_from_source( + config_module_path, + config_source, + f"{KIMI_K2_PACKAGE_NAME}.configuration_deepseek", + KIMI_K2_PACKAGE_NAME, ) - kimi_k2_module = _import_module_from_path( - model_module_path, f"{KIMI_K2_PACKAGE_NAME}.modeling_deepseek", KIMI_K2_PACKAGE_NAME + kimi_k2_module = _import_module_from_source( + model_module_path, + model_source, + f"{KIMI_K2_PACKAGE_NAME}.modeling_deepseek", + KIMI_K2_PACKAGE_NAME, ) + # Patch deprecated usage in Kimi implementation + _patch_dynamic_cache_compatibility() + # Patch Kimi Attention to init rope lazily to avoid save/load meta tensor error original_init_rope = kimi_k2_module.DeepseekV3Attention._init_rope original_forward = kimi_k2_module.DeepseekV3Attention.forward diff --git a/tests/unit/torch/speculative/test_speculative_utils.py b/tests/unit/torch/speculative/test_speculative_utils.py index 3de4999a69e..3ac7e51bbbd 100644 --- a/tests/unit/torch/speculative/test_speculative_utils.py +++ b/tests/unit/torch/speculative/test_speculative_utils.py @@ -19,12 +19,22 @@ was 2 (field count) and every conversation got dropped by the num_input_tokens <= 10 filter. """ +import hashlib +import sys +from unittest.mock import patch + import pytest import torch from _test_utils.torch.transformers_models import get_tiny_tokenizer from transformers import BatchEncoding -from modelopt.torch.speculative.utils import get_conversation_input_ids +import modelopt.torch.speculative.utils as spec_utils +from modelopt.torch.speculative.utils import ( + KIMI_K2_FILE_SHA256, + KIMI_K2_PACKAGE_NAME, + _setup_kimi_k2_decoder, + get_conversation_input_ids, +) CONVERSATIONS = [ {"role": "user", "content": "Explain why the sky is blue in a few sentences."}, @@ -73,3 +83,65 @@ def apply_chat_template(self, conversations, **kwargs): assert input_ids == expected # The pre-fix code saw len() in {1, 2} here and silently dropped the conversation. assert len(input_ids) > 10 + + +def _serve_kimi_k2_files(tmp_path, monkeypatch, sources): + """Serve ``sources`` as the downloaded Kimi-K2 files, with their hashes pinned.""" + for filename, source in sources.items(): + (tmp_path / filename).write_text(source) + digest = hashlib.sha256(source.encode()).hexdigest() + monkeypatch.setitem(KIMI_K2_FILE_SHA256, filename, digest) + monkeypatch.setattr( + spec_utils, "hf_hub_download", lambda repo_id, filename, **kwargs: str(tmp_path / filename) + ) + # Keep the global DynamicCache patch out of the rest of the test session. + monkeypatch.setattr(spec_utils, "_patch_dynamic_cache_compatibility", lambda: None) + + +def test_kimi_k2_decoder_refuses_unverified_code(tmp_path, monkeypatch): + """No downloaded Kimi-K2 code runs unless every file matches its pinned hash.""" + marker = tmp_path / "executed" + run_marker = f"open({str(marker)!r}, 'w').close()\n" + _serve_kimi_k2_files(tmp_path, monkeypatch, dict.fromkeys(KIMI_K2_FILE_SHA256, run_marker)) + # Only the modeling file fails its check, so the valid config file must not run either. + (tmp_path / "modeling_deepseek.py").write_text(run_marker + "# tampered\n") + + with pytest.raises(RuntimeError): + _setup_kimi_k2_decoder() + assert not marker.exists() + + +def test_kimi_k2_decoder_imports_verified_code(tmp_path, monkeypatch): + """Verified files load as one package, so the modeling file's relative import resolves.""" + modeling = ( + "from .configuration_deepseek import DeepseekV3Config\n" + "class DeepseekV3Attention:\n" + " _init_rope = forward = None\n" + "class DeepseekV3DecoderLayer:\n" + " config_class = DeepseekV3Config\n" + " forward = None\n" + ) + config = "class DeepseekV3Config:\n pass\n" + _serve_kimi_k2_files( + tmp_path, + monkeypatch, + {"configuration_deepseek.py": config, "modeling_deepseek.py": modeling}, + ) + + with patch.dict(sys.modules): + decoder_cls = _setup_kimi_k2_decoder() + assert decoder_cls.config_class.__module__ == f"{KIMI_K2_PACKAGE_NAME}.configuration_deepseek" + + +def test_kimi_k2_decoder_drops_failed_module(tmp_path, monkeypatch): + """A module that raises while importing is not left half-initialized in sys.modules.""" + _serve_kimi_k2_files( + tmp_path, + monkeypatch, + {"configuration_deepseek.py": "", "modeling_deepseek.py": "raise ImportError\n"}, + ) + + with patch.dict(sys.modules): + with pytest.raises(ImportError): + _setup_kimi_k2_decoder() + assert f"{KIMI_K2_PACKAGE_NAME}.modeling_deepseek" not in sys.modules From 164e15331dab008248d2594bb25b3606562913e5 Mon Sep 17 00:00:00 2001 From: Shengliang Xu <106840466+shengliangxu@users.noreply.github.com> Date: Tue, 27 Oct 2026 14:03:46 -0700 Subject: [PATCH 43/43] [NVBug 6938144] Stop making /workspace world-writable in the vLLM serve example image (#2997) 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 `examples/vllm_serve/Dockerfile` ran `chmod -R 777 /workspace` before switching to the `vllm` user. The precompiled CUDA extensions (`TORCH_EXTENSIONS_DIR=/workspace/torch_extensions`) and the editable ModelOpt install both live there and are loaded into the vLLM serving process at import, so any other user in the container, or on a shared volume, could replace them and run code as the serving user. - `/workspace` is now owned by the `vllm` user (`chown -R vllm`) and made non-writable for group and others (`chmod -R go-w`); root can still write. - README note on running the image as a different user (see the behavior change below). - 0.48.0 changelog entry, placed at the top of the Bug Fixes list so it does not collide with entries appended at the end. The base image's `vllm` user has UID 2000 and GID 0, and no `vllm` group exists (upstream `docker/Dockerfile` at v0.30.0), so only the user is chowned; `chown -R vllm:vllm` would fail the build. Group-0 write access (the OpenShift `g=u` convention) is deliberately not granted: every process in such a cluster runs with GID 0, which would leave the shared-volume case open. Not included: digest-pinning the base image and dropping the `dev-test` extras from the serving image. Behavior change: running the image as a user other than `vllm` no longer allows writing to `/workspace`. torch's JIT loader needs to write a lock file and `build.ninja` in the extension directory even when the `.so` is current, so such a run skips the precompiled CUDA extensions: ModelOpt warns and falls back to slower eager fake quantization. To run as another user, point `TORCH_EXTENSIONS_DIR` at a directory that user can write; the extensions are then built there on first use. ### Usage Building the image is unchanged (see `examples/vllm_serve/README.md`). Running it as a user other than `vllm`: ```bash mkdir -p ext_cache docker run --user $(id -u) -v "$PWD/ext_cache:/ext_cache" -e TORCH_EXTENSIONS_DIR=/ext_cache \ vllm-modelopt:v0.30.0 ``` ### Testing | Check | Result | |---|---| | pre-commit (rst, markdown, license, whitespace hooks) | pass | | Reduced image: the real base image `vllm/vllm-openai:v0.30.0`, stand-in files left world-writable as the old `chmod -R 777` left them, and the exact `RUN chown -R vllm /workspace && chmod -R go-w /workspace` line; then run as several users (no GPU, no network) | `/workspace` ends up `755 vllm:root` everywhere. As `vllm`, as root and as `2000:0`, creating a file, editing a source file, overwriting the `.so` and taking the lock file that torch's extension loader takes in the build directory all succeed. As UID 1234 with GID 0 or with its own GID, all four fail with `Permission denied`. With the README snippet (the caller's UID, a writable `/ext_cache`, `TORCH_EXTENSIONS_DIR=/ext_cache`) the three `/workspace` writes still fail and the lock is taken in `/ext_cache` | | Full `docker build` of the image (ModelOpt install) and a GPU serving run | not run: CI never builds this Dockerfile and the base image is several GB | | New test | none: a Dockerfile-text check would only restate the change | ### 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?: ❌ Deliberate: the image no longer lets arbitrary users write to `/workspace` (see the behavior change above). - 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 (see Testing) - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ - Did you get Claude approval on this PR?: ✅ `/claude review` approved with one suggestion, addressed in the follow-up commit. ### Additional Information NVBug 6938144. Targets the 0.48.0 release. ## Summary by CodeRabbit * **Bug Fixes** * The Docker image’s `/workspace` directory is writable only by `vllm`, rather than by all users. Running the image as another user may require configuring a writable extensions cache. * **Documentation** * Added instructions for running the image as another user by setting `TORCH_EXTENSIONS_DIR` to a writable directory. Without it, ModelOpt cannot load its CUDA extensions and falls back to slower eager fake quantization. --------- Signed-off-by: Shengliang Xu Signed-off-by: Rohan Joshi --- CHANGELOG.rst | 1 + examples/vllm_serve/Dockerfile | 8 ++++---- examples/vllm_serve/README.md | 11 +++++++++++ 3 files changed, 16 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index b492b77c6e8..4058290db06 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -138,6 +138,7 @@ Changelog - Harden the default locations of the TensorRT deployment runtime's build artifacts and the ONNX AutoTune timing cache: they no longer sit at fixed paths under the shared system temp directory, where another local user could pre-create them to plant symlinks or swap a built engine. Both now live in a private temporary directory that is created once per process; pass ``timing_cache_path`` (runtime) or ``--timing_cache`` (AutoTune) to keep a timing cache across processes. - Fix ``import modelopt.torch`` failing with ``ModuleNotFoundError`` for ``requests`` or ``huggingface_hub`` when ``nvidia-modelopt`` is installed without the ``[hf]`` extra. - 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. +- Harden ``examples/vllm_serve/Dockerfile``: ``/workspace`` is no longer world-writable, so no other unprivileged user in the container can replace the CUDA extensions or the ModelOpt sources that the vLLM server loads. The directory is now owned by the ``vllm`` user; rebuild the image, and to run it as a different user point ``TORCH_EXTENSIONS_DIR`` at a directory that user can write, since ModelOpt otherwise cannot load its CUDA extensions and falls back to slower eager fake quantization. - 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/examples/vllm_serve/Dockerfile b/examples/vllm_serve/Dockerfile index 8bb3f8c052d..2f6af8e62ac 100644 --- a/examples/vllm_serve/Dockerfile +++ b/examples/vllm_serve/Dockerfile @@ -27,13 +27,13 @@ RUN cd Model-Optimizer && \ # Llama4 requires this RUN pip install flash-attn==2.7.4.post1 --no-build-isolation -# Pre-compile CUDA extensions into a world-accessible directory so the vllm -# user can use the cache at runtime. +# Pre-compile CUDA extensions so the vllm user can use the cache at runtime. ENV TORCH_EXTENSIONS_DIR=/workspace/torch_extensions RUN python3 -c "import modelopt.torch.quantization.extensions as ext; ext.precompile()" || true -# Allow the non-root vllm user to access the workspace -RUN chmod -R 777 /workspace +# Only the vllm user (and root) may write here, since this code is loaded into the serving process. +# The base image's vllm user has GID 0 and no vllm group exists, so chown the user only. +RUN chown -R vllm /workspace && chmod -R go-w /workspace USER vllm diff --git a/examples/vllm_serve/README.md b/examples/vllm_serve/README.md index 8eaea64700d..05ebd4f4e1b 100644 --- a/examples/vllm_serve/README.md +++ b/examples/vllm_serve/README.md @@ -23,6 +23,17 @@ docker build -f examples/vllm_serve/Dockerfile -t vllm-modelopt:v0.30.0 . The Dockerfile defaults to vLLM 0.30.0. Override `VLLM_VERSION` only when building for a different tested release. +The image runs as the `vllm` user, which owns `/workspace`. To run it as a different user, point +`TORCH_EXTENSIONS_DIR` at a directory that user can write, for example a mounted volume. ModelOpt +then builds its CUDA extensions there on first use rather than under `/workspace`; without it, +ModelOpt warns and falls back to slower eager fake quantization. + +```bash +mkdir -p ext_cache +docker run --user $(id -u) -v "$PWD/ext_cache:/ext_cache" -e TORCH_EXTENSIONS_DIR=/ext_cache \ + vllm-modelopt:v0.30.0 +``` + For a direct installation from the ModelOpt repository root, install the tested vLLM release and the ModelOpt extras used by this example: