mirror of
https://github.com/NVIDIA/Model-Optimizer.git
synced 2026-10-22 06:21:37 +08:00
0.45.0dev
106
Commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
 kinjalpatel27
|
010b220dc0 |
vLLM fakequant export update for AWQ checkpoint (#1242)
### What does this PR do?
Type of change: Bug
Enables end-to-end AWQ checkpoint export and reload in the vLLM
fake-quant serving path (`MODELOPT_STATE_PATH`). Previously, the
`input_quantizer` was using incorrect `pre_quant_scale` especially with
grouped quantizers like `qkv_proj`, using simply the first
`input_quantizer.pre_quant_scale`. This MR adds
`_resmooth_experts_for_export` that non-mutatively averages
`pre_quant_scale` across MoE experts and unifies input `_amax`, required
because vLLM uses a single input quantizer per expert group. Adds
`merge_amax_tensors_for_group` (element-wise max for same-shape, `cat`
for GQA, scalar-max fallback) replacing the scalar-collapsing
`torch.stack().max()` that dropped per-channel `_amax` structure.
### Usage
```python
# Export AWQ checkpoint from HF model
from modelopt.torch.export.plugins.vllm_fakequant_hf import export_hf_vllm_fq_checkpoint
export_hf_vllm_fq_checkpoint(model, export_dir="./awq_vllm_checkpoint")
```
### Testing
**Step 1 — Export the quantized checkpoint:**
```bash
python examples/llm_ptq/hf_ptq.py \
--pyt_ckpt_path <MODEL_PATH> \
--recipe <AWQ_RECIPE> \
--calib_size 512 \
--export_path <EXPORT_DIR> \
--vllm_fakequant_export
```
This produces `<EXPORT_DIR>/vllm_fq_modelopt_state.pth` with the averaged per-expert
pre_quant_scale and unified _amax now included.
Step 2 — Serve via vLLM fakequant worker:
```bash
MODELOPT_STATE_PATH=<EXPORT_DIR>/vllm_fq_modelopt_state.pth \
python examples/vllm_serve/vllm_serve_fakequant.py \
<EXPORT_DIR> --tensor-parallel-size <TP>
```
Tested for quantization configurations:
```
FP8_DEFAULT_CFG
FP8_DEFAULT_CFG (input_q disabled)
INT8_SMOOTHQUANT_CFG
INT8_WEIGHT_ONLY_CFG
NVFP4_DEFAULT_CFG
NVFP4_AWQ_LITE_CFG
INT4_AWQ_CFG
NVFP4_AWQ_CFG
NVFP4_DEFAULT_CFG (input_q disabled)
```
### Before your PR is "*Ready for review*"
Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`).
Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors) (e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(..., weights_only=False)`, `pickle`, etc.).
- Is this change backward compatible?: ✅
- If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A
- Did you write any new necessary tests?: N/A
- Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A
### Additional Information
<!-- E.g. related issue. -->
<!-- This is an auto-generated comment: release notes by coderabbit.ai -->
## Summary by CodeRabbit
* **New Features**
* Added Nemotron-style MoE export support and group-aware AWQ resmoothing with optional requantization during export.
* Improved handling for shared-input / expert groups and tensor-parallel sharding of pre-quantization scales.
* **Bug Fixes**
* Removed AWQ reload limitation from known issues; improved checkpoint validation and safer save/load behavior.
* Better detection and handling of enabled weight-quantizers and clearer warnings for mismatched checkpoint keys.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---------
Signed-off-by: Kinjal Patel <kinjalpravin@nvidia.com>
|
||
|
jingyu-ml
|
26ae8da517 |
[2/3] Implicit Gemm NVFP4 (#1227)
### What does this PR do? Type of change: new feature <!-- Use one of the following: Bug fix, new feature, new example, new tests, documentation. --> - Add Conv3D implicit GEMM kernel with BF16 WMMA tensor cores and fused NVFP4 activation quantization for video diffusion VAE layers - Integrate into _QuantConv3d via QuantModuleRegistry — automatically dispatched when NVFP4 quantization is applied to nn.Conv3d - Move kernel from `experimental/conv/ to modelopt/torch/kernels/conv/`; move tests to `tests/gpu/torch/quantization/kernels/` ### Testing <!-- Mention how have you tested your change if applicable. --> - Added test cases to measure the difference between cuDNN and our CUDA implicit GEMM kernel - Added an NVFP4 fake quantization test using CUDA 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`, `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`: ✅ <!--- Mandatory --> - Did you write any new necessary tests?: ✅ <!--- Mandatory for new features or examples. --> - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ <!--- Only for new features, API changes, critical bug fixes or backward incompatible changes. --> ### Additional Information <!-- E.g. related issue. --> <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Per-backbone quantization/export in a single run with per-backbone checkpoints and backbone-aware quant filters * Configurable NVFP4 block-size via CLI/config; improved NVFP4 Conv3D inference path and Wan 2.2 quantization support * **Bug Fixes** * Video-model calibration now respects extra params and forces video decoding during calibration * **Documentation** * Added comprehensive Conv3D implicit‑GEMM kernel documentation; removed experimental Conv3D prototype docs/benchmark * **Tests** * New Wan 2.2 quantization/export tests and expanded Conv3D/FP4 kernel test coverage <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Jingyu Xin <jingyux@nvidia.com> |
||
|
jingyu-ml
|
feec81ad2b |
Add the Skip softmax for diffusion (#1166)
### What does this PR do?
Type of change: new feature, new example <!-- Use one of the following:
Bug fix, new feature, new example, new tests, documentation. -->
<!-- Details about the change. -->
## Summary
- Add skip-softmax sparse attention (BLASST) for diffusion models via
dedicated Triton kernels — an inference kernel with tile skipping and a
calibration kernel with vectorized multi-threshold sparsity measurement
- Add `triton_skip_softmax` method with exponential model calibration
(`scale_factor = a * exp(b * sparsity)`) and log-space fitting for
diffusion models
- Add Triton kernel backends for diffusers and LTX attention dispatch
- Fix calibration to skip RULER dataset generation when user provides
their own `forward_loop` (required for non-LLM models)
## Changes
### Triton kernels (`modelopt/torch/kernels/triton_fa.py`)
- **`_attn_fwd`**: Forward kernel with optional tile skipping — tiles
whose max attention score is far below the running softmax max are
skipped entirely (no V load, no softmax, no accumulation). Runtime
sparsity measurement via atomic counters.
- **`_attn_fwd_calibrate`**: Calibration kernel that computes full
attention while measuring how many tiles would be skipped at each of N
thresholds simultaneously. Uses per-program output buffers (zero atomic
contention) and vectorized multi-threshold comparison.
- **`attention()`** / **`attention_calibrate()`**: Python wrappers for
inference and calibration kernels.
### Kernel backends
(`modelopt/torch/sparsity/attention_sparsity/kernels/`)
- **`diffusers_triton_attention.py`**: Registers `modelopt_triton`
backend in diffusers' attention dispatch. Handles [B, S, H, D] → varlen
layout conversion, calibration/inference mode switching, thread-local
configuration, and counter accumulation.
- **`ltx_triton_attention.py`**: Patches `ltx_core.Attention` modules
for Triton dispatch with the same calibration/inference modes.
### Method
(`modelopt/torch/sparsity/attention_sparsity/methods/triton_skip_softmax.py`)
- `TritonSkipSoftmaxMethod`: Context managers for calibration (→
calibration kernel) and inference (→ forward kernel with tile skipping).
Three threshold priority levels: raw threshold > calibrated scale_factor
> static threshold.
### Calibration
(`modelopt/torch/sparsity/attention_sparsity/calibration/`)
- **`calibrator.py`**: `DynamicThresholdCalibrator` with `fit_logspace`
option — fits exponential model in log space (minimizes relative error)
for diffusion models where scale_factors span many orders of magnitude.
Records observed sparsity range for extrapolation warnings.
- **`calibrate.py`**: Skips RULER dataset when `forward_loop` is
provided; passes `fit_logspace` through from config.
### Config & conversion
- **`config.py`**: `CalibrationConfig.fit_logspace` field (default
False, recommended True for diffusion models).
`skip_softmax_raw_threshold` field for direct threshold mode.
- **`conversion.py`**: Auto-registers diffusers/LTX Triton backends on
`sparsify()`. Updated summary display.
### Example
- **`wan22_skip_softmax.py`**: End-to-end example for WAN 2.2 5B/14B
with baseline, raw-threshold, and calibrated modes. Supports runtime
sparsity reporting.
## Threshold modes
| Mode | How it works | Use case |
|------|-------------|----------|
| **Raw threshold** (`--raw-threshold -0.7`) | Passed directly to kernel
as `skip_threshold_log2` | Quick testing, sweeps |
| **Calibrated** (`--calibrate --target-sparsity 0.5`) | `scale_factor =
a * exp(b * target)`, then `threshold = scale_factor / seq_k` at runtime
| Production use with seqlen adaptation |
| **Static** (default `skip_softmax_threshold=0.1`) | `log2(lambda) *
sm_scale` | Fallback |
## Usage
```bash
# Fixed raw threshold (no calibration)
python examples/diffusers/sparsity/wan22_skip_softmax.py \
--model-path /path/to/Wan2.2-T2V-A14B-Diffusers \
--raw-threshold -0.7 \
--prompt "A cat playing piano" --output out.mp4
# With calibration (log-space fit for diffusion models)
python examples/diffusers/sparsity/wan22_skip_softmax.py \
--model-path /path/to/Wan2.2-T2V-A14B-Diffusers \
--calibrate --target-sparsity 0.5 \
--prompt "A cat playing piano" --output out.mp4
# Dense baseline for comparison
python examples/diffusers/sparsity/wan22_skip_softmax.py \
--model-path /path/to/Wan2.2-T2V-A14B-Diffusers \
--baseline \
--prompt "A cat playing piano" --output baseline.mp4
```
### 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`: ✅ <!---
Mandatory -->
- Did you write any new necessary tests?: ✅ <!--- Mandatory for new
features or examples. -->
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
❌ <!--- Only for new features, API changes, critical bug fixes or
backward incompatible changes. -->
### 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 skip-softmax sparse attention support for Diffusers models,
enabling efficient video generation
* Added support for both eager and Triton attention backends for sparse
attention
* Added new example script for Wan 2.2 text-to-video generation with
sparse attention optimization
* **Documentation**
* Updated documentation with sparse attention configuration guide and
usage examples
* **Tests**
* Added comprehensive unit tests for kernel backend registration and
skip-softmax functionality
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---------
Signed-off-by: Jingyu Xin <jingyux@nvidia.com>
|
||
|
realAsma
|
2d868d3f1f |
Performant layerwise calibration for large models (#1251)
## Summary
Adds **performant layerwise calibration** for quantizing large models
(e.g. DeepSeek-R1 671B) that don't fit entirely on GPU. ([Example
commands](#example-commands))
1. **Performant calibration for large models** — Each decoder layer is
moved from CPU/disk to GPU (accelerate) or unsharded (FSDP2) **only
once** and kept on GPU for the entire calibration step. Previously,
every calibration batch triggered weight transfer for every layer —
O(num_batches) weight movements per layer. Now it is O(1) per layer.
This also means you can **increase batch size** since only one layer's
weights occupy GPU at a time — e.g. DeepSeek-R1 on a single node
(8×80GB) with `batch_size=16` and `gpu_max_mem_percentage=0.5`.
2. **Checkpoint save/resume** — Saves progress after each layer, so jobs
that exceed cluster time limits (e.g. 4-hour Slurm windows for 100+
layer MoE models) can resume from the last completed layer.
3. **Rename** `sequential_calibrate` → `layerwise_calibrate` for
clarity.
### Design details
The existing layerwise state machine (skip/run/capture) already
processes one layer at a time, but skip-mode layers still kept their
parameters in the ModuleList — so frameworks transferred all weights
every forward pass. This PR adds:
- **`_SkipLayer`**: replaces fully-calibrated layers with a
parameter-free dummy in the ModuleList, so framework hooks have nothing
to transfer
- **`persistent_materialization`**: keeps the active layer on GPU for
the entire calibration step, avoiding repeated offload/reload cycles
Checkpoint save is per-layer; restore is bulk — quantizer state and
weights for layers 0..K-1 are restored once at the end of calibration,
keeping the hot path fast.
### Example commands
**Qwen3-8B** (NVFP4+GPTQ, single GPU):
```bash
python hf_ptq.py \
--pyt_ckpt_path Qwen/Qwen3-8B \
--recipe nvfp4_gptq_sequential.yaml \
--calib_size 64 \
--batch_size 16 \
--dataset cnn_dailymail \
--export_path outputs/qwen3_8b_nvfp4_gptq_seq \
--gpu_max_mem_percentage 0.5 \
--use_seq_device_map \
--vllm_fakequant_export
```
**DeepSeek-R1** (NVFP4 experts-only + FP8 KV, 8×80GB):
```bash
python hf_ptq.py \
--model unsloth/DeepSeek-R1-0528-BF16 \
--recipe ../../modelopt_recipes/general/ptq/nvfp4_experts_only-fp8_kv.yaml \
--dataset cnn_dailymail \
--batch_size 16 \
--calib_size 64 \
--calib_seq 512 \
--gpu_max_mem_percentage 0.5 \
--use_seq_device_map \
--trust_remote_code \
--export_path output/DeepSeek-R1-BF16-nvfp4-experts-only-fp8-kv \
--vllm_fakequant_export
```
### Example: NVFP4+GPTQ layerwise calibration on Qwen3-8B (36 layers,
single GPU — 20 GB peak)
**Initial run** (killed after layer 11):
```
Layerwise calibration: Found 36 transformer layers
Calibrating layer 1/36 | capture: [1]
Computing Hessians for 7 linear layers...
GPTQ time: 51.39s
Calibrating layer 2/36 | run: [1] | capture: [2]
Checkpoint: saved layer 0
GPTQ time: 50.06s
Calibrating layer 3/36 | skip: 1 | run: [2] | capture: [3]
Checkpoint: saved layer 1
...
Calibrating layer 12/36 | skip: 10 | run: [11] | capture: [12]
Checkpoint: saved layer 10
<killed>
```
**Resumed run** (picks up from layer 11, finishes all 36):
```
Layerwise calibration: Found 36 transformer layers
Checkpoint: resuming layerwise calibration from layer 11/36
Calibrating layer 12 (resumed)
GPTQ time: 51.45s
Calibrating layer 13/36 | skip: 11 | run: [12] | capture: [13]
Checkpoint: saved layer 11
...
Calibrating layer 36/36 | skip: 34 | run: [35] | capture: [36]
Checkpoint: saved layer 34
GPTQ time: 50.33s
Checkpoint: saved layer 35 (final)
Checkpoint: restored 11 previously calibrated layers
Layerwise calibration completed
Quantized model exported to: outputs/qwen3_8b_nvfp4_gptq_seq
GPU 0: Peak memory usage = 20.42 GB
```
## TODO
- [ ] Update CHANGELOG
## Test plan
- `tests/unit/torch/quantization/test_layerwise_calibrate.py` — unit
tests for skip/swap/restore
- `tests/unit/torch/quantization/test_sequential_checkpoint.py` —
checkpoint save/resume correctness
- `tests/gpu/torch/quantization/plugins/test_accelerate_gpu.py` —
CPU-offloaded layerwise + GPTQ + checkpoint resume
- `tests/gpu/torch/quantization/test_fsdp2.py` — FSDP2 layerwise
calibration
### Verified
- [x] Qwen3-8B: layerwise calibration + checkpoint save/restore +
fakequantized checkpoint export + vLLM serve
- [x] DeepSeek-R1: checkpoint resume tested
- [x] DeepSeek-R1: fakequantized checkpoint export verified
---------
Signed-off-by: realAsma <akuriparambi@nvidia.com>
|
||
|
|
361f7e391b |
Merge puzzletron compression algorithm (#1121)
### What does this PR do? Implement puzzletron compression algorithm based on Puzzle paper (https://arxiv.org/abs/2411.19146) <details> <summary> Th list of reviewed and merged MRs that resulted in the feature/puzzletron branch</summary> Merging dkorzekwa/any_model to feature/puzzletron [Add anymodel directories to feature/puzzletron by danielkorzekwa · Pull Request #974 · NVIDIA/Model-Optimizer](https://github.com/NVIDIA/Model-Optimizer/pull/974) - merged [Draft: anymodel activation scoring by danielkorzekwa · Pull Request #989 · NVIDIA/Model-Optimizer](https://github.com/NVIDIA/Model-Optimizer/pull/989) - merged [Draft: Merge anymodel pruning by danielkorzekwa · Pull Request #990 · NVIDIA/Model-Optimizer](https://github.com/NVIDIA/Model-Optimizer/pull/990/) - merged [Draft: Merging anymodel:build_library_and_stats by danielkorzekwa · Pull Request #993 · NVIDIA/Model-Optimizer](https://github.com/NVIDIA/Model-Optimizer/pull/993) - merged [Dkorzekwa/any model calc one block scores by danielkorzekwa · Pull Request #994 · NVIDIA/Model-Optimizer](https://github.com/NVIDIA/Model-Optimizer/pull/994) - merged [Draft: merge any_model: mip_and_realize_models by danielkorzekwa · Pull Request #995 · NVIDIA/Model-Optimizer](https://github.com/NVIDIA/Model-Optimizer/pull/995) - merged [Dkorzekwa/any model other modeqls by danielkorztiekwa · Pull Request #1007 · NVIDIA/Model-Optimizer](https://github.com/NVIDIA/Model-Optimizer/pull/1007/) - merged PR to 1007: https://github.com/NVIDIA/Model-Optimizer/pull/1039 - merged [Dkorzekwa/anymodel gptoss by danielkorzekwa · Pull Request #1020 · NVIDIA/Model-Optimizer](https://github.com/NVIDIA/Model-Optimizer/pull/1020) - merged [Merge any_model tutorial by danielkorzekwa · Pull Request #1035 · NVIDIA/Model-Optimizer](https://github.com/NVIDIA/Model-Optimizer/pull/1035) - merged [Merge mbridge distillation for any_model by danielkorzekwa · Pull Request #1036 · NVIDIA/Model-Optimizer](https://github.com/NVIDIA/Model-Optimizer/pull/1036) - merged [MR branch for the remaining difference between dkorzekwa/any_model an… by danielkorzekwa · Pull Request #1047 · NVIDIA/Model-Optimizer](https://github.com/NVIDIA/Model-Optimizer/pull/1047) - merged [Dkorzekwa/decilm hf code cleanup by danielkorzekwa · Pull Request #1071 · NVIDIA/Model-Optimizer](https://github.com/NVIDIA/Model-Optimizer/pull/1071) - merged [Dkorzekwa/decilm hf code cleanup 2 by danielkorzekwa · Pull Request #1073 · NVIDIA/Model-Optimizer](https://github.com/NVIDIA/Model-Optimizer/pull/1073) - merged [Dkorzekwa/anymodel subblock stats by danielkorzekwa · Pull Request #1085 · NVIDIA/Model-Optimizer](https://github.com/NVIDIA/Model-Optimizer/pull/1085) - merged [Dkorzekwa/anymodel subblock stats nodecilm by danielkorzekwa · Pull Request #1102 · NVIDIA/Model-Optimizer](https://github.com/NVIDIA/Model-Optimizer/pull/1102) - merged [Dkorzekwa/decilm cleanup post subblockstats by danielkorzekwa · Pull Request #1103 · NVIDIA/Model-Optimizer](https://github.com/NVIDIA/Model-Optimizer/pull/1103) - merged [code clean up by danielkorzekwa · Pull Request #1110 · NVIDIA/Model-Optimizer](https://github.com/NVIDIA/Model-Optimizer/pull/1110) - merged Merging into main: [Activation hooks redesign (reuse hooks component across both minitron and puzzletron) by danielkorzekwa · Pull Request #1022 · NVIDIA/Model-Optimizer](https://github.com/NVIDIA/Model-Optimizer/pull/1022) - merged [Dkorzekwa/puzzletron use importance hooks from prune by danielkorzekwa · Pull Request #1115 · NVIDIA/Model-Optimizer](https://github.com/NVIDIA/Model-Optimizer/pull/1115) - merged </details> <!-- Details about the change. --> ### Usage Puzzletron tutorial: https://github.com/NVIDIA/Model-Optimizer/tree/feature/puzzletron/examples/puzzletron ### Testing The main e2e test for compressing 9 models with Puzzletron: https://github.com/NVIDIA/Model-Optimizer/blob/feature/puzzletron/tests/gpu/torch/puzzletron/test_puzzletron.py 2-gpu nightly tests: - https://github.com/NVIDIA/Model-Optimizer/actions/runs/24468209205/job/71501061203 - https://github.com/NVIDIA/Model-Optimizer/actions/runs/24470214159/job/71508152952 ### 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?: ✅ - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added Puzzletron: end-to-end heterogeneous pruning & NAS workflow with AnyModel support, example pipelines, deployment and evaluation utilities, and tools for converting/pruning and exporting compressed checkpoints. * **Documentation** * Comprehensive Puzzletron tutorials, model-specific guides, evaluator instructions, example configs, and changelog entry. * **Chores** * CI/workflow updates (extras installation, longer GPU test timeout), pre-commit hook exclusion updated, and CODEOWNERS entries added. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com> Signed-off-by: Liana Mikaelyan <lmikaelyan@nvidia.com> Signed-off-by: Liana Mikaelyan <45925959+LianaMikael@users.noreply.github.com> Signed-off-by: Daniel Korzekwa <daniel.korzekwa@gmail.com> Signed-off-by: jrausch <jrausch@nvidia.com> Signed-off-by: root <root@pool0-00848.cm.cluster> Co-authored-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Co-authored-by: Liana Mikaelyan <lmikaelyan@nvidia.com> Co-authored-by: Liana Mikaelyan <45925959+LianaMikael@users.noreply.github.com> Co-authored-by: J Rausch <38429553+j-rausch@users.noreply.github.com> Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com> |
||
|
Chenhan D. YuandClaude Opus 4.6
|
3131195241 |
add: DFlash block diffusion speculative decoding (#1211)
DFlash (Block Diffusion for Flash Speculative Decoding) predicts an entire block of tokens in a single forward pass using masked parallel prediction with KV injection from the target model's hidden states. Key features: - Feature fusion (multi-layer hidden states -> FC + RMSNorm) - KV injection (fused features as K/V in every draft layer with QK-norm) - Random anchor sampling with bidirectional intra-block attention - Logit distillation with exponential loss decay (gamma weighting) - Multi-node DDP training with checkpoint resume - Export to z-lab compatible HF format - Online validation (context-dependent ground truth) Training recipe: modelopt_recipes/general/speculative_decoding/dflash.yaml Results: examples/speculative_decoding/doc/dflash_results.md ### ModelOpt Eval (online validation, osl=512) | Dataset | z-lab | ModelOpt (306K) | Diff | |---------|-------|-----------------|------| | gsm8k | 4.10 | **5.19** | **+1.09** | | MT-Bench | 3.58 | **4.36** | **+0.78** | ### z-lab Official Eval (dflash.benchmark, osl=512) | Dataset | z-lab | ModelOpt (306K) | Diff | |---------|-------|-----------------|------| | gsm8k | **5.00** | 4.08 | -0.92 | | MT-Bench | **3.28** | 2.99 | -0.29 | > z-lab model trained with block_size=16. ModelOpt trained with block_size=8. ## Evaluation Method Impact (gsm8k) | Eval Method | z-lab checkpoint | ModelOpt (306K) | |-------------|-----------------|-----------------| | Fixed GT (ModelOpt eval) | 2.95 | 4.23 | | Online GT (ModelOpt eval) | 4.10 | **5.19** | | z-lab official eval | **5.00** | 4.08 | ### What does this PR do? Type of change: ? <!-- Use one of the following: Bug fix, new feature, new example, new tests, documentation. --> <!-- 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`, `torch.load(..., weights_only=False)`, `pickle`, etc.). - Is this change backward compatible?: ✅ / ❌ / N/A <!--- If ❌, explain why. --> - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: ✅ / ❌ / N/A <!--- Mandatory --> - Did you write any new necessary tests?: ✅ / ❌ / N/A <!--- Mandatory for new features or examples. --> - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ / ❌ / N/A <!--- Only for new features, API changes, critical bug fixes or backward incompatible changes. --> ### Additional Information <!-- E.g. related issue. --> <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added DFlash speculative decoding mode with parallel block prediction support. * Included training launchers and MT-Bench evaluation scripts for DFlash models. * Added online acceptance rate validation for improved inference verification. * **Documentation** * DFlash quick start guide with configuration parameters and training examples. * Performance results and benchmarks for DFlash-trained models. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Chenhan Yu <chenhany@nvidia.com> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> |
||
|
ZhiyuandClaude Opus 4.6
|
d4279920e0 |
Generic Fused MoE Quantization + Export for transformers 5.0+ (#1187)
## What does this PR do? Add generic quantization and export support for **fused MoE expert modules** in HuggingFace transformers 5.0+. In transformers 5.0+, all major MoE models switched from sequential per-expert `nn.ModuleList` to **fused 3D tensor parameters** (`gate_up_proj`, `down_proj`). This breaks ModelOpt's existing per-expert quantization and export pipeline, which assumes iterable expert submodules. **Affected models (verified against transformers v5.5.0 source):** - `MixtralExperts` (Mixtral) - `Qwen2MoeExperts` (Qwen2-MoE) - `Qwen3MoeExperts` (Qwen3-MoE) - `Qwen3_5MoeExperts` (Qwen3.5-MoE) - `DeepseekV3NaiveMoe` (DeepSeek-V3) - `JambaExperts`, `OlmoeExperts`, and any future model following the same HF standard pattern **Key insight:** All these models share an identical fused expert structure and forward pattern. A single generic solution replaces N model-specific implementations. ### Context: relationship to PR #975 and PR #1170 - **PR #975** (`kmorabi/bump-transformers-5.0`): Adds experimental transformers 5.0 support but explicitly skips batched MoE experts. This PR fills that gap. - **PR #1170** (`chenjiel/refactor_qwen35`): Handles only Qwen3.5 using `_QuantFunctionalMixin`. This PR generalizes that approach to all fused MoE models. ### Changes **Quantization** (`modelopt/torch/quantization/plugins/huggingface.py`): - `_QuantFusedExperts(_QuantFunctionalMixin)` -- Generic wrapper that intercepts `F.linear` calls and applies per-expert quantization via storage-offset-based expert index recovery. Each expert gets its own weight and input quantizers (`nn.ModuleList`). - `_is_fused_experts_module()` -- Structural detector: `gate_up_proj` (3D) + `down_proj` (3D) + `num_experts` + `act_fn`. - `register_fused_experts_on_the_fly()` -- Auto-registration callback, added to `CUSTOM_MODEL_PLUGINS` before `register_sparse_moe_on_the_fly` so explicit registrations (Llama4, GptOss, etc.) take priority. - `_get_fused_expert_intermediate_dim()` -- Helper for cross-version attribute name resolution (`intermediate_dim` / `intermediate_size` / fallback to shape). **Export** (`modelopt/torch/export/moe_utils.py`, `unified_export_hf.py`, `layer_utils.py`): - `_export_fused_experts()` -- Splits fused 3D weights into per-expert 2D projections (`gate_proj`, `up_proj`, `down_proj`), handles amax fallback for uncalibrated experts, proportionally slices per-channel amax, and registers results under the standard `experts.{E}.gate_proj.weight` naming convention. - Integration in `_process_quantized_modules` and `_export_transformers_checkpoint` to dispatch to `_export_fused_experts` for fused expert modules. - Structural detection in `get_expert_linear_names` for fused experts. Added `MixtralSparseMoeBlock` to the `gate_proj`/`down_proj`/`up_proj` group (transformers 5.0 naming). **Tests** (`tests/unit/torch/quantization/plugins/test_fused_experts.py`): - Synthetic fused expert model matching the exact HF 5.0+ pattern. - Tests for structural detection, auto-registration, two-level registration (block + expert), quantizer creation, forward pass-through correctness, expert index recovery, and export output structure. ### Two-level registration design SparseMoeBlock --> _QuantSparseMoe (calibration control, token counting, top_k override) .experts --> _QuantFusedExperts (per-expert F.linear interception + quantization) `register_fused_experts_on_the_fly` runs first to register the inner expert module; `register_sparse_moe_on_the_fly` then registers the outer block. `_QuantSparseMoe.layer_sync_moe_local_experts_amax` skips fused experts (they are not iterable), as per-expert amax is managed internally by `_QuantFusedExperts`. ### Known limitations - **`@use_experts_implementation` backends**: The `F.linear` interception only works with `experts_implementation="eager"` (default). `batched_mm` / `grouped_mm` use `torch.bmm` / `torch._grouped_mm` instead and are not intercepted. - **Storage offset fragility**: Expert index recovery via `storage_offset()` breaks under `.contiguous()`, FSDP2 redistribution, or `torch.compile` materialization. Runtime assertions are included. - **Toggle state machine**: Assumes exactly 2 `F.linear` calls per expert. Documented in docstrings. - **Non-standard MoE models**: DBRX, GptOss, Llama4, Step3p5 have different layouts and are already explicitly handled. The generic solution does not attempt to cover these. ### Testing - [x] Unit tests with synthetic fused expert model: detection, registration, quantization, export - [x] Verify existing sequential MoE tests still pass (`test_sparse_moe.py`) - [ ] GPU test with a real MoE model on transformers 5.x ### Before your PR is "Ready for review" - Is this change backward compatible?: Yes -- existing explicit registrations take priority; sequential MoE models are unaffected. - Did you write any new necessary tests?: Yes - Did you update Changelog?: No (pending) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added quantization support for fused Mixture-of-Experts (MoE) modules with automatic detection, per-expert quantization handling, and export to per-expert submodules; unified checkpoint export now supports fused MoE experts. * **Tests** * Added end-to-end tests covering fused-experts detection, conversion, forward correctness, expert index recovery, and export. * **Changelog** * Updated release notes to announce fused MoE expert support for Hugging Face exports. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Zhiyu Cheng <zhiyuc@nvidia.com> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> |
||
|
|
3a177f6d6b |
GPTQ official (#853)
## What does this PR do? Implements official version of GPTQ with decoder level sequential calibration flow. Ref: [https://github.com/IST-DASLab/FP-Quant/tree/master](https://github.com/IST-DASLab/FP-Quant/tree/master) **Type of change:** New feature <!-- Use one of the following: Bug fix, new feature, new example, new tests, documentation. --> **Overview:** 1. Deletes gptq_lite configuration. gptq() is a more generic implementation and will resolve to gptq_lite when we set use_sequential:False 2. Introduce GPTQHelper class to handle hessian collection patching/unpatching, blockwise weight update, hessian initialization, hessian_inverse computation etc ## Usage <!-- You can potentially add a usage example below. --> ```python # Add a code snippet demonstrating how to use this python hf_ptq.py --pyt_ckpt_path Qwen/Qwen3-8B --qformat nvfp4_gptq --kv_cache_qformat none --dataset cnn_dailymail --batch_size 32 --calib_seq 512 --calib_size 512 --export_path exported_model ``` ## Testing <!-- Mention how have you tested your change if applicable. --> Measure perplexity and activation MSE on the following models ## 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 a new NVFP4 GPTQ quantization option for CLI/configuration. * **Improvements** * Replaced GPTQ-lite with a full GPTQ calibration pipeline for more accurate, Hessian-aware blockwise updates and token-based sampling. * Added sequential layer-by-layer calibration support and automatic promotion of NVFP4 static quantizers. * Improved logging and timing for calibration runs. * **Tests** * Expanded GPTQ tests, including export/roundtrip validation for quantized weights. <!-- 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: Suguna Velury <178320438+sugunav14@users.noreply.github.com> Signed-off-by: realAsma <akuriparambi@nvidia.com> Co-authored-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Co-authored-by: realAsma <86726418+realAsma@users.noreply.github.com> |
||
|
Keval MorabiaandRinZ27
|
5dc17dfd15 |
[Security] Enable torch.load(weights_only=True) for secure checkpoint loading + trust_remote_code fix (#1181)
### What does this PR do? - Add secure checkpoint loading support using `torch.serialization.add_safe_globals([cls])`. This also removes 1 existing pickle usage. - Remove hard-coded `trust_remote_code=True` - Replaces https://github.com/NVIDIA/Model-Optimizer/pull/1056 by @RinZ27 ### Testing <!-- Mention how have you tested your change if applicable. --> CICD tests ran ### 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`: ✅ <!--- Mandatory --> - Did you write any new necessary tests?: ✅ <!--- Mandatory for new features or examples. --> - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ <!--- Only for new features, API changes, critical bug fixes or backward incompatible changes. --> ### Additional Information NVBug: 5999336 <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added safe checkpoint save/load helpers and a --trust_remote_code CLI flag in examples to control remote-code loading. * **Bug Fixes** * Checkpoint loading now defaults to safer, weights-only semantics to reduce arbitrary-code exposure. * **Documentation** * CHANGELOG updated with security guidance and opt-in procedure for unsafe checkpoint loading. * **Tests** * New unit tests validating the safe-load behavior. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: RinZ27 <222222878+RinZ27@users.noreply.github.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Co-authored-by: RinZ27 <222222878+RinZ27@users.noreply.github.com> |
||
|
Keval Morabiaandh-guo18
|
80d2f02a2d |
Fix spec dec example tests (#1183)
### What does this PR do? Type of change: Test fix <!-- Use one of the following: Bug fix, new feature, new example, new tests, documentation. --> - Fix `tests/examples/speculative_decoding` - previously silently skipped - Avoid pulling nemotron-post-training-dataset-v2 in tests to reduce chances of HF loading timeout in CICD - Make slow and redundant tests manual to speed up CICD ### Testing <!-- Mention how have you tested your change if applicable. --> - Tests passing ### 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 * **Chores** * Removed git‑LFS install step from CI and deleted an automated branch‑cleanup workflow * Trimmed example environment dependencies and relaxed transformers compatibility; added an optional tokenization dependency * **Tests** * Switched tests to generate datasets dynamically and improved fixture handling * Standardized PTQ test parameters (explicit calibration dataset) and refined GPU/test selection * **Bug Fixes** * Improved device-awareness and numeric handling in speculative decoding attention paths <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Co-authored-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> |
||
|
Shengliang Xu
|
1cceb950d6 |
[OMNIML-3689] PTQ quant_cfg semantic correction. Design in doc _quant_cfg.rst (#1094)
### What does this PR do?
#### Summary
Redesigns the `quant_cfg` configuration format in ModelOpt's PyTorch
quantization stack, replacing the previous dict-based format with an
**ordered list of typed `QuantizerCfgEntry` dicts**.
##### Motivation
The old `quant_cfg` dict had several pain points:
- **Ambiguous precedence**: no explicit way to reason about which entry
wins when multiple keys match a quantizer
- **Mixed key namespaces**: wildcard paths and PyTorch class names lived
in the same dict level, requiring ad-hoc dispatch
- **Magic `"default"` key**: an implicit, undocumented catch-all that
was easy to misuse
- **Poor composability**: merging two configs required dict updates that
silently discarded keys
- **No YAML round-trip fidelity**: the nested structure couldn't be
expressed cleanly in YAML
##### New format
`quant_cfg` is now an ordered list of `QuantizerCfgEntry` TypedDicts.
Each entry has:
- `quantizer_name` *(required)*: `fnmatch` wildcard matched against
quantizer module names
- `cfg` *(optional)*: dict (or list of dicts) of
`QuantizerAttributeConfig` fields
- `enable` *(optional)*: toggles quantizer on/off independently of `cfg`
- `parent_class` *(optional)*: restricts match to quantizers whose
parent module is of the given PyTorch class (e.g. `"nn.BatchNorm2d"`)
Entries are applied in list order; later entries override earlier ones.
The canonical pattern is deny-all first (`_base_disable_all`), then
selectively re-enable and configure, then apply standard exclusions
(`_default_disabled_quantizer_cfg`).
##### Changes
**Core library (`modelopt/torch/quantization/`)**
- **`config.py`**:
- Added `QuantizerCfgEntry` TypedDict (line 163) and
`find_quant_cfg_entry_by_path()` helper for exact-match lookup of
entries by path.
- Added `normalize_quant_cfg_list()` (line 1539) that converts legacy
formats (flat dict, single-key dicts, `nn.*`-scoped dicts, `"default"`
key) to canonical `QuantizerCfgEntry` lists. After normalization every
entry is guaranteed to have explicit `quantizer_name`, `enable`, and
`cfg` keys.
- Converted `_default_disabled_quantizer_cfg` and
`_mamba_moe_disabled_quantizer_cfg` from dicts to lists of
`QuantizerCfgEntry`.
- Added `_base_disable_all` (line 205): canonical deny-all entry
(`[{"quantizer_name": "*", "enable": False}]`).
- Converted all ~30 built-in config constants (`INT8_DEFAULT_CFG`,
`FP8_DEFAULT_CFG`, `NVFP4_DEFAULT_CFG`, etc.) to list format using
`*_base_disable_all` and `*_default_disabled_quantizer_cfg` unpacking.
- KV-cache configs (`FP8_KV_CFG`, `NVFP4_KV_CFG`, etc.) are now minimal
lists designed to be concatenated with a primary config — they
intentionally omit `_base_disable_all` and `"algorithm"`.
- Added two `QuantizeConfig` Pydantic field validators: a
`mode="before"` validator that calls `normalize_quant_cfg_list()`, and a
`mode="after"` validator that validates `cfg` dicts against
`QuantizerAttributeConfig`.
- Updated `need_calibration()` to iterate the normalized list instead of
the old dict.
- Changed `QuantizeQuantCfgType` alias from `dict[str | Callable, ...]`
to `list[QuantizerCfgEntry]`.
- **`conversion.py`**:
- Rewrote `set_quantizer_by_cfg()` (line 217) to iterate the list
directly. Each entry's `parent_class` is resolved via
`QuantModuleRegistry[parent_class_name]` (the existing `_DMRegistryCls`
registry).
- Added `set_quantizer_attributes_full()` (line 314): full replacement
of quantizer attributes from a `QuantizerAttributeConfig`. Unspecified
fields revert to defaults, enforcing entry atomicity. Can also upgrade
`TensorQuantizer` → `SequentialQuantizer` or downgrade the reverse.
- Added `set_quantizer_attributes_partial()` (line 384): merges a
partial `dict` of attributes into existing quantizer state. Does NOT
change quantizer structure. Used for enable-only entries.
- Added `set_quantizer_by_cfg_context()` context manager (line 447) that
temporarily applies a `quant_cfg` list and restores original quantizer
state on exit.
- Deprecated `set_quantizer_attribute()` (line 525) with a
`DeprecationWarning` pointing to the new functions.
- **`tensor_quantizer.py`**:
- `TensorQuantizer.set_from_attribute_config()`: narrowed type hint from
`dict` to `dict[str, Any]`.
- Added `_axis_setter` and `_block_sizes_setter` custom setters so that
`axis` and `block_sizes` changes properly propagate to the calibrator
and maintain mutual exclusivity.
- `SequentialQuantizer.set_from_attribute_config()`: narrowed signature
to `list[QuantizerAttributeConfig] | list[dict[str, Any]]` (removed the
old union with single values).
- **`algorithms.py`**:
- Updated `_match_quantizer_cfg()` to iterate the list and return
`(matched_cfg, matched_enable)` tuple with last-match-wins.
- Updated `_cfg_to_dict()`, `estimate_quant_compression()`, and
`QuantRecipe` to work with the list-based format.
- Updated `get_auto_quantize_config()` to emit list-format `quant_cfg`.
- **`model_quant.py`**: `disable_quantizer()` / `enable_quantizer()` now
call `set_quantizer_attributes_partial()` directly instead of the
deprecated `set_quantizer_attribute()`. Updated docstrings and code
examples to show the list format.
- **`utils/core_utils.py`**: `disable_lora_quantizers_in_config()` and
`update_quant_cfg_with_kv_cache_quant()` updated to append
`QuantizerCfgEntry` dicts to the list.
- **Other**: minor updates to `backends/fp8_per_tensor_gemm.py`,
`backends/nvfp4_gemm.py`, `compress.py`, `model_calib.py`,
`export/unified_export_hf.py`, and
`sparsity/attention_sparsity/conversion.py` to use the list format.
- **`onnx/llm_export_utils/quantization_utils.py`**: Updated
quantization config construction to use list format.
**YAML recipes (`modelopt_recipes/`)**
- Converted all 5 general PTQ recipes to the new list format:
- `general/ptq/fp8_default-fp8_kv.yml`
- `general/ptq/nvfp4_default-fp8_kv.yml`
- `general/ptq/nvfp4_experts_only-fp8_kv.yml`
- `general/ptq/nvfp4_mlp_only-fp8_kv.yml`
- `general/ptq/nvfp4_omlp_only-fp8_kv.yml`
- Converted model-specific recipe:
`models/Step3.5-Flash/nvfp4-mlp-only.yaml`
**Documentation (`docs/`)**
- New guide: `docs/source/guides/_quant_cfg.rst` — comprehensive
reference covering entry format, ordering semantics, entry atomicity,
`enable` vs `cfg` independence, `parent_class` filtering, and common
patterns (deny-all-then-enable, customizing a built-in config, building
from scratch).
- Updated `_pytorch_quantization.rst` code examples to show the list
format with `copy.deepcopy` and `.append()`.
- Added `_quant_cfg.rst` to the quantization guide table of contents.
**Examples**
- Updated all quantization examples to use the list format:
`deepseek/ptq.py`, `diffusers/quantization/config.py`,
`llm_ptq/hf_ptq.py`, `llm_qat/main.py`, `vllm_serve/vllm_ptq_utils.py`,
`llm_autodeploy/run_auto_quantize.py`, `llm_eval/quantization_utils.py`,
`llm_ptq/example_utils.py`,
`windows/torch_onnx/diffusers/qad_example/sample_example_qad_diffusers.py`,
and 2 notebooks.
**Tests**
- New test file:
`tests/unit/torch/quantization/test_config_validation.py` — unit tests
for `need_calibration()`, `normalize_quant_cfg_list()` (new format,
legacy format conversions, error cases),
`find_quant_cfg_entry_by_path()`, `_match_quantizer_cfg()`, and
`QuantizeConfig` Pydantic validators.
- Extended `tests/unit/torch/quantization/test_quantize_cpu.py` with
tests for `set_quantizer_attributes_full()` (atomicity, parent_class
filtering, SequentialQuantizer creation), list ordering, enable-only
entry behavior, and end-to-end legacy dict format.
- Updated 20+ existing test files across `tests/unit/`, `tests/gpu/`,
`tests/gpu_megatron/`, and `tests/_test_utils/` to use the list format.
##### Backward compatibility
`normalize_quant_cfg_list()` is called automatically by the
`QuantizeConfig` Pydantic `mode="before"` validator, so existing code
passing the old dict-based format (flat dict like `{"*weight_quantizer":
{"num_bits": 8}}`, single-key dict lists, or `nn.*`-scoped dicts with
`parent_class` semantics) continues to work without modification. The
legacy `"default"` key is converted to `quantizer_name: "*"`.
`set_quantizer_attribute()` is preserved as a deprecated wrapper around
`set_quantizer_attributes_partial()`.
#### Test coverage
- **Unit tests**: new `test_config_validation.py` with tests for
normalization, validation, path lookup, and cfg matching. Extended
`test_quantize_cpu.py` with tests for full/partial attribute setting,
ordering, atomicity, and legacy backward compatibility.
- **System testing**:
```
python examples/llm_ptq/hf_ptq.py \
--model Qwen/Qwen3-8B \
--recipe general/ptq/fp8_default-fp8_kv \
--export_path=build/fp8_default-fp8_kv42 \
--calib_size=16 \
--batch_size=0 \
--trust_remote_code \
--export_fmt=hf
```
### Additional Information
<!-- E.g. related issue. -->
---------
Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
|
||
|
kinjalpatel27andFrida Hou
|
48da2812c7 |
vllm fakequant reload with modelopt state for HF (#805)
## What does this PR do? **Type of change:** new feature **Overview:** - Added support to reload HF exported checkpoint with modelopt state in vLLM fakequant - Added support to reload all quantizer parameters using `QUANT_FILE_PATH` instead of only amax ## Usage ```python cd $PWD/examples/vllm_serve python hf_ptq_export.py --pyt_ckpt_path meta-llama/Llama-3.2-3B-Instruct --quant_cfg NVFP4_DEFAULT_CFG --dataset cnn_dailymail --export_path llama3.2-3b --trust_remote_code --batch_size 1 --calib_size 512 --kv_cache_quant_cfg NVFP4_KV_CFG cd $PWD/examples/vllm_serve MODELOPT_STATE_PATH=llama3.2-3b/vllm_fq_modelopt_state.pth python vllm_serve_fakequant.py ../llm_ptq/llama3.2-3b/ -tp 1 --served-model-name llama3.2-3b--host 0.0.0.0 --port 8001 --trust-remote-code --enforce-eager --disable-custom-all-reduce --gpu-memory-utilization 0.8 ``` ## Testing - Exported the checkpoint using hf_ptq, reloaded to vllm fakequant example, manually checked the quantizer values - Repeated above with mixed quantization by disabling quantization for few layers - exported checkpoint with megatron-lm with nvfp4, reloaded the checkpoint with vllm ## Before your PR is "*Ready for review*" <!-- If you haven't finished some of the above items you can still open `Draft` PR. --> - **Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)** and your commits are signed. - **Is this change backward compatible?**: Yes - **Did you write any new necessary tests?**: No - **Did you add or update any necessary documentation?**: Yes - **Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**: Yes ## Additional Information <!-- E.g. related issue. --> <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * vLLM can reload fakequant from ModelOpt state, load external quantizer files, perform TP-aware conversion/sharding, and includes a HF PTQ export utility. New env vars: QUANT_FILE_PATH, MODELOPT_STATE_PATH, CALIB_BATCH_SIZE. * **Documentation** * Serve README updated with export/load workflows, example commands, env var usage, and known limitations. * **Bug Fixes** * Improved restore compatibility and device-fallback behavior across quantization plugins to tolerate extra restore arguments. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Kinjal Patel <kinjalpravin@nvidia.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Co-authored-by: Frida Hou <201670829+Fridah-nv@users.noreply.github.com> |
||
|
kaix-nv
|
2bad66c8b2 |
[3/n] Add skip-softmax to Triton flash attention kernel (#1081)
### What does this PR do?
Type of change: ? <!-- Use one of the following: Bug fix, new feature,
new example, new tests, documentation. -->
<!-- Details about the change. -->
New feature. Add skip-softmax tile skipping to the Triton flash
attention kernel.
### Usage
```python
# Add a code snippet demonstrating how to use this
from modelopt.torch.kernels import attention
# Skip-softmax with threshold 0.1 (tiles contributing < 10% are skipped)
out = attention(q, k, v, b_start_loc, b_seq_len, max_len,
skip_softmax_threshold=0.1)
# Via mtsa.sparsify() on HuggingFace models
import modelopt.torch.sparsity.attention_sparsity as mtsa
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.1-8B",
torch_dtype=torch.bfloat16,
device_map="cuda")
# Default config
mtsa.sparsify(model, mtsa.SKIP_SOFTMAX_TRITON_DEFAULT)
```
### Testing
<!-- Mention how have you tested your change if applicable. -->
Performance (TFLOPS at seq_len=16384, RTX 6000 Pro):
| SEQ_LEN | ModelOpt Triton | PyTorch SDPA | Flash Attention 2 |
Skip-Softmax t=0.01 | Skip-Softmax t=0.1 |
|---:|---:|---:|---:|---:|---:|
| 16384.0 | 188.849922 | 211.718193 | 224.242843 | 172.901804 |
279.861684 |
| 32768.0 | 175.321787 | 212.815740 | 224.833553 | 146.150702 |
262.490463 |
| 65536.0 | 167.302839 | 214.932407 | 226.456141 | 145.082937 |
243.344791 |
</body></html>
### Before your PR is "*Ready for review*"
Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)
and your commits are signed (`git commit -s -S`).
Make sure you read and follow the [Security Best
Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors)
(e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(...,
weights_only=False)`, `pickle`, etc.).
- Is this change backward compatible?: ✅ / ❌ / N/A <!--- If ❌, explain
why. -->
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: ✅ / ❌ / N/A
<!--- Mandatory -->
- Did you write any new necessary tests?: ✅ / ❌ / N/A <!--- Mandatory
for new features or examples. -->
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
✅ / ❌ / N/A <!--- Only for new features, API changes, critical bug fixes
or backward incompatible changes. -->
### Additional Information
<!-- E.g. related issue. -->
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit
* **New Features**
* Added a Triton "skip-softmax" tile-skipping option for flash attention
with a new attention keyword and configurable threshold (default 0.1).
* Added a new sparse attention method and a default sparse configuration
that enables the Triton skip-softmax method.
* **Tests**
* Added GPU tests covering threshold behavior, numerical fidelity vs
dense, shape preservation, decode-mode, and integration with sparsify.
* **Documentation**
* Updated changelog for the new feature and removed two prior listed
entries.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---------
Signed-off-by: Kai Xu <kaix@nvidia.com>
|
||
|
kaix-nv
|
b1f9f0121b |
[2/n] Add sparse softmax to the Triton flash attention kernel (#1078)
### What does this PR do?
Type of change: ? <!-- Use one of the following: Bug fix, new feature,
new example, new tests, documentation. -->
Type of change: New feature
Add N:M structured sparsity support to the Triton flash attention kernel
(`modelopt/torch/kernels/triton_fa.py`). For every M consecutive key
positions in the attention score tile, keeps the top-N values and sets
the rest to -inf before softmax. This is applied during prefill only.
**Supported patterns:** Any N:M where M=4 (N=1,2,3) or M=8 (N=1..4).
- Sink tokens and dense window blocks for preserving local attention and
attention sinks
**Performance (TFLOPS at seq_len=16384, RTX 6000):**
| Pattern | TFLOPS | % of Dense |
|---------|--------|------------|
| Dense | 89.3 | 100% |
| 2:4 (M=4) | 69.5 | 78% |
| 4:8 (M=8) | 57.3 | 64% |
### Usage
```python
# Add a code snippet demonstrating how to use this
from modelopt.torch.kernels import attention
# 2:4 sparsity (keep top 2 of every 4 K positions)
out = attention(q, k, v, b_start_loc, b_seq_len, max_len,
sparsity_n=2, sparsity_m=4)
# 4:8 sparsity with sink tokens and dense window
out = attention(q, k, v, b_start_loc, b_seq_len, max_len,
sparsity_n=4, sparsity_m=8,
num_sink_tokens=4, dense_window_blocks=2)
# Dense (default, zero overhead)
out = attention(q, k, v, b_start_loc, b_seq_len, max_len)
# Via mtsa.sparsify() on HuggingFace models
import modelopt.torch.sparsity.attention_sparsity as mtsa
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.1-8B",
torch_dtype=torch.bfloat16,
device_map="cuda")
# Default config
mtsa.sparsify(model, mtsa.SPARSE_SOFTMAX_DEFAULT)
```
### Testing
<!-- Mention how have you tested your change if applicable. -->
### Before your PR is "*Ready for review*"
Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)
and your commits are signed (`git commit -s -S`).
Make sure you read and follow the [Security Best
Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors)
(e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(...,
weights_only=False)`, `pickle`, etc.).
- Is this change backward compatible?: ✅ / ❌ / N/A <!--- If ❌, explain
why. -->
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: ✅ / ❌ / N/A
<!--- Mandatory -->
- Did you write any new necessary tests?: ✅ / ❌ / N/A <!--- Mandatory
for new features or examples. -->
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
✅ / ❌ / N/A <!--- Only for new features, API changes, critical bug fixes
or backward incompatible changes. -->
### Additional Information
<!-- E.g. related issue. -->
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit
* **New Features**
* N:M structured sparse softmax for Triton flash-attention prefill with
configurable dense-window and sink-token handling.
* **API**
* attention(...) accepts sparsity_n, sparsity_m, num_sink_tokens,
dense_window_size; HF/Triton prefill path propagates them.
* **Configuration**
* New config fields and exported preset to enable/configure Triton N:M
sparse softmax with validation.
* **Tests**
* Added GPU tests covering N:M behavior, tile structure,
forward/backward correctness.
* **Documentation**
* CHANGELOG and example docs updated with usage and CLI options.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---------
Signed-off-by: Kai Xu <kaix@nvidia.com>
|
||
|
Keval Morabia
|
d698864137 |
Add MBridge pruning distillation CICD test + min transformers bumped to 4.56 (#1109)
### What does this PR do? Type of change: new test <!-- Use one of the following: Bug fix, new feature, new example, new tests, documentation. --> - Add GitHub CICD for MBridge pruning and distillation for tiny Qwen3 - Make default dtype for test hf models bf16 instead of fp32. Also bumped min transformers from 4.53 to 4.56 since Json serialization of config with torch.bfloat16 object is not supported in older version. - Reference: nemo:26.02.01 has transformers 4.56 and trt-llm 1.3.0rc5 has transformers 4.57 <!-- Details about the change. --> ### Testing From nemo:26.02.01 container: ```bash pytest tests/examples/megatron_bridge ``` ### 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`: ✅ <!--- 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 * **New Features** * Added NeMo example tests to CI and made pruning score-format and search-space divisors configurable. * **Tests** * Added end-to-end pruning/distillation tests and expanded tiny-model test utilities with bfloat16 support. * Relaxed numerical tolerances, added a GPU capability gating for a heavy test, and standardized dtype handling across tests. * **Chores** * Standardized pip invocations in CI, updated minimum supported Transformers version in configs, and reformatted some project metadata. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> |
||
|
Daniel KorzekwaandKeval Morabia
|
714dbd2c0d |
Activation hooks redesign (reuse hooks component across both minitron and puzzletron) (#1022)
### What does this PR do? Type of change: Redesign of existing feature This PR introduces a shared activation hooks infrastructure for minitron and puzzletron. The activation hooks framework provides a reusable component for collecting and analyzing activations during forward passes, which is used by both minitron pruning and puzzletron algorithms. Note! Minitron megatron.py/mcore_minitron.py:ImportanceEstimatorRegistry code does not use this component yet - will be refactored in a separate MR. **Key changes:** - Added `modelopt/torch/prune/importance_hooks` module with base hooks framework: - `base_hooks.py`: Core hook infrastructure for registering and managing forward hooks - `base_hooks_analysis.py`: Analysis utilities for processing collected activations - `megatron_hooks.py`: Megatron-specific hook implementations - `compare_module_outputs.py`: Utilities for comparing module outputs - Added unit tests in `tests/gpu/torch/prune/importance_hooks`: - `test_base_hooks.py`: Tests for base hooks functionality - `test_base_hooks_analysis.py`: Tests for activation analysis utilities - Updated `test_mcore_gpt_minitron_pruning.py` to validate activation collection - Updated test utilities for distributed testing support ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ Yes - This is a new module that doesn't affect existing functionality - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: ✅ No - Did you write any new necessary tests?: ✅ Yes - Added comprehensive tests for the activation hooks infrastructure - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ N/A - This is infrastructure code that will be used by subsequent PRs <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Robust JSON utilities for complex objects. * Flushed-print helper for synchronized output. * Activation-based importance estimation framework with multiple hook implementations, Megatron plugin support, and layer-output comparison tools. * Project-root test fixture for test suites. * **Tests** * Many new end-to-end tests validating hooks, evaluation metrics, and multi-layer output comparisons. * Improved pruning tests with deterministic initialization and additional statistical assertions. * **Chores** * Tests: mirror rank/size into environment for test init and disable external telemetry. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Co-authored-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> |
||
|
kaix-nv
|
7c33d85607 |
[1/n] Add a Triton attention kernel with HF integration (#1034)
### What does this PR do?
Type of change: ? <!-- Use one of the following: Bug fix, new feature,
new example, new tests, documentation. -->
<!-- Details about the change. -->
- Adds a Triton flash attention kernel (triton_fa.py) with HF
integration for use in sparse attention and quantization workflows. The
kernel implements Flash Attention with varlen support, GQA, causal
masking, and forward/backward.
- Update the sparse attention to support backend="triton".
Key components:
- modelopt/torch/kernels/triton_fa.py -- Core Triton kernel
- modelopt/torch/kernels/hf_triton_attention.py -- HF adapter,
registered as attn_implementation="modelopt_triton"
- modelopt/torch/kernels/__init__.py -- Shared kernel registry
- modelopt/torch/sparsity/attention_sparsity/conversion.py -- Backend
selection (backend="triton" or "pytorch")
- modelopt/torch/sparsity/attention_sparsity/config.py -- Added "triton"
as valid backend option
- examples/llm_sparsity/attention_sparsity/hf_sa.py -- Updated example
to support --backend triton
### Usage
```python
# Direct kernel API (varlen packed format)
from modelopt.torch.kernels import attention
o = attention(
q, k, v, # [total_tokens, heads, head_dim]
b_start_loc=b_start_loc, # [batch] per-sequence start offsets
b_seq_len=b_seq_len, # [batch] per-sequence lengths
max_input_len=max_seq_len,
is_causal=True,
)
# HuggingFace integration (automatic via sparsify)
import modelopt.torch.sparsity.attention_sparsity as mtsa
config = {"sparse_cfg": {"*attn*": {"method": "flash_skip_softmax", "backend": "triton", "enable": True}}}
model = mtsa.sparsify(model, config=config)
# model now uses the Triton kernel for attention
# Or load directly with attn_implementation
model = AutoModelForCausalLM.from_pretrained(path, attn_implementation="modelopt_triton")
```
### Testing
<!-- Mention how have you tested your change if applicable. -->
`tests/gpu/torch/sparsity/attention_sparsity/test_triton_fa.py`
#### Kernel benchmark on RTX 6000
| SEQ_LEN | ModelOpt Triton | PyTorch SDPA | Flash Attention 2 |
|--------:|----------------:|-------------:|------------------:|
| 256.0 | 34.435353 | 26.199215 | 47.927293 |
| 512.0 | 60.216998 | 47.408218 | 80.736116 |
| 1024.0 | 81.209990 | 82.673526 | 94.197181 |
| 2048.0 | 88.800973 | 89.239451 | 94.822496 |
| 4096.0 | 88.302953 | 89.192071 | 96.826178 |
| 8192.0 | 89.538177 | 89.115835 | 91.563461 |
| 16384.0 | 85.457533 | 80.509254 | 81.391092 |
### Before your PR is "*Ready for review*"
Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)
and your commits are signed (`git commit -s -S`).
Make sure you read and follow the [Security Best
Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors)
(e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(...,
weights_only=False)`, `pickle`, etc.).
- Is this change backward compatible?: ✅ / ❌ / N/A <!--- If ❌, explain
why. -->
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: ✅ / ❌ / N/A
<!--- Mandatory -->
- Did you write any new necessary tests?: ✅ / ❌ / N/A <!--- Mandatory
for new features or examples. -->
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
✅ / ❌ / N/A <!--- Only for new features, API changes, critical bug fixes
or backward incompatible changes. -->
### Additional Information
<!-- E.g. related issue. -->
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit
## Release Notes
* **New Features**
* Added Triton flash attention backend support for sparse attention
operations, alongside the existing PyTorch backend, enabling improved
performance for compatible hardware.
* **Documentation**
* Updated README to document both available attention backends and their
configurations.
* **Tests**
* Added comprehensive test coverage for the new Triton backend,
including forward and backward pass validation.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---------
Signed-off-by: Kai Xu <kaix@nvidia.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>
|
||
|
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>
|
||
|
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
|
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>
|
||
|
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> |
||
|
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> |
||
|
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> |
||
|
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>
|
||
|
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>
|
||
|
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> |
||
|
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> |
||
|
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>
|
||
|
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> |
||
|
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> |
||
|
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> |
||
|
realAsma
|
ac30686c82 |
Track global_amax for weight FP4 MSE sweep; Refactor to NVFP4StaticQantizer, NVFP4MSECalibrator (#849)
## What does this PR do? **Type of change:** ? <!-- Use one of the following: Bug fix, new feature, new example, new tests, documentation. --> **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 ## Release Notes * **New Features** * Added NVFP4StaticQuantizer for improved 4-bit quantization with enhanced precision control * Introduced NVFP4MSECalibrator with flexible candidate generation for calibration optimization * **Improvements** * Optimized GPU kernels for Hopper+ graphics cards with better performance * Extended Triton support to broader GPU compatibility * Enhanced backward compatibility for restoring previously quantized models * **Tests** * Added comprehensive test coverage for new quantizers and calibration methods <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: realAsma <akuriparambi@nvidia.com> |
||
|
Jenny Chen
|
8a8c250d12 |
Latent MOE & Repeated MTP support for NemotronH; fix KV cache quant export (#830)
## What does this PR do? **Type of change:** New feature and bug fix **Overview:** Support Latent MOE and Repeated MTP for NemotronH models - Enable latent MOE modules during megatron import/export - Fix KV cache quantization export: remove old `qkv_layer.output_quantizer` export & replace with proper `k/v_bmm_quantizer` logic - Improvements to EP amax sync - Support repeated MTP import/export for NemotronH models (only BF16 export for MTP for now) ## Usage <!-- You can potentially add a usage example below. --> ```python # Add a code snippet demonstrating how to use this ``` ## Testing <!-- Mention how have you tested your change if applicable. --> ## Before your PR is "*Ready for review*" <!-- If you haven't finished some of the above items you can still open `Draft` PR. --> - **Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)** and your commits are signed. - **Is this change backward compatible?**: Yes/No <!--- If No, explain why. --> - **Did you write any new necessary tests?**: Yes/No - **Did you add or update any necessary documentation?**: Yes/No - **Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**: Yes/No <!--- Only for new features, API changes, critical bug fixes or bw breaking changes. --> ## Additional Information <!-- E.g. related issue. --> <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit ## Release Notes * **New Features** * Added support for grouped MLP and self-attention scaling operations in model export workflows * Enhanced model parallel training capabilities with improved component mapping * Expanded quantization configuration handling with dynamic module exclusion across distributed ranks * Improved support for additional transformer engine components * **Refactor** * Reorganized internal export and import logic for improved maintainability and specialist model architecture support <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: jenchen13 <jennifchen@nvidia.com> Signed-off-by: Jennifer Chen <jennifchen@nvidia.com> Signed-off-by: Jenny Chen <jennifchen@nvidia.com> |
||
|
sugunav14
|
02c5f292f0 |
GPTQ Lite implementation (#555)
## 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:** Adds support for GPTQ algorithm. This PR implements a modified version of the official GPTQ algorithm; the key difference is that updated activations from each layer are not used for hessian computation ## Usage <!-- You can potentially add a usage example below. --> Modify "algorithm" field in quant_cfg to "gptq_lite". Note: Does not currently work with AWQ ```python # Add a code snippet demonstrating how to use this ``` ## Testing <!-- Mention how have you tested your change if applicable. --> - [x] Added unit tests to test helper functions + e2e flow - [x] Perplexity and GPQA results | Model | Qformat | Perplexity wikitext2 | GPQA | |-------------|------------------------|------------------------|------| | Qwen3-8B | INT4 weight only (modelopt + amax/7) no GPTQ | 10.75 | n/a | | Qwen3-8B | INT4 weight only (modelopt + amax/7) | **10.56** | 0.388 | | Qwen3-8B | INT4 weight only + FP-Quant hessians + amax/7.5 | 10.25 | **0.449** | | Qwen3-8B | INT4 weight only (FP-Quant) | 10.24 | 0.46 | | Qwen3-8B | NVFP4 static weight only | 10.25 | n/a | | Qwen3-8B | NVFP4 static weight only no GPTQ | 10.25 | n/a | | Qwen3-0.6B | NVFP4 static weight only | **22.75** | n/a | | Qwen3-0.6B | NVFP4 dynamic weight only | 23.50 | n/a | | Qwen3-0.6B | NVFP4 static weight only with FP-Quant hessians | 22.0 | n/a | | Qwen3-0.6B | NVFP4 static weight only no GPTQ | 24.25 | n/a | Conclusions from results - Perplexity remains the same or shows improvement with Modelopt implementation. The magnitude of improvement is lesser in modelopt when compared to FP-Quant - GPQA shows no improvement with modelopt, but shows improvement with FP-Quant ## 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 <!--- 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/TensorRT-Model-Optimizer/blob/main/CHANGELOG.rst)?**: No <!--- Only for new features, API changes, critical bug fixes or bw breaking changes. --> ## Additional Information <!-- E.g. related issue. --> <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit ## Release Notes * **New Features** * GPTQ Lite quantization mode now available for efficient model calibration * GPU memory usage monitoring utility added * Quantization configuration extended to support complex nested structures and lists * **Tests** * Comprehensive test coverage added for GPTQ quantization 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: Suguna Velury <178320438+sugunav14@users.noreply.github.com> |
||
|
danisereb
|
4227bb7366 |
Add support for MXFP8 PTQ (#736)
## 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 MXFP8 PTQ, enabling MXFP8 hardware
acceleration during inference on Blackwell GPUs.
## Usage
<!-- You can potentially add a usage example below. -->
```bash
export MODEL_PATH=/my_home/hf_models/nvidia/OpenMath2-Llama3.1-8B
export OUTPUT_PATH=/my_home/hf_models/nvidia/OpenMath2-Llama3.1-8B-MXFP8
mkdir -p $OUTPUT_PATH
python examples/llm_ptq/hf_ptq.py \
--export_fmt hf \
--dataset cnn_dailymail \
--pyt_ckpt_path $MODEL_PATH \
--export_path $OUTPUT_PATH \
--qformat mxfp8
```
The `hf_quant_config.json` of the output checkpoint:
```json
{
"producer": {
"name": "modelopt",
"version": "0.41.0.dev50+g7a796a875"
},
"quantization": {
"quant_algo": "MXFP8",
"kv_cache_quant_algo": "FP8",
"group_size": 32,
"exclude_modules": [
"lm_head"
]
}
}
```
And `config.json` (only the `quantization_config`):
```json
...
"quantization_config": {
"ignore": [
"lm_head"
],
"quant_algo": "MXFP8",
"kv_cache_scheme": {
"dynamic": false,
"num_bits": 8,
"type": "float"
},
"producer": {
"name": "modelopt",
"version": "0.41.0.dev50+g7a796a875"
},
"quant_method": "modelopt"
}
```
## Testing
<!-- Mention how have you tested your change if applicable. -->
Used `hf_ptq.py` to quantize the model `nvidia/OpenMath2-Llama3.1-8B`
([available in
hugging-face](https://huggingface.co/nvidia/OpenMath2-Llama3.1-8B)), see
the example command above.
Checked that the generated MXFP8 checkpoint can be loaded with vLLM
(required changes in vLLM, not merged to main).
Added tests for `MXFP8QTensor` in
`tests/gpu/torch/quantization/test_qtensor_cuda.py`.
Added "mxfp8" in `tests/examples/llm_ptq/test_llm_ptq.py`
#### Support for Nemotron Models
Verify that Nemotron Nano V3 BF16 can be converted to MXFP8 using
`hf_ptq.py`:
https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16
## 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 MXFP8 quantization format support with new scaling mechanisms
and quantization utilities.
* Updated configuration options, example scripts, and utilities to
recognize and process MXFP8 quantization workflows.
* Extended quantization export pipelines to handle MXFP8 quantized
models.
* **Tests**
* Expanded test coverage for MXFP8 quantization across various tensor
shapes, data types, and device configurations.
<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: Daniel Serebrenik <daserebrenik@nvidia.com>
|
||
|
Frida Hou
|
aafd388394 |
add FP8 sweep option for static NVFP4 MSE (#758)
## 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:** ?
Adds fp8_scale_sweep mode to MSE calibrator for optimizing FP8-quantized
per-block scales in NVFP4 format.
## Usage
<!-- You can potentially add a usage example below. -->
Tested with this config
```python
NVFP4_WEIGHT_MSE_FP8_SWEEP_CFG = {
"quant_cfg": {
"*weight_quantizer": {
"num_bits": (2, 1),
"block_sizes": {-1: 16, "type": "static", "scale_bits": (4, 3)},
"axis": None,
"enable": True,
},
"*input_quantizer": {
"enable": False,
},
**_default_disabled_quantizer_cfg,
},
"algorithm": {
"method": "mse",
"fp8_scale_sweep": True,
},
}
```
## 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 FP8 scale sweep option for quantization calibration, enabling
optimized scale value sweeping for NVFP4 per-block quantization.
- Introduced new NVFP4_WEIGHT_MSE_CFG configuration preset for improved
weight quantization workflows.
**Tests**
- Added test coverage validating FP8 scale sweep functionality and reset
behavior.
<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: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
|
||
|
sychen52
|
b44c60ad34 |
Svdquant huggingface checkpoint export support (#754)
## What does this PR do?
**Type of change:** new feature
**Overview:**
## Usage
```bash
cd ./examples/llm_ptq/
python hf_ptq.py \
--pyt_ckpt_path Qwen/Qwen3-4B \
--export_path /home/scratch.shiychen_coreai/quantized_models/Qwen3-4B-svdq \
--qformat nvfp4_awq_svdquant --kv_cache_qformat none --sparsity_fmt dense --calib_size 8
```
## Testing
exported checkpoint and loaded.
## 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 nvfp4_svdquant as a new quantization format option for LLM model
quantization workflows.
* **Limitations**
* Multi-GPU export configurations using tensor or pipeline parallelism
are not supported with nvfp4_svdquant quantization.
<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: Shiyang Chen <shiychen@nvidia.com>
|
||
|
Keval Morabia
|
563a1e09c6 |
Add NAS to Minitron pruning for parameter based auto-pruning (#720)
## What does this PR do?
**Type of change:** New feature
- So far users we didnt have the NAS step from the Minitron paper so
users had to manually find a pruned architecture which fits that params
constraints and for that they have to tweak different combinations of
width and depth pruning
- This PR adds a simplified version of the NAS from paper. We first find
all candidate subnets that fit the user's params constraint and sort
them by parameter count.
- [Paper] Then we pick top K candidates, do distillation for ~2B tokens
and then select the one with best score (LM Loss / MMLU / other metric
we care about). Note that the one among these top K with highest params
is often not the best pruned model
- [ModelOpt] Then we pick top K candidates, and select the one with best
score (LM Loss / MMLU / other metric we care about). While doing KD
gives better indication on which one to pick, skipping it makes the
pruning much faster, much less compute intense, and finish everything in
single prune API instead of first exporting top K models, doing KD and
eval for all K models separately. We do print a Note in pruning step to
let users know this so they can do KD if they want slightly better
pruned model.
- Further full KD is still needed as usual
- We also restrict the search space choices (e.g. `hidden_size` multiple
of 256, `ffn_hidden_size` multiple of 512) to make the process
efficient. Users can configure this if they want to.
## Usage
<!-- You can potentially add a usage example below. -->
Pruning API is same as before:
```python
import modelopt.torch.prune as mtp
mtp.prune(
model,
mode="mcore_minitron",
constraints=constraints,
dummy_input=None, # Not used
config=config,
)
```
1. Manual Pruning (Existing):
```python
constraints = {"export_config": {"hidden_size: 3072", "ffn_hidden_size": 9216}}
config = {"forward_loop": forward_loop, "checkpoint": "/path/to/cache/pruning/scores.pth"}
mtp.prune(...)
```
2. NAS-based Auto Pruning (New):
```python
constraints = {"params": 6e9}. # prune to 6B params
config = {"forward_loop": forward_loop, "checkpoint": "/path/to/cache/pruning/scores.pth"}
# define the score_func to maximize (e.g MMLU, negative val loss, etc.)
from modelopt.torch.utils.plugins.megatron_mmlu import megatron_mmlu
def score_func(m):
return megatron_mmlu(m, tokenizer, percentage=0.05) # 5% sampled data for faster eval
config["score_func"] = score_func
# overwrite search space choices (showing defaults):
config["max_width_pruning"] = 0.4
config["max_depth_pruning"] = 0.2
config["hparams_to_skip"] = [] # can be used to disable pruning some hparams e.g. ["num_attention_heads"]
config["top_k"] = 10 # might be better to use 20 at the cost of longer time to prune
mtp.prune(...)
```
To configure search space (shows defaults):
```python
ss_config = mtp.mcore_minitron.get_mcore_minitron_config(
hidden_size_divisor=256,
ffn_hidden_size_divisor=512,
mamba_head_dim_divisor=8,
num_moe_experts_divisor=8,
num_layers_divisor=2,
)
mtp.prune(model, mode=[("mcore_minitron", ss_config)], ....)
```
## Testing
**Qwen3-8B -> 6B (~2 hours on 8xA5000)**
```python
0.4350 score -> {'num_layers': 34, 'hidden_size': 3328, 'ffn_hidden_size': 11264}
BEST 0.5705 score -> {'num_layers': 30, 'hidden_size': 3584, 'ffn_hidden_size': 11776}
0.4051 score -> {'num_layers': 36, 'hidden_size': 3840, 'ffn_hidden_size': 8192}
0.4593 score -> {'num_layers': 36, 'hidden_size': 3584, 'ffn_hidden_size': 9216}
0.2737 score -> {'num_layers': 36, 'hidden_size': 3072, 'ffn_hidden_size': 11776}
0.5556 score -> {'num_layers': 32, 'hidden_size': 3584, 'ffn_hidden_size': 10752}
0.3198 score -> {'num_layers': 28, 'hidden_size': 4096, 'ffn_hidden_size': 10240}
0.4119 score -> {'num_layers': 36, 'hidden_size': 4096, 'ffn_hidden_size': 7168}
0.3808 score -> {'num_layers': 36, 'hidden_size': 3328, 'ffn_hidden_size': 10240}
0.4783 score -> {'num_layers': 34, 'hidden_size': 3840, 'ffn_hidden_size': 8704}
```
**Nemotron-Nano-9B-v2 -> 7B (~2.5 hours on 8xA5000)**
```python
0.2629 score -> {'num_layers': 54, 'hidden_size': 4352, 'mamba_num_heads': 88, 'mamba_head_dim': 72, 'ffn_hidden_size': 15360}
0.2778 score -> {'num_layers': 50, 'hidden_size': 4480, 'mamba_num_heads': 128, 'mamba_head_dim': 56, 'ffn_hidden_size': 15680}
0.5041 score -> {'num_layers': 56, 'hidden_size': 4096, 'mamba_num_heads': 96, 'mamba_head_dim': 80, 'ffn_hidden_size': 14336}
BEST 0.6043 score -> {'num_layers': 48, 'hidden_size': 4352, 'mamba_num_heads': 120, 'mamba_head_dim': 80, 'ffn_hidden_size': 13824}
0.0772 score -> {'num_layers': 56, 'hidden_size': 4352, 'mamba_num_heads': 112, 'mamba_head_dim': 80, 'ffn_hidden_size': 10240}
0.3550 score -> {'num_layers': 50, 'hidden_size': 4480, 'mamba_num_heads': 112, 'mamba_head_dim': 64, 'ffn_hidden_size': 15680}
0.1016 score -> {'num_layers': 56, 'hidden_size': 4352, 'mamba_num_heads': 120, 'mamba_head_dim': 72, 'ffn_hidden_size': 10752}
0.5461 score -> {'num_layers': 46, 'hidden_size': 4480, 'mamba_num_heads': 128, 'mamba_head_dim': 72, 'ffn_hidden_size': 14848}
0.1992 score -> {'num_layers': 54, 'hidden_size': 4480, 'mamba_num_heads': 80, 'mamba_head_dim': 80, 'ffn_hidden_size': 14336}
0.5881 score -> {'num_layers': 48, 'hidden_size': 4480, 'mamba_num_heads': 112, 'mamba_head_dim': 80, 'ffn_hidden_size': 13824}
```
## 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 <!--- If No, explain why.
-->
- **Did you write any new necessary tests?**: Yes
- **Did you add or update any necessary documentation?**: Yes
- **Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**:
Yes <!--- Only for new features, API changes, critical bug fixes or bw
breaking changes. -->
## Additional Information
OMNIML-3043
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit
## Release Notes
* **New Features**
* Added NAS-based Auto Pruning for Minitron models as an alternative to
manual pruning using parameter constraints
* Introduced parameter counting capabilities for architecture search
* **Documentation**
* Expanded pruning guides with detailed examples and workflows for both
manual and automatic pruning approaches
* Updated configuration documentation with granular divisor parameters
* **Improvements**
* Enhanced parameter counting support for models with dynamic or
mixture-of-experts modules
<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: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
|
||
|
Frida Hou
|
18d9b1eea4 |
Add static per block MSE for NVFP4 weight (#613)
## 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:** ?
Support static block-wise MSE for NVFP4 weight quantization.
Add a FP4 triton kernel that take in scales for each block. It also
quantizes the scales to FP8.
This PR does the following:
1. Enable static NVFP4 implementation, i.e. block scales for weights are
calculated during calibration and feed into fake quant kernels
2.Extend mse_calibrate to support static NVFP4 with block scales
searching by MSE and global scale set as MAX
3.Refinements: calibrate weight quantizers only once during MSE
calibration
## Usage
<!-- You can potentially add a usage example below. -->
Example config:
```python
NVFP4_WEIGHT_MSE_CFG = {
"quant_cfg": {
"*weight_quantizer": {
"num_bits": (2, 1),
"block_sizes": {-1: 16, "type": "static", "scale_bits": (4, 3)},
"axis": None,
"enable": True,
},
"*input_quantizer": {
"enable": False,
},
**_default_disabled_quantizer_cfg,
},
"algorithm": {
"method": "mse",
"step_size": 0.25,
"start_multiplier": 0.25,
"stop_multiplier": 2.0,
},
}
NVFP4_WEIGHT_ACT_MSE_CFG = {
"quant_cfg": {
"*weight_quantizer": {
"num_bits": (2, 1),
"block_sizes": {-1: 16, "type": "static", "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": "mse",
"step_size": 0.25,
"start_multiplier": 0.25,
"stop_multiplier": 2.0,
},
}
```
## 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/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: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
|
||
|
Asha Anoosheh
|
b813ab548b |
Top-K KL Divergence loss (#747)
## What does this PR do? **Type of change:** New feature **Overview:** Writes a new KLDiv Logits loss which only uses top-k vocab values ## 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 * **New Features** * Added Top-K logit filtering capability for knowledge distillation workflows, enabling selective focus on high-probability tokens. * **Improvements** * Enhanced distributed tensor model-parallel operations with improved awareness for gradient computation and reduction. * Simplified legacy distributed operation constructs. * **Tests** * Introduced comprehensive test coverage for Megatron-based distillation, validating both standard and Top-K filtering variants. <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: Asha Anoosheh <aanoosheh@nvidia.com> |
||
|
realAsmaandAsma Thekkumpate
|
81c509c643 |
Fixes & Simplifications for MCore KVCache QAT/QAD; Unittests; Distributed Sync of KVCache Quantizer params (#727)
## What does this PR do? **Type of change:** Fix MCore KV Cache Quantization: Amax Device Placement Bug; Code clean up; Distributed Sync of KVCache Quantizer params; unittest expansion to hybrid models **Overview:** Fixes bugs preventing MCore KV Cache quantization from working during checkpoint restore. ### Bug Chain **Bug 1:** `is_enabled = self.weight_quantizer.is_enabled if hasattr(self, "weight_quantizer") else False` No `weight_quantizer` for KV-cache-only quant → `is_enabled=False` → metadata not saved → `modelopt_post_restore()` never called. *(Thanks to @jenchen13 )* **Bug 2:** After fixing Bug 1, `_amax` restored on CPU (via `_reset_pytorch_state_from_metadata`). Fallback `_calibrate_quantizers()` never called because `_amax` exists. **Bug 3:** Even if called, `_calibrate_quantizers()` fails — `core_attention` has no parameters → can't determine device/dtype. ### The Fix 1. Remove `is_enabled` check entirely — disabled modules may still need metadata restore. Explicitly skip `output_layer` from extra state callbacks (never quantized) 2. Set `dtype`/`device` on `core_attention` from parent Attention module, `modelopt_post_restore()` calls `self.to(device, dtype)` 3. Remove dead `_calibrate_quantizers()` code (will bring back similar logic for KV cache affine quantization) ### Previous Unit Test Was Wrong `model_test` was `mtq.quantize()`'d, not `mto.restore()`'d. Never tested actual restore path. ### Additional Fixes - Amax sync across DP/TP for KV cache quantizers - `flash_decode` auto-disabled ### Code Cleanup Removed ~100 lines of dead code. ## Testing 1. MCore KV Cache QAD with Nano V3 + Context Parallel works 2. Unit tests: hybrid models, KV+GEMM configs, correct restore workflow, backward pass validation ## 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?**: Yes --------- Signed-off-by: realAsma <akuriparambi@nvidia.com> Co-authored-by: Asma Thekkumpate <akuriparambi@cw-dfw-cs-001-vscode-02.cm.cluster> |
||
|
realAsma
|
cb343352ca |
Registry interface for custom quantization functional backend (#683)
## What does this PR do? **Type of change:** ? <!-- Use one of the following: Bug fix, new feature, new example, new tests, documentation. --> **Overview:** Add registry interface for custom quantization functional backend ## Usage <!-- You can potentially add a usage example below. --> see `tests/unit/torch/quantization/test_custom_backend.py` for usage example. ## 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/No - **Did you add or update any necessary documentation?**: Yes - **Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**: Yes ## Additional Information <!-- E.g. related issue. --> Signed-off-by: realAsma <akuriparambi@nvidia.com> |
||
|
Keval Morabia
|
b3ec529638 |
Make Minitron search space choices configurable for smaller test models (#709)
- Test update, no code change --------- Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> |