Commit Graph
225 Commits
Author SHA1 Message Date
jingyu-ml 1070d895dc Flux2-Dev Quantization (#947)
## What does this PR do?

**Type of change:** new example <!-- Use one of the following: Bug fix,
new feature, new example, new tests, documentation. -->

**Overview:** 

- Register Flux2Attention and Flux2ParallelSelfAttention in the
quantization plugin so bmm quantizers are patched (enables
--quantize-mha).
- Add Flux2-specific dummy input generation for HF checkpoint export.
- Guard check_conv_and_mha with hasattr for bmm quantizer attributes 

## Usage
<!-- You can potentially add a usage example below. -->

```bash
python quantize.py \
    --model flux2-dev \
    --model-dtype BFloat16 \
    --format fp4 --batch-size 2 --calib-size 1 \
    --n-steps 20 --quantized-torch-ckpt-save-path ./flux2-dev-fp4.pt --collect-method default \
    --hf-ckpt-dir ./flux2-dev-fp4
```

## Testing
<!-- Mention how have you tested your change if applicable. -->

## Before your PR is "*Ready for review*"
<!-- If you haven't finished some of the above items you can still open
`Draft` PR. -->

- **Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)**
and your commits are signed.
- **Is this change backward compatible?**: Yes<!--- If No, explain why.
-->
- **Did you write any new necessary tests?**: No
- **Did you add or update any necessary documentation?**: No
- **Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**:
No <!--- Only for new features, API changes, critical bug fixes or bw
breaking changes. -->

## Additional Information
<!-- E.g. related issue. -->


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Added Flux2-dev model support with Flux2-compatible dummy input
generation and default inference params (768×1024, guidance scale 4.0).

* **Refactor**
* Made attention quantization disabling more robust by iterating
available quantizers before disabling.

* **Infrastructure**
* Flux2 attention components are now optional and registered only when
present to avoid import issues.

* **Tests**
* Added Flux2 test helpers and coverage validating Flux2 dummy input
shapes.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
2026-03-14 11:42:02 -05:00
ZhiyuandClaude Opus 4.6 5b417377a8 Support Kimi-K2.5 PTQ (#820)
## What does this PR do?

**Type of change:** New model support <!-- Use one of the following: Bug
fix, new feature, new example, new tests, documentation. -->

**Overview:** Support Kimi-K2.5 PTQ.

## Usage
<!-- You can potentially add a usage example below. -->

```python
python3 hf_ptq.py   --pyt_ckpt_path moonshotai/Kimi-K2.5   --qformat nvfp4_mlp_only   --export_path ./kimi-k2.5-nvfp4   --trust_remote_code
```

## Testing
<!-- Mention how have you tested your change if applicable. -->

You may need `pip install transformers==4.57.1` and the model file here:
https://huggingface.co/nvidia/Kimi-K2.5-NVFP4/blob/main/modeling_kimi_k25.py

## Before your PR is "*Ready for review*"
<!-- If you haven't finished some of the above items you can still open
`Draft` PR. -->

- **Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)**
and your commits are signed.
- **Is this change backward compatible?**: Yes/No <!--- If No, explain
why. -->
- **Did you write any new necessary tests?**: Yes/No
- **Did you add or update any necessary documentation?**: Yes/No
- **Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**:
Yes/No <!--- Only for new features, API changes, critical bug fixes or
bw breaking changes. -->

## Additional Information
<!-- E.g. related issue. -->


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **Bug Fixes**
* Fixes model loading so mixed-precision (BF16) and expert weights are
correctly restored and stale placeholders removed.
* Adds error handling to avoid failures when optional decompression
components are missing.

* **New Features**
* Adds conditional patch/restore around pack‑quantized model loads,
final unpacking of weights after load, and on‑the‑fly decompression
during inference.
  * Improves logging for load and decompression events.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Zhiyu <zhiyuc@nvidia.com>
Signed-off-by: Zhiyu Cheng <zhiyuc@nvidia.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-03-14 06:56:24 +00:00
Wei-Ming Chen bc96f1ce39 [OMNIML-3277] Update kv cache behavior (#1012)
### What does this PR do?

Type of change: New feature <!-- Use one of the following: Bug fix, new
feature, new example, new tests, documentation. -->

<!-- Details about the change. -->
By default, FP8 KV cache quantization in hf_ptq.py now uses a constant
scale of 1.0 (amax=448.0) without a data-driven calibration pass. No KV
scales are written to the exported checkpoint — inference engines
(TRT-LLM, vLLM) use scale=1.0 when no scale is
present, so this is lossless. Pass --calibrate_kv_cache to opt into
data-driven per-tensor KV scale calibration (previous default behavior).
To support this cleanly in the quantization stack, a constant_amax field
is added to QuantizerAttributeConfig. Quantizers configured with
constant_amax skip calibration entirely (no forward pass needed), use
the fixed amax during fake-quant, and produce no
  _amax buffer in the state dict. 

### Usage

```python
  # Quantize with default constant KV scale (no calibration pass for KV)
  python hf_ptq.py --model ... --qformat fp8 --kv_cache_qformat fp8

  # Opt into data-driven KV calibration
  python hf_ptq.py --model ... --qformat fp8 --kv_cache_qformat fp8 --calibrate_kv_cache

  # Use constant_amax in a custom quant config
  quant_cfg = {
      "quant_cfg": {
          "*[kv]_bmm_quantizer": {"num_bits": (4, 3), "enable": True, "constant_amax": 448.0},
      },
      "algorithm": "max",
  }
  model = mtq.quantize(model, quant_cfg, forward_loop=calibrate_loop)
```

### Testing
<!-- Mention how have you tested your change if applicable. -->


- All 8 existing GPU HF export tests pass
(tests/gpu/torch/export/test_unified_hf_export_and_check_safetensors.py)
- Two new CPU unit tests added to TensorQuantizerTester:
test_constant_amax and test_constant_amax_skips_calibration

### 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?: ✅ (--calibrate_kv_cache flag
defaults to False; existing scripts without the flag now skip KV
calibration)
- 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?: ✅

### Additional Information
<!-- E.g. related issue. -->


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* New CLI flag --calibrate_kv_cache to optionally enable data-driven KV
cache calibration; default uses a fixed KV scale and omits KV scales
from exported checkpoints.
* Added constant_amax option to set fixed quantizer scales and skip
dynamic calibration for configured quantizers.

* **Bug Fixes**
* Removed forced flooring/clamp of KV cache scales; out-of-range
activations now emit a shorter warning.

* **Tests**
  * Added tests for constant_amax behavior and calibration interaction.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
2026-03-14 05:50:04 +00:00
jingyu-ml 812e8c60a2 Minor update on the LTX2 NVFP4 recipe (#1010)
### What does this PR do?

Type of change: minor code change <!-- Use one of the following: Bug
fix, new feature, new example, new tests, documentation. -->

1. Update the default calibration dataset for LTX_VIDEO_DEV and LTX2
from Gustavosta/Stable-Diffusion-Prompts to nkp37/OpenVid-1M, which
provides video-specific captions better suited for video model
calibration.
2. update the default recipe for ltx2: first 3 and last 3 layers stays
at higher precision.

### Usage

### Testing
<!-- Mention how have you tested your change if applicable. -->

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

### Additional Information
<!-- E.g. related issue. -->


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

* **Chores**
  * Updated default dataset configuration for LTX-Video and LTX2 models.
* Refined model filtering pattern for LTX-Video to support additional
model components.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
2026-03-13 22:42:48 -05:00
Rohan Joshi 2d7d1ec345 Skip softmax calibration with list of thresholds (#987)
Modify skip softmax calibration to use a list of thresholds instead of a
single threshold. Sparsity during inference is unchanged, but during
calibration we can use the list to gather statistics about many
thresholds in a single forward pass. Makes calibration 20x faster


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

* **New Features**
- Multi-threshold sparsity configuration for fine-grained control over
attention sparsity levels.

* **Improvements**
- Calibration efficiency: Single forward pass for collecting all
threshold data instead of iterating per threshold.
- Configuration format updated to support threshold lists for prefill
and decode phases.

* **Breaking Changes**
- Configuration API: `threshold` field renamed to `thresholds` and now
expects lists of values instead of scalars.
  - Sparsity statistics output updated to return per-threshold values.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Signed-off-by: Rohan Joshi <rohjoshi@nvidia.com>
2026-03-13 22:28:01 +00:00
Gwena Cunha 69c0d47946 [OMNIML-3252][ONNX] MOQ + Autotune moq integration docs (#1026)
### What does this PR do?

**Type of change**: documentation

**Overview**: This PR updates the documentation and does some folder
re-structuring and file re-naming related to
https://github.com/NVIDIA/Model-Optimizer/pull/951.

### Usage

Documentation

### Testing

Documentation

### 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`: ✅
- 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)?:
✅ (renamed `AutoQDQ` to `Autotune`)

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **Documentation**
  * Renamed AutoQDQ to Autotune across guides and changelog.
  * Updated Autotune guide descriptions and wording.
* Added a new section on optimizing Q/DQ node placement with Autotune,
including CLI usage and API links (appears twice in one README).
  * Applied minor grammar and capitalization corrections.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: gcunhase <4861122+gcunhase@users.noreply.github.com>
2026-03-12 10:47:07 -07:00
Wei-Ming Chen 34a9fc7924 Add support matrix for Nemotron-3 (#1023)
### What does this PR do?

Type of change: documentation <!-- Use one of the following: Bug fix,
new feature, new example, new tests, documentation. -->

<!-- Details about the change. -->
Add support matrix for Nemotron-3

### Usage

```python
# Add a code snippet demonstrating how to use this
```

### Testing
<!-- Mention how have you tested your change if applicable. -->

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

### Additional Information
<!-- E.g. related issue. -->


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Nemotron-3 model quantization support added; Mixture-of-Experts expert
support retained in auto-quantize scoring/grouping

* **Documentation**
* Hugging Face model support matrix updated to include Nemotron-3 and
adjust referenced models
* Quantization examples revised: nvfp4_mse with fp8 and effective
precision ~4.75; example configs clarified
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
Signed-off-by: Wei-Ming Chen <17592131+meenchen@users.noreply.github.com>
2026-03-11 18:12:39 +00:00
Keval Morabia cbab377983 Fix Mcore test utils for M-LM main (#1008)
### What does this PR do?

- M-LM main removed some functions we were using for running mcore
inference in unit tests - replace with alternative
- Future-proof `hybrid_override_pattern` -> `hybrid_layer_pattern`
rename for Mamba models

### Testing

Ran megatron tests with M-LM main branch and previous 0.16 release
version

### 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. -->
- 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)?:
N/A <!--- Only for new features, API changes, critical bug fixes or
backward incompatible changes. -->

### Additional Information
<!-- E.g. related issue. -->


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

* **Bug Fixes**
* Improved Mamba model pruning to dynamically select the appropriate
hybrid pattern configuration, ensuring correct handling across different
model variants.
* Simplified inference logic for pipeline-parallel models, improving
robustness and reliability.

* **Tests**
* Updated test infrastructure for Mamba and GPT model inference
validation with explicit hidden size parameters.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
2026-03-10 14:34:57 +05:30
yeyu-nvidia 5d0e012751 inplement mix hidden_states for eagle3; deprecate eagle1 (#946)
## What does this PR do?

new feature

**Overview:** 
Enable mix hidden_states in eagle3 training. Deprecate eagle1

## Usage
Add --mix_hidden_states True to launch_train.sh

```python
# Add a code snippet demonstrating how to use this
```

## Testing
<!-- Mention how have you tested your change if applicable. -->

## Before your PR is "*Ready for review*"
<!-- If you haven't finished some of the above items you can still open
`Draft` PR. -->

- **Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)**
and your commits are signed.
- **Is this change backward compatible?**: Yes/No <!--- If No, explain
why. -->
- **Did you write any new necessary tests?**: Yes/No
- **Did you add or update any necessary documentation?**: Yes/No
- **Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**:
Yes/No <!--- Only for new features, API changes, critical bug fixes or
bw breaking changes. -->

## Additional Information
<!-- E.g. related issue. -->


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Added --mix_hidden_states option to enable optional hidden-state
mixing during training.
* Added eagle_ttt_steps setting to control speculative multi-step
iterations.

* **Chores**
* Consolidated speculative decoding to EAGLE3 only; legacy Medusa/EAGLE1
paths removed.
* Unified configuration handling so models and plugins accept a single
config object.

* **Tests**
* Updated and expanded tests for hidden-state mixing and EAGLE3-only
scenarios.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Ye Yu <yeyu@nvidia.com>
2026-03-09 11:14:10 -07:00
d3748c2c60 Allow basename of dataset paths to match registered names (#997)
### What does this PR do?  Allow local dataset paths to match registered dataset configs

Type of change: Bug fix

<!-- Details about the change. -->

### Usage

```python
# Add a code snippet demonstrating how to use this
```

### Testing
<!-- Mention how have you tested your change if applicable. -->

### 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`, using
`torch.load(..., weights_only=True)`, avoiding `pickle`, etc.).

- Is this change backward compatible?: ✅ / ❌ / N/A <!--- If ❌, explain
why. -->
- If you copied code from any other source, did you follow IP policy in
[CONTRIBUTING.md](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md#-copying-code-from-other-sources)?:
✅ / ❌ / 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. -->

### Additional Information
<!-- E.g. related issue. -->


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Added a small sample dataset entry (minipile_100_samples) and support
for loading datasets from local filesystem paths with automatic
detection and config override.

* **Chores**
* Improved local-path resolution and substring-based matching against
registered dataset keys for consistent behavior.

* **Tests**
* Added a unit test to verify loading samples from a local dataset
snapshot.

* **Documentation**
* Updated docs to describe local-path support, matching behavior, and
updated function docstring.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Chenhan Yu <chenhany@nvidia.com>
Signed-off-by: Asha Anoosheh <aanoosheh@nvidia.com>
Co-authored-by: Chenhan Yu <chenhany@nvidia.com>
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-03-09 15:24:44 +00:00
Keval Morabia 1ccd945a51 Remove unused diffusers/cache_diffusion/pipeline and cuda-python dependency (#996)
`cuda-python` has mixed license and needs EStaff approval for usage. And
till 0.42, it was only used in
`examples/diffusers/cache_diffusion/pipeline` which has not been updated
in 9 months and not used anymore hence removing.

Also cherry-picked to `release/0.42.0` branch:
https://github.com/NVIDIA/Model-Optimizer/pull/984

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **Chores**
* Removed TensorRT/ONNX deployment and inference tooling, related model
export/configuration, and runtime helpers from the cache-optimized
diffusion examples; removed the cuda-python example dependency.
* **Tests**
* Removed the example benchmarking script and its associated benchmark
test.
* **Documentation**
* Strengthened dependency-review, security, and PR guidance; updated PR
template and contributing documentation.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
2026-03-07 00:37:19 +05:30
jingyu-ml 37d3f10cbd To support LTX2 ComfyUI format (#972)
### What does this PR do?

Type of change: Bug fix <!-- Use one of the following: Bug fix, new
feature, new example, new tests, documentation. -->

<!-- Details about the change. -->

- Added a flag merged_base_safetensor_path to the example code so that
user can export the ComfyUI style ckpt.

### Usage

```bash
python quantize.py \
    --model ltx-2 --format fp4 --batch-size 1 --calib-size 32 --n-steps 40 \
    --extra-param checkpoint_path=./ltx-2-19b-dev-fp8.safetensors \
    --extra-param distilled_lora_path=./ltx-2-19b-distilled-lora-384.safetensors \
    --extra-param spatial_upsampler_path=./ltx-2-spatial-upscaler-x2-1.0.safetensors \
    --extra-param gemma_root=./gemma-3-12b-it-qat-q4_0-unquantized \
    --extra-param fp8transformer=true \
    --quantized-torch-ckpt-save-path ./ltx-2-transformer.pt \
    --hf-ckpt-dir ./LTX2-NVFP4/ \
    --extra-param merged_base_safetensor_path=./ltx-2-19b-dev-fp8.safetensors
```

### Testing
<!-- Mention how have you tested your change if applicable. -->

### 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`, using
`torch.load(..., weights_only=True)`, avoiding `pickle`, etc.).

- Is this change backward compatible?: ✅ / ❌ / N/A <!--- If ❌, explain
why. -->
- If you copied code from any other source, did you follow IP policy in
[CONTRIBUTING.md](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md#-copying-code-from-other-sources)?:
✅ / ❌ / 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. -->

### Additional Information
<!-- E.g. related issue. -->


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

* **Documentation**
* Added new command-line parameters documentation for LTX-2 FP4
quantization examples (--hf-ckpt-dir and merged_base_safetensor_path
configuration options)

* **Improvements**
* Enhanced quantization pipeline to support conditional export behavior
based on model type
* Expanded LTX-Video model filtering patterns for more comprehensive
block detection

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
2026-03-06 05:06:36 +00:00
ynankani-nv 296a865c70 sample QAD example script (#933)
## What does this PR do?
sample QAD example script

**Type of change:** ?  new example
Example script for QAD on diffusion model like ltx-2

**Overview:** ?
1) Model loading 
2) NvFP4 fake quant PTQ using mtq.quantize
3) Distillation class wrapping using mtd.convert 
4) Using ltx-2 trainer code for training 
5) Checkpoint save in bf16 . post process bf16 model using Comfy-kitchen
to produce real quantized model for ComyUI inference

## Usage
<!-- You can potentially add a usage example below. -->

```python
accelerate launch --config_file fsdp_custom.yaml sample_example_qad_diffusers.py train  --config ltx2_qad.yaml 
```

## Testing
1) Tested improvement in Vbench score for PTQ and QAD checkpoint.

## Before your PR is "*Ready for review*"
<!-- If you haven't finished some of the above items you can still open
`Draft` PR. -->

- **Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)**
and your commits are signed.
- **Is this change backward compatible?**: NA <!--- If No, explain why.
-->
- **Did you write any new necessary tests?**: NA
- **Did you add or update any necessary documentation?**: NA
- **Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**:
NA <!--- Only for new features, API changes, critical bug fixes or bw
breaking changes. -->

## Additional Information
<!-- E.g. related issue. -->


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **Documentation**
* Added comprehensive README documenting the new Windows QAD example:
setup, usage, project layout, and workflow.

* **New Features**
* Added a complete Quantization-Aware Distillation training example with
distributed training config, PTQ calibration, teacher-student
distillation, CLI for training/inference, and an inference-checkpoint
creation utility.
  * Added requirements file listing needed Python packages and tooling.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: ynankani <ynankani@nvidia.com>
2026-03-06 04:15:54 +00:00
Chenjie Luo 42482b1b0f Add nvfp4_omlp_only config and simplify the config.py (#973)
### What does this PR do?

Type of change: ? new feature

1) Add vfp4_omlp_only config == nvfp4_mlp_only + o_proj quant
2) Add block sparse MOE to mlp only config
3) Simplfiy config.py
4) Update readme in llm_ptq mention these two configs for better
accuracy.

### Usage

huggingface_script.sh ... --quant nvfp4_omlp_only


### 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`, using
`torch.load(..., weights_only=True)`, avoiding `pickle`, etc.).

- Is this change backward compatible?: ✅ / ❌ / N/A <!--- If ❌, explain
why. -->
- If you copied code from any other source, did you follow IP policy in
[CONTRIBUTING.md](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md#-copying-code-from-other-sources)?:
✅ / ❌ / 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. -->

### Additional Information
<!-- E.g. related issue. -->


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Added nvfp4_omlp_only quantization format for NVFP4, enabling
selective quantization of MLP and output projection layers while
preserving attention QKV projection accuracy.

* **Changed**
* pass_through_bwd now defaults to True; set to False if using STE with
zeroed outlier gradients for better QAT accuracy.

* **Documentation**
* Updated post-training quantization guidance with NVFP4-specific
configuration recommendations and usage examples.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
2026-03-06 01:03:25 +00:00
skierat dd16a96fda example demonstrating how to train CosmosReason2 Eagle3 (#965)
## What does this PR do?

**Type of change:** New example

**Overview:**
Adds
examples/speculative_decoding/guides/train_eagle_head_cosmos_reason2.ipynb,
a step-by-step Jupyter notebook that walks through the full EAGLE3
draft-head training workflow for nvidia/Cosmos-Reason2-8B. The notebook
covers:

1. Installing dependencies
2. Authenticating with Hugging Face
3. Preparing training data from the Nemotron-Post-Training-Dataset-v2
(chat split) using a curated row-selection mapping
(guides/nemotron_mapping.csv)
4. Inspecting the bundled EAGLE3 config (guides/CR2_eagle_config.json)
tuned for Cosmos-Reason2 (YaRN RoPE, FlexAttention, reduced draft
vocabulary)
5. (Optional) Calibrating the draft vocabulary to 32k tokens for faster
training and inference
6. Launching training via launch_train.sh with FSDP2 multi-GPU support
7. Exporting the checkpoint to HF format and serving with vLLM

Also includes guides/nemotron_mapping.csv and
guides/CR2_eagle_config.json as companion files.

## Usage
Open and run
examples/speculative_decoding/guides/train_eagle_head_cosmos_reason2.ipynb
cell by cell. After training, serve the exported checkpoint with:
```
vllm serve nvidia/Cosmos-Reason2-8B \    --host 0.0.0.0 \    --port 8000 \    --speculative-model export/cosmos-reason2-8b-eagle3 \    --num-speculative-tokens 3 \    --dtype bfloat16
```

## Testing
Tested end-to-end on a 4xB100 GPUs. The exported checkpoint was
validated with specdec_bench

## Before your PR is "*Ready for review*"
<!-- If you haven't finished some of the above items you can still open
`Draft` PR. -->

- **Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)**
and your commits are signed.
- **Is this change backward compatible?**: Yes
- **Did you write any new necessary tests?**: No
- **Did you add or update any necessary documentation?**: Yes (the
notebook is self-documenting)
- **Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**:
No

## Additional Information
Cosmos-Reason2-8B requires at least one 80 GB GPU (H100/A100)

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Added a speculative-decoding configuration for draft model and
rotary/attention behavior.
* Added an end-to-end training notebook demonstrating training, export,
and deployment of a speculative-decoding draft head on Cosmos-Reason2.
* Added a data-preparation tool to download, normalize, and convert
Nemotron chat conversations into a standardized conversation format.

* **Documentation**
* Notebook documents environment setup, data prep, training/validation
cadence, export, and deployment steps.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Slawek Kierat <skierat@nvidia.com>
2026-03-05 23:53:49 +00:00
willg-nv e8f9687a60 Add Automated QDQ placement example - Part 4.1 (#841)
## What does this PR do?

**Type of change:** ?
This change implement a simple example to illustrate the usage of
Automated QDQ placement tool.

**Overview:** ?

## Usage
<!-- You can potentially add a usage example below. -->

```python
python3 -m modelopt.onnx.quantization.autotune \
    --model resnet50.bs128.onnx \
    --output ./resnet50_autotuned \
    --qdq-baseline resnet50_quantized.onnx \
    --schemes-per-region 50
```

## Testing
<!-- Mention how have you tested your change if applicable. -->

## Before your PR is "*Ready for review*"
<!-- If you haven't finished some of the above items you can still open
`Draft` PR. -->

- **Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)**
and your commits are signed.
- **Is this change backward compatible?**: Yes
- **Did you write any new necessary tests?**: No
- **Did you add or update any necessary documentation?**: Yes
- **Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**:
No

## Additional Information
<!-- E.g. related issue. -->


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Log output now includes timestamps for improved debugging and
traceability.

* **Documentation**
* Added comprehensive guide for QDQ placement optimization with
TensorRT, covering setup, usage examples, and deployment instructions.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Will Guo <willg@nvidia.com>
2026-03-04 18:28:24 +00:00
h-guo18 a34d613d3c Feat: Speculatice Decoding export with quantization support (#913)
## What does this PR do?

**Type of change:** ? <!-- Use one of the following: Bug fix, new
feature, new example, new tests, documentation. -->

**Overview:** 


Main changes:  
- Refactored speculative decoding export logics into `class
EagleExporter` to improve cohesion;

 
- Separated speculative decoding export entrance with quantization
export (`export_hf_checkpoint()`) due to their fundamental differences:
- Quantization export base model's state_dict and config, while
speculative decoding only export drafter's.
- Most of the model-specific logics of quantization export (e.g.
diffusers, vlms) are not needed for speculative decoding export.
- Quantization export produce different format than speculative decoding
checkpoint. (The former produce tokenizer config, generation config,
e.t.c, while the later does not need. )

## Usage
<!-- You can potentially add a usage example below. -->

To export an regular bf16 eagle checkpoint without quantization, the
commands are the same:
```python
python scripts/export_hf_checkpoint.py --model_path <x> --export_path <x>
```

To run PTQ on online-trained eagle checkpoint and export it:
```python
python hf_ptq.py --pyt_ckpt_path <x> --qformat fp8 --export_path <x>
```

The above two commands will produce drafter ckpt for deployment, in the
same foramt.

## Testing
<!-- Mention how have you tested your change if applicable. -->

Tested setting:
- Base model: llama3.1-8b
- Algorithms: eagle
- Export path tested: 
- (Unquantized online ckpt) `python scripts/export_hf_checkpoint.py
--model_path <x> --export_path <x>`
- (PTQ) export `python hf_ptq.py --pyt_ckpt_path <x> --qformat fp8
--export_path <x>`
- Tested deployment on vllm. Got normal AR. 

## Before your PR is "*Ready for review*"
<!-- If you haven't finished some of the above items you can still open
`Draft` PR. -->

- **Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)**
and your commits are signed.
- **Is this change backward compatible?**: Yes/No <!--- If No, explain
why. -->
- **Did you write any new necessary tests?**: Yes/No
- **Did you add or update any necessary documentation?**: Yes/No
- **Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**:
Yes/No <!--- Only for new features, API changes, critical bug fixes or
bw breaking changes. -->

## Additional Information
<!-- E.g. related issue. -->


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

## Release Notes

* **New Features**
  * Added export functionality for speculative decoding-optimized models
* Support for multiple speculative decoding architectures with
pre-configured deployment templates
* Enhanced model export detection and automatic routing for optimized
models

* **Tests**
  * Updated export validation tests for speculative decoding models

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
2026-03-04 01:46:02 +00:00
Chenjie Luo a4fde491cc Update MOE block detection logic and enable in huggingface_script.sh (#962)
### What does this PR do?

Type of change: Bug fix

Add moe expert calib ratio in huggingface_script.sh
Also fix minimax2.5 MOE detection which does not follow other HF MOE
layer convention

### Usage

scripts/huggingface_example.sh --model <MiniMax-M2.5> --quant nvfp4
--moe_calib_experts_ratio 1.0 --trust_remote_code

### 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`, using
`torch.load(..., weights_only=True)`, avoiding `pickle`, etc.).

- Is this change backward compatible?: ✅ / ❌ / N/A <!--- If ❌, explain
why. -->
- If you copied code from any other source, did you follow IP policy in
[CONTRIBUTING.md](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md#-copying-code-from-other-sources)?:
✅ / ❌ / 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. -->

### Additional Information
<!-- E.g. related issue. -->


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Configure MOE calibration experts ratio for quantization via an
environment/option, enabling finer control over calibration.

* **Bug Fixes**
* Improved detection of sparse MOE blocks to handle varying
expert/topology layouts, inferring expert counts when needed for more
reliable processing.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
2026-03-03 21:30:35 +00:00
Keval Morabia 860d0b4a70 Merge Linux and Windows Changelog (#954)
Since we no longer have any compiled packages, all releases are for all
platforms so we dont have separate windows releases hence merging
changelog as well.

Going forward, windows can test on latest linux version and if fixes
needed, they can go in next usual monthly linux release

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

* **Documentation**
* Reorganized changelog structure to consolidate Windows and Linux
release information in a unified view.
* Expanded Windows Support documentation across recent release versions.

* **Chores**
* Updated Windows example release badge to dynamically reflect the
latest PyPI release version.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
2026-03-03 11:31:59 +05:30
Chenhan D. Yu edde087bbd Adding a special list to handle HF models that require (#950)
## What does this PR do?

**Type of change:** ? <!-- Use one of the following: Bug fix, new
feature, new example, new tests, documentation. --> Bug fix

**Overview:** ?

1. Some models check `layer_types` and our offline overwrites
`num_hidden_layers=0` which will requires passing `layer_types=[]`. It
is also possible that the model will raise unexpected arguments when
passing `layer_types=[]`. So the current WAR is to use a list to handle
by checking the architectures.
2. Fix relative path issue in `launch_train.sh`



## Usage
<!-- You can potentially add a usage example below. -->

```python
# Add a code snippet demonstrating how to use this
```

## Testing
<!-- Mention how have you tested your change if applicable. -->

## Before your PR is "*Ready for review*"
<!-- If you haven't finished some of the above items you can still open
`Draft` PR. -->

- **Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)**
and your commits are signed.
- **Is this change backward compatible?**: Yes/No <!--- If No, explain
why. -->
- **Did you write any new necessary tests?**: Yes/No
- **Did you add or update any necessary documentation?**: Yes/No
- **Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**:
Yes/No <!--- Only for new features, API changes, critical bug fixes or
bw breaking changes. -->

## Additional Information
<!-- E.g. related issue. -->


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

* **Bug Fixes**
  * Enhanced model loading robustness for specific layer configurations.

* **Chores**
* Improved example script path handling and GPU initialization tracking.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Chenhan Yu <chenhany@nvidia.com>
2026-03-03 00:07:10 +00:00
Keval Morabia 82f1d216d1 Add Security and IP related contributing guide and configure coderabbit to catch such issues (#935)
### What does this PR do?

- Add Security related coding practices in `SECURITY.md` and merge with
`2_security.rst`
- Update `CONTRIBUTING.md` for instructions to follow if copying code
from other repositories
- Update PR template
- Cleanup dependency files
- New API `mto.load_modelopt_state` doing the insecure `torch.load(f,
weights_only=False)` instead of doing it separately everywhere. This
also allows us to later improve the input validation for
`modelopt_state_path` or use safer alternatives to `torch.load`

### Testing
<!-- Mention how have you tested your change if applicable. -->

N/A

### 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. hardcoded `trust_remote_code=True`, `torch.load(...,
weights_only=True)`, `pickle`, etc.).

- Is this change backward compatible?: ✅ <!--- If ❌, explain why. -->
- If you copied code from any other source, did you follow IP policy in
[CONTRIBUTING.md](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md#-copying-code-from-other-sources)?:
✅ <!--- Mandatory -->
- Did you write any new necessary tests?: NA <!--- Mandatory for new
features or examples. -->
- Did you add or update any necessary documentation and update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
NA <!--- Only for new features, API changes, critical bug fixes or
backward incompatible changes. -->

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **Documentation**
* Expanded and reorganized security guidance and contributor procedures;
updated PR template and several READMEs with clearer security,
submission, and installation instructions
* Replaced an older security document with an enhanced, centralized
security guidance

* **Chores**
* Adjusted example dependency lists and optional extras (adds, removals,
and version constraints)
* Enabled automated incremental reviews, added pre-merge security
checks, and introduced a knowledge-base of coding/security guidelines
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
2026-03-03 03:49:15 +05:30
Adi Renduchintalaandcoderabbitai[bot] f26b9c3a2f skip generate option for large models and mxfp8 (#942)
## What does this PR do?

**Type of change:** New feature

**Overview:** Adds a `--skip_generate` flag to `hf_ptq.py` that skips
the pre/post-quantization generation preview calls. These calls run
`model.generate()` which crashes for very large models (500B+) that are
split across GPU and CPU via `device_map="auto"` (e.g., models with
Mamba/Triton kernels that cannot handle CPU-offloaded tensors).

## Usage
```
python examples/llm_ptq/hf_ptq.py \
    --pyt_ckpt_path /path/to/model \
    --export_path /path/to/output \
    --qformat mxfp8 \
    --trust_remote_code \
    --export_fmt hf \
    --batch_size 1 \
    --skip_generate \
    --kv_cache_qformat none
```
## Testing
Tested with a 500B parameter NemotronH hybrid Mamba/attention model on
4x GB200 GPUs. Without --skip_generate, the script crashes at
model.generate() due to Mamba Triton kernels failing on CPU-offloaded
tensors. With --skip_generate, the generation preview is skipped and
quantization proceeds normally.

## Before your PR is "*Ready for review*"
<!-- If you haven't finished some of the above items you can still open
`Draft` PR. -->

- **Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)**
and your commits are signed.
- **Is this change backward compatible?**: Yes/No <!--- If No, explain
why. -->
- **Did you write any new necessary tests?**: Yes/No
- **Did you add or update any necessary documentation?**: Yes/No
- **Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**:
Yes/No <!--- Only for new features, API changes, critical bug fixes or
bw breaking changes. -->

## Additional Information
The --skip_generate flag sets generated_ids_before_ptq = None early,
which also causes the post-quantization generate to be skipped via the
existing if generated_ids_before_ptq is None: pass guard. Combined with
--batch_size 1 (to skip the get_max_batch_size forward-pass probe), this
eliminates all forward passes that can crash for device-map-split
models.

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

* **New Features**
* Introduced `--skip_generate` CLI option to skip pre-quantization text
and image generation, reducing processing time for very large models.
Useful when generation previews are computationally expensive.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: adithyare <adithyare@nvidia.com>
Signed-off-by: Adi Renduchintala <adithya.r@gmail.com>
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
2026-03-02 13:41:25 -08:00
jingyu-ml 2905cb0f2e Updated the diffusion config issue and more test cases (#937)
## What does this PR do?

**Type of change:** new tests, Bug fix <!-- Use one of the following:
Bug fix, new feature, new example, new tests, documentation. -->

**Overview:** 

- **Fixed the INT8 config issue**

- **Add HF checkpoint export test coverage**

1. The `--hf-ckpt-dir` export path had zero test coverage. This MR adds
tests at two levels:
2. Unit tests (tests/unit/torch/export/test_export_diffusers.py):
- Extended test_export_diffusers_real_quantized to parametrize over
INT8, INT8 SmoothQuant, FP8, and FP4 configs
- (previously only FP8). This gives 3 models x 4 configs = 12 test
cases.

3. GPU integration tests
(tests/gpu/torch/export/test_export_diffusers_hf_ckpt.py)
- New file testing the full quantize.py --hf-ckpt-dir pipeline via
subprocess with 4 combos:
- SDXL INT8 smoothquant min-mean (the exact scenario that triggered the
bug)
    - Flux INT8 smoothquant min-mean
    - SDXL FP8
    - Flux FP4

## Usage
<!-- You can potentially add a usage example below. -->

```python
# Add a code snippet demonstrating how to use this
```

## Testing
<!-- Mention how have you tested your change if applicable. -->

## Before your PR is "*Ready for review*"
<!-- If you haven't finished some of the above items you can still open
`Draft` PR. -->

- **Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)**
and your commits are signed.
- **Is this change backward compatible?**:No <!--- If No, explain why.
-->
- **Did you write any new necessary tests?**: Yes
- **Did you add or update any necessary documentation?**:No
- **Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**:No
<!--- Only for new features, API changes, critical bug fixes or bw
breaking changes. -->

## Additional Information
<!-- E.g. related issue. -->


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

## Release Notes

* **Tests**
* Added test coverage for exporting Diffusers models with Hugging Face
checkpoints across multiple quantization formats (INT8, FP8, FP4)
* Extended quantization export testing to validate multiple
configuration scenarios

* **Chores**
* Refined INT8 quantization configuration with improved calibrator
support for convolution layers

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
2026-02-28 23:54:02 +05:30
sugunav14 56e97c8096 Bug fix 5875873 (#865)
## What does this PR do?

**Type of change:** Bug fix <!-- Use one of the following: Bug fix, new
feature, new example, new tests, documentation. -->

**Overview:** Newer version of trl uses dtype instead of torch_dtype.
Modified code to set float32 as default for older versions of trl that
you torch_dtype.

## Usage
<!-- You can potentially add a usage example below. -->

```python
# Add a code snippet demonstrating how to use this
```

## Testing
<!-- Mention how have you tested your change if applicable. -->

## Before your PR is "*Ready for review*"
<!-- If you haven't finished some of the above items you can still open
`Draft` PR. -->

- **Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)**
and your commits are signed.
- **Is this change backward compatible?**: Yes/No <!--- If No, explain
why. -->
- **Did you write any new necessary tests?**: Yes/No
- **Did you add or update any necessary documentation?**: Yes/No
- **Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**:
Yes/No <!--- Only for new features, API changes, critical bug fixes or
bw breaking changes. -->

## Additional Information
<!-- E.g. related issue. -->


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

* **Bug Fixes**
* Enhanced error handling in model training examples to safely manage
missing dtype attributes, preventing crashes during initialization when
torch_dtype is not configured.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
2026-02-28 23:52:44 +05:30
Zhiyu a415667992 Enable Qwen3.5-MoE PTQ (#897)
## What does this PR do?

**Type of change:**  New model support

**Overview:** Add ModelOpt PTQ support for
https://huggingface.co/Qwen/Qwen3.5-397B-A17B

## Usage
<!-- You can potentially add a usage example below. -->

```python
python3 hf_ptq.py --pyt_ckpt_path /home/omniml_data_3/models/Qwen3.5-397B-A17B --qformat nvfp4_mlp_only --export_path /home/omniml_data_3/zhiyuc/checkpoints/Qwen3.5-397B-A17B-NVFP4 --trust_remote_code
```

## Testing
<!-- Mention how have you tested your change if applicable. -->

## Before your PR is "*Ready for review*"
<!-- If you haven't finished some of the above items you can still open
`Draft` PR. -->

- **Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)**
and your commits are signed.
- **Is this change backward compatible?**: Yes <!--- If No, explain why.
-->
- **Did you write any new necessary tests?**: Yes/No
- **Did you add or update any necessary documentation?**: Yes/No
- **Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**:
Not yet <!--- Only for new features, API changes, critical bug fixes or
bw breaking changes. -->

## Additional Information
<!-- E.g. related issue. -->


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

* **New Features**
* Added Qwen3.5 Mixture-of-Experts model support in quantization
workflows.

* **Bug Fixes**
* Enhanced error diagnostics during model export with detailed module
information.
* Improved dataset tokenizer processing with proper truncation and
length handling.
  * Fixed model export stability issue related to framework integration.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Zhiyu Cheng <zhiyuc@nvidia.com>
2026-02-27 00:06:07 +00:00
Frida Hou dfe705aab5 support static NVFP4 HF export (#858)
## What does this PR do?

**Type of change:** ? new feature

**Overview:** ?
Supports export `NVFP4StaticQantizer` in unified huggingface checkpoint,
as a deployment path for PTQ algorithms such as MSE

## Usage
<!-- You can potentially add a usage example below. -->

```python
# checkpoint generation
python examples/llm_ptq/hf_ptq.py --pyt_ckpt_path Qwen/Qwen3-8B  --qformat nvfp4_mse --export_path test-Qwen3-8B-Instruct-MSE-FP8-sweep-FP4 --kv_cache_qformat none --trust_remote_code 

```

## Testing
Tested generated Qwen3 8B checkpoint with trtllm serve and nv_eval
example in `Model-Optimizer-Internal/examples/nv_eval`.

NV eval results:
```
| Groups |Version|Filter|n-shot|  Metric   |   |Value |   |Stderr|
|--------|-------|------|------|-----------|---|-----:|---|-----:|
|mmlu_str|       |none  |      |exact_match|↑  |0.7186|±  |0.0036|
```


## Before your PR is "*Ready for review*"
<!-- If you haven't finished some of the above items you can still open
`Draft` PR. -->

- **Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)**
and your commits are signed.
- **Is this change backward compatible?**: Yes/No <!--- If No, explain
why. -->
- **Did you write any new necessary tests?**: Yes/No
- **Did you add or update any necessary documentation?**: Yes/No
- **Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**:
Yes/No <!--- Only for new features, API changes, critical bug fixes or
bw breaking changes. -->

## Additional Information
<!-- E.g. related issue. -->


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

* **New Features**
* Added support for static NVFP4 quantizers that utilize pre-computed
calibration scales.
* Introduced new NVFP4 W4A4 quantization configuration with optional FP8
scale sweep.

* **Performance Improvements**
* Static quantizers now skip unnecessary dynamic scaling factor
recalculation.
* Unified quantization handling for improved consistency and efficiency.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
Signed-off-by: Frida Hou <201670829+Fridah-nv@users.noreply.github.com>
2026-02-26 10:41:29 -08:00
yeyu-nvidia 4a11486f6c Enable multinode training for HF speculative decoding (#922)
## What does this PR do?

**Type of change:** 
New example

**Overview:** 
Modify launch_train.sh script to enable multi-node training.
Provide a slurm template script.

## Usage
Add required fields in slurm.sh.
Then use the command below to submit multi-node job:

```bash
bash slurm.sh
```

## Testing
<!-- Mention how have you tested your change if applicable. -->

## Before your PR is "*Ready for review*"
<!-- If you haven't finished some of the above items you can still open
`Draft` PR. -->

- **Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)**
and your commits are signed.
- **Is this change backward compatible?**: Yes/No <!--- If No, explain
why. -->
- **Did you write any new necessary tests?**: Yes/No
- **Did you add or update any necessary documentation?**: Yes/No
- **Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**:
Yes/No <!--- Only for new features, API changes, critical bug fixes or
bw breaking changes. -->

## Additional Information
<!-- E.g. related issue. -->


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

* **New Features**
* Added multi-node training support with new configuration options for
distributed setups.
* Introduced Slurm batch script for streamlined job submission to
cluster environments.

* **Improvements**
* Enhanced GPU resource management for multi-GPU and distributed
training configurations.
  * Updated speculative decoding model support and validation handling.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Ye Yu <yeyu@nvidia.com>
2026-02-26 18:23:43 +00:00
Keval Morabia 4eacb0da72 Add trust_remote_code cli option for mbridge distillation (#934)
## What does this PR do?

Nemotron2/3 need `AutoBridge(..., trust_remote_code=True)` which was
missing previously

## Testing
<!-- Mention how have you tested your change if applicable. -->

Nemotron-nano-v2 can be distilled using tokenized
Nemotron-Pretraining-SFT-v1 data

Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
2026-02-25 22:30:53 +01:00
jingyu-ml d78797b466 Update the LTX2 API calls during the calibration (#926)
## What does this PR do?

**Type of change:** Bug fix <!-- Use one of the following: Bug fix, new
feature, new example, new tests, documentation. -->

**Overview:** 

Update LTX-2 integration to match latest upstream API

1. The LTX-2 codebase removed/replaced several APIs. This MR updates all
affected files:
2. Replace cfg_guidance_scale with MultiModalGuiderParams: The pipeline
__call__ no longer accepts a single cfg_guidance_scale float. It now
requires two MultiModalGuiderParams objects (video_guider_params and
audio_guider_params) that control CFG, STG, rescale, cross-modality
guidance, and skip-step settings. Updated in ltx-2.py, ltx-2-fp8.py,
ltx-2-onestage.py, calibration.py, and models_utils.py.
3. Replace fp8transformer with QuantizationPolicy: The
TI2VidTwoStagesPipeline constructor no longer accepts the fp8transformer
boolean flag. FP8 quantization is now configured via
quantization=QuantizationPolicy.fp8_cast(). Updated in ltx-2-fp8.py and
pipeline_manager.py (with backwards-compatible support for the old
--extra-param fp8transformer=true CLI flag).
4. Remove DEFAULT_CFG_GUIDANCE_SCALE constant: Replaced by
DEFAULT_VIDEO_GUIDER_PARAMS and DEFAULT_AUDIO_GUIDER_PARAMS in all
import sites.

## Usage
<!-- You can potentially add a usage example below. -->

```bash
python quantize.py --model ltx-2 --format fp4 --batch-size 1 --calib-size 1 --n-steps 40 --extra-param checkpoint_path=./ltx-2-19b-dev-fp8.safetensors --extra-param distilled_lora_path=./ltx-2-19b-distilled-lora-384.safetensors --extra-param spatial_upsampler_path=./ltx-2-spatial-upscaler-x2-1.0.safetensors --extra-param gemma_root=./gemma-3-12b-it-qat-q4_0-unquantized --extra-param fp8transformer=true --hf-ckpt-dir ./ltx2-nvfp4
```

## Testing
<!-- Mention how have you tested your change if applicable. -->

## Before your PR is "*Ready for review*"
<!-- If you haven't finished some of the above items you can still open
`Draft` PR. -->

- **Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)**
and your commits are signed.
- **Is this change backward compatible?**: Yes/No <!--- If No, explain
why. -->
- **Did you write any new necessary tests?**: Yes/No
- **Did you add or update any necessary documentation?**: Yes/No
- **Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**:
Yes/No <!--- Only for new features, API changes, critical bug fixes or
bw breaking changes. -->

## Additional Information
<!-- E.g. related issue. -->


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

## Release Notes

* **Updates**
  * Default resolution for LTX2 models adjusted to 768x1280
* Guidance parameter configuration updated for video and audio pipelines
  * FP8 quantization parameter handling refined

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
2026-02-24 14:23:28 -08:00
03a1899dda Support force tokens to % of total experts during calibration (#910)
## What does this PR do?

**Type of change:** New feature

**Overview:** Adds a configurable `moe_calib_experts_ratio` parameter
that controls the percentage of experts to calibrate during the forward
pass in MoE (Mixture of Experts) models. Previously, the calibration
forward always routed tokens to **all** experts, which is expensive.
This PR allows the user to specify a ratio (default: still all experts
so no behavior change) to improve expert calibration coverage without
the cost of a full-expert forward. The token counting for the expert
coverage table now tracks the calibration routing and runs on CUDA for
efficiency.

**Changes include:**
- New `moe_calib_experts_ratio` field in `QuantizeAlgorithmConfig`
(`config.py`)
- Propagation of the ratio from the algorithm config to MoE modules
during calibration (`mode.py`)
- Updated `_QuantSparseMoe.forward` to use the configurable ratio
instead of hard-coding all experts (`huggingface.py`)
- New `--moe_calib_experts_ratio` CLI flag in `hf_ptq.py` (default
`0.25`)
- Moved `expert_token_count` tensor to CUDA and updated the HTML table
title in `moe_utils.py`

## Usage

Via hf_ptq.py CLI — calibrate 50% of experts during MoE calibration
python hf_ptq.py --model <model> --qformat int4_awq
--moe_calib_experts_ratio 0.5

Via Python API — pass the ratio through the algorithm config
import modelopt.torch.quantization as mtq

quant_cfg = {
    "quant_cfg": { ... },
    "algorithm": {
        "method": "awq_lite",
        "moe_calib_experts_ratio": 0.25,  # calibrate 1/4 of experts
    },
}
mtq.quantize(model, quant_cfg, forward_loop=calib_loop)

## Testing
Test with Qwen3 30B A3B calibration and check the tokens per expert.

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

## Release Notes

* **New Features**
* Added support for configurable expert calibration during Mixture of
Experts (MOE) model quantization. Users can now specify the percentage
of experts to include during calibration, enabling better expert
coverage and improved quantization accuracy for MOE models. Default: 25%
of all experts.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
Signed-off-by: Chenjie Luo <108829653+cjluo-nv@users.noreply.github.com>
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
Co-authored-by: realAsma <86726418+realAsma@users.noreply.github.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
2026-02-24 12:35:19 -08:00
Izzy Putterman 2802302bf9 SpecDec Bench: February Update (#875)
## What does this PR do?

**Type of change:** ? <!-- Use one of the following: Bug fix, new
feature, new example, new tests, documentation. -->

**Overview:** 
Addition of SpecBench Dataset
Addition of NVIDID SPEED-Bench dataset, preproc scripts, and custom
metrics aggregator
Addition of example of converting SpecBench Medusa to this FW
Addition of Initial TRTLLM AutoDeploy Specdec support

Updates to all frameworks for better perf (overlap/async scheduling etc)

## Usage
<!-- You can potentially add a usage example below. -->

```python
# Add a code snippet demonstrating how to use this
```

## Testing
<!-- Mention how have you tested your change if applicable. -->

## Before your PR is "*Ready for review*"
<!-- If you haven't finished some of the above items you can still open
`Draft` PR. -->

- **Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)**
and your commits are signed.
- **Is this change backward compatible?**: Yes/No <!--- If No, explain
why. -->
- **Did you write any new necessary tests?**: Yes/No
- **Did you add or update any necessary documentation?**: Yes/No
- **Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**:
Yes/No <!--- Only for new features, API changes, critical bug fixes or
bw breaking changes. -->

## Additional Information
<!-- E.g. related issue. -->


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

## Release Notes

* **New Features**
* Added SPEED-Bench dataset support with configurable throughput and
qualitative configurations
* Introduced SpecBench metrics with acceptance rate analysis and
visualizations
  * Added progress bar during benchmark execution
* New model implementations for auto-deployment and Medusa-style
speculative decoding
  * Data preparation utility for benchmark datasets
  * Enhanced metrics with per-category analysis and performance charts

* **Documentation**
  * Updated README with SPEED-Bench workflow and examples
  * New porting guide for integrating custom benchmark runners

* **Refactor**
  * Streamlined model and runner interfaces for improved flexibility
* Consolidated dataset implementations and removed deprecated base
classes

* **Chores**
  * Added required dependencies for data handling and visualizations

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Signed-off-by: Izzy Putterman <iputterman@nvidia.com>
2026-02-24 22:24:13 +05:30
Keval Morabia c689ea1176 flush print megatron tokenization stats and update readme (#927)
## What does this PR do?

When running the script, I often see the print stats for tokenization
(every `log_interval`) not showing up or showing up very very delayed.
Hence using `print(..., flush=True)` to fix this.

Also update README that the example shown for tokenization takes too
long to run, split into multiple .jsonl files for efficiently running
the tokenization; and try out a smaller dataset first to test the script

## Testing
<!-- Mention how have you tested your change if applicable. -->

Split Nemotron-pretraining-SFT-v1 dataset into multiple .jsonl splits
and then tokenize them parallelly in different slurm jobs.

Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
2026-02-24 17:41:38 +01:00
Keval Morabia f78385e70e Improve megatron dataset preprocessing script and update docs (#918)
## What does this PR do?

Improve megatron dataset preprocessing script and update docs

## Usage
<!-- You can potentially add a usage example below. -->

```python
python -m modelopt.torch.utils.plugins.megatron_preprocess_data \
    --hf_dataset nvidia/Nemotron-Pretraining-SFT-v1 \
    --hf_name Nemotron-SFT-General \
    --hf_split train \
    --hf_max_samples_per_split 10_000_000 \
    --json_keys text \
    --tokenizer Qwen/Qwen3-0.6B \
    --output_dir /path/to/tokenized/data/qwen3 \
    --workers 32 \
    --max_sequence_length 256_000
```

```python
python -m modelopt.torch.utils.plugins.megatron_preprocess_data \
    --jsonl_paths /path/to/data1.jsonl /path/to/data2.jsonl ... \
    --json_keys text \
    --tokenizer Qwen/Qwen3-0.6B \
    --output_dir /path/to/tokenized/data/qwen3 \
    --workers 32 \
    --max_sequence_length 256_000
```

## Testing
<!-- Mention how have you tested your change if applicable. -->

- Downloaded and tokenized Nemotron-Pretraining-SFT-v1 with
Nemotron-Nano-v2 tokenizer

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

* **Documentation**
* Updated data preparation guides with new CLI patterns and Hugging Face
Hub integration instructions.

* **New Features**
* Added batch tokenization via directory input and direct Hugging Face
dataset downloads with flexible subset/split filtering.

* **Configuration Updates**
* Optimized distillation settings: adjusted optimizer parameters and
increased checkpoint retention.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
2026-02-24 01:12:55 +05:30
Chenjie Luo 9975ba1065 Fix DeepSeek PTQ script (#912)
## What does this PR do?

**Type of change:** ? Bug fix

**Overview:** ?

Fix two bugs in the PTQ script

## Testing

Run DeepseekV3.2 PTQ and export


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

* **Refactor**
* Enhanced data type handling in quantization examples for bf16
operations
* Updated internal dependencies for quantization utilities to improve
modularity

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
2026-02-20 19:45:37 +00:00
Chenjie Luo ac7c985d96 [NVBUG: 5804406] Auto detect MOE layers (#900)
## What does this PR do?

**Type of change:** New feature, new tests

**Overview:** Replace hardcoded per-model MoE class registrations
(Mixtral, Qwen2Moe, Qwen3Moe, Qwen3Next, Llama4TextMoe, Qwen3VLMoe,
MiniMaxM2, etc.) with a single generic auto-detection mechanism
(`register_sparse_moe_on_the_fly`) that walks the model tree and
identifies MoE blocks by their structural attributes (`gate` + `experts`
with `top_k`/`num_experts`). This makes MoE quantization
forward-compatible with new HuggingFace MoE architectures without
requiring explicit registration for each model family.

Additionally, this PR:
- Tracks per-expert token routing counts during calibration via a gate
forward hook, enabling visibility into expert utilization.
- Saves an HTML report of expert token counts during export
(`save_expert_token_count_table`), highlighting under-utilized experts.
- Fixes the `topk` -> `top_k` attribute name for transformers >= 5.0
compatibility.
- Also move the ptq summary prints to a file in hf_ptq.py to reduce the
prints

## Usage

Auto-detection is transparent -- no user-facing API changes are needed.
Any HuggingFace MoE model with the standard `gate`/`experts` pattern is
automatically detected and quantized:

import modelopt.torch.quantization as mtq

# Any HuggingFace MoE model (Mixtral, Qwen3Moe, DeepSeek, etc.)
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-30B-A3B")

mtq.quantize(model, mtq.INT8_DEFAULT_CFG, forward_loop)

# During export, an .moe.html report with per-expert token counts is
saved automatically

## Testing
unittest, also test exporting qwen MOE

## Before your PR is "*Ready for review*"
<!-- If you haven't finished some of the above items you can still open
`Draft` PR. -->

- **Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)**
and your commits are signed.
- **Is this change backward compatible?**: Yes/No <!--- If No, explain
why. -->
- **Did you write any new necessary tests?**: Yes/No
- **Did you add or update any necessary documentation?**: Yes/No
- **Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**:
Yes/No <!--- Only for new features, API changes, critical bug fixes or
bw breaking changes. -->

## Additional Information
<!-- E.g. related issue. -->


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

* **New Features**
* Added expert token count visualization for Mixture of Experts models,
exported as HTML reports during model export.
* Enhanced sparse MoE quantization with improved calibration-aware
routing and automatic model block detection.

* **Tests**
* Added comprehensive test suite for sparse MoE quantization validation.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
2026-02-19 19:46:58 +00:00
h-guo18 eb99488da1 Fix: restore requires_grad in transformers5 reloading (#907)
## What does this PR do?

**Type of change:** ? <!-- Use one of the following: Bug fix, new
feature, new example, new tests, documentation. -->

**Overview:** 

Patch transformers 5.x parameter loading to preserve original
`requires_grad` settings.

In transformers v5.x, loading a checkpoint forcibly sets parameters'
requires_grad,
which unintentionally unfreeze frozen parameters (e.g. Base model in
eagle training).

This leads to optimizer initialization error since the restored
optimizer expected more parameter than the checkpoint.

This monkey-patch restores the original`requires_grad` after loading
parameters.

Reference:

https://github.com/huggingface/transformers/blob/v5.0.0.rc1-release/src/transformers/core_model_loading.py#L640

## Usage
<!-- You can potentially add a usage example below. -->

```python
# Add a code snippet demonstrating how to use this
```

## Testing
<!-- Mention how have you tested your change if applicable. -->

## Before your PR is "*Ready for review*"
<!-- If you haven't finished some of the above items you can still open
`Draft` PR. -->

- **Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)**
and your commits are signed.
- **Is this change backward compatible?**: Yes/No <!--- If No, explain
why. -->
- **Did you write any new necessary tests?**: Yes/No
- **Did you add or update any necessary documentation?**: Yes/No
- **Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**:
Yes/No <!--- Only for new features, API changes, critical bug fixes or
bw breaking changes. -->

## Additional Information
<!-- E.g. related issue. -->


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

* **Bug Fixes**
* Fixed model parameter loading in speculative decoding to properly
preserve gradient requirements for each parameter when using HuggingFace
Transformers 5.x, ensuring correct behavior during checkpoint resumption
and model initialization.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
2026-02-19 01:25:47 +00:00
h-guo18 b8a4586702 Refactor: Eagle data loading (#668)
## What does this PR do?

**Type of change:** Refactor <!-- Use one of the following: Bug fix, new
feature, new example, new tests, documentation. -->

**Overview:** 
Jira ticket: https://jirasw.nvidia.com/browse/OMNIML-2955

Main changes :

- Consolidate Eagle data loading with @ChenhanYu 's implementation of
`transformers_dataset.py`

- Refactor: baked the following logics from `example/main.py` to
`modelopt/torch` for cleaner example entrance:
  - default config selecting and merging with custom config
  - tokenizer post-processor (chat template and pad_tok_id)
  - d2t loading 
- Implementation refactor: In HF workflow, reuse base modfel's input
hidden states as input_embedding, instead of calculating from input_ids.
This has two main benefits:
    - Easier VLM support, which has various embedding processing logics.
    - Training effieicy.
    
- Deprecating eagle1 from the example. It is still available by setting
custom config.
  
 - Other minor fixes and readme updates.

## Usage
<!-- You can potentially add a usage example below. -->

```python
# Add a code snippet demonstrating how to use this
```

## Testing
<!-- Mention how have you tested your change if applicable. -->

Tested that training curves after changes (both online&offline) is
identical with original branch:
<img width="1073" height="634" alt="image"
src="https://github.com/user-attachments/assets/abfd7bea-c82c-48a7-8181-68c5a9e4da8d"
/>


## Before your PR is "*Ready for review*"
<!-- If you haven't finished some of the above items you can still open
`Draft` PR. -->

- **Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)**
and your commits are signed.
- **Is this change backward compatible?**: Yes/No <!--- If No, explain
why. -->
- **Did you write any new necessary tests?**: Yes/No
- **Did you add or update any necessary documentation?**: Yes/No
- **Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**:
Yes/No <!--- Only for new features, API changes, critical bug fixes or
bw breaking changes. -->

## Additional Information
<!-- E.g. related issue. -->


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

## Release Notes

* **New Features**
* Added draft vocabulary cache support for EAGLE model training,
enabling runtime vocabulary customization via `--draft_vocab_cache`
parameter
* Introduced new data loading utilities with sharding, streaming, and
tokenization support for large-scale training
  * Added optional `--log_steps` configuration to training launcher

* **Documentation**
* Updated EAGLE configuration guides with draft vocabulary cache setup
instructions and examples

* **Refactor**
* Restructured data pipeline for offline training with improved dataset
handling and batching
* Updated command-line arguments across training scripts (`--input-data`
replaces `--input-file`)

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
2026-02-18 14:04:06 -08:00
kaix-nv 9e38041d34 [OMNIML-2850] [3/n] Adds sparse attention calibration (#538)
## What does this PR do?

**Type of change:** ? <!-- Use one of the following: Bug fix, new
feature, new example, new tests, documentation. -->
new feature

**Overview:** ?
- This PR adds the sparse attention calibration algorithm
- Chunked prefill to support long ctx_len
- Separated calibration for prefill and decode

## Usage
<!-- You can potentially add a usage example below. -->

```python
import modelopt.torch.sparsity.attention_sparsity as mtsa

# Apply sparse attention with calibration
model = mtsa.sparsify(model, config=SKIP_SOFTMAX_CALIB)

# Print summary - now shows actual thresholds
mtsa.print_sparse_attention_summary(model)
# Output:
# Method: flash_skip_softmax, Threshold: Dynamic (λ=437.395926)

# Or llm_eval integration
# HuggingFace sparse attention example
python examples/llm_sparsity/attention_sparsity/hf_sa.py \
    --pyt_ckpt_path Qwen/Qwen3-4B \
    --sparse_attn skip_softmax_calib 
```

# The calibration method
## Calibration Algorithm
- Implemented the Inverse Power model: scale_factor = k / (1 -
sparsity)^p
- Fit model parameters (k, p) per phase using scipy.optimize.curve_fit
- At inference: threshold = k / (1 - target_sparsity)^p / seqlen

## Why Choosing the Inverse Power model?
The inverse power model better fits the relationship between sparsity
ratio and threshold_scale_factor.
<img width="2388" height="1082" alt="sparsity_model_analysis"
src="https://github.com/user-attachments/assets/4dfb45d4-8c16-4f15-a878-c8e08a9b6128"
/>

## Runtime Flexibility
- Target sparsity can be changed at inference time without recalibration
- Users can adjust module._sparse_method_instance.target_sparse_ratio
dynamically
- Threshold automatically adapts to sequence length

## Testing
<!-- Mention how have you tested your change if applicable. -->
The calibration results for `Qwen/Qwen3-30B-A3B-Thinking-2507` are shown
below and are mostly consistent with the ground-truth numbers collected
from the kernel side.
```
Prefill Calibration Results:
  Model: scale_factor = k / (1 - sparsity)^p
  Fitted k: 1003.3990
  Fitted p: 1.2589
  R-squared: 0.827549

Scale factors for different target sparsities:
  Target     Scale Factor
  ---------- ---------------
  50%        2401.35
  70%        4568.26
  80%        7610.98
  90%        18214.70
  95%        43591.65
```

## Before your PR is "*Ready for review*"
<!-- If you haven't finished some of the above items you can still open
`Draft` PR. -->

- **Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/CONTRIBUTING.md)**
and your commits are signed.
- **Is this change backward compatible?**: Yes/No <!--- If No, explain
why. -->
- **Did you write any new necessary tests?**: Yes/No
- **Did you add or update any necessary documentation?**: Yes/No
- **Did you update
[Changelog](https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/CHANGELOG.rst)?**:
Yes/No <!--- Only for new features, API changes, critical bug fixes or
bw breaking changes. -->

## Additional Information
<!-- E.g. related issue. -->

---------

Signed-off-by: Kai Xu <kaix@nvidia.com>
2026-02-18 02:18:12 +00:00
Chenjie Luo 3801923e9d Support MiniMax M2.1 (FP8 checkpoint) (#817)
## What does this PR do?

**Type of change:** ? new feature

**Overview:** ?

Support loading the MiniMax M2.1 (FP8) checkpoint for PTQ.

## Usage
scripts/huggingface_example.sh --model <minimax checkpoint> --quant
nvfp4 --trust_remote_code


## Before your PR is "*Ready for review*"
<!-- If you haven't finished some of the above items you can still open
`Draft` PR. -->

- **Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)**
and your commits are signed.
- **Is this change backward compatible?**: Yes/No <!--- If No, explain
why. -->
- **Did you write any new necessary tests?**: Yes/No
- **Did you add or update any necessary documentation?**: Yes/No
- **Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**:
Yes/No <!--- Only for new features, API changes, critical bug fixes or
bw breaking changes. -->

## Additional Information
<!-- E.g. related issue. -->


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

* **New Features**
  * Added MiniMax M2.1 model quantization support with nvfp4 format.
* Extended FP8 quantization capabilities with configurable dtype
parameter for enhanced precision control.

* **Improvements**
  * Enhanced detection of quantized linear module variants.
  * Improved weight unpacking for FP8-based linear modules.

* **Documentation**
  * Updated supported models table to include MiniMax M2.1.

<sub>✏️ Tip: You can customize this high-level summary in your review
settings.</sub>

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
Signed-off-by: Chenjie Luo <108829653+cjluo-nv@users.noreply.github.com>
2026-02-17 18:47:22 +00:00
mxinO ca1f9687bd [OMNIML-3505] LTX-2 Distillation Trainer (#892)
## What does this PR do?

**Type of change:** new example <!-- Use one of the following: Bug fix,
new feature, new example, new tests, documentation. -->

**Overview:** 
Adding LTX-2 distillation trainer.

## Usage
<!-- You can potentially add a usage example below. -->

```bash
accelerate launch \
    --config_file configs/accelerate/fsdp.yaml \
    --num_processes 8 \
    distillation_trainer.py --config configs/distillation_example.yaml
```

See readme for more details.

## Testing
Run training with single/multiple nodes.

## Before your PR is "*Ready for review*"
<!-- If you haven't finished some of the above items you can still open
`Draft` PR. -->

- **Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)**
and your commits are signed.
- **Is this change backward compatible?**: Yes <!--- If No, explain why.
-->
- **Did you write any new necessary tests?**: NA
- **Did you add or update any necessary documentation?**: Yes
- **Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**:
Yes <!--- Only for new features, API changes, critical bug fixes or bw
breaking changes. -->

## Additional Information
<!-- E.g. related issue. -->


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

## New Features
* Added distillation training support for LTX-2 models with quantization
integration.
* Introduced comprehensive documentation and example configurations for
distillation workflows.
* Includes multi-GPU and multi-node training setup with distributed
training support and customizable configuration templates.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Meng Xin <mxin@nvidia.com>
2026-02-14 13:28:21 -08:00
Keval Morabia b43ae641a4 MBridge pruning minor fix for saving pruned NemotronH (#887)
## What does this PR do?

**Type of change:** Bug fix <!-- Use one of the following: Bug fix, new
feature, new example, new tests, documentation. -->

## Testing
<!-- Mention how have you tested your change if applicable. -->

Nemotron Nano v2 pruned can be saved


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

* **Bug Fixes**
* Fixed Hugging Face model loading to properly respect the
`trust_remote_code` parameter during model instantiation.

* **Improvements**
* Enhanced distributed training logging with rank-0 aware warning and
logging mechanisms for cleaner, non-redundant output in multi-GPU and
multi-node scenarios.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
2026-02-13 11:58:56 -08:00
Zhiyu 95511a0bc9 Add Nemotron parse PTQ support (#786)
## What does this PR do?

**Type of change:** New model support <!-- Use one of the following: Bug
fix, new feature, new example, new tests, documentation. -->

**Overview:** Add PTQ support for
https://huggingface.co/nvidia/NVIDIA-Nemotron-Parse-v1.1

## Usage
<!-- You can potentially add a usage example below. -->

```python
python3 hf_ptq.py --pyt_ckpt_path /home/omniml_data_3/models/NVIDIA-Nemotron-Parse-v1.1 --qformat fp8 --export_path /home/omniml_data_3/zhiyuc/checkpoints/NVIDIA-Nemotron-Parse-v1.1-FP8 --trust_remote_code --kv_cache_qformat none --attn_implementation eager
```
By default, image-text data will be used in calibration for VLMs.

## Testing
<!-- Mention how have you tested your change if applicable. -->

## Before your PR is "*Ready for review*"
<!-- If you haven't finished some of the above items you can still open
`Draft` PR. -->

- **Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)**
and your commits are signed.
- **Is this change backward compatible?**: Yes <!--- If No, explain why.
-->
- **Did you write any new necessary tests?**: Yes/No
- **Did you add or update any necessary documentation?**: Yes/No
- **Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**:
Not yet <!--- Only for new features, API changes, critical bug fixes or
bw breaking changes. -->

## Additional Information
<!-- E.g. related issue. -->


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

* **New Features**
* Added support for Nemotron-Parse multimodal models, including proper
device mapping, processor loading, and generation handling.

* **Improvements**
* Enhanced quantization robustness with safer handling of quantization
attributes and fallback logic.
* Improved model loading with better device placement and encoder buffer
management for vision-language models.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Zhiyu Cheng <zhiyuc@nvidia.com>
2026-02-13 02:18:08 +00:00
Keval Morabia 8d67cb08d5 Add Megatron-Bridge recipe-free distillation example script (#861)
## What does this PR do?

**Type of change:** New example script <!-- Use one of the following:
Bug fix, new feature, new example, new tests, documentation. -->

- [x] M-Bridge recipe-free distillation script so its more easier to run
and can support pruned models
- [x] Fix resuming distillation run

## Usage
<!-- You can potentially add a usage example below. -->

```python
torchrun --nproc_per_node 8 distill.py \
    --teacher_hf_path Qwen/Qwen3-8B \
    --student_hf_path Qwen3-8B-NAS-Pruned-6B \
    --tp_size 8 \
    --data_paths <climbmix 25% tokenized (~90B tokens)> \
    --data_path_to_cache /path/to/cache/climbmix_dataset_indices_qwen3 \
    --seq_length 4096 \
    --mbs 8 \
    --gbs 768 \
    --train_iters 28500 \
    --lr 1e-4 \
    --min_lr 1e-5 \
    --lr_warmup_iters 100 \
    --eval_interval 500 \
    --eval_iters 32 \
    --log_interval 10 \
    --output_dir qwen3_8b_6b_mbridge_distill
```

## Testing
<!-- Mention how have you tested your change if applicable. -->

- [x] Re-ran Qwen3 8B -> 6B experiments and compare with Nemo2 results
from blog

Best subnet from NAS: `{'num_layers': 30, 'hidden_size': 3584,
'ffn_hidden_size': 11776} -> 5.99B params, 0.5718 score`

| Model | MMLU | GSM8K - flexible, strict | MBPP (coding) |
| ------- | ------ | ------- | ------- |
| Qwen3-8B | 74.9 | 87.5, 84.6 | 65.4 |
| Qwen3-8B-Pruned-6B | 57.6 | 11.6, 10.0 | 4.8 |
| Qwen3-8B-Pruned-6B (Distilled for 16k steps i.e. 50B tokens ~3k GPU
hours) | 71.6 | 78.0, 64.7 | 43.4 |
| Qwen3-8B-Pruned-6B (Distilled for 28.5k steps i.e. 90B tokens ~5.2k
GPU hours) | 71.9 | 78.1, 64.8 | 44.2 |
| Qwen3-4B | 70.0 | 81.1, 84.7 | 62.8 |

Previous Nemo2 experiments on depth pruned Qwen3 8B -> 6B (24 layers)
had MMLU ~72.0 so more or less similar. No hparam tuning done for
current M-Bridge distillation run

- [ ] (Separate PR) GitHub CI/CD test for example script with NeMo 26.02
container

## Before your PR is "*Ready for review*"
<!-- If you haven't finished some of the above items you can still open
`Draft` PR. -->

- **Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)**
and your commits are signed.
- **Is this change backward compatible?**: Yes <!--- If No, explain why.
-->
- **Did you write any new necessary tests?**: N/A
- **Did you add or update any necessary documentation?**: Yes
- **Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**:
Yes <!--- Only for new features, API changes, critical bug fixes or bw
breaking changes. -->


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

## Release Notes

* **New Features**
* Added complete distillation workflow and example for Megatron-Bridge
optimization.

* **Documentation**
* Enhanced setup guide with Docker workflows, data preparation steps,
and detailed distillation instructions.
  * Improved usage documentation and help references.

* **Improvements**
* Better data preprocessing output with human-readable formatting for
metrics.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
2026-02-11 22:29:51 +00:00
Frida Hou eb3e6edb50 [fix][5875912] Fix autoquant-autodeploy example (#878)
## What does this PR do?

**Type of change:** Bug fix <!-- Use one of the following: Bug fix, new
feature, new example, new tests, documentation. -->

**Overview:** ?
Please check Bug ticket

## Usage
<!-- You can potentially add a usage example below. -->

```python
# Add a code snippet demonstrating how to use this
```

## Testing
Tested with
```
./scripts/run_auto_quant_and_deploy.sh     --hf_ckpt ./models/Qwen/Qwen3-8B     --save_quantized_ckpt ./qwen3_8B_autoquant     --quant fp8     --effective_bits 10.0
```

## Before your PR is "*Ready for review*"
<!-- If you haven't finished some of the above items you can still open
`Draft` PR. -->

- **Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)**
and your commits are signed.
- **Is this change backward compatible?**: Yes/No <!--- If No, explain
why. -->
- **Did you write any new necessary tests?**: Yes/No
- **Did you add or update any necessary documentation?**: Yes/No
- **Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**:
Yes/No <!--- Only for new features, API changes, critical bug fixes or
bw breaking changes. -->

## Additional Information
<!-- E.g. related issue. -->


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

* **Refactor**
* Simplified LLM initialization by removing intermediate configuration
layer
  * Updated attention backend from triton to flashinfer

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
2026-02-11 20:33:08 +00:00
jingyu-ml 10efcb65f6 [3.1/4] Diffusion Quantized ckpt export - WAN 2.2 14B (#855)
## What does this PR do?

**Type of change:** documentation <!-- Use one of the following: Bug
fix, new feature, new example, new tests, documentation. -->

**Overview:** 

1. Added multi‑backbone support for quantization: --backbone now accepts
space- or comma-separated lists and resolves to a list of backbone
modules.
2. Introduced PipelineManager.iter_backbones() to iterate named backbone
modules and updated get_backbone() to return a single module or a
ModuleList for multi‑backbone.
3. Updated ExportManager to save/restore per‑backbone checkpoints when a
directory is provided, with {backbone_name}.pt files, and to create
target directories when missing.
4. Simplified save_checkpoint() calls to rely on the registered
pipeline_manager by default.


**Usage: **

```bash
python quantize.py --model wan2.2-t2v-14b --format fp4 --batch-size 1 --calib-size 32 \
    --n-steps 30 --backbone transformer transformer_2 --model-dtype BFloat16 \
    --quantized-torch-ckpt-save-path ./wan22_mo_ckpts \
    --hf-ckpt-dir ./wan2.2-t2v-14b 
```

Plans

- [x] [1/4] Add the basic functionalities to support limited image
models with NVFP4 + FP8, with some refactoring on the previous LLM code
and the diffusers example. PIC: @jingyu-ml
- [x] [2/4] Add support to more video gen models. PIC: @jingyu-ml 
- [x] [3/4] Add test cases, refactor on the doc, and all related README.
PIC: @jingyu-ml
- [ ] [4/4] Add the final support to ComfyUI. PIC @jingyu-ml 

## Testing
<!-- Mention how have you tested your change if applicable. -->

## Before your PR is "*Ready for review*"
<!-- If you haven't finished some of the above items you can still open
`Draft` PR. -->

- **Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)**
and your commits are signed.
- **Is this change backward compatible?**: No <!--- If No, explain why.
-->
- **Did you write any new necessary tests?**:No
- **Did you add or update any necessary documentation?**: Yes
- **Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**:
Yes <!--- Only for new features, API changes, critical bug fixes or bw
breaking changes. -->

## Additional Information
<!-- E.g. related issue. -->


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

## Release Notes

* **New Features**
* Unified Hugging Face export support for diffusers pipelines and
components
  * LTX-2 and Wan2.2 (T2V) support in diffusers quantization workflow
* Comprehensive ONNX export and TensorRT engine build documentation for
diffusion models

* **Documentation**
* Updated to clarify support for both transformers and diffusers models
in unified export API
* Expanded diffusers examples with LoRA fusion guidance and additional
model options (Flux, SD3, SDXL variants)
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
2026-02-10 22:34:12 -08:00
Chenjie LuoandZhiyu 5e43b2a5f5 Support Qwen3 Next MTP load and export (#860)
## What does this PR do?

Fix MTP export for Qwen3 Next

**Overview:** ?

For Qwen3 next, the MTP weights are not stored separately in
safetensors. So we use "mtp" weights key to decide if the weights are
for MTP or not.


## Testing
Qwen3 Next PTQ and check if MTP is in the exported checkpoint.

scripts/huggingface_example.sh --model
<Qwen3-Next-80B-A3B-Instruct/Thinking> --quant nvfp4 --trust_remote_code

## Before your PR is "*Ready for review*"
<!-- If you haven't finished some of the above items you can still open
`Draft` PR. -->

- **Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)**
and your commits are signed.
- **Is this change backward compatible?**: Yes/No <!--- If No, explain
why. -->
- **Did you write any new necessary tests?**: Yes/No
- **Did you add or update any necessary documentation?**: Yes/No
- **Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**:
Yes/No <!--- Only for new features, API changes, critical bug fixes or
bw breaking changes. -->

## Additional Information
<!-- E.g. related issue. -->


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

* **Refactor**
* Optimized Multi-Token Prediction weight loading with improved layer
detection and handling.

* **Chores**
* Simplified status reporting to display total loaded weights and
detected layers.
  * Removed verbose per-file warnings for cleaner console output.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
Co-authored-by: Zhiyu <zhiyuc@nvidia.com>
2026-02-09 22:48:15 +00:00
yeyu-nvidia a8f5314c93 fix the path change in torch v2.10 for spec dec (#863)
## What does this PR do?

**Type of change:** 
bug fix

**Overview:** 
torch v2.10 changes the path for _SDPAMerger. will need to use the new
path for import

## Usage
<!-- You can potentially add a usage example below. -->

```python
# Add a code snippet demonstrating how to use this
```

## Testing
<!-- Mention how have you tested your change if applicable. -->

## Before your PR is "*Ready for review*"
<!-- If you haven't finished some of the above items you can still open
`Draft` PR. -->

- **Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)**
and your commits are signed.
- **Is this change backward compatible?**: Yes/No <!--- If No, explain
why. -->
- **Did you write any new necessary tests?**: Yes/No
- **Did you add or update any necessary documentation?**: Yes/No
- **Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**:
Yes/No <!--- Only for new features, API changes, critical bug fixes or
bw breaking changes. -->

## Additional Information
<!-- E.g. related issue. -->


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

* **Chores**
* Updated internal import references to reflect organizational changes
in dependencies.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Ye Yu <yeyu@nvidia.com>
2026-02-09 17:56:53 +00:00
Zhiyu 452c5a09b0 GLM-4.7 MTP support (#792)
## What does this PR do?

**Type of change:** ? <!-- Use one of the following: Bug fix, new
feature, new example, new tests, documentation. -->

**Overview:** Enable GLM-4.7 PTQ workflow, including loading the
standalone MTP modules and export as-is.

## Usage
<!-- You can potentially add a usage example below. -->

```python
python3 hf_ptq.py --pyt_ckpt_path /home/omniml_data_3/models/GLM-4.7 --qformat nvfp4_mlp_only --export_path /home/omniml_data_3/zhiyuc/checkpoints/GLM-4.7-NVFP4-0203 --trust_remote_code
```

## Testing
<!-- Mention how have you tested your change if applicable. -->

## Before your PR is "*Ready for review*"
<!-- If you haven't finished some of the above items you can still open
`Draft` PR. -->

- **Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)**
and your commits are signed.
- **Is this change backward compatible?**: Yes <!--- If No, explain why.
-->
- **Did you write any new necessary tests?**: Yes/No
- **Did you add or update any necessary documentation?**: Yes/No
- **Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**:
Yes <!--- Only for new features, API changes, critical bug fixes or bw
breaking changes. -->

## Additional Information
<!-- E.g. related issue. -->


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

* **New Features**
* Added quantization support for GLM-4.7 model with automatic handling
of specialized layer architecture.
* Added image-text data calibration capabilities for Nemotron VL model
quantization.

* **Documentation**
* Updated support matrix to reflect newly supported models and
quantization features.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Zhiyu Cheng <zhiyuc@nvidia.com>
2026-02-04 15:43:24 -08:00
Keval Morabia 944dd1a284 Move parallel_state init and warnings to Quant DynamicModule + MBridge pruning doc update (#854)
## What does this PR do?

**Type of change:** Minor improvement <!-- Use one of the following: Bug
fix, new feature, new example, new tests, documentation. -->

Only quantization DynamicModules use the parallel_state attribute so for
all other model opt methods, we see a parallel state not initialized
warning which could be confusing hence moving it to QuantModule class
instead

Minor update to MBridge pruning docs

## Testing
<!-- Mention how have you tested your change if applicable. -->

N/A

## Before your PR is "*Ready for review*"
<!-- If you haven't finished some of the above items you can still open
`Draft` PR. -->

- **Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)**
and your commits are signed.
- **Is this change backward compatible?**: Yes <!--- If No, explain why.
-->
- **Did you write any new necessary tests?**: Yes
- **Did you add or update any necessary documentation?**: Yes
- **Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**:
Yes <!--- Only for new features, API changes, critical bug fixes or bw
breaking changes. -->

## Additional Information
<!-- E.g. related issue. -->


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

* **New Features**
* Distributed parallel state support is now available in quantization
workflows for multi-GPU training.

* **Bug Fixes**
* Improved resource cleanup in distributed training to ensure proper
environment finalization.

* **Documentation**
* Updated example paths and added new manual pruning configuration
examples.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
2026-02-04 22:23:21 +00:00
jingyu-ml 2e43c80609 [2/4] Diffusion Quantized ckpt export (#810)
## What does this PR do?

**Type of change:** New feature <!-- Use one of the following: Bug fix,
new feature, new example, new tests, documentation. -->

**Overview:** 

This MR adds HuggingFace checkpoint export support for LTX‑2 by treating
TI2VidTwoStagesPipeline as a diffusion-like pipeline, exporting only the
stage‑1 transformer (with QKV-fusion-enabled dummy inputs) and falling
back to writing model.safetensors when save_pretrained isn’t available.
It also preserves the original forward in DynamicModule patching
(_forward_pre_dm) so downstream callers can still invoke the pre-patched
forward implementation.

**Changes**

1. Added the calibration & quantization support of the LTX2, even with
FP8 precision.
2. Preserve original forward before `DynamicModule` patching: when
patching forward, we now stash the pre-patched implementation in
`self._forward_pre_dm` (once) so downstream code can still call the
original forward, then re-bind forward to the class implementation. This
is needed for the LTX2 FP8 calibration.
3. Added LTX‑2 HF export path: `export_hf_checkpoint()` now also treats
ltx_pipelines.ti2vid_two_stages.TI2VidTwoStagesPipeline as a
“diffusion-like” object and routes it through
_export_diffusers_checkpoint() (import guarded; no hard dependency).
4. Generalized component discovery: introduced
get_diffusion_components() (aliasing the old get_diffusers_components)
to support non-diffusers pipelines; for LTX‑2 it returns only
stage_1_transformer.
5. Enabled QKV fusion for LTX‑2 backbone: added a model-aware dummy
forward generator (generate_diffusion_dummy_forward_fn) that builds
minimal LTX Modality inputs (including correct timesteps broadcasting)
so shared-input hooks can run and fuse QKV when applicable.
6. Export fallback for non-save_pretrained modules: when a component
lacks save_pretrained (LTX‑2 transformer), export now writes
model.safetensors + minimal config.json instead of pytorch_model.bin.

Plans

- [x] [1/4] Add the basic functionalities to support limited image
models with NVFP4 + FP8, with some refactoring on the previous LLM code
and the diffusers example. PIC: @jingyu-ml
- [x] [2/4] Add support to more video gen models. PIC: @jingyu-ml 
- [ ] [3/4] Add test cases, refactor on the doc, and all related README.
PIC: @jingyu-ml
- [ ] [4/4] Add the final support to ComfyUI. PIC @jingyu-ml 

## Usage
<!-- You can potentially add a usage example below. -->
```bash
python quantize.py --model ltx-2 --format fp4 --batch-size 64 --calib-size 1 --n-steps 40 --extra-param checkpoint_path=/home/scratch.omniml_data_2/jingyux/models/LTX-2/ltx-2-19b-dev-fp8.safetensors --extra-param distilled_lora_path=/home/scratch.omniml_data_2/jingyux/models/LTX-2/ltx-2-19b-distilled-lora-384.safetensors --extra-param spatial_upsampler_path=/home/scratch.omniml_data_2/jingyux/models/LTX-2/ltx-2-spatial-upscaler-x2-1.0.safetensors --extra-param gemma_root=/home/scratch.omniml_data_2/jingyux/models/LTX-2/gemma-3-12b-it-qat-q4_0-unquantized --extra-param fp8transformer=true --hf-ckpt-dir ./ltx2-nvfp4
```

## Testing
<!-- Mention how have you tested your change if applicable. -->

## Before your PR is "*Ready for review*"
<!-- If you haven't finished some of the above items you can still open
`Draft` PR. -->

- **Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)**
and your commits are signed.
- **Is this change backward compatible?**: Yes <!--- If No, explain why.
-->
- **Did you write any new necessary tests?**:No
- **Did you add or update any necessary documentation?**:No
- **Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**:No
<!--- Only for new features, API changes, critical bug fixes or bw
breaking changes. -->

## Additional Information
<!-- E.g. related issue. -->

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

## Release Notes

* **New Features**
* Added LTX-2 video model support with complete quantization and export
pipeline integration
* Introduced `--extra-param` CLI option for flexible model configuration
and parameter passing
* Enhanced export capabilities with broader diffusion model
compatibility

* **Chores**
* Changed default model data type from Half to BFloat16 for improved
numerical stability

<sub>✏️ Tip: You can customize this high-level summary in your review
settings.</sub>

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
2026-02-04 10:41:33 +00:00