Puzzletron v2 sync 20260730 - #2037
Draft
grzegorz-k-karch wants to merge 183 commits into
Draft
Conversation
) ### What does this PR do? Type of change: Bug fix Fix `--quant_cfg` CLI parsing by typing `quant_cfg` as `str | None` instead of `str | QuantizeConfig | None` ### Testing ``` accelerate launch --config_file examples/gpt-oss/configs/zero3.yaml examples/gpt-oss/sft.py --config examples/gpt-oss/configs/sft_full.yaml --model_name_or_path openai/gpt-oss-20b --quant_cfg MXFP4_MLP_WEIGHT_ONLY_CFG --output_dir gpt-oss-20b-qa ``` ### 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?: 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 <!--- Run `/claude review`. NVIDIA org members can self-trigger for complex changes; orthogonal to CodeRabbit. --> ### Additional Information <!-- E.g. related issue. --> <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Refactor** * Quantization config parameter now accepts string identifiers or none; resolution behavior for named presets remains unchanged. * **Documentation** * Updated argument reference to reflect the parameter type change while preserving the deprecation note and usage guidance. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Kinjal Patel <kinjalpravin@nvidia.com>
## Summary - Moves `tools/launcher/examples/gemma-4/gemma-4-E4B-it/specdec_bench_mtp_vllm.yaml` to `tools/launcher/examples/google/gemma-4-E4B-it/specdec_bench_mtp_vllm.yaml` to align with the `google/` vendor-prefix convention used by other models in the launcher examples directory. - Updates the embedded `--yaml` path in the YAML comment to match the new location. ## Test plan - [ ] Confirm the file appears at the new path - [ ] Confirm the old path is gone <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Chores** * Updated an example benchmark invocation reference to reflect the project’s reorganized directory layout. This edits only the example comment in the benchmark YAML; no benchmark settings, model paths, task arguments, or runtime configuration were modified. Low-risk documentation alignment. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: Chenhan D. Yu <5185878+ChenhanYu@users.noreply.github.com>
### What does this PR do? Type of change: new tests / tooling Adds pre-commit validation for speculative-decoding recipes (the existing `check-modelopt-recipes` hook only ran on PTQ) and for launcher YAML references into the recipe library. - `tools/precommit/check_modelopt_recipes.py`: accept `speculative_eagle` / `speculative_dflash` / `speculative_medusa` in addition to `ptq`, so per-model spec-dec recipes (e.g. `modelopt_recipes/models/Qwen3-8B/dflash.yaml`) get full Pydantic validation via `load_recipe()` at commit time. - `tools/precommit/check_launcher_yaml.py` (new): scans every `tools/launcher/examples/**/*.yaml` for `--config <path>` and `data.chat_template=<path>` references, verifies the resolved files exist, and runs `load_recipe()` on any path under `modelopt_recipes/`. Skips `<<global_vars.x>>` interpolation. `pass_filenames: false` so recipe-side edits also re-validate all launcher references. ### Usage ```bash pre-commit run check-modelopt-recipes --all-files pre-commit run check-launcher-yaml --all-files ``` ### Testing Smoke-tested both hooks manually: | Scenario | Result | |---|---| | spec-dec recipe with `dflash_block_size: not_an_int` | exit 1, Pydantic int_parsing error | | launcher YAML with non-existent `--config` path | exit 1, source file + resolved path reported | | launcher YAML with non-existent `data.chat_template` path | exit 1 | | Current repo state (all valid) | exit 0 | ### 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?: ✅ (hooks themselves are the 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?: ❌ pending ### Additional Information Motivated by the per-model recipe migration in #TBD — without these hooks, broken `--config` paths and recipe schema typos surface only at CI or runtime. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Chores** * Added an automated pre-commit check that validates launcher example YAMLs, reporting parse errors and missing or invalid references. * Expanded recipe validation to cover additional recipe types beyond PTQ, improving detection of invalid recipe formats and metadata. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
…1687) ### What does this PR do? Type of change: Bug fix Exclude Qwen visual and vision_tower modules from NVFP4 quantization and keep the Qwen linear attention projection exclusions. These modules can produce matrix dimensions that are incompatible with vLLM 0.22.1's ModelOpt FP4 Marlin fallback path, causing checkpoint load or profiling failures such as `size_n = 4304 is not divisible by tile_n_size = 64`. ### Usage N/A. This is a recipe configuration fix. ### Testing - `python -m pytest tests/unit/recipe/test_presets.py tests/unit/recipe/test_loader.py -q` - `python -m pre_commit run --files modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml tests/unit/recipe/test_loader.py tests/unit/recipe/test_presets.py` - E2E validation with `vllm/vllm-openai:v0.22.1`: PTQ export validation passed with zero Marlin-incompatible quantized layers, and vLLM `/health`, `/v1/models`, and `/v1/completions` passed. The final PR broadens the validated visual MLP exclusions to the full `*visual*` subtree and adds the common `*vision_tower*` naming pattern. ### 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?: Yes - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: Yes - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A - Did you get Claude approval on this PR?: N/A ### Additional Information N/A <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Tests** * Added unit tests that verify the built-in PTQ recipe and preset correctly disable incompatible projection and visual components for certain quantization modes. * Ensures quantization settings are validated across recipes and presets. * **Chores** * Updated quantization configuration to disable quantizers for select projection and vision-related model layers. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
…vllm.ai) (#1685) ### What does this PR do? Type of change: documentation (agent skill) The `deployment` and `evaluation` skills cross-check vLLM launch commands against `recipes.vllm.ai`, but the SGLang reference had **no** equivalent authoritative recipe source. This adds the **SGLang cookbook** (`docs.sglang.io/cookbook`) — the direct SGLang analog of `recipes.vllm.ai`, using the same `(hw, variant, quant, strategy, nodes)` variant-tuple command generator — as the cross-check source for SGLang deployments. Changes (canonical `.agents/` source; `.claude/` are symlinks): - **`SKILL.md`**: cookbook cross-check note under the SGLang quick-start, mirroring the existing vLLM cu130/sm_103 block — covers the JS-rendered (Mintlify) page, the raw-markdown fallback in `sgl-project/sglang` `docs_new/`, the URL-fragment variant selection, and the SM120 (RTX PRO 6000) nightly-image gotcha. - **`references/sglang.md`**: a new authoritative-recipe section, a MoE/FP4 backend flag matrix (`flashinfer_mxfp4` vs `marlin`, `deepep` vs `megamoe`), and per-hardware notes (Blackwell / Hopper / RTX PRO 6000). ### Usage N/A — documentation only. Agents now cross-check SGLang launch flags against `docs.sglang.io/cookbook/<category>/<org>/<model>`, fetching the raw markdown at `github.com/sgl-project/sglang/blob/main/docs_new/cookbook/<category>/<org>/<model>.mdx` when the page is JS-rendered. ### Testing Markdown-only change to skill docs; no code paths affected. ### 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?: N/A (docs only) - Did you update Changelog?: N/A (agent skill docs, not a library feature/API change) - Did you get Claude approval on this PR?: ❌ (pending) 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Documentation** * Added authoritative cookbook cross-check guidance for SGLang deployments and how to select model variants via recipe parameters. * Explained how to fetch raw cookbook content when rendered pages omit commands and which flag areas the cookbook governs. * Expanded backend/strategy guidance for MoE, dispatch modes, and runtime tuning (batching, request sizing, CUDA-graph). * Added hardware-specific notes including RTX PRO 6000 nightly-image requirement and processor-specific recommendations. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Zhiyu Cheng <zhiyuc@nvidia.com>
…295242) (#1678) ### What does this PR do? Type of change: Bug fix Fixes the GPT-OSS MXFP4 → NVFP4 PTQ path (`examples/llm_ptq/hf_ptq.py` with `--cast_mxfp4_to_nvfp4`), which failed in three independent ways. The documented command now runs end-to-end and produces a bit-exact (100% lossless) NVFP4 checkpoint. Addresses **nvbug 6295279** (OMNIML-5046) and **nvbug 6295242** (OMNIML-5045). 1. **nvbug 6295242 — CUDA illegal memory access on load.** GPT-OSS ships native MXFP4 weights that Transformers dequantizes to BF16; the threaded weight loader trips an illegal-memory access when `device_map="auto"` shards the dequant across **multiple GPUs**. The missing optional `kernels` package only *forces* the dequant path — it is not the root cause. `get_model` now detects MXFP4 checkpoints and loads them with `Mxfp4Config(dequantize=True)` on a **sequential** device map so the dequant stays on a single device. `kernels` is no longer required. 2. **nvbug 6295279 #1 — `NotImplementedError: Mxfp4GptOssExperts` during unified HF export.** Forcing `dequantize=True` yields plain `GptOssExperts` (even when `kernels` is installed), which ModelOpt wraps and exports normally. 3. **nvbug 6295279 #2 — `FileNotFoundError` in the cast step.** `--cast_mxfp4_to_nvfp4` treated `--pyt_ckpt_path` as a local dir; a HF Hub ID now resolves to its cached snapshot dir via `_resolve_model_path`. Also fixes a **static-block NVFP4 regression** (surfaced by the cast's `force_weight_quantizers_static`, introduced by #1560's now-unconditional `weight_only_quantize`): `_QuantGptOssExperts` / `_QuantLlama4TextExperts` quantize their expert weights transposed in the forward (`_transposed_quantize`), but the inherited `iter_weights_for_calibration` fed the non-transposed weight, locking a mismatched block-quant `_original_shape` and raising `ValueError: Input shape has changed`. The override now calibrates on the transposed view, matching both the forward and the export's `_amax` orientation. ### Why this regressed (it worked when the cast was added) `get_model` never had explicit handling for a *natively pre-quantized MXFP4* checkpoint — GPT-OSS fell through the generic *unquantized-checkpoint* branch and relied on Transformers' **implicit** MXFP4 behavior, which is fragile across three axes. The cast was originally validated (#1372, 2026-05-01) in the "lucky" quadrant of each: - **GPU count:** `device_map="auto"` on a single GPU never shards, so the dequant stays on one device. On multiple GPUs `auto` balances the model and shards the MXFP4→BF16 dequant across devices → CUDA illegal-memory crash (6295242). - **`kernels` presence:** without `kernels`, Transformers auto-dequantizes to BF16 `GptOssExperts` (exportable). With `kernels` installed it keeps the packed `Mxfp4GptOssExperts` kernel path → export `NotImplementedError` (6295279 #1). - **Transformers version:** the kernel-backed experts wrapper and the threaded multi-GPU weight loader are newer-Transformers behavior (env here is 5.5.4). Earlier versions simply dequantized MXFP4 → BF16, which is what the old generic path happened to need. The QA env sat in the *breaking* quadrant (multi-GPU and/or `kernels` present, newer Transformers), so the implicit path failed. The new branch makes both decisions explicit and deterministic (`dequantize=True` + single-device load), regardless of environment — mirroring the existing `has_pack_quantized_config` branch for compressed-tensors checkpoints. The fourth issue (static-block `Input shape has changed`) is a separate regression: it was introduced by **#1560 (2026-06-02, "Make sure all weight quantizers have `_amax`")**, a month *after* the cast landed. #1560 made `weight_only_quantize` unconditional in `max_calibrate`; previously it ran only when no calibration `forward_loop` was supplied, and the cast always supplies one — so the non-transposed weight-quantizer call simply never happened before. The conflict only appears at the intersection of (a) transposed-quantize experts (GPT-OSS/Llama4), (b) static-block NVFP4 — which `--cast_mxfp4_to_nvfp4` forces via `force_weight_quantizers_static` — and (c) #1560. CI's GPT-OSS NVFP4 coverage uses the *dynamic*-block path, which never locks the block shape, so #1560 looked safe. ### Usage ```bash python hf_ptq.py \ --pyt_ckpt_path openai/gpt-oss-20b \ --qformat nvfp4_mlp_only \ --cast_mxfp4_to_nvfp4 \ --export_path ./gpt-oss-20b-nvfp4 ``` ### Testing - Ran the documented command end-to-end on 2xB200 (`openai/gpt-oss-20b`): cast overrode **48/48** expert weight quantizers, **100% lossless** layers/blocks, exported a valid packed-NVFP4 HF checkpoint (uint8 weights + FP8 per-block `weight_scale` + per-tensor `weight_scale_2` + `hf_quant_config.json`). - Verified plain `--qformat nvfp4_mlp_only` (no cast) still works end-to-end. - **Independently verified the export is bit-exact:** dequantized the exported NVFP4 weights (ModelOpt's E2M1 LUT + pack layout) and compared against Transformers' canonical MXFP4→BF16 dequant (`Mxfp4Config(dequantize=True)`) over all 24 layers × both expert weights — `max_abs_err = 0`, 100% bitwise-equal in bf16. So `dequant(exported NVFP4) == dequant(original MXFP4)` exactly. - New unit tests: `test_get_original_hf_quant_method_*` (load detection) and `test_gpt_oss_experts_iter_weights_for_calibration_transposed` (the transpose regression). Existing `test_cast_mxfp4_to_nvfp4.py` (8 tests) still pass. `pre-commit` clean. **Known limitation:** verified for gpt-oss-20b (fits one GPU). gpt-oss-120b dequantized does not fit a single GPU, so `sequential` would still span GPUs — that case would need a CPU-dequant-then-dispatch path and is left as a follow-up. ### 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)?: ✅ (0.45 Bug Fixes) - Did you get Claude approval on this PR?: ❌ (not yet run) ### Additional Information nvbug 6295279, nvbug 6295242 / OMNIML-5046, OMNIML-5045. 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Bug Fixes** * Prevented CUDA illegal-memory access during MXFP4→NVFP4 casting. * Fixed expert-weight calibration orientation to avoid shape mismatches. * **New Features** * Support loading native MXFP4 checkpoints with automatic dequantization. * Resolve remote model identifiers to local checkpoints when casting MXFP4→NVFP4, improving reliability. * **Tests** * Added unit and GPU regression tests covering quant-method detection, casting, and expert-weight calibration. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
…6293731, 6293762) (#1691) ### What does this PR do? Type of change: Bug fix Fixes two sglang deployment failures on multimodal Gemma (`gemma-4-31B-it`) caused by general PTQ presets leaking quantization into the SigLIP vision branch via broad wildcards: - **NVBug 6293731** — `general/ptq/fp8_default-kv_fp8`: the `w8a8_fp8_fp8` unit enables bare `*weight_quantizer` / `*input_quantizer`, which also match the vision tower (`model.vision_tower.*`, `model.visual.*`) and the vision embedding projection (`model.embed_vision.*`). The exported checkpoint deploys but emits **garbled text** in sglang. - **NVBug 6293762** — `general/ptq/nvfp4_mlp_only-kv_fp8`: the `*mlp*` enables also match the vision tower's block MLPs (`model.vision_tower.encoder.layers.*.mlp`), and an image request **crashes** the FP4 kernel at decode: `ValueError: too many values to unpack (expected 2)` in sglang's `modelopt_quant.py` `apply`. ### Fix Add `*embed_vision*` / `*vision_tower*` / `*visual*` disable rules to the shared `configs/ptq/units/default_disabled_quantizers` unit, alongside the existing `*router*` / `*lm_head*` entries. Both the composed `general/ptq/*` recipes **and** the `configs/ptq/presets/model/*` presets import this unit, so: - every general recipe (`fp8_default`, `nvfp4_default`, `nvfp4_mlp_only`, `nvfp4_omlp_only`, …) keeps the vision branch in BF16 by default — fixing the whole vision-overreach class, not just the two reported recipes; - the `test_general_ptq_yaml_matches_config_dicts` YAML↔preset parity test stays satisfied (both sides pick up the new entries from the one shared unit). The rules are **no-ops on text-only models** (nothing matches). A recipe that intentionally wants to quantize the vision branch can re-enable these after importing the unit. Files changed: - `modelopt_recipes/configs/ptq/units/default_disabled_quantizers.yaml` (+14) ### Testing Re-export of `gemma-4-31B-it` with the affected recipes and re-deploy in sglang (the env from the bug reports: `lmsysorg/sglang:v0.5.12.post1`, GB200) to confirm fp8_default no longer garbles text and nvfp4_mlp_only no longer crashes on image requests. _(Results to be appended.)_ Unit-level: `tests/unit/recipe/test_loader.py::test_general_ptq_yaml_matches_config_dicts` (parity) passes for all four general presets. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ (text-only checkpoints unaffected; new rules only match vision modules that should never have been quantized by a general recipe) - If you copied code from any other sources or added a new PIP dependency: N/A - Did you write any new necessary tests?: N/A (recipe data fix; covered by the existing parity test + verified by real PTQ export + sglang deploy) - Did you update Changelog?: N/A - Did you get Claude approval on this PR?: ❌ (pending) ### Additional Information NVBug 6293731 and 6293762. Reported on modelopt 0.45.0rc0, GB200, gemma-4-31B-it, sglang 0.5.12.post1. Tracked under OMNIML-5034. Companion to PR #1690 (same vision-overreach class on the gemma-specific `w4a8_awq` recipe, NVBug 6294017). 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Chores** * Updated quantization configuration to preserve BF16 precision for vision encoder components in multimodal models. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Zhiyu Cheng <zhiyuc@nvidia.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
### What does this PR do? Type of change: Bug fix Fixes the `llm_ptq` example-test failure on current `main`: ```text tests/examples/llm_ptq/test_hf_ptq_args.py::test_qwen_autoquant_disabled_layers_are_scoped_to_qwen_models ``` `*linear_attn.in_proj_a*` and `*linear_attn.in_proj_b*` were promoted into the global default disabled quantizer list, so they are no longer Qwen-only AutoQuant exclusions. This PR removes the redundant entries from the example-specific Qwen AutoQuant exclusion tuple and narrows the test to the remaining Qwen-only pattern, `*shared_expert_gate*`. ### Usage N/A ### Testing - Reproduced the failure on `origin/main` (`cc17f2c45`) with: - `python -m pytest tests/examples/llm_ptq/test_hf_ptq_args.py::test_qwen_autoquant_disabled_layers_are_scoped_to_qwen_models -vv` - Verified the fix with: - `python -m pytest tests/examples/llm_ptq/test_hf_ptq_args.py::test_qwen_autoquant_disabled_layers_are_scoped_to_qwen_models tests/unit/recipe/test_presets.py::test_w4a16_nvfp4_preset_disables_vllm_marlin_incompatible_projections tests/unit/recipe/test_loader.py::test_nvfp4_weight_only_recipe_disables_vllm_marlin_incompatible_projections -vv` - `python -m py_compile examples/llm_ptq/example_utils.py tests/examples/llm_ptq/test_hf_ptq_args.py` - `git diff --check` ### 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 Observed while checking the failing `trtllm-pr (llm_ptq) / run-test` job on PR #1697. The failure reproduces on `main`, so this PR is independent of #1697. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Chores** * Adjusted quantization exclusion configuration for the Qwen model family to narrow which layers are excluded. * **Tests** * Updated test expectations to align with the revised quantization exclusion behavior for Qwen models. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
…fo in clear_stale_value_info (#1697) ### What does this PR do? Type of change: Bug fix INT4 quantization upgrades the model to opset >= 21, at which point ONNX Runtime runs type inference while building the AWQ calibration `InferenceSession`. Custom ops backed by TensorRT plugins (domain `trt.plugins`) have no ORT type-inference function, so their output types are only known from the `value_info` that TensorRT type/shape inference populated earlier in preprocessing. `clear_stale_value_info` cleared `value_info` wholesale, dropping those types, so ORT failed output type inference for the custom op at model load, e.g.: ``` Node (Conv-2) Op (IdentityConv) output arg (X2) type inference failed ``` - `modelopt/onnx/utils.py`: in `clear_stale_value_info`, preserve `value_info` entries for outputs of `trt.plugins`-domain nodes (which ORT cannot re-derive); clear the rest as before. - `tests/gpu/onnx/quantization/test_plugin.py`: add a regression test quantizing a model with the built-in `CustomSkipLayerNormPluginDynamic` plugin at INT4 + awq_clip (the opset >= 21 path), asserting the quantized model is produced and the custom op survives. ### Usage ```python python -m modelopt.onnx.quantization \ --onnx_path=model.onnx \ --quantize_mode=int4 \ --calibration_method=awq_clip \ --trt_plugins=/path/to/plugin.so ``` ### Testing - `pytest tests/gpu/onnx/quantization/test_plugin.py -k int4_awq` — fails before the fix (ORT type-inference error at calibration-session load) and passes after. The full `test_plugin.py` (including the existing INT8 quantization and autocast cases) passes. - The example [here](https://github.com/NVIDIA/Model-Optimizer/blob/main/examples/onnx_ptq/README.md#quantize-an-onnx-model-with-custom-op) also failed before this fix, now passes. ### Before your PR is "*Ready for review*" Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors). - Is this change backward compatible?: ✅ - 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 ### Additional info Fixing regression inserted by #1565 <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Bug Fixes** * Preserve metadata for TensorRT plugin outputs during cleanup and correctly reconcile output data types so custom plugin ops remain intact after optimization/quantization. * **Tests** * Added a GPU ONNX regression test covering int4 quantization with AWQ calibration to ensure TensorRT plugins are retained. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: Gwenaelle Cunha Sergio <gcunhasergio@nvidia.com>
…#1702) ### What does this PR do? Type of change: Bug fix Fixes nvbug **6311147** (OMNIML-5103). `examples/deepseek/deepseek_v3/ptq.py` resolved the cloned DeepSeek-V3 / DeepSeek-V3.2-Exp inference repos relative to its own directory (`deepseek_v3/`) via `Path(__file__).resolve().parent`. But the [README](https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/deepseek) clones those repos into the parent `examples/deepseek/` directory and runs the script from there, so the lookup landed one level too deep and raised `ValueError: DeepSeek-V3 or DeepSeek-V3.2-Exp not found` (the error message also printed the wrong directory). The fix resolves from `parent.parent` via a single `DEEPSEEK_DIR` base shared by both repo paths and the error message. ### Usage ```bash # Run from examples/deepseek/ as documented in the README, after cloning # DeepSeek-V3 (or DeepSeek-V3.2-Exp) into that directory: torchrun --nproc-per-node 8 --master_port=12346 deepseek_v3/ptq.py \ --model_path $DS_CKPT \ --config DeepSeek-V3/inference/configs/config_671B.json \ --quant_cfg NVFP4_DEFAULT_CFG \ --output_path $FP4_QUANT_PATH ``` ### Testing - Confirmed against the repro path: with the file at `examples/deepseek/deepseek_v3/ptq.py` and the repos cloned into `examples/deepseek/`, `Path(__file__).resolve().parent.parent` now points at `examples/deepseek/` so `DeepSeek-V3/inference` resolves correctly. - Verified the sibling `examples/deepseek/deepseek_v4/` does not share the bug (it takes an explicit `--dsv4_inference_dir` argument instead). - `pre-commit` clean. ### 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?: N/A (one-line path fix in an example script that requires the DeepSeek repos + multi-GPU checkpoint to exercise) - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A (bug is in a 0.45-cycle example, not a regression from a released version) - Did you get Claude approval on this PR?: ❌ (not yet run) ### Additional Information nvbug 6311147 / OMNIML-5103. 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Bug Fixes** * Improved path resolution in the example script to more reliably locate the required inference repository. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: Chenjie Luo <chenjiel@nvidia.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
… + checkpoint I/O knobs (#1571) ### What does this PR do? Type of change: new feature <!-- Use one of the following: Bug fix, new feature, new example, new tests, documentation. --> Groups all layerwise-calibration options under a nested `LayerwiseConfig` and adds three new behavior knobs to it. All changes are backward compatible. #### 1. Nested `layerwise` config `QuantizeAlgorithmConfig.layerwise` changes from `bool` to a Pydantic submodel: ```python class LayerwiseConfig(ModeloptBaseConfig): enable: bool = False get_qdq_activations_from_prev_layer: bool = False checkpoint_dir: str | None = None save_every: int = 1 save_quantizers_only: bool = False ``` Backward compatibility: - `layerwise: True/False` still accepted (emits `DeprecationWarning`). - Flat `layerwise_checkpoint_dir` silently migrated into `layerwise.checkpoint_dir`. - Legacy `use_sequential` alias preserved (and resolved during flat-key migration so it can't be dropped). - Conflicting flat+nested `checkpoint_dir` values raise. - All 7 shipped PTQ recipes (`modelopt_recipes/general/ptq/*.yaml`, `huggingface/qwen3_5*/ptq/*.yaml`) migrated to the canonical nested shape — no semantic change. #### 2. `get_qdq_activations_from_prev_layer` — correct GPTQ vs max-calib semantics Controls what layer N's calibration sees: - **True** (GPTQ default): activations carry the quantize-dequantize error of layers 0..N-1 — GPTQ's Hessian-compensation goal. - **False** (max/mse/local_hessian default): full-precision activations, matching the non-layerwise pass exactly. The False branch wraps the next-layer input capture forward with the existing `set_quantizer_by_cfg_context` deny-all idiom (`{"quantizer_name": "*", "enable": False}`). GPTQ's per-algorithm default is enforced via `@model_validator(mode="after")` that reads `LayerwiseConfig.model_fields_set` — works for every input shape (empty constructor, bool, partial dict, full dict) and lets explicit user values override. #### 3. `save_every` — gate the large activation-cache writes `save_every: int = 1` (`ge=1`). With N > 1, the per-layer `next_inputs.pt` (cached activation tensors, the largest checkpoint artifact for most models) is only written for the boundary layer of each N-layer window. Per-layer weight/quantizer/output_meta files are still written every layer (resume needs them to replay skip layers correctly). Interrupting mid-window re-calibrates that window on resume. #### **UPDATE: drop this in current PR and moved to #1640 #### 4. `save_quantizers_only` — algorithm-aware weight-blob skipping `save_quantizers_only: bool = False`. When True, skip `weights.pt` entirely and persist just the per-quantizer `state_dict` slice (carries `_amax`) to a new `quantizer_buffers.pt`. On resume, `full_restore` reloads only the quantizer slice and trusts that algorithm semantics didn't mutate `layer.weight`. Safety is enforced by a **whitelist**: `_supports_save_quantizers_only: ClassVar[bool] = False` on `QuantizeAlgorithmConfig`, overridden to `True` only on `MaxCalibConfig`, `MseCalibConfig`, `LocalHessianCalibConfig` (audited — these only touch `_amax`). Weight-mutating algorithms (GPTQ folds Hessian updates, AWQ/SmoothQuant fold pre-quant scales) reject the flag at config-construction time so in-place weight updates can't be silently lost on resume. ### Usage ```python import modelopt.torch.quantization as mtq # GPTQ — `get_qdq_activations_from_prev_layer` defaults to True (Hessian semantics). # save_every reduces activation-cache I/O. mtq.quantize( model, { "quant_cfg": [...], "algorithm": { "method": "gptq", "layerwise": { "enable": True, "checkpoint_dir": "/path/to/ckpts", "save_every": 4, }, }, }, forward_loop=forward_loop, ) # Max-calibration — `get_qdq_activations_from_prev_layer` defaults to False (FP from prior layers). # save_quantizers_only skips the weights blob since max only updates _amax. mtq.quantize( model, { "quant_cfg": [...], "algorithm": { "method": "max", "layerwise": { "enable": True, "checkpoint_dir": "/path/to/ckpts", "save_quantizers_only": True, }, }, }, forward_loop=forward_loop, ) ``` ### Testing New / updated unit tests in `tests/unit/torch/quantization/`: - **`test_config_validation.py`** — `TestLayerwiseNestedConfig` covers nested-form acceptance, bool-form `DeprecationWarning`, flat `layerwise_checkpoint_dir` migration, conflicting flat+nested checkpoint_dir, `use_sequential` alias survival under migration, per-algorithm qdq defaults (parametrized Max/GPTQ), `save_every` `ge=1` validation, and the `save_quantizers_only` whitelist — parametrized rejection on `[GPTQ, AWQLite, SmoothQuant]` and acceptance on `[Max, Mse, LocalHessian]`. - **`test_layerwise_calibrate.py`** — - `test_layerwise_no_qdq_matches_sequential_amax` — behavioral equivalence: layerwise + `qdq=False` produces the same per-quantizer `_amax` as the non-layerwise (sequential) max-calibration flow (verified via `torch.testing.assert_close`). - `test_layerwise_save_every_writes_next_inputs_only_at_window_boundaries` — window-save layout (all layer dirs present, `next_inputs.pt` only at boundaries). - `test_layerwise_save_quantizers_only_resume_matches_one_shot_amax` — end-to-end resume: full run → manifest rewound → fresh model resumes → final `_amax` matches the one-shot baseline; also pins the on-disk shape (no `weights.pt`, `quantizer_buffers.pt` present). ## End-to-end correctness verification Ran 4 PTQ jobs on **Qwen3-8B** with NVFP4 W4A16 quant_cfg and `--calib_size 16`, one GPU each on 4 GPUs | Run | Algorithm | Layerwise config | Purpose | |---|---|---|---| | **A** | GPTQ | `enable=true, get_qdq_activations_from_prev_layer=true, save_every=5` | New nested form + the new`save_every` knob | | **B** | MSE | `enable=true, get_qdq_activations_from_prev_layer=false, save_quantizers_only=true` | New nested form + the new `save_quantizers_only` knob | | **C** | GPTQ | Legacy flat form: `layerwise: true, layerwise_checkpoint_dir: ...` | Backward-compat baseline for A (exercises the migration validator) | | **D** | MSE | `enable=false` | Non-layerwise baseline for B | Across two pairwise comparisons (A vs C ; B vs D), all 905 tensors in the exported HF checkpoints are bit-identical with hf_quant_config.json matching exactly — confirming the new layerwise knobs preserve correctness and the flat-form backward-compat path is intact. ### 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?: ✅ — `layerwise: True/False` still accepted with a `DeprecationWarning`; flat `layerwise_checkpoint_dir` silently migrated; `use_sequential` alias preserved. The two new knobs default to no-op behavior (`save_every=1`, `save_quantizers_only=False`). - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A — no new dependencies. - Did you write any new necessary tests?: ✅ — see Testing section. - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ ready for review. - Did you get Claude approval on this PR?: ✅ — `/claude review` consulted iteratively; review findings (GPTQ default survival, save_quantizers_only whitelist scope, docstring accuracy, dead `layer` param) addressed in-PR. ### Additional Information Notes on design choices that came out of internal review: - GPTQ's `qdq=True` default uses a `model_validator(mode="after")` + `model_fields_set` check rather than a `default_factory` — a `default_factory` is only fired when the field is absent, so any user-supplied dict (the natural way to enable layerwise) would silently lose the GPTQ default. - `save_quantizers_only` is enforced as a whitelist (`_supports_save_quantizers_only`) rather than a per-algorithm blacklist, which keeps future weight-mutating algorithms safe by default. - `set_quantizer_by_cfg_context` is reused for the qdq-disable scope instead of a bespoke helper, keeping `model_calib.py` aligned with the existing "deny-all" idiom documented at `conversion.py:240`. - Pre-validation recipe helpers in `examples/llm_ptq/example_utils.py` (`needs_checkpoint_path_update` / `resolve_checkpoint_dir`) accept both flat and nested shapes since they run before Pydantic validation; `resolve_checkpoint_dir` now also returns the resolved path so the caller doesn't re-derive it. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Layerwise calibration now uses a nested config (enable, get_qdq_activations_from_prev_layer, checkpoint_dir, save_every, save_quantizers_only); checkpointing supports quantizer-only saves and windowed commits with resume support. * **Behavior** * GPTQ calibration defaults get_qdq_activations_from_prev_layer=True when unspecified; save_every controls when next-inputs/manifests are written and must be positive. * **Deprecations** * Legacy flat-style layerwise keys are still accepted but emit DeprecationWarning and are auto-migrated (conflicts detected). * **Examples/UX** * Tools auto-detect legacy vs nested checkpoint layouts and report the resolved path. * **Tests** * Expanded coverage for nested configs, validation, checkpoint/resume semantics, and windowed saves. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Frida Hou <201670829+Fridah-nv@users.noreply.github.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
## Summary `tools/mcp/` — a new MCP server exposing the existing `tools/launcher/core.py` orchestration as **typed MCP tools** that codex / Claude Code agents can call directly, instead of shelling out to `uv run launch.py --yaml ...` and parsing prose output. Tracked under [OMNIML-5123](https://jirasw.nvidia.com/browse/OMNIML-5123) (Epic). Ships **Phase 1 + Phase 1.5** together: the core launcher surface plus the four highest-leverage helpers from the `cell.md` simplification loop ([OMNIML-5128](https://jirasw.nvidia.com/browse/OMNIML-5128) partial, [OMNIML-5132](https://jirasw.nvidia.com/browse/OMNIML-5132) full). ## Nine tools **Phase 1 — core launcher surface:** | Tool | Description | |---|---| | `list_examples` | Enumerate `tools/launcher/examples/` with model + description metadata extracted from each YAML | | `verify_setup` | Fail-fast probe for the named executor. Docker: `docker info` (daemon up) + `docker info --format` runtime-registry check for the `nvidia` runtime — no image pull, daemon-fast. Slurm: `ssh -o BatchMode=yes -o ConnectTimeout=5` to the cluster login node. ~1 s probe saves 30+ s of wasted submission on bad config | | `submit_job` | Submit a launcher YAML. Mode is determined by mutually-exclusive args: `hf_local` → Docker (local GPU), `cluster_host` → Slurm (remote SSH). Returns experiment_id immediately; the actual job runs detached | | `job_status` | Filesystem-based status from nemo_run's experiment dir (`_DONE`, `status_*.out`) — no in-memory registry, survives MCP server restarts | | `job_logs` | Read `log_<task>.out` from experiment dir; per-task filtering + optional tail | **Phase 1.5 — `cell.md` simplification (OMNIML-5128 / 5132):** | Tool | Description | |---|---| | `wait_for_experiment` | Replaces the agent's `while True: status; sleep` poll loop with one tool call. Reuses `job_status_impl` so terminal-state semantics stay identical. Returns final status plus `waited_seconds`; on timeout returns structured `{ok: False, reason: "wait_timeout", last_status: …}` | | `provision_passwordless_ssh_dry_run` | No-side-effect inspection of `~/.ssh/` that emits the exact `ssh-keygen` / `ssh-copy-id` commands the operator should run to make `verify_setup(executor='slurm')` pass. Closes the verify_setup "ssh_auth_failed → now what?" gap | | `read_cluster_artifact` | Uses nemo_run's tunnel primitives, not a reinvented SSH layer. `path=None` wraps `nemo experiment logs <id> <job_idx>` (built-in log fetch); with a `path`, uses the experiment's `Tunnel` to read the file. Structured failure on subprocess error / timeout | | `open_draft_pr` | `git push -u origin HEAD` + `gh pr create --draft …`. Validates cwd is a git repo first; on gh failure after push succeeds, reports `branch_pushed=True` so the operator can retry just the PR-open step | ## Design constants 1. **Single `submit_job` with mode by args** (not separate `submit_docker` / `submit_slurm` tools). Keeps the LLM tool catalog compact; mutual-exclusion is a runtime check. 2. **Filesystem is the source of truth** for status + logs. No in-memory registry. Survives MCP server restarts cleanly — important because operators / agents kill + restart their hosts often. 3. **`verify_setup` is auto-called by `submit_job`** by default (skippable when caller just probed). The probe is ~1 s; the cost of a misconfigured submission is 30+ s of cluster timeout or container-pull. Always-on verify pays back immediately. 4. **Delegate to nemo_run for tunnels.** `read_cluster_artifact` and `wait_for_experiment` use nemo_run's existing `Experiment` / `Tunnel` / `nemo experiment logs` primitives rather than reinventing SSH/rsync. One source of truth for cluster I/O. ## Layout ``` tools/mcp/ ├── pyproject.toml # name: modelopt-mcp, console_script ├── modelopt_mcp/ │ ├── __init__.py │ ├── server.py # FastMCP entry; 9 tool definitions │ └── bridge.py # thin wrapper over launcher's core.py │ # + filesystem status/log helpers │ # + tunnel/PR helpers (Phase 1.5) └── tests/ └── test_bridge.py # 34 unit tests, fully hermetic # (mocked subprocess + tmp_path fixtures) ``` ## Install Two paths, both **from source via uv**. No PyPI wheel exists; OMNIML-5123 opted for the uvx-from-git pattern to skip publication overhead. ### End-user install (recommended) `uvx` from the git subdirectory — single command, no manual clone: ```bash # Claude Code claude mcp add modelopt -- uvx --from \ "git+https://github.com/NVIDIA/Model-Optimizer.git#subdirectory=tools/mcp" \ modelopt-mcp # Codex codex mcp add modelopt -- uvx --from \ "git+https://github.com/NVIDIA/Model-Optimizer.git#subdirectory=tools/mcp" \ modelopt-mcp ``` Under the hood `uvx` clones the whole repo to its cache, installs `tools/mcp/` as the entry, and resolves the sibling `modelopt-launcher` dep via `[tool.uv.sources]` (`path = "../launcher"`) inside the cloned tree. ### Dev install (local checkout) ```bash uv pip install -e tools/launcher # sibling dep first uv pip install -e tools/mcp # then this package modelopt-mcp # entry on PATH ``` ### Why no plain `pip install` today Two specific reasons, worth flagging so reviewers know what's intentional vs missing: 1. **Nothing on PyPI yet.** Neither `modelopt-mcp` nor `modelopt-launcher` are published — this PR introduces the package but doesn't add release machinery. 2. **`pip` doesn't read `[tool.uv.sources]`.** Even from a local checkout, plain `pip install -e tools/mcp` fails because `modelopt-launcher` is a bare name (no URL) and pip can't find it. Sticking with `uv` / `uvx` is the practical path while we're git-only. If we later want plain-pip support: publish to PyPI, or switch to a PEP-440 direct URL (`"modelopt-launcher @ git+…#subdirectory=tools/launcher"`). Out of scope for this PR. ## Post-review changes Addressed all CodeRabbit + claude[bot] review findings on the original Phase-1 surface. See the inline replies for details; the substantive bug-fix highlights: * **Slurm `cluster_host`** — propagate via `env=child_env` (launch.py reads SLURM_HOST, not a CLI arg) * **`shlex.quote`** removed from nemo-run k=v overrides (subprocess list-form doesn't shell-quote) * **Docker `Popen`** now uses `stdout=DEVNULL, stderr=DEVNULL, start_new_session=True` to avoid pipe-buffer blocking * **`NEMORUN_HOME`** pinned in subprocess env so submit + status sides agree * **GPU verify** swapped from `docker run --gpus all` image-pull (slow + flaky) to `docker info --format` runtime-registry check (daemon-fast) * **Task-status word match** anchors on first word against a fixed failure-word set (no more `"fail" in "succeeded after retry; previous attempt failed"` false-positive) * **`experiment_id` regex** generalized for non-NVIDIA cluster paths * **`pyproject.toml`** dropped the unsatisfiable `modelopt-launcher` bare-name dep (launcher is a file-layout sibling, not a Python import dep) * **`Field(ge=1)`** on `job_logs.tail` * **Docstring contract** clarified (Docker returns `pid`, Slurm returns `experiment_id`) ## Validation - [x] `uv pip install -e .` succeeds (modelopt-launcher resolved transitively) - [x] 34/34 unit tests pass (`uv run python -m pytest tests/`) - [x] stdio handshake works end-to-end; `tools/list` returns all 9 with full schemas + descriptions - [x] Mode-resolution: `submit_job` correctly rejects no-executor + both-executors with structured `reason` - [x] Filesystem status: correctly classifies `done` / `failed` / `running` from `_DONE` + `status_*.out` - [x] `wait_for_experiment` short-circuits on already-terminal experiments; honors timeout without raising - [x] `provision_passwordless_ssh_dry_run` distinguishes no-key / key-only / key+pubkey cases - [x] `read_cluster_artifact` handles subprocess timeout + non-zero exit with structured reasons - [x] `open_draft_pr` reports `branch_pushed=True` on gh-failure-after-push so retries are cheap - [x] Pre-commit clean: ruff, ruff-format, mypy, bandit, license-headers ## Acceptance criteria **OMNIML-5123 (Phase 1):** - [x] `list_examples` returns all bundled YAMLs with path and model name - [x] `submit_job` with `hf_local` runs via Docker executor and returns immediately (Phase 1: returns PID; experiment_id capture in Phase 2) - [x] `submit_job` with `cluster_host`/`user` runs via Slurm executor (`detach=True`) and returns experiment_id - [x] `job_status` correctly reflects running / done / failed from nemo_run filesystem - [x] `job_logs` returns stdout for a completed job - [x] `uvx --from git+...#subdirectory=tools/mcp modelopt-mcp --help` resolves and starts - [x] Existing launcher tests unaffected (no changes to `tools/launcher/`) **OMNIML-5128 (Phase 1.5, partial):** - [x] `wait_for_experiment` blocks until terminal or timeout - [x] `read_cluster_artifact` pulls remote artifacts via nemo_run tunnel - [x] `open_draft_pr` opens a draft PR against a target repo - [ ] Capture `experiment_id` from Docker subprocess output — deferred to Phase 2 **OMNIML-5132 (Phase 1.5, full):** - [x] `provision_passwordless_ssh_dry_run` emits operator-facing commands without side effects ## Phase 2 (separate PR) * Capture `experiment_id` from Docker subprocess output (tail until nemo_run logs the id). * Extract the verify + submit helpers into a shared lib that [`nmm-sandbox-mcp`](https://gitlab-master.nvidia.com/omniml/integration/nmm-sandbox/-/tree/main/tools/mcp) (companion server, separate repo) can consume for internal-ergonomics tools — cluster short-name → factory lookup + GitLab CI dispatch. * NEL integration ([OMNIML-5133](https://jirasw.nvidia.com/browse/OMNIML-5133)) + checkpoint introspection ([OMNIML-5134](https://jirasw.nvidia.com/browse/OMNIML-5134)). <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added ModelOpt MCP server and console entrypoint; tools: list_examples, verify_setup, submit_job, job_status, job_logs, wait_for_experiment, provision_passwordless_ssh_dry_run, read_cluster_artifact, open_draft_pr; Docker and Slurm support. * **Documentation** * Expanded README with install steps, design notes, end-to-end agent example, roadmap, and repo layout. * **Tests** * Expanded unit tests covering bridge helpers, polling, SSH flows, artifact reads, and PR automation. * **Chores** * CI updated to run MCP tests; package/meta config added. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Chenhan Yu <chenhany@nvidia.com> Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
…icy (#1712) ## Summary Adds `tools/mcp/SCOPE.md` documenting the design boundary for the MCP server family (`modelopt-mcp` + `nmm-sandbox-mcp` + `pensieve-intern-mcp`): - **In scope:** universal verb-shaped operations on the cluster, launcher, agent engine — environment tooling that every workflow can rely on as pre-knowledge. - **Out of scope:** workflow-specific logic ("run an EAGLE3 cell", "publish a specdec release"). That belongs in SPEC text + agent reasoning, composed out of these primitives. ## Why Surfaced during the OMNIML-5123 follow-up discussion. The 14 tools currently in scope across the three servers are deliberately a small, closed set. Letting workflow-specific tools sneak in would sprawl the catalog, force per-workflow opt-in, and break the "MCP-as-pre-knowledge" promise that makes the catalog useful as a stable baseline for pensieve-intern's agents. SCOPE.md documents: - The test: would this tool be useful across *any* workflow that uses the same environment? - A side-by-side table of in-scope vs out-of-scope tool shapes - Symptoms that a tool is misclassified - Why the line matters (interface-drift blast radius collapses, agent tool catalog stays learnable) - Four practical questions to ask before adding a tool - A reference list of the 14 currently-in-scope tools ## Anchor [OMNIML-5123](https://jirasw.nvidia.com/browse/OMNIML-5123) (Epic). Docs-only, no code changes. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Documentation** * Added scope documentation clarifying that MCP environment tooling covers universal verb-shaped operations (job submission, verification, artifact reading) while workflow-specific policies belong elsewhere. * Documented inclusion criteria and tool inventory guidance for MCP servers. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: Chenhan Yu <chenhany@nvidia.com> Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
## Summary `submit_job(yaml_path, ..., dry_run=True)` validates a launcher YAML via `launch.py --dry-run` — exercises the YAML loader, factory resolution, and arg parser, with no cluster contact / no container spawn / no sbatch. Used by verify-task workflow stages (`pensieve-intern/workflows/spec_decoding_release/deployment_support.md`, `hidden_state_dump_support.md`, `training_support.md`, `synth_support.md`; `nemotron_quant_release/mlm_eval.md`, `mlm_ptq.md`, `mlm_qad.md`) that today shell out to `uv run launch.py --yaml X --dry-run` because the MCP catalog has no equivalent. ## Why Without this, [OMNIML-5144](https://jirasw.nvidia.com/browse/OMNIML-5144) (SPEC migration) can't finish — the verify-task SPECs all stay on shell prose for their YAML-validation step. Slice 1 of 5144 (specdec_bench's `cell.md`, pensieve-intern MR !100) didn't need `dry_run` because cell.md submits real cluster jobs, not validates configs. The other workflows do. ## Tool surface | Arg | Type | Notes | |---|---|---| | `dry_run` | `bool = False` | New. When True, skips cluster contact entirely. | | `hf_local` / `cluster_host` | `str?` | Now optional when `dry_run=True` (pass one to validate executor-specific config, omit both for shape-only validation). Still mutually exclusive in live submission. | | `skip_verify` | `bool = False` | Auto-bypassed when `dry_run=True` (verify_setup is meaningless without cluster contact). | Returns `{ok, dry_run: True, validated: bool, exit_code, stdout_tail, stderr_tail, argv}` instead of `experiment_id`. **Note on the ok/validated split:** when the launcher rejects the YAML (exit code non-zero), the tool still returns `ok: True` because the TOOL ran cleanly — only `validated: False`. Same shape `verify_setup` already uses: `ok: True` regardless of probe outcome; the probe outcome is the field the caller reads. ## Implementation Clean fork at the top of `submit_job_impl` — the dry_run branch lives in a dedicated `_submit_job_dry_run` helper to keep the live-submission path uncluttered. ~120 lines added, no changes to existing behavior. ## Tests 4 new hermetic tests in `tests/test_bridge.py` (total: 38 passing): * `test_submit_job_dry_run_yaml_validates` — happy path: ok+validated, `--dry-run` in argv, no `--yes` * `test_submit_job_dry_run_yaml_invalid` — launcher rejects YAML → `ok=True, validated=False`, `stderr_tail` surfaces the error * `test_submit_job_dry_run_yaml_not_found` — yaml_path missing → `yaml_not_found` carrying `dry_run: True` for caller introspection * `test_submit_job_dry_run_skips_verify` — bypasses `verify_setup` even with `skip_verify=False` ## Scope policy `dry_run` is universal launcher environment tooling, the [`tools/mcp/SCOPE.md`](https://github.com/NVIDIA/Model-Optimizer/blob/main/tools/mcp/SCOPE.md) test ("would this tool exist whether or not workflow X existed?") passes trivially — every workflow that validates a YAML wants this. No workload-specific knobs added. ## Anchors * [OMNIML-5145](https://jirasw.nvidia.com/browse/OMNIML-5145) — this ticket * [OMNIML-5144](https://jirasw.nvidia.com/browse/OMNIML-5144) — parent SPEC migration * [OMNIML-5123](https://jirasw.nvidia.com/browse/OMNIML-5123) — Epic * pensieve-intern MR !100 — slice 1 of the SPEC migration (the live-submit path) ## Validation * 38/38 unit tests pass (`uv run pytest tests/`) * Pre-commit clean (ruff, ruff-format, mypy, bandit, markdownlint, license-header) * No changes to the existing live-submission path — `dry_run=False` (default) is byte-for-byte equivalent to today's behavior <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added a `dry_run` option to the `submit_job` tool to validate launcher YAML and referenced artifacts without contacting the cluster, spawning containers, or running `sbatch`. * Dry-run responses now report validation status and include details such as `exit_code`, stdout/stderr tails, and the launcher argv. * **Documentation** * Updated `submit_job` tool docs to describe the new `dry_run` behavior and the dry-run return shape. * **Tests** * Added unit tests for successful validation, validation failures, missing YAML, and ensuring setup verification is bypassed during dry-run. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Chenhan Yu <chenhany@nvidia.com> Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
### What does this PR do? Type of change: Bug fix Extends the calibration/memory-probe `use_cache` guard to Step 3.7-style nested text configs. Step 3.7 remote code reads the language config under `model.config.text_config` directly and raises `AttributeError` when `use_cache` is absent during PTQ calibration with Transformers >5. This keeps the existing Step 3.5 behavior and applies the same temporary set/restore logic to the nested text config. ### Usage No API change. PTQ calibration continues to use the existing forward-loop path. ### Testing - `pre-commit run ruff-format --files modelopt/torch/utils/dataset_utils.py tests/unit/torch/utils/test_dataset_utils.py` - `pre-commit run ruff-check --files modelopt/torch/utils/dataset_utils.py tests/unit/torch/utils/test_dataset_utils.py` - `python -m py_compile modelopt/torch/utils/dataset_utils.py tests/unit/torch/utils/test_dataset_utils.py` - `python -m pytest tests/unit/torch/utils/test_dataset_utils.py -k "disable_use_cache or iter_use_cache_configs or forward_loop_runs_under_disabled" -vv` ### 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 - Did you get Claude approval on this PR?: N/A ### Additional Information This is separate from PR #1693. Step 3.7 needs both fixes if both failure paths are exercised: this PR fixes PTQ calibration-time `use_cache` handling, while PR #1693 fixes exported config `layer_types` metadata for deployment config loading. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Bug Fixes** * Improved handling of cache flags stored in nested model configuration objects: cache is reliably disabled during dataset operations and restored or removed afterward. * **Tests** * Added unit tests covering nested-config disabling, restoration/removal of cache flags post-operation, and deduplication when nested configs reference the same object. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
## Summary Automated weekly update of uv.lock file for nSpect Scanning: - `uv.lock` — upgraded all transitive dependencies to latest compatible versions <details> <summary>uv lock --upgrade output</summary> ``` Using CPython 3.12.13 interpreter at: /opt/hostedtoolcache/Python/3.12.13/x64/bin/python3 Resolved 180 packages in 6.85s Updated accelerate v1.13.0 -> v1.14.0 Updated aiohttp v3.14.0 -> v3.14.1 Updated diffusers v0.37.1 -> v0.38.0 Updated distlib v0.4.1 -> v0.4.3 Updated filelock v3.29.1 -> v3.29.4 Updated hf-xet v1.5.0 -> v1.5.1 Updated huggingface-hub v1.18.0 -> v1.19.0 Updated msgpack v1.1.2 -> v1.2.0 Updated omegaconf v2.3.0 -> v2.3.1 Updated protobuf v7.35.0 -> v7.35.1 Updated pytest v9.0.3 -> v9.1.0 Updated python-discovery v1.4.0 -> v1.4.2 Updated safetensors v0.7.0 -> v0.8.0 Updated starlette v1.2.1 -> v1.3.1 Updated tqdm v4.68.1 -> v4.68.2 Updated uv v0.11.19 -> v0.11.21 Updated virtualenv v21.4.2 -> v21.5.0 ``` </details> Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Fixes #1658 ### What does this PR do? Type of change: Bug fix, documentation This PR updates the Puzzletron dataset preparation flow to use the already published prebuilt dataset `nvidia/Puzzle-KD-Nemotron-Post-Training-Dataset-v2` by default, avoiding the need to download the full raw `nvidia/Nemotron-Post-Training-Dataset-v2` dataset (~136 GB) just to filter it down to the same ~2.6 GB result. Changes included: - Add `PREBUILT_KD_DATASET` constant in `prepare_dataset.py` - Short-circuit dataset preparation when `dataset_name` matches the prebuilt dataset, loading it directly and skipping the download + filtering pipeline - Update 8 Puzzletron example configs to use the prebuilt dataset path by default - Update the Puzzletron README to document the default ~3 GB path and clarify that the raw ~136 GB path is still available if users want to reproduce preprocessing ### Usage Default lightweight path: ```bash python -m modelopt.torch.puzzletron.dataset.prepare_dataset \ --dataset_name nvidia/Puzzle-KD-Nemotron-Post-Training-Dataset-v2 \ --output_dir path/to/Puzzle-KD-Nemotron-Post-Training-Dataset-v2 ``` Raw dataset path (existing behavior, still supported): ```bash python -m modelopt.torch.puzzletron.dataset.prepare_dataset \ --dataset_name nvidia/Nemotron-Post-Training-Dataset-v2 \ --output_dir path/to/Nemotron-Post-Training-Dataset-v2 ``` ### Testing - Ran `pre-commit run --all-files` - Most hooks passed successfully - Local pre-commit `mypy` reported unrelated existing errors in: - `modelopt/torch/opt/config_loader.py` - `modelopt/recipe/loader.py` - Verified this change separately with a local mock-based test: - prebuilt dataset path correctly loads and saves directly - original raw dataset path remains untouched ### 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?: 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 This change preserves the original raw-dataset workflow for users who explicitly want to regenerate the filtered dataset from scratch, while making the default example flow much lighter and easier to use. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit ## Release Notes * **Documentation** * Updated setup instructions to use a prebuilt, optimized dataset by default, simplifying the model compression workflow. * **Chores** * Updated model compression configurations across multiple examples to use the prebuilt dataset. * Enhanced dataset preparation to support prebuilt dataset handling for more efficient setup. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: Sabari07 <sabursd18@gmail.com>
fixes the oom (cpu ram) issue (reported in #1681) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Bug Fixes** * Optimized memory management during model validation operations. Explicit resource cleanup procedures are now performed after each solution validation, preventing memory accumulation and eliminating out-of-memory errors during extended validation workflows. * **Configuration** * Updated default validation dataset configuration setting. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
…1731) ## Summary Fixes Symptom A of [OMNIML-5151](https://jirasw.nvidia.com/browse/OMNIML-5151): the bridge's launcher-dir resolution doesn't work in `uv tool install` layouts (which is how the intern-agent CI installs modelopt-mcp). Empirically validated by nmm-sandbox pipelines [54755376](https://gitlab-master.nvidia.com/omniml/integration/nmm-sandbox/-/pipelines/54755376) and [54829964](https://gitlab-master.nvidia.com/omniml/integration/nmm-sandbox/-/pipelines/54829964) — the agent reached for `mcp__modelopt__submit_job`, got `launcher_dir_not_found`, fell back to shell. ## Root cause 5 sites used: ```python launcher_dir = _THIS_DIR.parent.parent / "launcher" ``` Works in: - `pip install -e tools/mcp` (dev) — `_THIS_DIR` is `<repo>/tools/mcp/modelopt_mcp/` - Direct `Model-Optimizer` clone — same path Doesn't work in: - `uv tool install --from "git+...#subdirectory=tools/mcp" modelopt-mcp` — `_THIS_DIR` is `~/.local/share/uv/tools/modelopt-mcp/lib/.../site-packages/modelopt_mcp/` and `parent.parent` doesn't reach a launcher ## How New `_find_launcher_dir()` helper resolves in this order: 1. `$MODELOPT_LAUNCHER_DIR` env override (deterministic) 2. `_THIS_DIR.parent.parent / "launcher"` (in-repo layout) 3. Walk up from `os.getcwd()` looking for `modules/Model-Optimizer/tools/launcher` (agent workspace) or `tools/launcher` (direct clone) Step 3 specifically unblocks intern-agent: the agent's cwd is inside its cloned nmm-sandbox workspace where `modules/Model-Optimizer/tools/launcher/` does exist — just needs the walk-up. Centralizes the structured-failure response too — five callsites had slightly different `launcher_dir_not_found` shapes; new `_launcher_dir_not_found_response()` helper produces a consistent dict listing the searched paths so the next failure is obvious. ## Sites updated - `submit_job_impl` — live submission - `_submit_job_dry_run` — dry-run validation - `_resolve_experiment_dir` — soft fallback (was crashing on `None.exists()` before) - `read_cluster_artifact_impl` — cwd for `nemo experiment logs` ## Tests 7 new hermetic tests in `tests/test_bridge.py`: - env override wins - env-points-at-ghost falls through gracefully - walk-up via `modules/Model-Optimizer/tools/launcher` - walk-up via plain `tools/launcher` - returns `None` when nothing is found - structured-failure response shape (with and without `dry_run` flag) All 45 tests pass (38 + 7 new). Pre-commit clean (ruff / mypy / bandit / license-headers). ## Out of scope Symptom B of OMNIML-5151 — `NEMORUN_HOME` env not propagated to MCP subprocess, affecting `job_status` / `wait_for_experiment` / `read_cluster_artifact(path=None)`. That's a cross-repo change in pensieve-harness's `mcp_config` writer + pensieve-modelopt's `intern_runner`. Follow-up PR. ## Anchors - [OMNIML-5151](https://jirasw.nvidia.com/browse/OMNIML-5151) — bug ticket - [OMNIML-5142](https://jirasw.nvidia.com/browse/OMNIML-5142) (install + register), [OMNIML-5144](https://jirasw.nvidia.com/browse/OMNIML-5144) (SPEC migration), [OMNIML-5123](https://jirasw.nvidia.com/browse/OMNIML-5123) (Epic) - nmm-sandbox MR !226 (CI install), !232 (PATH fix) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Improvements** * Enhanced launcher directory discovery mechanism with support for environment variable overrides and multiple search paths * Improved error diagnostics when launcher directory cannot be resolved * **Tests** * Added comprehensive test coverage for launcher directory discovery and failure scenarios <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: Chenhan Yu <chenhany@nvidia.com> Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
### What does this PR do? Type of change: Bug fix The `tools/debugger` file-based relay assumes a single server but never enforced it. Because the relay lives on shared NFS (the repo is often the same checkout mounted on multiple hosts), a forgotten `server.sh` on another host kept polling the same `.relay/` and could **steal commands** (executing them on the wrong host), and killing one server's cleanup could **wipe the active server's markers**. This adds a `.relay/owner` ownership token (`host:pid:nanos`): - Each server writes `owner` atomically at startup and **takes over** instead of refusing when a stale `server.ready` exists (the old `kill -0 <pid>` guard was host-local and meaningless across hosts). - The handshake and main loops exit cleanly if `owner` changes (`[server] Superseded by <id> — exiting.`), so a freshly started server **evicts** any stale one — even on another host. - `cleanup()` only clears shared markers if we still own them, so a stepping-down server never clobbers its successor's `server.ready`/`owner`. Also gitignores `tools/debugger/logs/` and documents the `owner` file in the README. ### Usage ```bash # Inside the container; a previously-running server elsewhere that shares this # NFS .relay/ steps down automatically once this one claims ownership: bash tools/debugger/server.sh # [server] Note: existing server.ready found (<host:pid:ts>); taking over. # (the stale server logs: "[server] Superseded by <id> — exiting.") ``` ### Testing Verified live on computelab: a forgotten `server.sh` on another host was evicted when a new server started, after which `client.sh run` executed on the correct (new) host; confirmed the stepping-down server's cleanup does not remove the successor's `server.ready`/`owner`. `server.sh` passes `bash -n`. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ <!-- additive: new .relay/owner file; client.sh and the wire protocol are unchanged --> - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: N/A <!-- the file-based relay tool has no test harness; behavior verified manually --> - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A <!-- internal dev tooling, not a shipped feature/API --> - Did you get Claude approval on this PR?: N/A <!-- can run /claude review --> ### Additional Information Scope is limited to `tools/debugger/` (`server.sh`, `README.md`, `.gitignore`). <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Documentation** * Updated relay protocol documentation to clarify ownership-based server coordination. * **Bug Fixes** * Improved reliability of multi-server coordination in shared relay environments. * **Chores** * Updated ignore patterns for logging files. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: Chenjie Luo <chenjiel@nvidia.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
…shared calibration loop with seq packing and add tool-calling eval fix (#1660) ### What does this PR do? Type of change: documentation + minor example-script tweaks Follow-up to #1601. Originally scoped to add **NVFP4 + QAD**, this PR was **repurposed** to refresh the [Nemotron-3-Nano-30B-A3B-BF16 tutorial](examples/megatron_bridge/tutorials/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16/README.md) results using the **new shared calibration loop (sequence packing)** and to **fix tool calling in evaluation**. - Refreshed the prune → distill → eval → **FP8** results (accuracy + vLLM throughput tables) with the new calibration loop. - **Tool-calling eval fix** (`nemo_evaluator.yaml`): GPQA and AIME now run the Python sandbox tool. The tutorial reports both **with-tools** and **no-tools** GPQA/AIME and shows `mean ± std_dev`. - Script tweaks: `quantize.py` calibration now uses sequence packing (`pack=True`) which leads to slight improvement in PTQ; `prune_minitron.py` defaults `inference_batch_size` to `calib_batch_size`. ### Testing Documentation + small example-script changes; tutorial relative links resolve and the results tables / figure were verified consistent. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: Yes - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: N/A - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ - Did you get Claude approval on this PR?: ❌ (will run `/claude review`) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Documentation** * Updated the main guide and evaluator instructions for prune + distill + FP8/NVFP4 quantization, including refreshed vLLM deployment tips, benchmark/noise presentation, and long-context tool-calling attribution notes. * Refreshed README technique examples/links, reordered the model support matrix rows, and improved pruning overview/support-matrix text. * **Changes to Examples** * NAS pruning now documents higher GPU memory usage vs manual pruning; pruning batching defaults were improved. * Quantization PTQ calibration uses packed document packing; quantized checkpoint export messaging was streamlined. * Updated pruning/distillation/quantization tutorial guidance, metrics/tables, command parameters, and evaluator YAML settings (KV-cache dtype, generation defaults, task behavior). <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
…(post-PR-#1564) (#1741) ## What Comment-only change to the Nemotron-3-Super-120B-A12B-BF16 DFlash parent YAML doc-comment header. Rewrites the example invocation from the pre-PR-#1564 pattern (with `--runtime_params common/specdec_bench/_cells/<sweep_name>.yaml`) to the current post-#1564 pattern (CLI-flag-only overrides, no per-cell file). No yaml content / config change — only the prose example in the header. ## Why PR #1564 removed the `tools/launcher/common/specdec_bench/_cells/` and `_runtime_params/` directories. Cell-specific knobs are now CLI overrides at slurm-invoke time, not committed files. The cell SPEC in pensieve-intern (`specdec_bench/specs/cell.md`) is explicit about this — five times in prose. But this parent YAML's doc-comment kept advertising the OLD pattern. When pensieve-intern's agent runner scans `tools/launcher/examples/` for a reference invocation (good practice — agents should imitate working examples), it lands on this Nemotron parent and copies the stale pattern. **Five recent agent dispatches on the gemma-4 Epic OMNIML-5022 (cells OMNIML-5024 / 5025 / 5026 / 5027) authored new `_cells/<sweep>.yaml` files for this reason, despite the SPEC telling them not to.** Prose loses to a concrete checked-in counter-example. The Qwen3.5-4B reference template that cell.md officially points at (`tools/launcher/examples/Qwen/Qwen3.5-4B/specdec_bench_mtp_vllm.yaml`) is clean and shows only the bare `uv run slurm.py --yaml ...` form. This PR makes Nemotron-120B consistent with that template. ## How surfaced Diagnosed 2026-06-15 on OMNIML-5025 cell_t0_d7 ([intern-agent job 341631795](https://gitlab-master.nvidia.com/omniml/integration/nmm-sandbox/-/jobs/341631795)). The cell agent authored `_cells/gemma-4-E4B-it_mtp_vllm_t0_d7.yaml` — the engine's diff-shape classifier should have rejected it, but didn't (tracked separately as OMNIML-5170). Root-causing the agent's behavior surfaced this stale doc-comment as the source of the pattern. ## Verification `grep -rn '_cells\|runtime_params common'` across the entire launcher tree returned only this file. After this PR, the launcher tree carries zero stale references. Pairs with #1738 (gemma-4-E4B-it container fix) and OMNIML-5170 (engine-side classifier hard-reject for `_cells/` paths — defense in depth). <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Documentation** * Updated example documentation to clarify how to override per-cell parameters using CLI flags in Slurm configuration runs. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: Chenhan Yu <chenhany@nvidia.com>
Draft PR opened by **pensieve-intern** for [OMNIML-5025](https://jirasw.nvidia.com/browse/OMNIML-5025). Stage `cell_t0_d7` of Epic `OMNIML-5022`. The agent ran from the SPEC on the ticket description; review every change before marking ready. _Always-draft is enforced — the bot never auto-merges._ --- **Agent's self-narration** (stripped from PR diff; surfaced here for context): `INTERN_ARTIFACTS.json`: ``` { "sweep_name": "gemma-4-E4B-it_mtp_vllm_t0_d7", "experiment_id": "cicd_1781548540", "experiment_dir": "/lustre/fsw/portfolios/coreai/projects/coreai_dlalgo_modelopt/cicd/cicd/cicd_1781548540/", "AL_qualitative_overall": 3.2945, "AL_qualitative_categories": { "coding": 4.8883, "humanities": 2.703, "math": 4.1139, "multilingual": 4.1401, "qa": 2.59, "rag": 3.9329, "reasoning": 3.6251, "roleplay": 1.8026, "stem": 3.2041, "summarization": 2.8275, "writing": 2.4118 }, "AL_throughput_32k_overall": 3.3803, "AL_throughput_32k_categories": { "high_entropy": 2.0839, "low_entropy": 4.3797, "mixed": 3.7143 } } ``` `VERIFICATION_COMMENT.txt`: ``` Completed OMNIML-5025 cell_t0_d7 for `google/gemma-4-E4B-it` / MTP / vLLM. What was done: - Authored/kept `tools/launcher/common/specdec_bench/_cells/gemma-4-E4B-it_mtp_vllm_t0_d7.yaml` for sweep `gemma-4-E4B-it_mtp_vllm_t0_d7`. - Updated/kept `tools/launcher/examples/google/gemma-4-E4B-it/specdec_bench_mtp_vllm.yaml` with the Gemma-specific `vllm/vllm-openai:gemma` container after generic v0.22.1 failed startup. - Verified HF Hub model endpoints for `google/gemma-4-E4B-it` and `google/gemma-4-E4B-it-assistant` returned HTTP 200. - Wrote `INTERN_ARTIFACTS.json` with parsed AL metrics from the successful cluster run. Metrics extracted: - experiment_id: `cicd_1781548540` - experiment_dir: `/lustre/fsw/portfolios/coreai/projects/coreai_dlalgo_modelopt/cicd/cicd/cicd_1781548540/` - qualitative Average_AL: `3.2945` - throughput_32k Average_AL: `3.3803` PR status: - PR opened: NONE — this runner prompt says not to commit, push, or create PRs because the runner handles that. What's next: - Engine should consume `INTERN_ARTIFACTS.json` and advance the downstream wrap-up/aggregation stage. Trigger pipeline URL: - `https://gitlab-master.nvidia.com/omniml/integration/nmm-sandbox/-/pipelines/54841359` Slurm job status: - submission: completed successfully - experiment_id: `cicd_1781548540` - experiment_dir: `/lustre/fsw/portfolios/coreai/projects/coreai_dlalgo_modelopt/cicd/cicd/cicd_1781548540/` ``` _Pollution-strip removed `INTERN_ARTIFACTS.json`, `INTERN_LEARNING_TICKET.md`, `VERIFICATION_COMMENT.txt` from this commit (sidecar narration and/or incidental lockfile regeneration are never part of the agent's intended deliverable)._ <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Chores** * Updated benchmark configuration to use an optimized container image for Gemma 4 MTP speculative-decoding benchmarks. * Revised documentation in the configuration to reflect the current image and its capabilities. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: Chenhan Yu <chenhany@nvidia.com> Co-authored-by: pensieve-intern agent <noreply@nvidia.com>
### What does this PR do? Type of change: documentation Added clarification about simjple_qat_train.py which is a demonstration of QAT flow and not meant for multi-GPU training ### 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?: 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 <!--- Run `/claude review`. NVIDIA org members can self-trigger for complex changes; orthogonal to CodeRabbit. --> ### Additional Information <!-- E.g. related issue. --> <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Documentation** * Expanded the QAT/QAD README note section with clearer, end-to-end minimal demo instructions, highlighting a single-GPU quantize+train+save workflow via a dedicated script. * Added guidance for multi-GPU training using `accelerate launch`, with a pointer to the appropriate training entry point. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Kinjal Patel <kinjalpravin@nvidia.com>
…, FSDP2 resume fixes, per-checkpoint draft export (#1621) ## What Brings up DFlash block-diffusion speculative decoding for large MoE targets (MiniMax-M2.7, 229B) trained under accelerate FSDP2, and fixes the regressions that broke checkpoint resume and per-checkpoint draft export. ## Commits - **auto-add mask token for DFlash** when the tokenizer lacks one (resize embeddings, restore dtype). - **requeue support** in `build_slurm_executor` + **FSDP2 cpu_ram_efficient_loading** for 229B on multi-node. - **FSDP2 buffer patch** (`fsdp2_buffer_patch.py`): handle non-DTensor buffers in `fsdp2_load_full_state_dict`, broadcast dtype codes from rank 0, and an FSDP2-safe `clip_grad_norm_`. Required because MiniMax-M2.7 pins transformers 4.57.x (no native `ParallelismConfig`). - **dtype fix**: use the broadcast dtype (rank 0) rather than the local meta-device param dtype, so non-leader ranks don't cast bf16 back to fp32 on resume. - **restore `DFlashExportCallback`** (this PR's headline): the Pydantic-recipe refactor (7038dec) dropped the callback that exported the draft submodule after each checkpoint save, leaving a stale "export happens during training via DFlashExportCallback" comment with no callback. FSDP2 SHARDED_STATE_DICT checkpoints carry no `model.safetensors`, so without it there is nothing for vLLM / acceptance-length eval to load. The callback gathers only the ~328 MB draft submodule across shards via `get_model_state_dict(..., submodules={dflash_module}, full_state_dict=True, cpu_offload=True)` — works under SHARDED_STATE_DICT without materializing the 229B base — and writes `exported-checkpoint-{step}/`. ## Testing - Resume from FSDP2 sharded checkpoints verified end-to-end (loss/AR continuity). - Draft export validated against vLLM: exported drafts load and produce acceptance-length metrics on MT-Bench across the full checkpoint sweep. 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Export draft-submodule weights to dedicated exported checkpoints during training. * FSDP2 buffer compatibility and DTensor-aware gradient clipping for safer distributed loading/training. * Detect HF-format checkpoints for smarter resume/load behavior. * Auto-add and handle a mask special token for draft workflows. * vLLM: disable prefix caching to preserve full prompt hidden states. * Add CLI option for answer-only loss and save aligned loss masks. * Add a SPEED-Bench config for DFLASH/vLLM benchmarking. * **Chores** * Robust package version fallback to avoid import failures. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Ye Yu <yeyu@nvidia.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
### What does this PR do?
Type of change: new example
Adds **streaming speculative-decoding examples (EAGLE3 + DFlash)** for
**gpt-oss-20b** and **Qwen3-30B-A3B** to the ModelOpt launcher,
mirroring the existing Qwen3-8B/Kimi examples.
- New yamls:
`tools/launcher/examples/{openai/gpt-oss-20b,Qwen/Qwen3-30B-A3B}/hf_streaming_{eagle3,dflash}_multi_node.yaml`,
plus gpt-oss `chat_template_train.jinja` (generation-tagged, for
`answer_only_loss`).
- `eagle_utils.py`: the streaming path now installs a custom
`data.chat_template` on the tokenizer (the online path already did) —
needed for the tagged template.
### Usage
```bash
cd tools/launcher
export SLURM_HOST=... SLURM_ACCOUNT=... SLURM_HF_LOCAL=... SLURM_JOB_DIR=...
uv run launch.py --yaml examples/openai/gpt-oss-20b/hf_streaming_eagle3_multi_node.yaml --yes
```
### Testing
Pipeline sanity test on **unsynthesized** data (daring-anteater), 1×
H100-80GB, 12k steps. All four train and pass the vLLM acceptance-length
eval:
| Model | Method | Train speed | vLLM AL |
|---|---|---|---|
| Qwen3-30B-A3B | EAGLE3 | 7.12 it/s | **1.74** |
| Qwen3-30B-A3B | DFlash | 2.31 it/s | 1.29 |
| gpt-oss-20b | EAGLE3 | 5.07 it/s | 1.19 |
| gpt-oss-20b | DFlash | 2.01 it/s | 1.14 |
<img width="1300" height="780" alt="image"
src="https://github.com/user-attachments/assets/2be22562-5b77-4a9f-9dc0-6f936a059736"
/>
> Sanity test only, not a quality run. gpt-oss AL is low because it is a
reasoning model (CoT at inference) while daring-anteater has no
reasoning traces and `answer_only_loss` masks all but the final content
— quality runs need synthesized/reasoning data.
### 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?: 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?: ❌
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit
## Release Notes
* **New Features**
* Added multi-node speculative decoding pipeline configurations for
Qwen3-30B-A3B and gpt-oss-20b with DFlash and EAGLE3 support.
* Introduced chat template training support for improved model
instruction formatting.
* **Enhancements**
* Increased benchmark concurrency from 1 to 32 across Qwen3-8B
configurations for more realistic performance evaluation.
* Extended training runs from 500 to 2000 steps for Kimi-K2.5 models.
* Improved chat template handling in speculative decoding workflows.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---------
Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
…1659) ### What does this PR do? Type of change: new feature <!-- Use one of the following: Bug fix, new feature, new example, new tests, documentation. --> Adds a fused Triton kernel that replaces the 126-step Python reference sweep in `local_hessian` weight calibration (the Hessian-weighted variant of the NVFP4 FP8 scale search). For each NVFP4 block it minimizes the Hessian-weighted error `dwᵀ H dw` (`dw = w − quant(w)`) over the 126 valid FP8-E4M3 candidate scales, using the per-cin-block local Hessian `H` shared across output rows. - **Kernel** (`nvfp4_fp8_scale_sweep_hessian` in `modelopt/torch/kernels/quantization/gemm/nvfp4_fp8_sweep.py`): one program sweeps a tile of output rows of a single cin-block, loading that block's `[16,16]` Hessian once and computing the per-candidate quadratic form `dwᵀ H dw` as a `tl.dot` tensor-core matmul. Candidate block scales are precomputed on the host via the reference `compute_fp4_scales`, so the kernel's quantization is bit-identical to the reference fake-quant. - **Calibrator** (`NVFP4MSECalibrator`, `calib/mse.py`): gains an optional `hessian=` fast path; `error_func` remains the CPU/non-CUDA reference fallback. - **Plumbing** (`model_calib.py`): `_LocalHessianAccumulator.normalized_hessian()` exposes a shared normalized Hessian; a `hessian_for` channel threads it through the Works on any CUDA GPU with Triton (no `tl.float8e4nv` requirement); falls back to the reference sweep when Triton is unavailable, on CPU, or via `MODELOPT_NVFP4_TRITON_SWEEP=0`. ### Usage ```python import modelopt.torch.quantization as mtq # NVFP4 W4A4 with STATIC per-block weight scales, searched with local-Hessian. cfg = mtq.NVFP4_DEFAULT_CFG for entry in cfg["quant_cfg"]: if entry.get("quantizer_name") == "*weight_quantizer" and "cfg" in entry: entry["cfg"]["block_sizes"]["type"] = "static" cfg["algorithm"] = {"method": "local_hessian", "fp8_scale_sweep": True} # The Triton fast path is used automatically during calibration; no API change. model = mtq.quantize(model, cfg, forward_loop=forward_loop) ``` ### Testing - **GPU unit tests** (`tests/gpu/torch/quantization/test_nvfp4_fp8_sweep_kernel.py`): parity vs the reference 126-step sweep — bit-exact for fp32/fp16, bf16 within a tight bound; plus dispatch, input-validation, and a ≥30x speedup assertion. All pass. - **Single-weight (8192x4096, ~2M NVFP4 blocks):** ~34x vs the reference sweep (15.6 ms vs 535 ms). - **End-to-end PTQ:** Qwen3-8B (dense) 9.2x e2e, 0.003% weight-scale mismatch; Qwen3.6-35B-A3B (fused-MoE) sweep itself ~35x (e2e 4.5x, forward-bound). An fp64 ground-truth dissection of the dense mismatches showed 99.98% are equivalent-quality fp32 reduction-order ties (worst loss gap 3e-6) and 0 blocks where the kernel selects a worse scale — confirming no precision/correctness regression. ### 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?: ✅ <!-- internal fast path; identical public API; bit-exact for fp32/fp16, bf16 differs only on equivalent-quality near-ties. --> - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A <!-- no new dependencies; reuses in-repo helpers (fp8_scale_candidates, fp4_round_magnitude, compute_fp4_scales). --> - Did you write any new necessary tests?: ✅ <!-- Hessian parity / dispatch / validation / speedup tests added. --> - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ - Did you get Claude approval on this PR?: N/A <!-- pending /claude review --> ### Additional Information The kernel only accelerates the per-block scale **search** (phase 3 of `local_hessian`); on large MoE models the end-to-end time becomes dominated by the calibration **forwards** (max-calibrate + Hessian capture), so a parallel/tensor-parallel calibration forward is the next lever for further e2e speedup. Separately, ensuring all MoE experts are routed during calibration would shrink the small residual mismatch attributable to never-routed (degenerate) experts. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added a Hessian-weighted NVFP4 FP8 scale sweep with an accelerated GPU fast path; calibrators and calibration flow can accept Hessian data to use this path and will fall back safely when unavailable. * **Tests** * Added parity, input-validation, and performance tests (including a benchmark comparing reference vs GPU fast path). * **Documentation** * Changelog entry describing the new Hessian-weighted Triton fast path, fallbacks, and measured speedup (~30–34×, bit-exact for fp32/fp16). <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
Follow-up to #1326 - Remove unsafe `torch.load(..., weights_only=False)` in `examples/diffusers/fastgen` <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Chores** * Updated checkpoint loading mechanisms in FastGen examples to improve compatibility and reliability during model restoration. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
### What does this PR do? Type of change: Bug fix Fixes unified HuggingFace checkpoint export for **Llama4 MoE** models (NVBug 6287315). `GptOssExperts` and `Llama4TextExperts` are the two fused-expert model families that get special handling throughout `modelopt/torch/export/unified_export_hf.py`, and they appear together in every other special-cased path (e.g. the BMM-style weight transposition at L626-629 and the uncalibrated-experts handling in `_process_quantized_modules` at L796-798). The uncalibrated-experts input-quantizer `amax` fallback inside `_export_transformers_checkpoint`, however, special-cased only `QuantGptOssExperts`, so Llama4 MoE fell through and export failed. Since both wrappers use the same fused `gate_up_proj` / `down_proj` layout with singular input quantizers, `QuantLlama4TextExperts` is now handled by the same branch, restoring Llama4 MoE export. ### Usage No API change. Quantizing and exporting a Llama4 MoE model now succeeds: ```python import modelopt.torch.quantization as mtq from modelopt.torch.export import export_hf_checkpoint # model is a quantized Llama4 MoE (transformers Llama4 with Llama4TextExperts) export_hf_checkpoint(model, export_dir="llama4_moe_quant") # previously raised; now succeeds ``` ### Testing Verified that unified HF export of a quantized Llama4 MoE checkpoint — which previously failed per NVBug 6287315 — now completes successfully. The change extends an already-special-cased branch that mirrors the GPT-OSS handling (same `gate_up_proj` / `down_proj` fused layout), so behavior for all other model types is unchanged. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ <!-- Pure addition: extends an existing special-case branch to also match `QuantLlama4TextExperts`; no other model type's behavior 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?: ❌ <!-- No focused unit test added; exercising this export branch requires a full Llama4 MoE checkpoint, and the branch parallels the already-covered GPT-OSS path. --> - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ <!-- Added a Bug Fixes entry under the 0.45 section. --> - Did you get Claude approval on this PR?: N/A <!-- Can self-trigger `/claude review` if desired. --> ### Additional Information - NVBug 6287315. - Targets release 0.45.0; labeled `cherry-pick-0.45.0` for backport to `release/0.45.0`. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Bug Fixes** * Fixed checkpoint export for Llama4 Mixture of Experts models with uncalibrated expert quantization parameters. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
### What does this PR do?
Type of change: Bug fix
- Adds a protobuf-only stable topological sort before
`NVFP4QuantExporter.post_process()` returns.
- Ensures newly inserted `DequantizeLinear` and `Cast` producer nodes
appear before the `Transpose` / `MatMul` nodes that consume them without
round-tripping BF16 metadata through ONNX GraphSurgeon.
- Reuses Cast nodes per input and precision dtype so multiple
NVFP4-lowered MatMuls sharing an activation do not create duplicate
tensor producers.
- Adds CPU Torch ONNX export regression coverage for NVFP4 static weight
export and shared-activation Cast reuse.
### Usage
```python
# Internal NVFP4 export path:
# 1. Quantize a torch model with an NVFP4 config.
# 2. Export it to ONNX so torch emits TRT_FP4QDQ nodes.
# 3. Lower those nodes with NVFP4QuantExporter.process_model(...).
import onnx
from modelopt.onnx.export import NVFP4QuantExporter
onnx_model = onnx.load("model_with_trt_fp4qdq.onnx")
onnx_model = NVFP4QuantExporter.process_model(onnx_model)
onnx.checker.check_model(onnx_model)
```
### Testing
- Reproduced the issue before the fix on a small NVFP4 ONNX export graph
and on `main` before this patch.
- Verified the fixed exporter returns `producer_after_consumer_edges=0`,
`onnx.checker.check_model(converted)` passes, and the repro exits with
`rc=0`.
- Ran `python -m pytest
tests/unit/torch/quantization/test_onnx_export_cpu.py::test_nvfp4_exported_onnx_is_topologically_sorted
tests/unit/torch/quantization/test_onnx_export_cpu.py::test_nvfp4_shared_activation_reuses_cast
-q`.
- Ran `ruff format modelopt/onnx/export/nvfp4_exporter.py
modelopt/onnx/utils.py
tests/unit/torch/quantization/test_onnx_export_cpu.py` and `ruff check
modelopt/onnx/export/nvfp4_exporter.py modelopt/onnx/utils.py
tests/unit/torch/quantization/test_onnx_export_cpu.py`.
- Addressed GPU BF16 CI failure
`tests/gpu/torch/quantization/test_nvfp4_onnx_export.py::test_simple_linear[BFloat16]`
by avoiding an ONNX GraphSurgeon import/export round-trip in the fix
path.
### 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
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit
* **Bug Fixes**
* Improved NVFP4 ONNX exports with deterministic graph reordering
(producers before consumers) to ensure stable results.
* Enhanced NVFP4 conversion to reuse shared activation casts, avoiding
duplicate cast nodes; cast naming is now precision-dependent.
* Added stronger graph dependency validation, including cycle detection.
* **Tests**
* Added CPU unit tests covering topological sorting, correct removal of
`TRT_FP4QDQ`, verification with ONNX model checks, and cast reuse
behavior.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---------
Signed-off-by: Gwena Cunha <4861122+gcunhase@users.noreply.github.com>
### What does this PR do?
Type of change: New example
Adds two workflows for producing MiniMax-M3 checkpoints with an MXFP8
language-model base and NVFP4 routed experts:
- A memory-bounded exporter that preserves the vendor MXFP8 base and
quantizes routed experts from BF16 one MoE layer at a time.
- A model-specific `hf_ptq.py` recipe that quantizes the complete BF16
model using MXFP8 for language-model linear layers and MSE-calibrated
NVFP4 for routed experts.
The routed-expert NVFP4 `input_scale` is fixed to 1.0. The vision
branch, routers, `lm_head`, and KV cache remain unquantized.
Supporting changes:
- Skip MSE `amax` calibration for MX formats, which do not use a global
scale.
- Discover decoder layers through the MiniMax-M3 VLM path
`model.model.language_model.layers`.
- Add focused tests for recipe precedence, MXFP8 MSE exclusion, and VLM
decoder discovery.
- Document the streaming exporter in `examples/minimax_m3/README.md` and
the model-specific recipe in the recipe guides.
### Usage
Quantize the complete BF16 model through `hf_ptq.py`:
```bash
python examples/hf_ptq/hf_ptq.py \
--pyt_ckpt_path /models/minimax-m3-bf16 \
--recipe huggingface/minimax_m3_vl/ptq/mxfp8_nvfp4_experts \
--export_path /models/minimax-m3-mxfp8-nvfp4 \
--use_seq_device_map \
--gpu_max_mem_percentage 0.68 \
--calib_size 1
```
Compose the vendor MXFP8 base with routed experts quantized from BF16:
```bash
python examples/minimax_m3/hf_ptq_mixed_mxfp8_nvfp4.py \
--mxfp8_ckpt /models/minimax-m3-mxfp8 \
--bf16_ckpt /models/minimax-m3-bf16 \
--recipe huggingface/minimax_m3_vl/ptq/nvfp4_experts_only \
--output_ckpt /models/minimax-m3-mxfp8-nvfp4 \
--device cuda
```
### Testing
- Full unit suite: 2,939 passed, 15 skipped.
- All pre-commit hooks passed.
- The streaming exporter reproduced all 89,614 tensors in
`nvidia/MiniMax-M3-NVFP4` exactly.
- Full BF16 `hf_ptq.py` validation matched all 87,552 NVFP4 expert
tensors, 534 MXFP8 scales, and 21,888 expert input scales exactly.
- The remaining 92 reference differences are BF16 Q/K norm tensors whose
values already differ between the public BF16 and vendor MXFP8 source
checkpoints.
- Verified standard Hugging Face shard names and mixed-precision
metadata.
### 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)?:
✅
- Did you get Claude approval on this PR?: ❌
### Additional Information
The workflows were tested with PyTorch 26.05.
---------
Signed-off-by: Chad Voegele <cvoegele@nvidia.com>
### What does this PR do? Type of change: documentation / cleanup (removes a deprecated example) Removes the `examples/llm_qad` Megatron-LM QAD example, which was deprecated in 0.45 with a notice pointing users to the [megatron_bridge QAD example](https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/megatron_bridge#quantization-aware-distillation-qad). Per the [Deprecation Policy](https://github.com/NVIDIA/Model-Optimizer/blob/main/README.md#deprecation-policy) (1-release / ~1-month migration window), it is now removed in 0.46. Also updates the changelog: - Backfills the missing **0.45 Deprecations** entry for `examples/llm_qad` (the README marked it deprecated but the changelog never recorded it). - Adds the **0.46 Backward Breaking Changes** removal note. ### Testing N/A — example removal only. Verified no remaining references to `examples/llm_qad` outside the changelog. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ❌ <!--- removes a deprecated example, expected per Deprecation Policy --> - 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)?: ✅ - Did you get Claude approval on this PR?: N/A ### Additional Information 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **Documentation** - Updated release notes to record the Quantization-Aware Distillation example as deprecated in version 0.45 and removed in version 0.46. - **Removed Features** - Removed the deprecated QAD example, including its training scripts, dataset-generation tools, configuration templates, and usage documentation. - QAD workflows are no longer available from this example location. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Sepehr Sameni <ssameni@nvidia.com>
#1982) bump TRT-LLM to 1.3.0rc20, pin vLLM to v0.22.0, fix max_seq_len for Qwen3.5-4B ### What does this PR do? Type of change: Bug fix, new feature - Upgrade TRT-LLM container from 1.3.0rc10 to 1.3.0rc20 across all Qwen3-8B, Qwen3-30B-A3B, Kimi-K2.5, and gpt-oss-20b launcher configs. - Replace Kimi-K2.5 aarch64-specific vLLM image (v0.22.0-aarch64) with the multi-arch v0.22.0 tag, which resolves to amd64/arm64 automatically. - Fix Qwen3.5-4B throughput_32k runs: raise max_seq_len from 40960 to 65536 to accommodate outlier prompts (~46.6K tokens) that caused VLLMValidationError and aborted the entire benchmark run. - Fix specdec_bench_mtp_vllm.yaml: remove reference to non-existent runtime_params_throughput_32k.yaml; use --max_seq_len 65536 instead ### Usage ``` cd Model-Optimizer/tools/launcher uv run launch.py --yaml examples/Qwen/Qwen3-8B/megatron_lm_ptq_local.yaml hf_local=/mnt/hf-local --yes ``` ### Testing 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 ### Additional Information N/A <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Improvements** * Updated launcher example pipelines and the default launcher Slurm container to the newer TensorRT-LLM `1.3.0rc20` image. * Updated Kimi-K2.5 workflows to use the multi-architecture vLLM `0.22.0` image (removing prior architecture-specific variants). * Increased the Qwen3.5-4B long-context benchmark maximum sequence length to 65,536 tokens and simplified the corresponding throughput configuration. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: noeyy-mino <174223378+noeyy-mino@users.noreply.github.com>
#1949) ### What does this PR do? Type of change: New feature. This PR adds ordered, module-name-specific search spaces to AutoQuant so different runtime decision groups can use different candidate formats. The user-facing design now separates fixed PTQ configuration from the AutoQuant search space: - A normal top-level **quantize** config defines the fixed baseline for modules that are not searched. - **auto_quantize.module_search_spaces** explicitly lists only the module families AutoQuant should search. - Top-level **auto_quantize.candidate_formats** remains available for the existing global-search mode, but it is mutually exclusive with a fixed **quantize** baseline. - Fixed and searched groups still participate in one integrated calibration, sensitivity-scoring, active-MoE cost, LP selection, checkpoint, and export flow. Additional safeguards: - Reject module rules that partially split a runtime-fused decision group. - Keep BF16/no-quant as an internal sensitivity baseline while allowing each search-space rule to control whether it is solver-selectable. - Isolate fixed groups from unrelated calibration algorithms. - Reject infeasible budgets from the resolved per-group choices before calibration. - Fingerprint the fixed PTQ baseline, runtime groups, candidate choices, scoring boundaries, replay attributes, and cost weights before checkpoint reuse. - Preserve the previous global candidate-format API and one-candidate module rules for backward compatibility. ### Design rationale The fixed configuration should use the same PTQ recipe system users already use for uniform NVFP4, W4A16, FP8, and other baseline configurations. AutoQuant then only needs to describe what is genuinely searched. For example, the Qwen3.6 recipe reuses the model-specific quantization configuration from `huggingface/qwen3_5_moe/ptq/w4a16_nvfp4-fp8_attn-kv_fp8_cast`. Unmatched modules retain that PTQ baseline, while shared experts, self-attention, linear-attention, and lm_head are explicitly searched. The original PTQ recipe is unchanged, and a loader test asserts that both recipes resolve to the same fixed `quantize` configuration. This remains one AutoQuant operation rather than staged PTQ plus AutoQuant: - NAS SearchSpace models generic architecture hyperparameters and is not connected to AutoQuant calibration snapshots, sensitivity scoring, runtime-fused groups, active-MoE accounting, or quant-config export. - Ordered QuantizeConfig wildcard rules can describe a final static assignment, but cannot express per-family candidate sets in one global budget solve. - Applying fixed PTQ before or after a separate AutoQuant pass would remove those modules from the shared sensitivity baseline and active-MoE numerator/denominator. The implementation therefore composes the existing PTQ recipe schema with the existing AutoQuant/PuLP path instead of introducing another search framework or solver. ### Usage ~~~yaml imports: model_quant_cfg: huggingface/qwen3_5/ptq/w4a16_nvfp4-fp8_attn-kv_fp8_cast.quant_cfg w4a16_nvfp4: configs/ptq/presets/model/w4a16_nvfp4 fp8: configs/ptq/presets/model/fp8 # Same model-specific baseline as the Qwen3.5-MoE PTQ recipe. quantize: algorithm: method: max layerwise: enable: false quant_cfg: - $import: model_quant_cfg auto_quantize: constraints: effective_bits: 6.0 cost_model: active_moe cost: active_moe_expert_ratio: 0.03125 # Only these modules are AutoQuant decision variables. module_search_spaces: - module_name_patterns: - "*mlp.shared_expert*" - "*linear_attn*" - "*self_attn*" - "*lm_head*" candidate_formats: - $import: w4a16_nvfp4 - $import: fp8 allow_no_quant: false ~~~ The legacy global-search form remains supported by omitting **quantize** and supplying top-level **auto_quantize.candidate_formats**. ### Qwen3.6 search-space comparison Compared the default W4A16 NVFP4/FP8 search against the historical two-format module-specific recipe that keeps routed experts on the W4A16 PTQ baseline while shared experts and attention search W4A16 versus FP8. The final PR recipe keeps every matched quantizable group in W4A16 or FP8 (`allow_no_quant: false`). The table below is retained as evidence for the fixed-baseline design; the final recipe validation follows it. Both lanes used local raw-gradient scoring, active-MoE accounting, batch size 1, code32k calibration, full parser-on LiveCodeBench with 8 repeats, and full SciCode. Active GiB/token includes scale storage. | Target bits | Active GiB/token, default -> fixed W4 routed | Attention BF16+FP8, default -> fixed W4 routed | LCB, default -> fixed W4 routed | SciCode, default -> fixed W4 routed | |---:|---:|---:|---:|---:| | 5.8 | 2.1289 -> 2.0135 | 49.2% -> 80.0% | 70.8150 -> 72.4945 (+1.6795 pp) | 40.0518 -> 40.4216 (+0.3698 pp) | | 6.1 | 2.2342 -> 2.1070 | 61.5% -> 93.8% | 72.5220 -> 73.5683 (+1.0463 pp) | 39.1642 -> 39.3861 (+0.2219 pp) | | 6.4 | 2.3387 -> 2.2181 | 67.7% -> 79.2% | 72.0540 -> 73.9813 (+1.9273 pp) | 38.2027 -> 39.9038 (+1.7011 pp) | | 6.7 | 2.4392 -> 2.3173 | 72.3% -> 93.8% | 72.5220 -> 73.4581 (+0.9361 pp) | 38.7944 -> 39.6820 (+0.8876 pp) | BF16 references are 74.3667 LCB and 40.7914 SciCode. With a 1.5 percentage-point tolerance, the fixed-routed-expert lane jointly passes at targets 6.1, 6.4, and 6.7; the default lane has no joint pass. This comparison is end-to-end rather than an isolated solver ablation because the historical default lane used full attention-family grouping while the module-specific lane used runtime-required grouping. Final no-BF16 recipe validation used the model-specific PTQ baseline, codeblend 16k calibration, local gradient scoring, active-MoE accounting, batch size 1, and full parser-on LiveCodeBench with 8 repeats: | Target bits | Active GiB/token | Linear attention FP8/W4 | Self attention FP8/W4 | Shared experts FP8/W4 | lm_head | LCB pass@1 avg-of-8 | Avg completion tokens | Delta vs BF16 | |---:|---:|---:|---:|---:|:---:|---:|---:|---:| | 6.0 | 2.0823 | 80/10 | 36/4 | 97/23 | W4 | 73.2930 | 39,319.1 | -1.0737 pp | | 6.3 | 2.1848 | 59/31 | 35/5 | 84/36 | FP8 | 73.8711 | 38,515.8 | -0.4956 pp | Both targets pass the BF16-minus-1.5pp threshold. The model-specific baseline intentionally leaves linear-attention A/B and non-linear helper modules unquantized; all searched quantizable projections are W4A16 or FP8, and the artifact audit found no missing quant tensors. ### Testing - Fresh four-GPU Qwen3.6 model-specific-baseline E2E runs completed calibration, gradient search, ModelOpt export, and runtime validation at 6.0 and 6.3 effective bits. Slurm jobs 2649968 and 2649969 completed with exit code 0. Both artifacts had complete shards, valid runtime metadata, and zero quantized modules missing quant tensors. - Pre-commit passed: Ruff, formatting, mypy, Bandit, Markdown/YAML checks, and recipe validation. - 306 focused AutoQuant, recipe-loader, and HF PTQ tests passed for the implementation. - All 212 recipe-loader and HF PTQ mapping tests passed for the model-specific PTQ baseline integration. After the final `allow_no_quant: false` update, the focused loader/HF mapping tests and full pre-commit recipe validation passed. The original PTQ recipe remains unchanged, and the resolved fixed baselines compare equal. - The distributed AutoQuant unit test passed separately outside the local socket sandbox. - Wider quantization validation reached 782 passing tests; remaining local failures were unrelated optional dependency/socket limitations. - Historical four-GPU Qwen3.6 E2E calibration, search, export, and metadata verification for the fixed-W4 baseline variant: - Slurm job 2645235 completed with exit code 0 in 9m53s. - Achieved 5.99 effective bits. - All 40 routed-expert groups remained fixed W4A16. - Exported valid searched allocations: linear attention FP8/W4 = 78/12, self-attention = 37/3, shared experts = 112/8, lm_head = W4. - Produced and verified a 21 GiB ModelOpt checkpoint. - Completed the four-target full parser-on LCB/SciCode evaluation summarized above. - Final no-BF16 parser-on LCB jobs 2650360 and 2650367 completed with exit code 0; their 8-repeat results are summarized above. ### Before this PR is ready for review - Is this change backward compatible?: ✅ - If code was copied or a new dependency was added, was the contribution guidance followed?: N/A - Were necessary tests added?: ✅ - Was the changelog updated?: ✅ - Claude approval after the latest update?: Pending ### Additional Information Related work: OMNIML-5477, OMNIML-5119. --------- Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
… meta placeholders) (#1640) ### What does this PR do? **Type of change:** New feature Adds two opt-in memory optimizations to layerwise calibration (built on the nested `LayerwiseConfig` foundation merged in #1571): 1. **Skip layer-weight checkpoint + in-memory writeback** (`LayerwiseConfig.calib_mutates_weights`). For amax-only algorithms (`max`/`mse`/`local_hessian`) — which only update `TensorQuantizer._amax` and never touch `layer.weight` — set `calib_mutates_weights=False` to (a) checkpoint only the quantizer-state slice (`quantizer_buffers.pt`) instead of the full layer weights blob, and (b) skip the offload/shard *writeback* in `persistent_materialization`. A model validator restricts `calib_mutates_weights=False` to those amax-only algorithms; weight-mutating algorithms (GPTQ/AWQ/SmoothQuant) reject it. Also fixes the FSDP2 `writeback=False` gather ordering so the layer is actually all-gathered (not left sharded) during calibration. 2. **Meta-device skip-layer placeholders.** Already-calibrated ("skip") layers now emit zero-filled `meta` tensors instead of real-device buffers, eliminating their activation memory. (Limitation: models whose parent `forward` runs real-device ops on the hidden state *between* decoder blocks are unsupported and should use non-layerwise calibration — this is unconditional, not gated by `calib_mutates_weights`.) ### Usage ```python import modelopt.torch.quantization as mtq cfg = mtq.INT8_DEFAULT_CFG # Amax-only algorithm: skip the per-layer weight checkpoint + writeback to save memory. cfg["algorithm"] = {"method": "max", "layerwise": {"enable": True, "calib_mutates_weights": False}} mtq.quantize(model, cfg, forward_loop) ``` ### Testing - **Unit** (`tests/unit/torch/quantization/`): `test_config_validation.py` (the `calib_mutates_weights` whitelist accept/reject), `test_layerwise_calibrate.py` (checkpoint save/resume, the meta-placeholder behavior, and incomplete-checkpoint guard). All passing on CPU (125 passed, 1 skipped — the meta inter-layer-ops case). - **GPU** (multi-GPU runner): `tests/gpu/torch/quantization/test_fsdp2.py::test_layerwise_calibrate_fsdp2` asserts every FSDP layer is actually all-gathered (params are **not** `DTensor`s) under `writeback=False` — direct coverage of the FSDP2 fix; plus `plugins/test_accelerate_gpu.py` and `test_gptq.py`. ### 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.). > Note: checkpoint resume uses `torch.load(..., weights_only=False)` on the per-layer `weights.pt` / `quantizer_buffers.pt` / `next_inputs.pt` files. These are written by the calibration run itself into the user-supplied `checkpoint_dir` (not third-party artifacts); the pre-existing pattern is unchanged by this PR. `full_restore` now raises on an incomplete checkpoint rather than silently restoring partial state. - Is this change backward compatible?: ✅ — `calib_mutates_weights` defaults to `True` (prior behavior); both optimizations are opt-in and the layerwise config is unreleased. - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A — no copied code or new dependencies. - Did you write any new necessary tests?: ✅ - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ - Did you get Claude approval on this PR?: ❌ <!-- run /claude review --> ### Additional Information Co-authored with @realAsma. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit ## Release Notes * **New Features** * Added `layerwise.calib_mutates_weights` to control whether layerwise calibration performs weight writeback, with safer support limited to amax-only algorithms. * Improved layerwise skip-mode to reduce activation memory by using meta-device placeholders for already-calibrated layers. * **Improvements** * Enhanced layerwise checkpointing/resume with stronger manifest validation, scenario-aware persistence behavior, and clearer progress reporting. * Extended writeback control for persistent weight materialization across supported backends. * Updated `gptq` documentation to clarify Layerwise vs Non-layerwise semantics. * **Tests** * Expanded GPU/CPU/FSDP2 and validation coverage for meta placeholders, checkpoint layout, and config deprecations. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: realAsma <akuriparambi@nvidia.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Co-authored-by: realAsma <akuriparambi@nvidia.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
### Description MBridge main (nemo:26.08) will initialize Nemotron-H as HybridModel instead of MambaModel (subclassed of HybridModel). Also make minimum nemo container 26.04 ### Testing Tested Nemotron-3-Nano PTQ with MBridge main (fails otherwise) Tested locally `tests/gpu_megatron` and `tests/examples/megatron_bridge` with `nemo:26.06.01` + Mount latest MBridge/Mcore GH CICD tests will be added with nemo:26.08 release <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added TE HybridModel stack-spec support and enabled Megatron-Bridge hybrid providers/models for export/import and runtime handling. * **Updates** * Dataset packing now oversamples raw text at **16x** and improves the packed-mode underflow warning. * Quantization: `--quant_cfg` now defaults to `None` unless explicitly set (or via `--recipe`). * Distillation example: validation settings are provided via a dedicated top-level validation configuration. * Improved plugin import warnings to report the originating call location; model stats now support HybridModel. * **Deprecations** * Megatron-Bridge / Megatron-LM optimization features now require NeMo container `nemo:26.04` or newer (`nemo:26.06` recommended). * The Mamba stack specification helper is deprecated. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
…#1897) ### What does this PR do? Save hf checkpoint at every valitation iteration during distillation. Addition functionality: - added `--validate_only` in `examples/megatron_bridge/distill.py` to enable computing validation losses for iter 0 - added `--reset_optimizer` in `examples/megatron_bridge/distill.py` to enable not using presaved optimizer, e.g., when changing the number of train iters. ### Usage - examples/megatron_bridge/distill.py - examples/megatron_bridge/README.md (line 228) ### Testing - tests/examples/megatron_bridge/test_distill.py ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ - Did you write any new necessary tests?: ✅ <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added `--validate_only` for student-only validation at iteration 0. * Added `--hf_validation_export_path` with `--hf_validation_export_interval` to export validated student HuggingFace artifacts during distillation. * Added `prepare_data_blend.py` for YAML-driven token-budgeted data blends. * Added `--max_tokens` to stop Megatron preprocessing after a token budget. * **Bug Fixes** * Validation exports avoid duplicate checkpoints, preserve the student architecture/config, and export only student artifacts. * **Documentation** * Expanded researcher and tutorial guides for iterative workflows and token-budgeted blends. * **Tests** * Updated distillation/blend/max_tokens test coverage, including validate-only and interval-based exports. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
### What does this PR do? Adds a new entry to the **Latest News** section of `README.md` for the NVIDIA Developer blog post on quantizing Nemotron 3 Ultra (550B) to NVFP4 with Model Optimizer: - **Blog:** [Creating the NVIDIA Nemotron 3 Ultra NVFP4 Checkpoint with NVIDIA Model Optimizer](https://developer.nvidia.com/blog/creating-the-nvidia-nemotron-3-ultra-nvfp4-checkpoint-with-nvidia-model-optimizer/) - **Checkpoint:** [NVIDIA-Nemotron-3-Ultra-550B-A55B-NVFP4](https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-NVFP4) ### 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?: 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 <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit ## Summary by CodeRabbit * **Documentation** * Updated the “Latest News” section with a new announcement for **2026/06/26** about the creation of the NVIDIA Nemotron 3 Ultra NVFP4 checkpoint using NVIDIA Model Optimizer. * Included throughput/accuracy highlights and a link to the related Hugging Face checkpoint. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Seonghee Lee <jessielee.shl@gmail.com>
### What does this PR do? Type of change: documentation This PR clarifies when users should use ONNX quantization with Autotune enabled versus the direct Autotune entry point. - Adds a warning to the Autotune guide explaining that direct Autotune is a lower-level Q/DQ placement tool and does not replace calibrated ONNX PTQ. - Updates the ONNX quantization guide to show `autotune=True` in the Python API and explain that it uses default Autotune settings. - Updates ONNX PTQ example documentation to prefer `python -m modelopt.onnx.quantization ... --autotune=<mode>` for accuracy-sensitive PTQ from an unquantized model. - Updates the direct Autotune CLI help text to point users back to the full ONNX quantization workflow when calibration data and accuracy validation are required. ### Usage ```python N/A — documentation/help text change. ``` ### Testing - Ran `python -m py_compile modelopt/onnx/quantization/autotune/__main__.py`. - Built the Sphinx documentation with `python -m sphinx -b html docs/source docs/build/html`; build succeeded. Remaining warnings are from optional documentation imports and existing cross-reference labels. ### 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?: N/A - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Documentation** * Clarified that Direct Autotune is an advanced tool for Q/DQ placement experiments, not a replacement for full calibrated ONNX quantization. * Added guidance on when to use Direct Autotune versus the end-to-end ONNX PTQ workflow. * Documented the optional `autotune=True` setting, expected calibration-time impact, and representative calibration data requirements. * Expanded links and guidance across ONNX PTQ examples and updated CLI help text with the recommended workflow. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Gwenaelle Cunha Sergio <gcunhasergio@nvidia.com>
Re-organize Changelog for 0.46.0 into subsections Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
…runing (#1955) ### What does this PR do? Type of change: New feature Add Minitron pruning support for MoE models loaded with the efficient fused **grouped GEMM** experts (`TEGroupedMLP`), in addition to the existing `TESequentialMLP` path. - New `_DynamicTEGroupedLinear` + `_DynamicTEGroupedMLP` dynamic modules slice the per-expert grouped weights (`num_moe_experts` / `moe_ffn_hidden_size` / `hidden_size`) and reorder/drop experts via `num_gemms`, mirroring the SequentialMLP path with a minimal DynamicModule diff. - Since Minitron prunes homogeneously, a single shared `moe_ffn_hidden_size` is pruned across experts (SequentialMLP keeps one per expert). Same pruned width and independent per-expert weights; only the kept-channel index set is shared. - `examples/megatron_bridge/prune_minitron.py` uses grouped GEMM by default (`--no_moe_grouped_gemm` to fall back to SequentialMLP). ### Usage ```bash torchrun --nproc_per_node 2 prune_minitron.py \ --hf_model_name_or_path <moe-model> \ --prune_target_active_params 3e9 \ --output_hf_path /tmp/pruned # add --no_moe_grouped_gemm to use SequentialMLP ``` ### Testing Verified in an `nvcr.io/nvidia/nemo:26.06` + MBridge main (as of 23 Jul) mounted so it mimics nemo:26.08 behavior. - GPT MoE dynamic-module + pruning + parameter-sorting tests parametrized over both expert impls. - Mamba hybrid NAS metric tests: params-based now covers grouped GEMM, memory-based stays SequentialMLP (exact param counts / top-k / search-space-size assertions hold identically). - NemotronH end-to-end `test_prune_minitron` exercises real-forward NAS without grouped GEMM (to be enabled in 26.08 container) - Compared Nemotron-3-Nano-30B-A3B pruning: SequentialMLP vs GroupedGEMM on 4x B300 (accuracy + speed) | Metric | TESequentialMLP (old) | TEGroupedLinear (new) | |---|---|---| | Calibration time | 6.5 mins | 3.5 mins| | Time to evaluate Top-10 pruned candidates | 23 mins | 11 mins | Top-10 Pruned Candidates — MMLU Scores | # | export_config | params | TESequentialMLP (old) | TEGroupedLinear (new) | |---|---|---|---|---| | 1 | L52, h2688, mamba 56×48, 96 experts, ffn 1536, shared 3072 | 20.09B | 0.2713 | 0.2846 | | 2 | L52, h2688, mamba 48×56, 104 experts, ffn 1536, shared 3072 | 21.61B | 0.2580 | 0.2594 | | 3 | L52, h2560, mamba 48×64, 96 experts, ffn 1536, shared 3712 | 19.28B | 0.3951 | 0.3895 | | 4 | L52, h2304, mamba 64×64, 104 experts, ffn 1856, shared 3072 | 22.28B | 0.4951 | 0.4944 | | 5 | L52, h2560, mamba 48×48, 96 experts, ffn 1792, shared 3328 | 21.99B | 0.2685 | 0.2580 | | 6 | L48, h2560, mamba 56×56, 104 experts, ffn 1792, shared 3072 | 23.68B | 0.4741 | 0.4657 | | 7 | L46, h2560, mamba 64×56, 104 experts, ffn 1792, shared 3072 | 23.68B | 0.2385 | 0.2371 | | 8 | L52, h2688, mamba 48×56, 96 experts, ffn 1536, shared 3072 | 20.09B | 0.2587 | 0.2622 | | 9 | L52, h2304, mamba 64×64, 96 experts, ffn 1856, shared 3072 | 20.70B | 0.4888 | 0.4860 | | 10 | L50, h2560, mamba 48×48, 104 experts, ffn 1792, shared 3712 | 23.68B | 0.2517 | 0.2531 | ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ - Did you write any new necessary tests?: ✅ - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ - Did you get Claude approval on this PR?: ✅ <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit ## Summary by CodeRabbit * **New Features** * Improved MoE pruning and NAS/search-space handling to support both grouped-GEMM and sequential expert layouts, including optimized grouped-GEMM execution for Minitron. * Added `--no_moe_grouped_gemm` to force the sequential expert path when required. * **Bug Fixes** * Tightened `--prune_score_func` parsing for MMLU to accept only `mmlu_<N>pct_bs<bs>`. * Added a safety check in Megatron prefill to prevent int32 indexing overflow by asking to reduce calibration batch size. * **Tests** * Expanded GPT and Mamba GPU pruning/search-space tests to cover both MoE modes. * **Documentation** * Updated release notes and bridge/pruning guidance for the grouped-GEMM behavior. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
### What does this PR do?
Type of change: new example + bug fixes
Adds a **MiniMax-M3 DSpark streaming training recipe** under
`tools/launcher/examples/MiniMaxAI/MiniMax-M3/`, plus the four fixes it
needs to actually run. Each fix addresses a failure mode that is silent
or misleading without it:
1. **Gemma-style final norm for the fake base**
(`modeling_final_norm.py`, `modeling_fakebase.py`): M3 uses a
gemma-style final RMSNorm (`(1 + weight)` scale, fp32
multiply-then-cast). Selecting the norm by `model_type` alone picks
plain `rmsnorm` — MiniMax's VL remote code coerces its `model_type`-less
`text_config` to **mixtral** — which silently drops the `+1` and
corrupts the distillation target. New `_FinalGemmaRMSNorm` + selection
by the explicit `use_gemma_norm` config flag (only MiniMax sets it).
2. **Loud failure for un-maskable `answer_only_loss`**
(`hf_streaming_dataset.py`): with a fast tokenizer whose chat template
has no `{% generation %}` tags,
`apply_chat_template(return_assistant_tokens_mask=True)` only warns and
returns an **all-zero mask** — training runs at zero loss on every
sample with no other symptom. Now raises at tokenization with an
actionable message.
3. **`SERVE_BLOCK_SIZE` knob** (`train_eagle_streaming.sh`): nemo_run
exports env values unquoted, so a multi-token
`SERVE_EXTRA_ARGS="--trust-remote-code --block-size 128"` loses
everything after the first token. M3's MSA sparse attention requires KV
block 128 (`ValueError: No common block size for 16` at engine init
otherwise), so `--block-size` gets a dedicated single-token knob.
4. **Relax the speculative_decoding `transformers` pin to `<5.13`**
(match `pyproject.toml`): the old `<5.4` pin downgrades recent vLLM
containers (e.g. transformers 5.12.1, which also provides in-tree
`minimax_m3_vl`) and breaks `vllm serve` (`ALLOWED_LAYER_TYPES` needs
>=5.5.3).
The example itself encodes the validated M3 specifics: generation-tagged
chat template copy (required for `answer_only_loss` — see fix 2), draft
dims + base `rope_theta=5e6` set explicitly (not inherited), mask token
200063 (reserved slot; added tokens end at 200060), `EAGLE_CAPTURE_IDS`
= draft default `target_layer_ids+1` + final layer, `trust_remote_code`
at serve/export, and AWS-EFA NIXL notes (UCX segfaults at agent init on
EFA nodes; LIBFABRIC required there).
### Usage
```bash
cd tools/launcher
export SLURM_HOST=localhost SLURM_ACCOUNT=<account> SLURM_PARTITION=<partition> \
SLURM_HF_LOCAL=<hf_models_dir> SLURM_JOB_DIR=<experiments_dir> NEMORUN_HOME=$PWD
uv run launch.py --yaml examples/MiniMaxAI/MiniMax-M3/hf_streaming_dspark_multi_node.yaml \
identity=$HOME/.ssh/id_ecdsa detach=True --yes
```
### Testing
- `tests/unit/torch/speculative/plugins/`: **129 passed** inside a
current vLLM x86_64 nightly container (transformers 5.12.1), including 3
new tests for the assistant-mask guard.
- Generation-tagged template verified against real corpus samples:
`input_ids` identical to the original template, mask covers exactly the
assistant turns (think prefix + content + eos), contiguous, no
user-prompt leak.
- The recipe is exercised end-to-end by a live M3 DSpark training run
(this yaml modulo cluster paths): streaming serve + NIXL transport +
resume all healthy; drafter MT-Bench AL exceeds our Kimi-K2.6 DSpark
reference by ~8k steps.
- `ruff check` / `ruff format` (0.15.20) clean.
### Before your PR is "*Ready for review*"
- Is this change backward compatible?: ✅ (norm selection only changes
models with `use_gemma_norm=True`, previously mis-normed; the guard
turns a silent zero-loss run into an error; `SERVE_BLOCK_SIZE` is
opt-in; the pin relax widens the allowed range)
- Did you write any new necessary tests?: ✅
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
❌ (can add if desired)
🤖 Generated with [Claude Code](https://claude.com/claude-code)
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit
- **New Features**
- Added Gemma-style RMS normalization support for speculative decoding.
- Added MiniMax-M3 multi-node training configuration, including a full
Jinja chat template for tool calls, multimodal content, and thinking
modes.
- Added optional `SERVE_BLOCK_SIZE` support for vLLM serve launches.
- **Bug Fixes**
- Improved `answer_only_loss` masking validation: now fails fast with
clear errors when required `{% generation %}` markers are missing or
when using a non-fast tokenizer.
- **Compatibility**
- Expanded the supported Transformers version range for speculative
decoding examples.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---------
Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
### What does this PR do?
Type of change: New feature (+ refactor/removal of the legacy multi-node
path) <!-- Use one of the following: Bug fix, new feature, new example,
new tests, documentation. -->
<!-- Details about the change. -->
Consolidates multi-node FSDP2 post-training quantization into the
standard hf_ptq.py entry point behind a --use_fsdp2 flag, and removes
the separate
Accelerate-based multinode_ptq.py script and fsdp2.yaml config. FSDP2
PTQ is now launched with torchrun and supports single-node multi-GPU and
multi-node out
of the box.
Highlights:
- --use_fsdp2 / --cpu_offload on hf_ptq.py — calibration runs under
PyTorch FSDP2; decoder layers are sharded (root stays replicated), with
optional CPU
offload of decoder shards between forwards.
- New modelopt/torch/utils/model_load_utils.py — parallel, round-robin
safetensors loading: each rank reads only the decoder layers it owns
from disk, then
broadcasts them so every rank can shard its slice. Non-decoder weights
(embed/lm_head/norm) are read on rank 0 and broadcast.
- New FSDP2 helpers in modelopt/torch/utils/distributed.py — fsdp2_wrap,
shard_dataloader, fsdp_aware_forward_loop, broadcast_state_dict.
- Distributed export (unified_export_hf.py) — replaces the Accelerate
get_state_dict gather with get_model_state_dict(full_state_dict,
cpu_offload,
broadcast_from_rank0); rank 0 writes files, other ranks sync on a
barrier.
- core_utils.py — relaxes the stale "root must be an FSDPModule" assert
(only decoder layers are wrapped), and adds a CPU↔GPU mirror so weight
access/writeback
works for CPU-offloaded shards.
- Docs — rewritten examples/llm_ptq/README.md FSDP2 section and the
SLURM PTQ reference, both using torchrun --use_fsdp2.
v1 scope: standard causal-LM checkpoints only. --use_fsdp2 raises
NotImplementedError for VILA, multimodal/VL,
pack-quantized/compressed-tensors, speculative decoding, MTP,
auto-quantize, and sparsity.
Design Notes
1. Why a custom parallel-safetensors loader instead of HF device_map /
accelerate
AutoModelForCausalLM.from_pretrained(device_map="auto" | "cpu") and
accelerate.load_checkpoint_in_model both load the full checkpoint on
every rank from disk (per-rank CPU peak ≈ full model size). For 70B
that's ~140 GiB/rank; for 200B+ it OOMs the node before sharding can
run. parallel_load_and_prepare_fsdp2 round-robins decoder layers across
ranks so each rank reads only ~model_size / world_size from disk, then
per-layer broadcasts to peers. Per-rank CPU peak is bounded by the
largest single layer + transient broadcast, not the full model. This is
what makes 200B+ FSDP2 PTQ feasible on commodity nodes; HF/accelerate's
existing entry points don't expose this composition.
2. Why a custom FSDP2 stack instead of keeping multinode_ptq.py +
accelerate launch
The deleted multinode_ptq.py ran on accelerate launch --config_file
fsdp2.yaml and duplicated hf_ptq.py's load → calibrate → export path.
Two consequences:
- Two divergent scripts for the same operation: CLI surface, calibration
loop, and export rewrites had to land twice. They had already drifted.
- Users had to know which script applied at which scale.
The new code path unifies under hf_ptq.py --use_fsdp2. Going direct to
FSDP2 primitives (fully_shard, CPUOffloadPolicy, MixedPrecisionPolicy)
instead of
routing through accelerate's wrappers buys:
- Direct control over mp_policy and offload_policy (accelerate's plugin
layer hides them).
- The parallel-read loader above (incompatible with accelerate's
per-rank full load).
- No YAML config file in the example dir.
-
The "custom stack" is intentionally thin: fsdp2_wrap is a 5-line wrapper
over fully_shard; the rest is pure torch.distributed. We're not
reimplementing FSDP2 —just composing the public PyTorch surface directly
rather than via accelerate's adapter.
3. fsdp_aware_forward_loop ↔ transformers_trainer.py:_quantize_model
duplication
Both implement the same trick: mtq.quantize unwraps the FSDP module
before calling the user's forward_loop, and forwarding through the
unwrapped module bypasses FSDP2's pre/post-forward hooks (no all-gather
→ broken calibration). Both capture the outer wrapped model and forward
through it instead. This PR extracts the pattern into
fsdp_aware_forward_loop (in distributed.py) as the canonical helper. The
QLoRA path (transformers_trainer.py:_quantize_model) keeps its inlined
version this PR because the QLoRA forward loop has training-specific
quirks (batch shape, loss accumulation, gradient flow) that need a
careful pass to share the helper cleanly. The TODO in the helper's
docstring marks the consolidation target.
### Usage
```python
# Single node, multiple GPUs
torchrun --standalone --nproc_per_node=<num_gpus> hf_ptq.py \
--pyt_ckpt_path <model> \
--qformat nvfp4 \
--export_path <out> \
--use_fsdp2
# Multi-node (run on each node)
torchrun \
--nnodes=<N> --node_rank=<rank> \
--master_addr=<node0_ip> --master_port=<port> \
--nproc_per_node=<gpus_per_node> \
hf_ptq.py \
--pyt_ckpt_path <model> --qformat nvfp4 \
--export_path <out> --use_fsdp2 --cpu_offload # --cpu_offload for very large models
```
### Testing
<!-- Mention how have you tested your change if applicable. -->
- tests/gpu/torch/quantization/test_fsdp2.py:
test_writeback_root_unwrapped (assert relaxation, root unwrapped) and
test_writeback_cpu_offload (CPU↔GPU mirror
round-trip under CPUOffloadPolicy).
- tests/unit/torch/utils/test_model_load_utils.py: pure-function tests
for weight_map_for (sharded / single-file / missing) and
read_safetensors_subset.
- Manual end-to-end PTQ runs on single- and multi-node torchrun (FP8 /
NVFP4 / NVFP4 layerwise), with and without --cpu_offload.
### 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 ❌, explain why. -->
hf_ptq.py is fully backward compatible (FSDP2 is opt-in via a new flag),
but the legacy examples/llm_ptq/multinode_ptq.py script and fsdp2.yaml
are removed. Users of the old multi-node entry point must migrate to
hf_ptq.py --use_fsdp2
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: N/A <!---
Mandatory -->
- Did you write any new necessary tests?: ✅ <!--- Mandatory for new
features or examples. -->
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
❌<!--- Only for new features, API changes, critical bug fixes or
backward incompatible changes. -->
- Did you get Claude approval on this PR?: ✅ / ❌ / N/A <!--- Run
`/claude review`. NVIDIA org members can self-trigger for complex
changes; orthogonal to CodeRabbit. -->
### Additional Information
<!-- E.g. related issue. -->
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit
* **New Features**
* FSDP2-enabled PTQ: torchrun single-/multi-node execution, CPU offload,
and NVFP4 layerwise calibration; distributed model loading and export
support.
* **Documentation**
* Rewritten PTQ guide and SLURM notes with explicit torchrun/FSDP2
instructions.
* **Removed**
* Legacy Accelerate-based multinode PTQ workflow and YAML config.
* **Tests**
* Added FSDP2-focused tests for quantization writeback and safetensors
distributed loading.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---------
Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
Signed-off-by: sugunav14 <178320438+sugunav14@users.noreply.github.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: realAsma <86726418+realAsma@users.noreply.github.com>
…2022) ### What does this PR do? Type of change: Documentation update - Update documentation guide for ONNX INT4 PTQ on Windows cuda13 host - mention about compatible onnxruntim-gpu and cupy-cuda13x packages. ### Testing - Windows's onnx_ptq\genai_llm INT4 PTQ example with a 1B genai-cuda-ep ONNX model + local doc building ### 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?: ✅ / ❌ / N/A <!--- If ❌, explain why. --> - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: ✅ / ❌ / N/A <!--- Mandatory --> - Did you write any new necessary tests?: ✅ / ❌ / N/A <!--- Mandatory for new features or examples. --> - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ / ❌ / N/A <!--- Only for new features, API changes, critical bug fixes or backward incompatible changes. --> - Did you get Claude approval on this PR?: ✅ / ❌ / N/A <!--- Run `/claude review`. NVIDIA org members can self-trigger for complex changes; orthogonal to CodeRabbit. --> ### Additional Information <!-- E.g. related issue. --> <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Documentation** * Clarified Windows CUDA prerequisites for calibration and GPU-accelerated quantization. * Added setup guidance for CUDA 12 and CUDA 13.x, including compatible packages and cuDNN requirements. * Expanded installation verification steps for CUDA, ONNX Runtime, and CuPy. * Updated the GenAI LLM example with CUDA version compatibility guidance. * **Enhancements** * Added runtime logging of detected CUDA environment paths and version details during quantization. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: vipandya <vipandya@nvidia.com>
add scripts and a skill to use flashinfer to do layerwise benchmark with
different backends.
### What does this PR do?
Type of change: ? new skill and scripts
### Usage
```bash
python .agents/skills/benchmark-model-kernels/scripts/benchmark_model.py \
openai/gpt-oss-120b \
--tp 4 --ep 4 --ms 8 16 \
--flashinfer_repo $HOME/flashinfer \
--workdir /tmp/gpt-oss-benchmark
```
or
```bash
python .agents/skills/benchmark-model-kernels/scripts/benchmark_via_builtin.py \
--flashinfer_repo $HOME/flashinfer \
--ms 8 16 \
--nks 1280,2880 2880,1024 \
--moe_hidden_size 2880 \
--moe_intermediate_size 2880 \
--moe_num_experts 32 \
--moe_top_k 4 \
--workdir /tmp/gpt-oss-via-builtin
```
Or example prompt to trigger the skill:
```
Benchmark the kernel shapes for NVIDIA-Nemotron-3-Nano-30B-A3B
```
example output (TP1, EP4):
```
flashinfer 0.6.13; checkout ~/flashinfer3 @ 453bc6e1c6c852a4dc8a35c672f271e1e0d216a5; NVIDIA Graphics Device (sm_100 / 148 SMs / 178 GiB); 600 W power limit
GEMM
module_name,M,N,K,backend,with_quant,runtime
model.layers.*.mixer.shared_experts.up_proj,1,3712,2688,bf16,False,8.384
model.layers.*.mixer.shared_experts.up_proj,16,3712,2688,bf16,False,8.896
model.layers.*.mixer.shared_experts.up_proj,1,3712,2688,fp8_cudnn,False,6.112
model.layers.*.mixer.shared_experts.up_proj,1,3712,2688,fp8_cudnn,True,8.832
model.layers.*.mixer.shared_experts.up_proj,16,3712,2688,fp8_cudnn,False,6.176
model.layers.*.mixer.shared_experts.up_proj,16,3712,2688,fp8_cudnn,True,8.944
model.layers.*.mixer.shared_experts.up_proj,1,3712,2688,fp8_cutlass,False,7.280
model.layers.*.mixer.shared_experts.up_proj,1,3712,2688,fp8_cutlass,True,10.000
model.layers.*.mixer.shared_experts.up_proj,16,3712,2688,fp8_cutlass,False,7.216
model.layers.*.mixer.shared_experts.up_proj,16,3712,2688,fp8_cutlass,True,9.984
model.layers.*.mixer.shared_experts.up_proj,1,3712,2688,fp8_trtllm,False,7.583
model.layers.*.mixer.shared_experts.up_proj,1,3712,2688,fp8_trtllm,True,10.303
model.layers.*.mixer.shared_experts.up_proj,16,3712,2688,fp8_trtllm,False,7.520
model.layers.*.mixer.shared_experts.up_proj,16,3712,2688,fp8_trtllm,True,10.288
model.layers.*.mixer.shared_experts.up_proj,1,3712,2688,nvfp4_cudnn,False,7.456
model.layers.*.mixer.shared_experts.up_proj,1,3712,2688,nvfp4_cudnn,True,10.016
model.layers.*.mixer.shared_experts.up_proj,16,3712,2688,nvfp4_cudnn,False,7.472
model.layers.*.mixer.shared_experts.up_proj,16,3712,2688,nvfp4_cudnn,True,10.192
model.layers.*.mixer.shared_experts.up_proj,1,3712,2688,nvfp4_cutedsl,False,6.928
model.layers.*.mixer.shared_experts.up_proj,1,3712,2688,nvfp4_cutedsl,True,9.488
model.layers.*.mixer.shared_experts.up_proj,16,3712,2688,nvfp4_cutedsl,False,6.160
model.layers.*.mixer.shared_experts.up_proj,16,3712,2688,nvfp4_cutedsl,True,8.880
model.layers.*.mixer.shared_experts.up_proj,1,3712,2688,nvfp4_cutlass,False,7.232
model.layers.*.mixer.shared_experts.up_proj,1,3712,2688,nvfp4_cutlass,True,9.792
model.layers.*.mixer.shared_experts.up_proj,16,3712,2688,nvfp4_cutlass,False,8.320
model.layers.*.mixer.shared_experts.up_proj,16,3712,2688,nvfp4_cutlass,True,11.040
model.layers.*.mixer.shared_experts.up_proj,1,3712,2688,nvfp4_trtllm,False,7.024
model.layers.*.mixer.shared_experts.up_proj,1,3712,2688,nvfp4_trtllm,True,9.488
model.layers.*.mixer.shared_experts.up_proj,16,3712,2688,nvfp4_trtllm,False,6.976
model.layers.*.mixer.shared_experts.up_proj,16,3712,2688,nvfp4_trtllm,True,9.584
model.layers.*.mixer.shared_experts.down_proj,1,2688,3712,bf16,False,9.504
model.layers.*.mixer.shared_experts.down_proj,16,2688,3712,bf16,False,9.791
model.layers.*.mixer.shared_experts.down_proj,1,2688,3712,fp8_cudnn,False,6.560
model.layers.*.mixer.shared_experts.down_proj,1,2688,3712,fp8_cudnn,True,9.344
model.layers.*.mixer.shared_experts.down_proj,16,2688,3712,fp8_cudnn,False,6.688
model.layers.*.mixer.shared_experts.down_proj,16,2688,3712,fp8_cudnn,True,9.488
model.layers.*.mixer.shared_experts.down_proj,1,2688,3712,fp8_cutlass,False,7.712
model.layers.*.mixer.shared_experts.down_proj,1,2688,3712,fp8_cutlass,True,10.496
model.layers.*.mixer.shared_experts.down_proj,16,2688,3712,fp8_cutlass,False,7.680
model.layers.*.mixer.shared_experts.down_proj,16,2688,3712,fp8_cutlass,True,10.480
model.layers.*.mixer.shared_experts.down_proj,1,2688,3712,fp8_trtllm,False,8.560
model.layers.*.mixer.shared_experts.down_proj,1,2688,3712,fp8_trtllm,True,11.344
model.layers.*.mixer.shared_experts.down_proj,16,2688,3712,fp8_trtllm,False,8.623
model.layers.*.mixer.shared_experts.down_proj,16,2688,3712,fp8_trtllm,True,11.424
model.layers.*.mixer.shared_experts.down_proj,1,2688,3712,nvfp4_cudnn,False,8.464
model.layers.*.mixer.shared_experts.down_proj,1,2688,3712,nvfp4_cudnn,True,11.216
model.layers.*.mixer.shared_experts.down_proj,16,2688,3712,nvfp4_cudnn,False,8.512
model.layers.*.mixer.shared_experts.down_proj,16,2688,3712,nvfp4_cudnn,True,11.360
model.layers.*.mixer.shared_experts.down_proj,1,2688,3712,nvfp4_cutedsl,False,7.584
model.layers.*.mixer.shared_experts.down_proj,1,2688,3712,nvfp4_cutedsl,True,10.336
model.layers.*.mixer.shared_experts.down_proj,16,2688,3712,nvfp4_cutedsl,False,6.416
model.layers.*.mixer.shared_experts.down_proj,16,2688,3712,nvfp4_cutedsl,True,9.264
model.layers.*.mixer.shared_experts.down_proj,1,2688,3712,nvfp4_cutlass,False,8.992
model.layers.*.mixer.shared_experts.down_proj,1,2688,3712,nvfp4_cutlass,True,11.744
model.layers.*.mixer.shared_experts.down_proj,16,2688,3712,nvfp4_cutlass,False,8.800
model.layers.*.mixer.shared_experts.down_proj,16,2688,3712,nvfp4_cutlass,True,11.648
model.layers.*.mixer.shared_experts.down_proj,1,2688,3712,nvfp4_trtllm,False,7.648
model.layers.*.mixer.shared_experts.down_proj,1,2688,3712,nvfp4_trtllm,True,10.208
model.layers.*.mixer.shared_experts.down_proj,16,2688,3712,nvfp4_trtllm,False,7.104
model.layers.*.mixer.shared_experts.down_proj,16,2688,3712,nvfp4_trtllm,True,9.887
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,1,4608,2688,bf16,False,8.704
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,16,4608,2688,bf16,False,8.736
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,1,4608,2688,fp8_cudnn,False,6.512
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,1,4608,2688,fp8_cudnn,True,9.232
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,16,4608,2688,fp8_cudnn,False,6.464
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,16,4608,2688,fp8_cudnn,True,9.232
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,1,4608,2688,fp8_cutlass,False,8.016
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,1,4608,2688,fp8_cutlass,True,10.736
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,16,4608,2688,fp8_cutlass,False,7.648
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,16,4608,2688,fp8_cutlass,True,10.416
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,1,4608,2688,fp8_trtllm,False,7.776
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,1,4608,2688,fp8_trtllm,True,10.496
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,16,4608,2688,fp8_trtllm,False,7.904
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,16,4608,2688,fp8_trtllm,True,10.672
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,1,4608,2688,nvfp4_cudnn,False,7.808
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,1,4608,2688,nvfp4_cudnn,True,10.367
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,16,4608,2688,nvfp4_cudnn,False,7.920
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,16,4608,2688,nvfp4_cudnn,True,10.640
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,1,4608,2688,nvfp4_cutedsl,False,7.104
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,1,4608,2688,nvfp4_cutedsl,True,9.664
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,16,4608,2688,nvfp4_cutedsl,False,6.304
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,16,4608,2688,nvfp4_cutedsl,True,9.024
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,1,4608,2688,nvfp4_cutlass,False,8.736
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,1,4608,2688,nvfp4_cutlass,True,11.296
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,16,4608,2688,nvfp4_cutlass,False,7.583
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,16,4608,2688,nvfp4_cutlass,True,10.303
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,1,4608,2688,nvfp4_trtllm,False,6.784
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,1,4608,2688,nvfp4_trtllm,True,9.248
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,16,4608,2688,nvfp4_trtllm,False,7.040
model.layers.*.mixer.q_proj|model.layers.*.mixer.k_proj|model.layers.*.mixer.v_proj,16,4608,2688,nvfp4_trtllm,True,9.647
model.layers.*.mixer.o_proj,1,2688,4096,bf16,False,10.368
model.layers.*.mixer.o_proj,16,2688,4096,bf16,False,10.880
model.layers.*.mixer.o_proj,1,2688,4096,fp8_cudnn,False,6.912
model.layers.*.mixer.o_proj,1,2688,4096,fp8_cudnn,True,9.663
model.layers.*.mixer.o_proj,16,2688,4096,fp8_cudnn,False,7.008
model.layers.*.mixer.o_proj,16,2688,4096,fp8_cudnn,True,9.888
model.layers.*.mixer.o_proj,1,2688,4096,fp8_cutlass,False,8.160
model.layers.*.mixer.o_proj,1,2688,4096,fp8_cutlass,True,10.912
model.layers.*.mixer.o_proj,16,2688,4096,fp8_cutlass,False,8.064
model.layers.*.mixer.o_proj,16,2688,4096,fp8_cutlass,True,10.944
model.layers.*.mixer.o_proj,1,2688,4096,fp8_trtllm,False,8.480
model.layers.*.mixer.o_proj,1,2688,4096,fp8_trtllm,True,11.232
model.layers.*.mixer.o_proj,16,2688,4096,fp8_trtllm,False,8.752
model.layers.*.mixer.o_proj,16,2688,4096,fp8_trtllm,True,11.632
model.layers.*.mixer.o_proj,1,2688,4096,nvfp4_cudnn,False,8.768
model.layers.*.mixer.o_proj,1,2688,4096,nvfp4_cudnn,True,11.520
model.layers.*.mixer.o_proj,16,2688,4096,nvfp4_cudnn,False,8.864
model.layers.*.mixer.o_proj,16,2688,4096,nvfp4_cudnn,True,11.680
model.layers.*.mixer.o_proj,1,2688,4096,nvfp4_cutedsl,False,7.904
model.layers.*.mixer.o_proj,1,2688,4096,nvfp4_cutedsl,True,10.656
model.layers.*.mixer.o_proj,16,2688,4096,nvfp4_cutedsl,False,6.784
model.layers.*.mixer.o_proj,16,2688,4096,nvfp4_cutedsl,True,9.600
model.layers.*.mixer.o_proj,1,2688,4096,nvfp4_cutlass,False,9.248
model.layers.*.mixer.o_proj,1,2688,4096,nvfp4_cutlass,True,12.000
model.layers.*.mixer.o_proj,16,2688,4096,nvfp4_cutlass,False,9.264
model.layers.*.mixer.o_proj,16,2688,4096,nvfp4_cutlass,True,12.080
model.layers.*.mixer.o_proj,1,2688,4096,nvfp4_trtllm,False,7.184
model.layers.*.mixer.o_proj,1,2688,4096,nvfp4_trtllm,True,9.744
model.layers.*.mixer.o_proj,16,2688,4096,nvfp4_trtllm,False,7.744
model.layers.*.mixer.o_proj,16,2688,4096,nvfp4_trtllm,True,10.432
model.layers.*.mixer.out_proj,1,2688,4096,bf16,False,10.368
model.layers.*.mixer.out_proj,16,2688,4096,bf16,False,10.880
model.layers.*.mixer.out_proj,1,2688,4096,fp8_cudnn,False,6.912
model.layers.*.mixer.out_proj,1,2688,4096,fp8_cudnn,True,9.663
model.layers.*.mixer.out_proj,16,2688,4096,fp8_cudnn,False,7.008
model.layers.*.mixer.out_proj,16,2688,4096,fp8_cudnn,True,9.888
model.layers.*.mixer.out_proj,1,2688,4096,fp8_cutlass,False,8.160
model.layers.*.mixer.out_proj,1,2688,4096,fp8_cutlass,True,10.912
model.layers.*.mixer.out_proj,16,2688,4096,fp8_cutlass,False,8.064
model.layers.*.mixer.out_proj,16,2688,4096,fp8_cutlass,True,10.944
model.layers.*.mixer.out_proj,1,2688,4096,fp8_trtllm,False,8.480
model.layers.*.mixer.out_proj,1,2688,4096,fp8_trtllm,True,11.232
model.layers.*.mixer.out_proj,16,2688,4096,fp8_trtllm,False,8.752
model.layers.*.mixer.out_proj,16,2688,4096,fp8_trtllm,True,11.632
model.layers.*.mixer.out_proj,1,2688,4096,nvfp4_cudnn,False,8.768
model.layers.*.mixer.out_proj,1,2688,4096,nvfp4_cudnn,True,11.520
model.layers.*.mixer.out_proj,16,2688,4096,nvfp4_cudnn,False,8.864
model.layers.*.mixer.out_proj,16,2688,4096,nvfp4_cudnn,True,11.680
model.layers.*.mixer.out_proj,1,2688,4096,nvfp4_cutedsl,False,7.904
model.layers.*.mixer.out_proj,1,2688,4096,nvfp4_cutedsl,True,10.656
model.layers.*.mixer.out_proj,16,2688,4096,nvfp4_cutedsl,False,6.784
model.layers.*.mixer.out_proj,16,2688,4096,nvfp4_cutedsl,True,9.600
model.layers.*.mixer.out_proj,1,2688,4096,nvfp4_cutlass,False,9.248
model.layers.*.mixer.out_proj,1,2688,4096,nvfp4_cutlass,True,12.000
model.layers.*.mixer.out_proj,16,2688,4096,nvfp4_cutlass,False,9.264
model.layers.*.mixer.out_proj,16,2688,4096,nvfp4_cutlass,True,12.080
model.layers.*.mixer.out_proj,1,2688,4096,nvfp4_trtllm,False,7.184
model.layers.*.mixer.out_proj,1,2688,4096,nvfp4_trtllm,True,9.744
model.layers.*.mixer.out_proj,16,2688,4096,nvfp4_trtllm,False,7.744
model.layers.*.mixer.out_proj,16,2688,4096,nvfp4_trtllm,True,10.432
model.layers.*.mixer.in_proj,1,10304,2688,bf16,False,17.520
model.layers.*.mixer.in_proj,16,10304,2688,bf16,False,16.032
model.layers.*.mixer.in_proj,1,10304,2688,fp8_cudnn,False,8.912
model.layers.*.mixer.in_proj,1,10304,2688,fp8_cudnn,True,11.632
model.layers.*.mixer.in_proj,16,10304,2688,fp8_cudnn,False,9.375
model.layers.*.mixer.in_proj,16,10304,2688,fp8_cudnn,True,12.143
model.layers.*.mixer.in_proj,1,10304,2688,fp8_cutlass,False,10.527
model.layers.*.mixer.in_proj,1,10304,2688,fp8_cutlass,True,13.248
model.layers.*.mixer.in_proj,16,10304,2688,fp8_cutlass,False,10.415
model.layers.*.mixer.in_proj,16,10304,2688,fp8_cutlass,True,13.183
model.layers.*.mixer.in_proj,1,10304,2688,fp8_trtllm,False,9.328
model.layers.*.mixer.in_proj,1,10304,2688,fp8_trtllm,True,12.048
model.layers.*.mixer.in_proj,16,10304,2688,fp8_trtllm,False,9.392
model.layers.*.mixer.in_proj,16,10304,2688,fp8_trtllm,True,12.160
model.layers.*.mixer.in_proj,1,10304,2688,nvfp4_cudnn,False,9.088
model.layers.*.mixer.in_proj,1,10304,2688,nvfp4_cudnn,True,11.648
model.layers.*.mixer.in_proj,16,10304,2688,nvfp4_cudnn,False,9.024
model.layers.*.mixer.in_proj,16,10304,2688,nvfp4_cudnn,True,11.744
model.layers.*.mixer.in_proj,1,10304,2688,nvfp4_cutedsl,False,9.424
model.layers.*.mixer.in_proj,1,10304,2688,nvfp4_cutedsl,True,11.984
model.layers.*.mixer.in_proj,16,10304,2688,nvfp4_cutedsl,False,7.616
model.layers.*.mixer.in_proj,16,10304,2688,nvfp4_cutedsl,True,10.336
model.layers.*.mixer.in_proj,1,10304,2688,nvfp4_cutlass,False,9.728
model.layers.*.mixer.in_proj,1,10304,2688,nvfp4_cutlass,True,12.288
model.layers.*.mixer.in_proj,16,10304,2688,nvfp4_cutlass,False,10.144
model.layers.*.mixer.in_proj,16,10304,2688,nvfp4_cutlass,True,12.864
model.layers.*.mixer.in_proj,1,10304,2688,nvfp4_trtllm,False,ERROR: FlashInfer produced no result row and no error message; see scratchpad/nemotron_ep4_v6/driver.log
model.layers.*.mixer.in_proj,1,10304,2688,nvfp4_trtllm,True,ERROR: FlashInfer produced no result row and no error message; see scratchpad/nemotron_ep4_v6/driver.log
model.layers.*.mixer.in_proj,16,10304,2688,nvfp4_trtllm,False,ERROR: FlashInfer produced no result row and no error message; see scratchpad/nemotron_ep4_v6/driver.log
model.layers.*.mixer.in_proj,16,10304,2688,nvfp4_trtllm,True,ERROR: FlashInfer produced no result row and no error message; see scratchpad/nemotron_ep4_v6/driver.log
MoE
H=2688 F=1856 E=32 top_k=6 activation=Relu2
module_name,M,N,K,backend,with_quant,runtime
model.layers.*.mixer.experts,1,,,bf16_cutlass,False,54.271
model.layers.*.mixer.experts,16,,,bf16_cutlass,False,175.552
model.layers.*.mixer.experts,1,,,fp8_cutlass,False,41.312
model.layers.*.mixer.experts,1,,,fp8_cutlass,True,44.032
model.layers.*.mixer.experts,16,,,fp8_cutlass,False,94.799
model.layers.*.mixer.experts,16,,,fp8_cutlass,True,97.567
model.layers.*.mixer.experts,1,,,fp8_trtllm,False,26.480
model.layers.*.mixer.experts,1,,,fp8_trtllm,True,29.200
model.layers.*.mixer.experts,16,,,fp8_trtllm,False,97.680
model.layers.*.mixer.experts,16,,,fp8_trtllm,True,100.448
model.layers.*.mixer.experts,1,,,nvfp4_cutlass,False,40.944
model.layers.*.mixer.experts,1,,,nvfp4_cutlass,True,40.623
model.layers.*.mixer.experts,16,,,nvfp4_cutlass,False,70.432
model.layers.*.mixer.experts,16,,,nvfp4_cutlass,True,71.072
model.layers.*.mixer.experts,1,,,nvfp4_trtllm,False,22.784
model.layers.*.mixer.experts,1,,,nvfp4_trtllm,True,25.119
model.layers.*.mixer.experts,16,,,nvfp4_trtllm,False,57.407
model.layers.*.mixer.experts,16,,,nvfp4_trtllm,True,59.999
```
### Testing
run locally
### 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?: ✅ / ❌ / N/A <!--- If ❌, explain
why. -->
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: ✅ / ❌ / N/A
<!--- Mandatory -->
- Did you write any new necessary tests?: ✅ / ❌ / N/A <!--- Mandatory
for new features or examples. -->
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
✅ / ❌ / N/A <!--- Only for new features, API changes, critical bug fixes
or backward incompatible changes. -->
- Did you get Claude approval on this PR?: ✅ / ❌ / N/A <!--- Run
`/claude review`. NVIDIA org members can self-trigger for complex
changes; orthogonal to CodeRabbit. -->
### Additional Information
<!-- E.g. related issue. -->
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit
- **New Features**
- Added a guided workflow for planning and running per-rank GEMM and
fused-MoE kernel benchmarks.
- Added automatic model inspection to derive supported kernel shapes,
tensor-parallel and expert-parallel settings, routing details, and
required benchmark parameters.
- Added FlashInfer built-in benchmark execution with BF16, FP8, NVFP4,
and activation-quantization timing.
- Added structured previews, logs, CSV reports, error details, and
dry-run validation.
- **Bug Fixes**
- Improved handling of unsupported layouts, padding, missing results,
invalid settings, and benchmark failures.
- **Tests**
- Added extensive coverage for model inspection, argument validation,
benchmark generation, quantization, routing, reporting, and error
handling.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---------
Signed-off-by: Shiyang Chen <shiychen@nvidia.com>
### What does this PR do? **Type of change:** Bug fix (CI) `mcp==2.0.0` was released and removed the `mcp.server.fastmcp` module (the 2.0 SDK replaces it with `mcp.server.mcpserver`). `tools/mcp/pyproject.toml` declared an unpinned `mcp>=1.0`, so CI now resolves `mcp==2.0.0`, and `tools/mcp/modelopt_mcp/server.py`'s `from mcp.server.fastmcp import FastMCP` fails at import: ``` ModuleNotFoundError: No module named 'mcp.server.fastmcp' ERROR collecting tools/mcp/tests/test_bridge.py ``` This breaks the `mcp` unit job — and thus the `unit-pr-required-check` gate — on **every** PR whose diff touches `pyproject.toml`, `noxfile.py`, or `.github/workflows/unit_tests.yml` (the changed-files paths that trigger the `mcp` job). This PR pins `mcp>=1.0,<2`, keeping the 1.x line that still ships `mcp.server.fastmcp`. Migrating the server to the mcp 2.0 API (`mcp.server.mcpserver`) is a larger change tracked separately. ### Usage ```python # N/A - dependency pin only ``` ### Testing - `uv pip install -e tools/mcp` now resolves an `mcp` 1.x wheel, so `from mcp.server.fastmcp import FastMCP` imports and `tools/mcp/tests/test_bridge.py` collects again. ### 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 — tightens an existing dependency's upper bound. - Did you write any new necessary tests?: N/A — restores collection of the existing `tools/mcp` tests. - Did you update Changelog?: N/A — CI/build fix, nothing user-facing in the wheel. - Did you get Claude approval on this PR?: ❌ ### Additional Information Unblocks the `unit-pr-required-check` gate for in-flight PRs (surfaced on #2000). Follow-up: migrate `modelopt_mcp/server.py` to the mcp 2.0 API and relax the pin. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Bug Fixes** * Prevented compatibility issues by limiting the MCP package to supported version 1.x releases. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
### What does this PR do? Type of change: Bug fix (CI) The `deploy-preview` job in the `Docs` workflow fails on **every pull request opened from a fork**, which blocks merging for all external contributors. **Root cause.** `deploy-preview` runs `rossjrw/pr-preview-action@v1`, which pushes the built HTML to the `gh-pages` branch. The workflow declares `permissions: contents: write`, but for a `pull_request` event originating from a forked repository GitHub caps the `GITHUB_TOKEN` at **read-only** — the `permissions:` block cannot elevate above that cap. The push is therefore rejected: ``` remote: Permission to NVIDIA/Model-Optimizer.git denied to github-actions[bot]. fatal: unable to access 'https://github.com/NVIDIA/Model-Optimizer.git/': The requested URL returned error: 403 ``` The job's `if:` condition gated on `github.event_name`, `github.event.action` and the `changes` path filter, but never on whether the PR came from a fork — so it always ran and always failed. Because the `changes` filter matches `docs/**`, `modelopt/**` and `.github/workflows/pages.yml`, essentially any substantive fork PR trips this. **Fix.** Restrict `deploy-preview` to PRs whose head branch lives in this repository: ```yaml github.event.pull_request.head.repo.full_name == github.repository ``` A skipped job is not a failed job, so fork PRs are no longer blocked by it. ### Usage N/A — CI-only change. ### Testing Behaviour by scenario: | Scenario | Before | After | | --- | --- | --- | | PR from a branch in this repo | preview deployed | preview deployed (**unchanged**) | | PR from a fork | ❌ fails with 403 | ⏭️ skipped | | Fork deleted (`head.repo` is `null`) | ❌ fails | ⏭️ skipped | - Confirmed against workflow history: recent `Docs` runs on in-repo branches (`main`, `chenjiel/nvfp4-act-headroom`, `mxin/qad-skill`, `haoguo/dspark-ptq-script`) all succeed, while fork-branch runs fail with the 403 above. - `build-docs` was already passing on the affected PRs — only the deploy step failed, so documentation builds are unaffected either way. - YAML parses; `pre-commit run --files .github/workflows/pages.yml` passes. (`yamlfmt` excludes `^.github/workflows/`, so this file is not auto-formatted.) - This PR edits `.github/workflows/pages.yml`, which is itself in the `changes` filter, so it exercises `deploy-preview` on the in-repo path — the preview deploy on this PR passing is a self-check that the unchanged path still works. The `closed` cleanup path is gated by the same condition. That is intentional: a fork PR never deployed a preview directory, so there is nothing to remove. ### 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?: N/A — workflow-condition change; not unit-testable - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A — CI infrastructure, not user-facing - Did you get Claude approval on this PR?: ❌ — not yet run ### Additional Information Currently blocking #1975 (`add Qwen3-VL support for DFlash training`), which is approved with every other check green and sits at `mergeStateStatus: BLOCKED` solely because of this job. Note that a PR only picks up this fix once its branch contains it, since workflows run from the PR branch's own definitions. A follow-up option, if doc previews for external contributors are wanted: build in the `pull_request` workflow and deploy from a separate `workflow_run`-triggered workflow, which executes in the base-repo context and does get a write token. Deliberately not using `pull_request_target` here — that would run unreviewed PR code with write permissions. Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
### What does this PR do? Type of change: Bug fix - Synchronizes existing ONNX tensor declarations when folding `Constant -> Cast` patterns, so the published `value_info` dtype matches the converted Constant payload. - Computes `GatherND` output shapes directly when autocast is running in custom-op mode, avoiding a generic input-0 shape copy for this shape-changing standard ONNX operator. - Restricts the previous input-0 shape fallback to actual custom-op outputs. - Adds regression tests for folded Constant dtype metadata and custom-op-mode `GatherND` shape propagation. ### Usage ```python # No user-facing API change. Existing Autocast usage remains: $ python -m modelopt.onnx.autocast --onnx=model.onnx ``` ### Testing - Reproduced the metadata issue with ModelOpt `0.44.0` and with main ToT `cba8a5c62a1a54fe89fb69bfd484ea0a653c633a` before the fix. - Verified the standalone reproduction no longer reports stale Constant dtype metadata or input-derived `GatherND` output shape after the fix. - Ran `pytest tests/unit/onnx/autocast/test_precisionconverter.py::test_folded_constant_cast_updates_value_info_type tests/unit/onnx/autocast/test_precisionconverter.py::test_custom_op_mode_uses_schema_shape_for_standard_gathernd -q` - Ran `pytest tests/unit/onnx/autocast/test_precisionconverter.py -q` ### 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)?: ❌ <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Bug Fixes** * Precision conversion now better preserves/restores full model input/output metadata when enabled, including after custom-op type/shape propagation. * Improved shape inference for standard ops during custom-op mode (including `GatherND`, `Gather`, `Unsqueeze`, and `Shape`), with correct scalar (rank-0) handling and stricter input-shape propagation. * When folding redundant `Cast` after `Constant`, element-type metadata is now updated consistently across the graph and nested subgraphs. * **Tests** * Added regression and custom-op mode tests for `Constant -> Cast -> Identity` folding and for `GatherND`/rank-change shape propagation. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Gwena Cunha <4861122+gcunhase@users.noreply.github.com> Signed-off-by: Gwenaelle Cunha Sergio <gcunhasergio@nvidia.com>
### What does this PR do? Type of change: documentation The container setup section uninstalls `nvidia-lm-eval` and the note above it states that it is "replaced with `lm-eval` from the repo", but none of the install commands actually install `lm-eval`. Following the README therefore leaves the environment without `lm-eval`, which the [Evaluation](https://github.com/NVIDIA/Model-Optimizer/blob/main/examples/puzzletron/README.md#evaluation) step requires. `lm-eval` ships in `examples/llm_eval/requirements.txt`, so this adds that one install line to the setup block. ### Usage N/A ### Testing Traced the dependency graph to confirm `lm-eval` is reachable from no install path the README has currently: - `puzzletron` extra: `fire`, `hydra-core`, `immutabledict`, `lru-dict`, `pandas`, `typeguard` - `dev-test` extra: pytest/coverage tooling, `timm`, `torchprofile`, `torchvision`, `torch-geometric` - `examples/puzzletron/requirements.txt`: `math-verify`, `ray`, `transformers<5.0` `lm_eval[api,ifeval]>=0.4.10` appears only in `examples/llm_eval/requirements.txt`. `pre-commit run --files examples/puzzletron/README.md` passes (markdownlint included). ### 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?: N/A - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A <!-- docs-only --> - Did you get Claude approval on this PR?: N/A <!-- not an NVIDIA org member, cannot self-trigger --> ### Additional Information Fixes #1786 <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Documentation** * Updated Puzzletron container setup instructions to include installing the LLM evaluation dependencies. * Clarified the setup sequence before running smoke tests. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: Rishi Khare <rishiskhare@gmail.com>
…) subgraphs during FP16/BF16 conversion (#1628) ### What does this PR do? Type of change: Bug fix Fixes [6058841] — `python -m modelopt.onnx.quantization --high_precision_dtype fp16` crashed with *"Inconsistent type on If node"* on models containing control-flow `If`/`Loop`/`Scan` subgraphs. **Root cause.** The FP16/BF16 `PrecisionConverter` blindly converted *every* subgraph initializer to the parent control-flow node's precision, without the activation-bracketing casts it uses in the main graph. This left inconsistent tensor types that crash ONNX shape inference / TensorRT strongly-typed parsing: - A `Gemm` inside an `If` branch reading an outer-scope activation (fp32) ended up with fp16 weights → `B has inconsistent type tensor(float16)`. - `Resize` `scales` (which must stay fp32 per the ONNX spec) was converted to fp16 → `ParseData type mismatch ... Expected:float Actual:float16`. **Fix.** - A subgraph node is converted to low precision **only when all of its float inputs are subgraph initializers** eligible for low precision. Any node consuming a float activation / outer-scope tensor, or an input that must stay high precision (e.g. `Resize` `scales`), stays high precision — so each node's inputs share a single precision. - Float **outer-scope captures** and **low→high precision boundaries** inside subgraphs are reconciled with `Cast` nodes; a captured tensor's preserved subgraph `value_info` is synced to its real main-graph precision; control-flow node **outputs** are treated as high precision (a low-precision `If` branch may still convert eligible nodes *inside* its body, but the parent control-flow node's outputs stay high precision). - `Constant`→`Cast` folding in `remove_redundant_casts` now refreshes the constant's `value_info` so a same-type-constrained consumer (e.g. `Greater`) isn't left with a stale, conflicting type. (Pre-existing main-graph bug surfaced once the `If` models completed conversion.) **Known limitation.** A low-precision `Loop`/`Scan` whose body carries a **float loop-carried or scan-input state variable** is not yet reconciled here — its body's formal inputs are treated as high precision while the main graph casts the parent's inputs to low precision. This PR fully covers `If` branches (both precisions) and high-precision `Loop`/`Scan` bodies (outer-scope captures reconciled); the low-precision `Loop`/`Scan` body case is tracked as a follow-up. ### Usage ```bash python -m modelopt.onnx.quantization \ --quantize_mode int8 --high_precision_dtype fp16 \ --onnx_path model.onnx \ --output_path model_strongType_int8+fp16.onnx ``` ### Testing Validated on both reported models: | graph pattern | convert | strict `infer_shapes(check_type=True)` | ORT load | numerics vs FP32 | |---|---|---|---|---| | `If` branch with `Gemm` reading an outer-scope input | ✅ | ✅ | ✅ | bit-exact (Gemms kept fp32) | | `If` branch with `Resize` (`scales` initializer) | ✅ | ✅ | ✅ | max abs err 9e-5 (fp16 tol) | Also verified with `keep_io_types=False`. Added 5 regression tests in `tests/unit/onnx/autocast/test_precisionconverter.py` — `If` Gemm-with-outer-scope-input, `If` Resize-scales-stay-FP32, chained-`If` capture, high-precision `Loop`-body capture, and `Constant`→`Cast` fold `value_info` refresh — that fail without the fix and pass with it. Full `tests/unit/onnx/autocast/` suite passes (214). `pre-commit` clean. ### 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)?: ✅ - Did you get Claude approval on this PR?: ❌ (draft) ### Additional Information Fixes bug [6058841]. Draft pending final review. 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **Bug Fixes** - Fixed FP16/BF16 conversion for ONNX models with control-flow subgraphs (for example `If`/`Loop`), especially when `--high_precision_dtype fp16` is used. - Preserved FP32 precision for branch weights that depend on outer-scope FP32 activations. - Added/adjusted casts for precision reconciliation across subgraph boundaries and updated value type metadata after `Constant`→`Cast` folding. - **Tests** - Added regression tests covering control-flow nesting/chaining, `Resize` inputs that must remain FP32, initializer precision selection, and strict shape/type inference. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
### What does this PR do?
Type of change: new feature
Adds online DFlash training support for Qwen3-VL–style vision-language
models.
Changes include:
- Load VLMs through the Transformers 5 `AutoModelForImageTextToText`
API, while retaining compatibility with the legacy VLM auto-model API.
- Run the base model through its top-level multimodal forward when
image/video inputs are present, ensuring vision embeddings are injected
before collecting DFlash target hidden states.
- Extend `VisionLanguageDataCollator` to:
- propagate `answer_only_loss`, chat-template, and DFlash
label-alignment settings;
- apply `VLM_MIN_PIXELS` / `VLM_MAX_PIXELS` processor limits;
- derive assistant-only masks from ChatML/Llama chat boundaries when
processor generation masks are unavailable;
- enforce the fixed `training_seq_len` required by DFlash block
training.
- Preserve the existing text-only DFlash path.
### Usage
```bash
python -m torch.distributed.run \
--nproc_per_node 4 \
examples/speculative_decoding/main.py \
--config modelopt_recipes/general/speculative_decoding/dflash.yaml \
model.model_name_or_path=/path/to/qwen3-vl-model \
model.trust_remote_code=true \
data.data_path=/path/to/train.jsonl \
data.vlm_processor=/path/to/qwen3-vl-model \
data.vlm_img_dir=/path/to/image/root \
training.training_seq_len=4096 \
training.answer_only_loss=true \
dflash.dflash_block_size=8 \
dflash.dflash_mask_token_id=151669
### Testing
- git diff --check
- Parsed all modified Python modules successfully.
- Ran iterative multi-node Slurm smoke tests with a Qwen3-VL-family model and mixed multimodal data:
- validated VLM model loading with Transformers 5;
- validated distributed initialization, DFlash conversion, and VLM collation paths;
- identified and addressed processor padding/truncation behavior required by fixed-size DFlash blocks.
This PR remains draft pending a completed end-to-end training smoke test and automated regression coverage.
### 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?: ❌ — automated Qwen3-VL/DFlash regression coverage still needs to be added before review.
- Did you update Changelog (https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ❌ — evaluate and add an entry before marking ready for review if this is considered user-facing speculative-decoding support.
- Did you get Claude approval on this PR?: N/A
### Additional Information
The PR intentionally excludes local Slurm launch scripts, logs, model paths, datasets, and environment-specific configuration.
<!-- This is an auto-generated comment: release notes by coderabbit.ai -->
## Summary by CodeRabbit
* **New Features**
* Expanded VLM data-collation controls, including `shift_labels` and more robust `answer_only_loss` masking.
* Improved Qwen3-VL speculative decoding for Transformers 5.3+ with correct video frame grouping and safer position-id handling.
* Improved DFlash RoPE export to reliably read `rope_theta` from newer config formats.
* **Bug Fixes**
* Hardened multimodal preprocessing and training loss masking to keep label/attention alignment consistent.
* Improved behavior when anchor sampling yields no valid blocks.
* More resilient VLM model loading when certain Transformers auto classes are unavailable.
* **Tests**
* Added coverage for RoPE export, Qwen3-VL position-id logic across Transformers versions, and VLM label-mode/collator options.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---------
Signed-off-by: Slawomir Kierat <skierat@nvidia.com>
Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
Co-authored-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
### What does this PR do?
**Type of change:** New feature (developer tool) + new tests
Adds a **standalone, cross-process resource monitor** —
`tools/resource_monitor.py` —
that samples a workload's GPU and CPU usage *from the outside* while it
runs, so
we can verify per-run memory budgets (e.g. the single-GPU layerwise PTQ
target of
≤80 GB GPU / ≤80 GB CPU for OMNIML-4947) without instrumenting
`hf_ptq.py` itself.
The monitor wraps any command, samples at a fixed interval, and on exit
writes a
CSV timeseries plus a peak/mean/min summary:
- **GPU** (per device, via NVML with an `nvidia-smi` fallback): used
memory,
utilization %, power draw (W), temperature (°C).
- **CPU** (via `psutil`): system total/used/free memory + utilization %,
and the
monitored process tree's RSS + CPU %.
It is opt-in from `examples/hf_ptq/scripts/huggingface_example.sh` via
`MODELOPT_MEM_MONITOR=1` (off by default → byte-for-byte identical
behavior). When
enabled it wraps the `hf_ptq.py` run and writes the trace/summary to a
**sibling**
`${SAVE_PATH}_mem_monitor/` directory, kept out of the exported
checkpoint that is
uploaded and consumed downstream.
**Files:**
- `tools/resource_monitor.py` — the sidecar (NVML + `nvidia-smi`
fallback; `psutil`).
- `tests/unit/tools/test_resource_monitor.py` — CPU-only unit tests (run
in the `unit` nox lane).
- `examples/hf_ptq/scripts/huggingface_example.sh` — opt-in
`MODELOPT_MEM_MONITOR=1` wrapper.
- `examples/hf_ptq/requirements.txt` — adds `psutil`.
- `pyproject.toml` — adds `psutil` to the `dev-test` extra
(deterministic import in the unit lane).
- `.github/workflows/unit_tests.yml` — adds `tools/resource_monitor.py`
to the unit-test path filters.
#### Why a new tool instead of extending
`modelopt/torch/utils/memory_monitor.py`?
The existing `GPUMemoryMonitor` is a fundamentally different tool and
cannot serve
this use case by extension:
| | `modelopt.torch.utils.memory_monitor.GPUMemoryMonitor` |
`tools/resource_monitor.py` (this PR) |
|---|---|---|
| Scope | **In-process** thread inside the workload | **Cross-process**
— wraps an external command |
| Survives workload OOM/SIGKILL | ❌ dies with the process | ✅ keeps
sampling, still writes the summary |
| Import cost | Pulls `torch` (~19 s) — lives in the workload |
Torch-free (`psutil`+`pynvml`, ~0.03 s) |
| Metrics | GPU device memory only | GPU mem/util/**power/temp** +
**CPU** mem/util + process-tree RSS |
| Output | In-memory / logs | CSV timeseries + peak/mean/min summary |
Merging the two would force `torch` into a standalone sidecar (defeating
the point)
or split the in-process monitor's threading model. A future refactor may
factor out
a **shared torch-free sampling core with two thin frontends**
(in-process + sidecar);
that is tracked as a follow-up rather than blocking this monitoring
harness, which
PR #2008 (single-GPU disk-offload PTQ) depends on.
### Usage
```bash
# Wrap mode (preferred): monitor exits with the workload's return code
python tools/resource_monitor.py --gpus 2,3 --out mem.csv --summary peak.txt -- \
python hf_ptq.py --pyt_ckpt_path=<model> --qformat=nvfp4 ...
# Opt-in from the HF PTQ example (off by default):
MODELOPT_MEM_MONITOR=1 CUDA_VISIBLE_DEVICES=2,3 CUDA_DEVICE_ORDER=PCI_BUS_ID \
bash examples/hf_ptq/scripts/huggingface_example.sh <args>
# -> writes ${SAVE_PATH}_mem_monitor/mem_trace.csv and mem_peak.txt
```
### Testing
- **Unit (CPU-only, in the `unit` nox lane):** `pytest
tests/unit/tools/test_resource_monitor.py`
— 11 tests covering `--gpus` parsing (CSV + space-separated, UUID/MIG
rejection),
the disabled/`nvidia-smi` sampling paths (including `[N/A]` → `None` and
the
smi-failure-yields-empty guard), CPU sampling, the accumulator, and
end-to-end
CSV/summary + exit-code propagation in wrap mode.
- **GPU-validated** on a B200 node (GPUs 2,3): confirmed the `gpu{i}_*`
memory /
utilization / power / temperature columns populate and the summary is
written.
### Before your PR is "*Ready for review*"
- Is this change backward compatible?: ✅ — new tool; the example wrapper
is off unless `MODELOPT_MEM_MONITOR=1`.
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: ✅ — `psutil`
added to `dev-test` + the hf_ptq example requirements;
`pynvml`/`nvidia-ml-py` is an optional runtime dep (graceful
`nvidia-smi` fallback).
- Did you write any new necessary tests?: ✅ —
`tests/unit/tools/test_resource_monitor.py`.
- Did you update Changelog?: N/A — repo-level `tools/` script, not
shipped in the wheel.
- Did you get Claude approval on this PR?: ❌ <!-- run /claude review -->
### Additional Information
Part of **OMNIML-4947** (single-GPU disk-offload PTQ). This is PR 1 of
the stack —
the monitoring harness that PR #2008 (disk-offload layerwise PTQ +
offload-aware
export) builds on.
---------
Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
### What does this PR do? Type of change: Bug fix Prevents recursively collected text-submodel reverse mappings from rewriting an already nested multimodal model namespace during unified Hugging Face export. Transformers reverses the Qwen3.5 text mapping into a broad `^model.` -> `model.language_model.` rename. ModelOpt previously applied that rule to every key in the full VLM state dict, moving `model.visual.*` under the language model and nesting `model.language_model.*` twice. This change drops the reverse rule only when its target child namespace is already registered. Standalone text models continue to use the conversion. ### Usage ```python # Existing export_hf_checkpoint usage is unchanged. ``` ### Testing - `python -m pytest -q tests/unit/torch/export` (`112 passed`, `1 skipped` because optional Diffusers is not installed) - Targeted pre-fix reproduction confirmed both malformed Qwen3.5 namespaces; both regression cases pass after the fix - Tiny `Qwen3_5MoeForConditionalGeneration` meta-device model-tree probe passed - `python -m pre_commit run --files modelopt/torch/export/quant_aware_conversion.py tests/unit/torch/export/test_quant_aware_conversion.py` ### 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 No API or dependency changes. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Bug Fixes** * Improved model conversion to prevent incorrect renaming of nested visual or sibling model components. * Fixed duplicate namespace prefixes in converted model weights. * Preserved correct reverse mapping for text-only model configurations. * **Tests** * Added coverage for nested multimodal and text-only model conversion scenarios. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
Signed-off-by: Grzegorz Karch <gkarch@nvidia.com> # Conflicts: # examples/llm_qad/configs/qwen3-30b-a3b-instruct-2507-moe_template.conf # examples/llm_qad/configs/qwen3-8b_template.conf # examples/pruning/minitron_vs_puzzletron/README.md # examples/pruning/minitron_vs_puzzletron/advanced_compression_experiments.md # examples/puzzletron/configs/gptoss-20b_remove_experts_memory/gptoss-20b.yaml # examples/puzzletron/configs/gptoss-20b_remove_experts_memory/gptoss-20b_remove_experts_memory.yaml # examples/puzzletron/configs/gptoss-20b_remove_experts_memory/validate_model_defaults.yaml # examples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/Llama-3_1-8B.yaml # examples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/llama-3_1-8B_pruneffn_memory.yaml # examples/puzzletron/configs/llama-3_1-8B_pruneffn_runtime/Llama-3_1-8B.yaml # examples/puzzletron/configs/llama-3_1-8B_pruneffn_runtime/llama-3_1-8B_pruneffn_runtime.yaml # examples/puzzletron/configs/llama-3_2-3B_pruneffn_memory/Llama-3_2-3B.yaml # examples/puzzletron/configs/llama-3_2-3B_pruneffn_memory/llama-3_2-3B_pruneffn_memory.yaml # examples/puzzletron/configs/llama-3_2-3B_pruneffn_memory/validate_model_defaults.yaml # examples/puzzletron/configs/mistral-small-24b-instruct-2501_pruneffn_memory/Mistral-Small-24B.yaml # examples/puzzletron/configs/mistral-small-24b-instruct-2501_pruneffn_memory/mistral-small-24b-instruct-2501_pruneffn_memory.yaml # examples/puzzletron/configs/mistral-small-24b-instruct-2501_pruneffn_memory/validate_model_defaults.yaml # examples/puzzletron/configs/nemotron-nano-12b-v2/nemotron_nano_12b_v2.yaml # examples/puzzletron/configs/nemotron-nano-12b-v2/nemotron_nano_12b_v2_pruneffn_memory.yaml # examples/puzzletron/configs/nemotron-nano-12b-v2/validate_model_defaults.yaml # examples/puzzletron/configs/nemotron-nano-30b-A3b-v3/nemotron_nano_v3.yaml # examples/puzzletron/configs/qwen2_5_7b_instruct_pruneffn_memory/qwen2_5_7b_instruct.yaml # examples/puzzletron/configs/qwen2_5_7b_instruct_pruneffn_memory/qwen2_5_7b_instruct_pruneffn_memory.yaml # examples/puzzletron/configs/qwen2_5_7b_instruct_pruneffn_memory/validate_model_defaults.yaml # examples/puzzletron/configs/qwen3-8b_pruneffn_memory/qwen3_8b.yaml # examples/puzzletron/configs/qwen3-8b_pruneffn_memory/qwen3_8b_pruneffn_memory.yaml # examples/puzzletron/configs/qwen3-8b_pruneffn_memory/validate_model_defaults.yaml # modelopt/torch/puzzletron/bypass_distillation/bypass_utils.py # modelopt/torch/puzzletron/bypass_distillation/data_classes.py # modelopt/torch/puzzletron/diagnostics/__init__.py # modelopt/torch/puzzletron/sewing_kit/utils.py # tests/_test_utils/examples/hf_ptq_example_utils.py # tests/unit/torch/puzzletron/test_bypass_dataloaders.py # tests/unit/torch/puzzletron/test_bypass_keys_to_learn.py # tests/unit/torch/puzzletron/test_bypass_losses.py # tests/unit/torch/puzzletron/test_bypass_lr_scheduler.py # tests/unit/torch/puzzletron/test_bypass_utils.py
Contributor
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
Comment |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Type of change: ?
Usage
# Add a code snippet demonstrating how to use thisTesting
Before your PR is "Ready for review"
Make sure you read and follow Contributor guidelines and your commits are signed (
git commit -s -S).Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded
trust_remote_code=True,torch.load(..., weights_only=False),pickle, etc.).CONTRIBUTING.md: ✅ / ❌ / N/AAdditional Information