Skip to content

Puzzletron v2 sync 20260730 - #2037

Draft
grzegorz-k-karch wants to merge 183 commits into
feature/puzzletron_v2from
gkarch/puzzletron_v2_sync_20260730
Draft

Puzzletron v2 sync 20260730#2037
grzegorz-k-karch wants to merge 183 commits into
feature/puzzletron_v2from
gkarch/puzzletron_v2_sync_20260730

Conversation

@grzegorz-k-karch

Copy link
Copy Markdown
Contributor

What does this PR do?

Type of change: ?

Usage

# Add a code snippet demonstrating how to use this

Testing

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.).

  • 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?: ✅ / ❌ / N/A
  • Did you get Claude approval on this PR?: ✅ / ❌ / N/A

Additional Information

kinjalpatel27 and others added 30 commits June 11, 2026 12:01
)

### 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>
gcunhase and others added 28 commits July 20, 2026 12:21
### 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
@copy-pr-bot

copy-pr-bot Bot commented Jul 30, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@coderabbitai

coderabbitai Bot commented Jul 30, 2026

Copy link
Copy Markdown
Contributor

Important

Review skipped

Draft detected.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: c93ef3d5-a982-4a2d-a085-e4f138d449d5

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

  • 🔍 Trigger review

Comment @coderabbitai help to get the list of available commands.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.