mirror of
https://github.com/NVIDIA/Model-Optimizer.git
synced 2026-10-19 05:49:19 +08:00
0.44.0dev
197
Commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
 Ajinkya Rasane
|
e4df91bf04 |
OMNIML-2663] Replace modelopt FP8 QDQ nodes with native ONNX QDQ nodes (#852)
## What does this PR do? **Type of change:** New feature **Overview:** - Updated FP8 quant exporter to replace modelopt custom QDQ nodes with native ONNX QDQ nodes - Updated get_onnx_bytes_and_metadata to make convert_float_to_float16() default instead of autocast - Created util functions to fix graph structure after conversion ## Testing ``` python torch_quant_to_onnx.py --quantize_mode=fp8 \ --onnx_save_path=<model_path> \ --calibration_data_size 64 \ --batch_size 128 python evaluate.py --onnx_path=<model_path> \ --model_name=vit_base_patch16_224 \ --results_path=./results.txt \ --batch_size 128 ``` Results: Before replacement: ``` The top1 accuracy of the model is 85.06% The top5 accuracy of the model is 97.558% Inference latency of the model is 5.27963 ms ``` After replacement: ``` The top1 accuracy of the model is 85.054% The top5 accuracy of the model is 97.542% Inference latency of the model is 5.74771 ms ``` ## 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 - Replaced modelopt QDQ nodes with native ONNX qdq nodes - **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 <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * ONNX utilities to remove redundant Casts, fold Constant→Cast patterns, and convert targeted Casts to FP16. * **Improvements** * FP8 QDQ nodes now converted to native ONNX QDQ/Dequantize nodes for improved compatibility. * Export pipeline streamlined: consistent FP16 handling, unified weight quantization, cast cleanup ordering, and added logging for better traceability. * **Tests** * Unit tests updated to use the new ONNX utilities. * **Changelog** * Entry added noting FP8 QDQ → native ONNX QDQ conversion. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> |
||
|
sugunav14andrealAsma
|
beac6e9faa |
Sequential calibrate refactor (#982)
### What does this PR do?
Type of change: New feature <!-- Use one of the following: Bug fix, new
feature, new example, new tests, documentation. -->
The current sequential calibration support has O(N^2) complexity for
collecting updated activations for a decoder layer. To solve this, we
adopted a modular/plugin based approach which involves hooks to capture
the updated activations by running forward on the previous decoder layer
using cached prev layer activations. This leads to an issue with nested
modules i.e. the logic in the parent module might need to be replicated
in the lower level modules to ensure equivalence. For example, in the
nemotron model, the parent module NemotronHModel has logic to create and
select appropriate mask based on the decoder layer type (mamba vs
attention).
This PR implements a more generic solution for sequential calibration,
by choosing to collect activations using model forward, thereby ensuring
that all the parent module logic is preserved. We use an attribute
"state"on the modules to indicate whether to perform recomputation/skip
the layer while running module forward. This can help us avoid redundant
computations for getting updated activations.
The overall flow is as follows
1. The user must register a get_decoder_layers() function that returns a
list of layers to be calibrated sequentially
2. LayerActivationCollector, goes through the list of layers and patches
module forward with a "state aware" module forward
3. When model.forward() is called, all the parent logic is recomputed as
expected (embeddings, residual connections, generating attention mask
etc).
4. Lets say we are currently calibrating layer N and we want to get
updated activations; we set layer N to capture and layer N-1 to run
(because this layer was processed previously and updated activations
need to be generated). Already processed layers are set to skip. When
model.forward() is called, all the previous decoder layer computations
are skipped. Layer N-1 uses the cached inputs to generate new
activations. Layer N inputs are captured using the same logic as before
and cached so that they can be used to get updated activations for Layer
N+1.
### Usage
```python
# Sequential calibrate config
NVFP4_SEQUENTIAL_CFG = {
"quant_cfg": {
"*weight_quantizer": _nvfp4_quantizer,
"*input_quantizer": _nvfp4_quantizer,
**_default_disabled_quantizer_cfg,
},
"algorithm": {"method": "max", "use_sequential": True},
}
```
### 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?: ✅
- 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)?:
✅
- Did you write any new necessary tests?: ✅
- 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**
* Public sequential, per-layer calibration API and an
activation-collection utility.
* Broader model discovery support including Nemotron-H and homogeneous
HuggingFace variants.
* **Improvements**
* Clearer validation/error messages and deterministic
patching/unpatching with guaranteed cleanup and resource handling.
* Consolidated discovery/registration flow for decoder-layer handling
and improved per-layer logging/progress.
* **Tests**
* Extensive new unit tests covering discovery, per-layer capture/replay,
inter-layer behavior, and edge cases.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---------
Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
Signed-off-by: realAsma <akuriparambi@nvidia.com>
Co-authored-by: realAsma <akuriparambi@nvidia.com>
|
||
|
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>
|
||
|
Frida Hou
|
0f8482a982 |
[minor]: Fix AutoQuant Megatron test (#1040)
### What does this PR do? Type of change: Bug fix <!-- Use one of the following: Bug fix, new feature, new example, new tests, documentation. --> The `quantizer_states` are on different rank, so the tests are failing, while I can move them to the same device and compare, I think it's sufficient to just compare the search outcome. ### 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. --> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> |
||
|
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>
|
||
|
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> |
||
|
kinjalpatel27
|
58417e5014 |
Added block wise RHT (#1014)
### What does this PR do?
Added support for RHT with non-power of 2. Rotate quantization
configuration can be used to specify block_size as well.
```
NVFP4_KV_ROTATE_BLOCK_32_CFG = {
"quant_cfg": {
"*q_bmm_quantizer": {
"enable": False,
"rotate": {"enable": True, "block_size": 32},
},
"*k_bmm_quantizer": {
**_nvfp4_quantizer,
"rotate": {"enable": True, "block_size": 32},
},
"*v_bmm_quantizer": _nvfp4_quantizer,
},
"algorithm": "max",
}
```
### Testing
```
pytest tests/gpu/torch/quantization/test_hadamard.py -k test_hadamard_transform_block
```
### Before your PR is "*Ready for review*"
Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)
and your commits are signed (`git commit -s -S`).
Make sure you read and follow the [Security Best
Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors)
(e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(...,
weights_only=False)`, `pickle`, etc.).
- Is this change backward compatible?: ✅
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: N/A
- Did you write any new necessary tests?: ✅
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
✅
### Additional Information
<!-- E.g. related issue. -->
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit
* **New Features**
* Block-granular randomized Hadamard transform (RHT) for non-power-of-2
dimensions.
* Rotation configuration expanded to accept block-size and an option to
perform rotation in FP32; rotation settings are now exposed via the
public API.
* **Tests**
* Added tests validating block-granular RHT across varied dimensions and
block sizes.
* **Documentation**
* Changelog updated to mention the new block-granular RHT capability.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---------
Signed-off-by: Kinjal Patel <kinjalpravin@nvidia.com>
|
||
|
Keval Morabia
|
d0abca7981 |
Support megatron tokenization for post training datasets (#1018)
### What does this PR do? Update megatron_preprocess_data.py to support applying chat template for tokenizing chat based post training datasets <!-- 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. --> - Tokenized Nemotron-Post-Training-Dataset-v2 (~2B tokens for stem + chat + math + code splits) - Doing distillation on pruned nano v2 7B ### 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. --> <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Bug Fixes** * Added runtime logging for data truncation operations to enhance processing visibility * Improved handling of chat-formatted conversation data in list format * Eliminated duplicate log messages during data encoding operations <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> |
||
|
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> |
||
|
Gwena Cunha
|
26cad678d2 |
[OMNIML-3252][ONNX] Add real Q/DQ scales in Autotune (#951)
## What does this PR do?
**Type of change:** New feature
**Overview:** ONNX Autotune (also called Auto Q/DQ) is currently and
standalone feature of ModelOpt that automatically adds Q/DQ where
relevant according to information obtained from TensorRT inference. One
issue is that the scales in those Q/DQ nodes are random.
This PR does 2 major things:
1. Integrates Auto Q/DQ into the ONNX quantization workflow; and
2. Enables calibration data to be used to obtain the correct scales for
the Q/DQ nodes.
## Usage
```python
$ python -m modelopt.onnx.quantization --onnx_path=model.onnx --autotune={quick,default,extensive}
```
> Please see `__main__.py` for other args.
## Testing
1. Added unittest for Q/DQ node placement validation:
`tests/gpu/onnx/quantization/test_autotune_quantization_integration.py`
2. Verified that accuracy was recovered by integrating MOQ with
Autotune. Results on RTX 3090 with TRT 10.12.0.36 (`--stronglyTyped`)
with ViT, as per `examples/onnx_ptq`:
| Model | Top-1 acc | Top-5 acc |
|--------------------------|---------------|----------------|
| FP32 | 85.1% | 97.5% |
| FP16 (FP32 with --fp16) | 85.1% | 97.5% |
| Quant (MOQ) | 82.4% | 96.4% |
| Quant (Autotune) | 0.1% | 0.5%|
| Quant (MOQ + Autotune) | 79.6% | 95.0% |
Notice that accuracy was mostly recovered from standalone Autotune to
MOQ + Autotune (real Q/DQ scales). The drop in accuracy between MOQ and
MOQ + Autotune is likely due to some sensitive nodes being quantized,
such as `BiasAdd` (see bug 5916898).
## 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?**: No (will be
done in a different PR)
- **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. -->
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit
* **New Features**
* Autotuning added to ONNX quantization: CLI flags, presets, per-region
tuning, and FP8/INT8 support; accepts in-memory models and optional
output dirs; node-filter loading and explicit-flag CLI behavior.
* Activation-operation accessor exposed and autotune helpers added to
the package API.
* **Bug Fixes**
* Safer graph rewiring to avoid corrupting quantized graphs when targets
are absent.
* **Tests**
* New integration test and model helper validating autotune quantization
consistency.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
## Additional information
To reproduce accuracy with ViT, call `download_example_onnx.py` and
`image_prep.py` without `--fp16`.
If `--fp16` is used here, quantizing this model with `--autotune`
results in the following error:
```
[modelopt][onnx] - ERROR - Benchmark failed: Converting dtype('float16') to a ctypes type
```
This is fixed in https://github.com/NVIDIA/Model-Optimizer/pull/978.
---------
Signed-off-by: gcunhase <4861122+gcunhase@users.noreply.github.com>
|
||
|
realAsmaandClaude Opus 4.6
|
fe83270139 |
Refactor HF _QuantSparseMoe: config-driven token counting, NemotronH detection (#970)
## What does this PR do? **Type of change:** New feature **Overview:** Extend `_QuantSparseMoe` to support NemotronH-style MoE blocks (which use `n_routed_experts` instead of `num_experts`) and refactor the MoE calibration features to be config-driven and lazy-initialized. Key changes: - `_is_sparse_moe_block` in `plugins/huggingface.py` now accepts `n_routed_experts` (NemotronH pattern) in addition to `num_experts` - `_QuantSparseMoe` is refactored: token counting and forced expert forwarding are now opt-in via config knobs (`moe_calib_experts_ratio`, `moe_count_expert_calib_tokens`). When both are off (default), forward is a zero-overhead pass-through. - Token counting buffer and gate hook are lazy-initialized on first use instead of eagerly in `_setup` - `_QuantSparseMoe` gets `layer_sync_moe_local_experts_amax` to sync input quantizer amax across experts (same as Megatron path) - Extract shared `sync_moe_experts_input_amax` utility into `utils.py`, also fixing missing weight amax for experts that received no tokens during calibration. Megatron's `_MegatronSequentialMLP` now calls this shared utility. - `SequentialQuantizer` delegates `amax` property ## Testing - Updated and added unit tests in `test_sparse_moe.py` covering default config, lazy init, token counting, top_k restoration, and end-to-end quantize with both features enabled. ## Before your PR is "*Ready for review*" - **Is this change backward compatible?**: Yes - **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 Signed-off-by: realAsma <akuriparambi@nvidia.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> |
||
|
realAsmaandClaude Opus 4.6
|
a5d46ff12b |
Auto Quantize improvements and bug fixes for large sparse MoEs (#953)
## What does this PR do?
**Type of change:** New feature + Bug fixes
**Overview:**
Enable AutoQuantize for NemotronH and large SparseMoE models, and update
the FP8 workflow split between `mtq.auto_quantize` and `mtq.quantize`.
`mtq.auto_quantize` is now positioned as the lightweight search phase
(lite calibration + scoring), while `mtq.quantize` is used for
heavier/final calibration workflows (longer calibration passes,
force-all-token style MoE calibration, and advanced recipes such as
GPTQ, MSE, etc.).
### Algorithm & feature changes
- **NemotronH / SparseMoE support**: Updated `quant_module` and
`score_module` rules (should eventually move to the proposed modeling
lib). In future, this should be the only change needed to support new
models — the bug fixes below were unearthed while enabling NemotronH
- **Config generation**: Added
`mtq.get_auto_quantize_config(search_state, constraints=None,
verbose=False)` to re-solve from `search_state` and produce plain-dict
configs (no redundant `output_quantizer`), with optional verbose summary
- **FP8 workflow split**: Use lite calibration in `mtq.auto_quantize`,
then run longer/final calibration with `mtq.quantize` using the
generated config
- **Performance**: Pass `name_to_module` to
`enable_weight_access_and_writeback` to avoid O(N^2) overhead on large
MoE models
- **Calibration caching in checkpoint**: Save/restore quantizer
calibration states (metadata + state_dict) per recipe in the
AutoQuantize checkpoint, so resuming a search skips redundant
calibration
- **Per-rank distributed checkpointing**: When `torch.distributed` is
initialized, each rank saves/loads its own checkpoint file
(`search_state{rank}.pt`), with backward-compatible fallback to the
single-file path
### API updates
- **Config API naming**: Use `mtq.get_auto_quantize_config(...)` for
exporting the searched recipe into a quantize-ready config
- **Recommended usage pattern**:
```python
# 1) Lightweight search + lite calibration
model, search_state = mtq.auto_quantize(
model,
constraints={"effective_bits": 6.0},
quantization_formats=[mtq.NVFP4_DEFAULT_CFG, mtq.FP8_DEFAULT_CFG],
data_loader=data_loader,
forward_step=forward_step,
loss_func=loss_func,
num_calib_steps=64, # lite calibration during search
num_score_steps=128,
)
# 2) Export searched config (optionally re-solve constraints)
auto_quantize_config = mtq.get_auto_quantize_config(
search_state,
constraints={"effective_bits": 6.0},
verbose=True,
)
# 3) Final / longer calibration pass with quantize
model = mtq.quantize(
model,
config=auto_quantize_config,
forward_loop=long_calibration_loop, # e.g. force-all-token style MoE calibration
)
```
### Bug fixes
- Fixed `disabled_layers` handling so fused kernels (e.g. Mamba blocks)
are properly skipped
- Fixed gradient checkpointing to keep all modules except the
checkpointed modules in eval
- Fixed FP8 fake quant NaN/inf when `amax ≈ 0`
- Fixed `SequentialQuantizer.convert_to_single_quantizer` to operate on
`module` instead of `model`, avoiding O(N^2) CPU iteration on SparseMoE
models with 1000s of submodules
- Switched to proper `F.kl_div` for KL divergence scoring
### Not yet exposed to `llm_ptq`
`mtq.get_auto_quantize_config` is not yet wired into `llm_ptq`. The
plain config records per-expert quantization settings for all MoE
experts, resulting in large JSON files. For my experiments I used a
quick workaround. A follow-up PR will add a better config representation
and expose it to `llm_ptq`.
## Testing
- Tested on NemotronH-tiny and Nemotron-Super-RL models
- Verified auto_quantize scoring + config generation end-to-end
- Unit test for checkpoint resume verifies calibration cache correctness
(metadata + tensor values)
- Existing unit tests pass
## Before your PR is "*Ready for review*"
- **Is this change backward compatible?**: Yes
- **Did you write any new necessary tests?**: AutoQuantize end-to-end
requires GPU + large MoE models; verified manually on NemotronH-tiny and
Nemotron-Super-RL. Unit test coverage to follow.
- **Did you add or update any necessary documentation?**: No
- **Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**:
Yes
## Additional Information
Follow-up planned: expose `mtq.get_auto_quantize_config` to `llm_ptq`
with a compact config format for MoE models. AWQ support in AutoQuantize
can also be removed in a future PR to keep it lightweight.
---------
Signed-off-by: realAsma <akuriparambi@nvidia.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
|
||
|
|
695c8e8522 |
Integrate Automated QDQ placement tool - part 4.3 (#843)
## What does this PR do? This PR upload user guide of Automated QDQ placement tool. This tool automatically search QDQ insertion points with better performance. **Overview:** ? ## 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 * **Documentation** * Added comprehensive guide for Automated Q/DQ Placement Optimization workflow, including quick start instructions, advanced usage patterns, configuration options, best practices, and troubleshooting. * **New Features** * Exposed public API for CLI parser programmatic access. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Will Guo <willg@nvidia.com> Signed-off-by: Gwena Cunha <4861122+gcunhase@users.noreply.github.com> Co-authored-by: Gwena Cunha <4861122+gcunhase@users.noreply.github.com> Co-authored-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> |
||
|
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> |
||
|
Hrishith Thadicherla
|
0a84bb2caf |
Graph Surgery Framework for TRT-RTX (#992)
### What does this PR do?
Type of change: new feature
Adds the **ONNX Graph Surgery** framework — graph-level transformations
on exported ONNX models for optimized inference with ONNX Runtime.
**GQA Attention Replacement** — Replaces native attention subgraphs
(Q/K/V projections, RoPE, softmax, KV cache) with a single fused
`GroupQueryAttention` operator. Handles RoPE cache computation,
attention mask reformatting, Q/K/V weight fusion, and KV cache I/O
automatically from a HuggingFace model ID. Supports FP16/BF16, INT4/AWQ
quantized, and combined QKV models.
**DequantizeLinear Weight Transpose** — Transposes quantized weights in
`DequantizeLinear` nodes to column-major layout for providers like
NvTensorRtRtx. Handles INT4/UINT4 packed formats.
**Whisper Encoder Cross-Attention KV** — Adds cross-attention K/V
projection outputs to the Whisper encoder for the ONNX Runtime GenAI
pipeline. Loads cross-attention weights from HuggingFace and generates
`genai_config.json`.
### Usage
```bash
# GQA attention replacement
python -m modelopt.onnx.graph_surgery replace-gqa \
-i model.onnx -o model_gqa.onnx \
-m meta-llama/Llama-3.2-1B \
--max-seq-len 4096 --dtype float16
# DequantizeLinear weight transpose
python -m modelopt.onnx.graph_surgery transpose-dq \
-i model_quantized.onnx -o model_transposed.onnx
# Whisper encoder cross-attention KV
python -m modelopt.onnx.graph_surgery add-cross-kv \
-i encoder_model.onnx -o encoder_with_kv.onnx \
-m openai/whisper-large-v3-turbo
```
### Testing
Added test case for GQA graph surgery at
`tests/unit/onnx/test_gqa_graph_surgery.py`. Builds a toy attention
subgraph with native ops matching the real Optimum export pattern and
applies GQA surgery on it. Compared end outputs of both models — outputs
match exactly.
```
tests/unit/onnx/test_gqa_graph_surgery.py::TestGQAGraphSurgery::test_gqa_node_exists PASSED
tests/unit/onnx/test_gqa_graph_surgery.py::TestGQAGraphSurgery::test_gqa_attributes PASSED
tests/unit/onnx/test_gqa_graph_surgery.py::TestGQAGraphSurgery::test_node_count_reduced PASSED
tests/unit/onnx/test_gqa_graph_surgery.py::TestGQAGraphSurgery::test_rotary_emb_nodes_removed PASSED
tests/unit/onnx/test_gqa_graph_surgery.py::TestGQAGraphSurgery::test_position_ids_removed PASSED
tests/unit/onnx/test_gqa_graph_surgery.py::TestGQAGraphSurgery::test_logits_match
Original nodes: 106 -> GQA nodes: 13
Logits shape: (1, 4, 64)
Original[0,:4]: [-2906. 7704. 15248. 8488.]
GQA [0,:4]: [-2906. 7704. 15248. 8488.]
Max abs diff: 0.000000
Mean abs diff: 0.000000
PASSED
```
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit
* **New Features**
* ONNX graph-surgery toolkit: new utilities to replace attention with
GQA, add encoder cross-attention KV outputs, convert FP16→BF16, and
transpose DequantizeLinear weights.
* New CLI with subcommands to run the above transformations.
* New helper utilities for graph manipulation, RoPE cache generation,
and Whisper GenAI config creation.
* **Tests**
* Added unit tests validating GQA surgery and DQ-transpose
transformations.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---------
Signed-off-by: Hrishith Thadicherla <hthadicherla@nvidia.com>
|
||
|
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> |
||
|
Gwena Cunha
|
0ad287ca7b |
[ONNX][Autotune] Replace CUDA memory management from CUDART to PyTorch (#998)
### What does this PR do? **Type of change**: Bug fix **Overview**: Replace CUDA memory management from CUDART to PyTorch (higher-level API). ### Usage ```python # Add a code snippet demonstrating how to use this ``` ### Testing 1. Added unittests. 2. Tested that this PR does not break https://github.com/NVIDIA/Model-Optimizer/pull/951 or https://github.com/NVIDIA/Model-Optimizer/pull/978 ### 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?: ✅ - 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?: ✅ <!--- 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 Summary of changes in `benchmark.py — TensorRTPyBenchmark`: | What changed | Before | After | |---|---|---| | Imports | `contextlib` + `from cuda.bindings import runtime as cudart` | `import torch` (conditional) | | Availability flag | `CUDART_AVAILABLE` | `TORCH_CUDA_AVAILABLE = torch.cuda.is_available()` | | `__init__` guard | checks `CUDART_AVAILABLE or cudart is None` | checks `TORCH_CUDA_AVAILABLE` | | `_alloc_pinned_host` | `cudaMallocHost` + ctypes address hack, returns `(ptr, arr, err)` | `torch.empty(...).pin_memory()`, returns `(tensor, tensor.numpy())` | | `_free_buffers` | `cudaFreeHost` + `cudaFree` per buffer | `bufs.clear()` — PyTorch GC handles deallocation | | `_allocate_buffers` | raw `device_ptr` integers, error-code returns | `torch.empty(..., device="cuda")`, `tensor.data_ptr()` for TRT address | | `_run_warmup` | `cudaMemcpyAsync` + `cudaStreamSynchronize` | `tensor.copy_(non_blocking=True)` inside `torch.cuda.stream()` | | `_run_timing` | same cudart pattern | same torch pattern | | `run` — stream lifecycle | `cudaStreamCreate()` / `cudaStreamDestroy()` | `torch.cuda.Stream()` / `del stream` | | `run` — stream arg to TRT | raw integer handle | `stream.cuda_stream` (integer property) | | Error handling | `cudaError_t` return codes | PyTorch raises `RuntimeError`, caught by existing `except Exception` | Related to https://github.com/NVIDIA/Model-Optimizer/pull/961 <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Refactor** * TensorRT benchmarking migrated from direct CUDA runtime calls to PyTorch CUDA tensors, pinned memory, and CUDA stream primitives — simplifying buffer management, transfers, and timing semantics. * **Tests** * Expanded GPU autotune benchmark tests with broader unit and integration coverage for CUDA/TensorRT paths, pinned-host/device buffering, stream behavior, warmup/timing, and end-to-end latency scenarios. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: gcunhase <4861122+gcunhase@users.noreply.github.com> |
||
|
|
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> |
||
|
willg-nvandKeval Morabia
|
6d77ce754f |
Integrate Automated QDQ placement tool - part 4.4 (#961)
### What does this PR do? Many minor changes: 1. Add preset mode to AutoQDQ. 2. Add pattern cache tests. 3. increase batch size for stable QDQ insertion 4. update LICENSE 2024 -> 2026. 5. add cuda-python to pyproject.toml for `[onnx]` ### 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 mode presets (quick, default, extensive) with a new --mode option for autotuning. * Introduced AutoQDQ: automated Q/DQ placement tool for ONNX quantization with pattern caching and checkpoint/resume. * **Documentation** * Updated CLI help and examples to show mode usage and override semantics. * **Tests** * Added comprehensive tests for pattern cache and mode-presets/explicit-override behavior; re-enabled a GPU autotuning workflow test; minor test updates. * **Chores** * Added "cuda-python" to optional dependencies. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Will Guo <willg@nvidia.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Co-authored-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> |
||
|
Wei-Ming Chen
|
22423041f7 |
API to measure MSE for target quantizers (#940)
## 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:** add an API to measure MSE for target quantizers given a
forward loop
## Usage
<!-- You can potentially add a usage example below. -->
```python
# 1. Quantize the model as usual
model = mtq.quantize(model, quant_cfg, forward_loop)
# 2. Compute MSE for all quantizers
mse = mtq.compute_quantization_mse(model, forward_loop)
# 3. Print the top-5 noisiest quantizers
for name, err in sorted(mse.items(), key=lambda x: -x[1])[:5]:
print(f"{name}: {err:.4e}")
```
## Testing
<!-- Mention how have you tested your change if applicable. -->
Unit test and test with HF PTQ
## 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 an API to measure per-quantizer mean-squared error (MSE) between
original and fake-quantized tensors; supports wildcard and callable
filtering, skips disabled/non-fake-quant quantizers, and runs safely
under no-grad.
* **Tests**
* Added comprehensive tests for MSE validity, pattern and callable
filtering, union behavior, exclusion of disabled quantizers,
preservation of model state, and forward-hook cleanup.
* **Documentation**
* Updated changelog to document the new MSE measurement API.
<!-- 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>
|
||
|
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> |
||
|
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> |
||
|
noeyy-mino
|
31f0783e35 |
add new cases for ckpts on HF (#938)
## What does this PR do? **Type of change:** new tests **Overview:** Add new deployment tests for newly added checkpoints on HF ## Usage pytest test_deploy.py --run-release ```python None ``` ## Testing None ## 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?**: 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 ## Additional Information N/A <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added support for new model deployments: Qwen3 variants (480B, 235B, 397B), Kimi K2.5, and additional Mixtral configurations. * **Improvements** * Expanded backend support for select models, including TensorRT-LLM integration for Nemotron models. * **Chores** * Renamed environment variable from MODELOPT_LOCAL_MODEL_ROOT to MODELOPT_LOCAL_EAGLE_MODEL for improved clarity. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: noeyy-mino <174223378+noeyy-mino@users.noreply.github.com> |
||
|
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> |
||
|
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> |
||
|
Keval Morabia
|
a076e6c218 |
Make GPU tests >2x faster by reusing spawn processes between tests (#958)
### What does this PR do? Type of change: Test infra improvement <!-- Use one of the following: Bug fix, new feature, new example, new tests, documentation. --> * Instead of spawning multiprocesses per test case (adds ~10s overhead x 150 tests), we spawn once per file (gpu_megatron) or session (gpu) to make the tests much much faster * Remove some unused / unnecessary tests * Fix and enable Minitron NAS test on Hybrid models for Mcore 0.16 Test time (`pytest tests/<test_type>`) on 2x RTXPro 6000 Blackwell: https://github.com/NVIDIA/Model-Optimizer/actions/runs/22625069275 | | gpu | gpu_megatron | |------------|-----|--------------| | **Before** | 27m | 40m | | **After** | 16m | 19m | Full PR-merge CI/CD now only takes ~45mins ### Testing <!-- Mention how have you tested your change if applicable. --> Tested on 1,2,4,8 GPU setup ### 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?: ✅ <!--- 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?: ✅ <!--- 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. --> <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Chores** * Reworked GPU test infra to use a persistent distributed worker pool with improved per-worker teardown, aggregated error reporting, Megatron-specific cleanup, reduced GPU CI timeouts, and a small build-system dependency change. * **Tests** * Converted many multi-process GPU tests to a fixture-driven distributed runner, updating test entry points and orchestration while preserving test logic; adjusted imports and conditional guards for several GPU tests. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> |
||
|
Keval Morabia
|
ba29ad706d |
Enable RTX Pro 6000 Blackwell runners for CI/CD (#944)
## What does this PR do? **Type of change:** CI/CD Improvement <!-- Use one of the following: Bug fix, new feature, new example, new tests, documentation. --> Updated CI/CD test matrix (new `cuda13-gpu-trtllm` dedicated job for gpu tests on trtllm container) | Workflow | Trigger | Test Matrix | GPU Runner | |---|---|---|---| | **GPU tests** | PR | `cuda13-gpu`, `cuda13-gpu-megatron`, `cuda13-gpu-trtllm` | 1x RTX Pro 6000 | | **GPU tests** | Nightly | `cuda13-gpu`, `cuda13-gpu-megatron`, `cuda13-gpu-trtllm` | 2x RTX Pro 6000 | | **Example tests (torch)** | PR | `llm_distill`, `llm_qat`, `llm_sparsity`, `speculative_decoding` | 1x H100 | | **Example tests (torch)** | Nightly | `llm_distill`, `llm_qat`, `llm_sparsity`, `speculative_decoding` | 2x RTX Pro 6000 | | **Example tests (trtllm)** | PR | `llm_ptq`, `vlm_ptq` | 1x RTX Pro 6000 | | **Example tests (trtllm)** | Nightly | `llm_autodeploy`, `llm_eval`, `llm_ptq`, `vlm_ptq` | 2x RTX Pro 6000 | | **Example tests (onnx)** | PR | `diffusers`, `torch_onnx` | 1x L4 | | **Example tests (onnx)** | Nightly | `diffusers`, `torch_onnx` | 2x RTX Pro 6000 | ## Testing <!-- Mention how have you tested your change if applicable. --> - [ ] Per PR tests pass in this PR - [ ] Nightly tests manually triggered: https://github.com/NVIDIA/Model-Optimizer/actions/runs/22495679199 (GPU tests), ? (Example tests) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Chores** * Unified and simplified CI test matrices and reduced duplicate workflow configuration. * Updated GPU runner targets and container images for test jobs to newer GPU types. * **Tests** * Added new GPU test variants and several new test modules. * Simplified test gating: megatron auto-skip removed; tests now use a single dependency check for the mamba provider. * **Tooling** * Added a new tox environment for an additional GPU test suite. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> |
||
|
willg-nv
|
94b2cb272d |
Integrate Automated QDQ placement tool - part 3.3 (#839)
## What does this PR do?
This PR implements QDQ autotuner CLI. This is the initial version of
CLI, it will be integrated to modelopt.onnx.quantization.autotune.
Usage:
```
python -m modelopt.onnx.quantization.autotune
--onnx_path model.onnx --schemes_per_region 50
--pattern_cache cache.yaml --qdq_baseline baseline.onnx
--quant_type int8 --verbose
```
PR 3.1: https://github.com/NVIDIA/Model-Optimizer/pull/837
PR 3.2 https://github.com/NVIDIA/Model-Optimizer/pull/838
PR 3.3: https://github.com/NVIDIA/Model-Optimizer/pull/839
**Overview:** ?
## 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?**: Document will
be added in part 4.
- **Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**:
CHANGE log will be added in part 4.
## 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 a command-line interface for ONNX quantization autotuning with
configurable parameters for models, output paths, quantization
strategies, and TensorRT benchmarking.
* Introduced an automated workflow for pattern-based region optimization
with state management, baseline comparison, and benchmarking
capabilities.
<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: Will Guo <willg@nvidia.com>
|
||
|
sugunav14
|
fcdaf6519d |
Support decoder block-level sequential calibration (#924)
## 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:** Add support for sequential calibration of layers (at
decoder level granularity) in ModelOpt.
Calibration flow
1. Get list of decoder blocks
2. For current block call get input activations (considering weight and
activation QDQ from all other previous blocks) and call specified
calibration function.
functions added
1. get_decoder_layers() -> to detect and get list of blocks to iterate
over
2. LayerActivationCollector class -> to get input activations to the
layer
3. sequential_calibrate() -> to perform the described calibration flow
4. use_sequential field in QuantizeAlgorithmConfig
## Usage
<!-- You can potentially add a usage example below. -->
```python
# Sample config
NVFP4_DEFAULT_CFG = {
"quant_cfg": {
"*weight_quantizer": {
"num_bits": (2, 1),
"block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)},
"axis": None,
"enable": True,
},
"*input_quantizer": {
"num_bits": (2, 1),
"block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)},
"axis": None,
"enable": True,
},
**_default_disabled_quantizer_cfg,
},
"algorithm": {
"method": "max",
"use_sequential": True,
}
```
Set use_sequential=True in QUANT_CFG's "algorithm" section.
## 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**
* Sequential layer-by-layer calibration: Quantization now supports
processing decoder layers sequentially to improve memory efficiency on
large models.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---------
Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
|
||
|
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>
|
||
|
willg-nv
|
35e60991c7 |
Integrate Automated QDQ autotuner - part 3.2 (#838)
## What does this PR do? This PR implements QDQAutotuner class. This class is used to drive the main Autotuner workflow. The workflow is: 1. uses RegionSearch to build regions 2. generate QDQ ONNX models and evaluate perf 3. save best model This PR is part 2/4 of https://github.com/NVIDIA/Model-Optimizer/pull/703. PR 3.1: https://github.com/NVIDIA/Model-Optimizer/pull/837 PR 3.2 https://github.com/NVIDIA/Model-Optimizer/pull/838 PR 3.3: https://github.com/NVIDIA/Model-Optimizer/pull/839 **Overview:** ? ## 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?**: Not in this part. - **Did you add or update any necessary documentation?**: No, document will be updated in part 4. - **Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**: No, change log will be updated when all changes are ready. ## Additional Information <!-- E.g. related issue. --> <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Introduced ONNX Q/DQ autotuning framework with automatic region discovery and pattern-based optimization. * Added model profiling and quantization scheme generation capabilities. * Enabled state persistence and quantization model export functionality. * Introduced configuration management for quantization parameters and profiling workflows. <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: Will Guo <willg@nvidia.com> |
||
|
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> |
||
|
Gwena CunhaandKeval Morabia
|
3fe7e65b70 |
[5615343,5597780,5371126] Upgrade ORT to 1.24 (#928)
## What does this PR do? **Type of change:** Bug fix **Overview:** Upgrade ORT to 1.24.x to fix various bugs (5615343, 5597780, 5371126). TODO: Verify no regressions by bumping ORT once @ajrasane is back ## Usage See each bug. ## Testing See each bug. ## 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)?**: 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 ## Bug Fixes * Upgraded ONNX Runtime to version 1.24.2, addressing multiple reported issues * Updated system requirements and installation documentation to reflect the new ONNX Runtime version <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: gcunhase <4861122+gcunhase@users.noreply.github.com> Signed-off-by: Gwena Cunha <4861122+gcunhase@users.noreply.github.com> Co-authored-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> |
||
|
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> |
||
|
Jenny Chen
|
6b0ea4d43d |
Fix MOE layer sync test (#936)
## What does this PR do? **Type of change:** Bug Fix **Overview:** ? Fix MOE layer sync test by initializing weights in MOE layer differently [Link to bug](https://github.com/NVIDIA/Model-Optimizer/actions/runs/22288772586/job/64472124733#step:7:851) ## 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. --> https://github.com/NVIDIA/Model-Optimizer/actions/runs/22414958311 ## 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 * **Tests** * Enhanced quantization testing with improved expert weight initialization patterns for expert-based models. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: Jennifer Chen <jennifchen@nvidia.com> |
||
|
willg-nv
|
e589ac80f6 |
Integrate Automated QDQ benchmark - part 3.1 (#837)
## What does this PR do? This PR integrates benchmark module to QDQ autotunner. This benchamrk module is used to evaluate ONNX model perf. This PR is 1/3 of https://github.com/NVIDIA/Model-Optimizer/pull/703. Once all small PRs are merged #703 could be closed. PR 3.1: https://github.com/NVIDIA/Model-Optimizer/pull/837 PR 3.2 https://github.com/NVIDIA/Model-Optimizer/pull/838 PR 3.3: https://github.com/NVIDIA/Model-Optimizer/pull/839 ## 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?**: No, document will be added in part 4. - **Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**: No, change log will be updated when all changes are merged. ## Additional Information <!-- E.g. related issue. --> <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added ONNX quantization autotuning capabilities with a consolidated module providing streamlined import paths for core components. * Introduced unified benchmarking framework supporting TensorRT-based model evaluation with both command-line and Python API implementations. * Added support for timing cache persistence, custom plugin libraries, shape validation, and dynamic input shape configuration for flexible model testing and optimization. <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: Will Guo <willg@nvidia.com> |
||
|
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>
|
||
|
kinjalpatel27
|
df47e816a7 |
Added support to rotate in fp32 (optional) (#885)
## What does this PR do?
**Type of change:** New Feature
**Overview:**
This MR adds support to perform rotation for RHT in float32 if enabled
by quantization configuration. It also makes rotate argument in
quantization configuration of type bool (for backward compatibility) or
dict (added option for float32 rotation)
## Usage
```
python hf_ptq.py --pyt_ckpt_path meta-llama/Llama-3.2-3B-Instruct --qformat nvfp4 --export_fmt hf --dataset cnn_dailymail --export_path test --trust_remote_code --inference_pipeline_parallel 1 --batch_size 1 --calib_size 4 --kv_cache_qformat nvfp4_rotate
```
Updated `NVFP4_KV_ROTATE_CFG` locally with `"rotate": {"enable": True,
"rotate_fp32": True}`
```
...
model.layers.27.self_attn.k_bmm_quantizer TensorQuantizer((2, 1) bit fake block_sizes={-1: 16, 'type': 'dynamic', 'scale_bits': (4, 3)}, amax=8.3750 rotated (fp32) calibrator
=MaxCalibrator quant)
...
```
Updated `NVFP4_KV_ROTATE_CFG` locally with `"rotate": {"enable": True,
"rotate_fp32": False}`
```
model.layers.27.self_attn.k_bmm_quantizer TensorQuantizer((2, 1) bit fake block_sizes={-1: 16, 'type': 'dynamic', 'scale_bits': (4, 3)}, amax=8.3750 rotated calibrator=MaxCalibrator quant)
```
## Testing
Updated unit test in `tests/gpu/torch/quantization/test_hadamard.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
- **Did you write any new necessary tests?**: No (updated existing test)
- **Did you add or update any necessary documentation?**: Yes
- **Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**:
Yes
## 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 rotational input capability prior to quantization for RHT
(Rotated Hyperplane Transform).
* Introduced granular rotation configuration options enabling FP32
casting for improved numerical stability during transforms.
* **Tests**
* Expanded test coverage for rotation functionality with parameterized
FP32 casting scenarios.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---------
Signed-off-by: Kinjal Patel <kinjalpravin@nvidia.com>
|
||
|
Jenny Chen
|
02fa3623d8 |
Sync MOE layer input quantizer only (#903)
## What does this PR do? **Type of change:** Bug fix **Overview:** in MOE layer we currently sync both the weight and input quantizers so that all experts have the same weight amaxes and activation amaxes. VLLM/TRTLLM actually support non-uniform weight amaxes in MOE so we only need to sync the activation amaxes. ## 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** * Improved input quantizer synchronization for Mixture of Experts models to ensure correct amax value handling across local experts. * **Documentation** * Fixed typos and clarified wording in quantization documentation. * **Tests** * Added test coverage for Mixture of Experts quantizer synchronization functionality. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Jennifer Chen <jennifchen@nvidia.com> |
||
|
kaix-nv
|
70ffb6f87c |
Remove test_llama_eval_sparse_attention (#914)
## 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:** ? Fix the example test. ## 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 * **Tests** * Removed a sparse attention test variant to streamline test coverage. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: Kai Xu <kaix@nvidia.com> |
||
|
Keval Morabia
|
9e23c6c312 |
Upgrade Dev containers for CICD to latest (#891)
## What does this PR do? - Upgrade CICD test containers to latest - Enable torch 2.10 testing in CICD ## Testing <!-- Mention how have you tested your change if applicable. --> CI/CD in this PR should pass <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added support for mixed-precision gradient handling with FSDP2. * **Documentation** * Updated Linux installation guide with CUDA 13.x support and cupy dependency guidance. * **Chores** * Updated CI/CD workflows and test infrastructure to support PyTorch 2.10 and CUDA 13. * Updated container image versions and test environment configurations. * Updated TensorRT-LLM version requirements. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> |
||
|
Tony@NV
|
7c4c9fdbc9 |
Support multiple-batch input for autocast calibration. (#760)
## What does this PR do? Add multi-batch calibration data support for autocast precision conversion. This enhancement allows users to provide multiple batches of calibration data (via a directory of NPZ files or Polygraphy JSON with multiple batches) to aggregate tensor statistics across batches, resulting in more robust precision conversion decisions. ## Usage ### Single NPZ file (existing behavior) ``` python -m modelopt.onnx.autocast --onnx_path model.onnx --calibration_data calibration_data.npz --output_path model_fp16.onnx ``` ### Directory containing multiple NPZ files for multi-batch calibration (new) ``` python -m modelopt.onnx.autocast --onnx_path model.onnx --calibration_data calibration_data_dir/ --output_path model_fp16.onnx ``` ## Testing - Tested with single NPZ file to ensure backward compatibility - Tested with directory containing multiple NPZ files for multi-batch calibration - Verified that aggregated statistics (absmax, min, max) are correctly computed across batches ## 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. - **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 Key changes: - Added `TensorStats` dataclass to store aggregated tensor statistics (absmax, min_val, max_val, shape) - Updated `ReferenceRunner` to: - Load multiple NPZ files from a directory (`_load_inputs_from_npz`) - Aggregate statistics across batches (`_aggregate_tensor_stats`) - Process multi-batch inference in `run()` method - Updated `IORangeRule` and `DepthOfReductionRule` to handle both raw numpy arrays and `TensorStats` objects - Enhanced `--calibration_data` CLI help text to document multi-batch support <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added multi-batch calibration support via directories of NPZ files or Polygraphy JSON files. * Implemented cross-batch statistics aggregation for more robust precision conversion decisions. * **Documentation** * Expanded calibration_data CLI option guidance with detailed support for multiple input formats and batch processing benefits. <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: Tony Yin <toyin@nvidia.com> |
||
|
Frida Hou
|
adcce614cb |
add local hessian calibration (#788)
## 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:** Add a new calibration method for weight scale search. It considers activation information by weighing scale candidates with local hessian matrix. Initial experiments with Qwen3 8B NVFP4 shows improvements. ## Usage <!-- You can potentially add a usage example below. --> Use `NVFP4_W4A4_WEIGHT_LOCAL_HESSIAN_CFG` quantization config for quantization and evaluation. ## 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 local Hessian-weighted MSE calibration pathway for NVFP4 per-block quantization with configurable amax search parameters and FP8 scale sweep support. * **Tests** * Added test coverage for the new local Hessian weight-only quantization configuration. <!-- 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> |
||
|
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>
|
||
|
realAsma
|
c4b662fbc8 |
[Bug fix] Fake quantized model save after HF accelerate hooks are added (#906)
## What does this PR do? **Type of change:** Bug fix **Overview:** Fix `AttributeError: Can't get local object 'add_hook_to_module.<locals>.new_forward'` when saving a quantized model a second time after restoring it with `device_map="auto"`. When a model is loaded with `device_map="auto"`, accelerate's `add_hook_to_module` patches every submodule (including `TensorQuantizer` instances) and injects three instance attributes: `_hf_hook`, `_old_forward`, and `forward` (a `functools.partial` wrapping a local function). These are not picklable and were leaking into the modelopt state dict collected by `get_modelopt_state()`, causing `torch.save` to fail. This PR adds the three accelerate-injected attributes to `TensorQuantizer._skip_properties_for_save_restore` so they are excluded from the serialized state, matching the existing pattern used for HuggingFace and DeepSpeed attributes. ## Usage ```python mto.enable_huggingface_checkpointing() # Quantize and save model = AutoModelForCausalLM.from_pretrained(name, device_map="auto") model = mtq.quantize(model, mtq.FP8_DEFAULT_CFG, forward_loop=forward_loop) model.save_pretrained(save_dir) # Restore and save again (this previously failed) model2 = AutoModelForCausalLM.from_pretrained(save_dir, device_map="auto") model2.save_pretrained(save_dir_round2) # now works ``` ## Testing - Added unit test `test_tensor_quantizer_modelopt_state_with_accelerate_hook` in `tests/unit/torch/quantization/plugins/test_accelerate.py` that verifies accelerate hook attributes are excluded from modelopt state and the state dict remains picklable. ## 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. - **Is this change backward compatible?**: Yes — only adds entries to a skip set; existing saved checkpoints are unaffected. - **Did you write any new necessary tests?**: Yes - **Did you add or update any necessary documentation?**: No (internal fix, no API change) - **Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**: No ## Additional Information The root cause is in accelerate's `add_hook_to_module`, which defines `new_forward` as a local function and binds it via `functools.partial` onto `module.forward`. Since local functions cannot be pickled, any `TensorQuantizer` that has been hooked by accelerate becomes unserializable unless these attributes are excluded. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Bug Fixes** * Enhanced compatibility with accelerate library by excluding framework-specific hooks and attributes from model state serialization, preventing issues during save/restore operations. * **Tests** * Added test to validate that accelerate-related attributes are properly excluded from model state and that the state remains picklable. * **Public API** * TensorQuantizer is now publicly exported. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: realAsma <akuriparambi@nvidia.com> |
||
|
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> |
||
|
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> |
||
|
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>
|
||
|
Keval Morabia
|
ae69d5d9b9 |
Separate CI job for Megatron GPU tests (#888)
## What does this PR do? [Short term]: Megatron based tests take a long time often resulting in CICD timeout. Splitting megatron tests into a dedicated CICD job for faster overall CI/CD run [Mid/Long term]: Run all megatron gpu tests using `torchrun` instead of `pytest` so all dist processes are already created and all individual tests no longer need to setup and destroy their processes which adds a lot of overhead per test ## Testing <!-- Mention how have you tested your change if applicable. --> - [x] 1-GPU CI/CD passing (on this PR) - [x] 2-GPU CI/CD passing (on nightly run - manually triggered): https://github.com/NVIDIA/Model-Optimizer/actions/runs/22000517688 --------- Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> |
||
|
willg-nvandAjinkya Rasane
|
3e95d9fd77 |
Integrate Automated QDQ placement tool - part 2.3 (#846)
## What does this PR do? This PR implement RegionInspect tool. This tool could be used to visualize the regions parititioned by RegionSearch classes. This tool could be used to analyze if the partitioned regions match the fusion patterns. **Overview:** ? ## 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 - **Did you write any new necessary tests?**: Yes - **Did you add or update any necessary documentation?**: No, document update is in Part 4. - **Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**: No, CHANGELOG will be updated when all changes are ready. ## Additional Information <!-- E.g. related issue. --> <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit ## New Features * Added a region inspection tool for ONNX models. Analyzes model structure and generates detailed reports including region statistics, hierarchical relationships, node coverage metrics, and size distribution analysis. Available through a command-line interface with configurable parameters. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Will Guo <willg@nvidia.com> Co-authored-by: Ajinkya Rasane <131806219+ajrasane@users.noreply.github.com> |