mirror of
https://github.com/NVIDIA/Model-Optimizer.git
synced 2026-10-25 06:43:41 +08:00
0.45.0dev
249
Commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
 Frida HouandKeval Morabia
|
97d153118e |
[minor] Add custom calibration backend registry (#1281)
### 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 a public backend-specific calibrator registration API to support FP8 scale-sweep calibration, allowing backends to supply custom calibrators used during FP8 tuning. * **Tests** * Added unit tests confirming registry insertion/overwrite, that registered calibrators are invoked when FP8 scale-sweep is enabled, are not invoked when disabled, and that calibration falls back to defaults when no backend is registered. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Co-authored-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> |
||
|
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> |
||
|
Keval MorabiaandClaude Sonnet 4.6
|
3d0f0db49e |
[CI] Replace tox with nox, use nemo:26.04 for megatron tests, and simplify CI workflows (#1286)
### What does this PR do? Type of change: New feature / infrastructure improvement Follow-up to #1285 for correct CI test environment for megatron based tests Replaces `tox` + `tox-current-env` with `nox` for all test, lint, docs, and wheel build sessions. The primary motivation was that `tox-current-env` is incompatible with uv venvs in NGC containers (e.g. NeMo's `/opt/venv`) — it picks the system Python via `sys._base_executable` instead of the container's venv Python which has megatron packages pre-installed. Key changes: - **`noxfile.py`** replaces `tox.ini` with GPU, CPU unit, partial-install, pre-commit, docs, and wheel sessions - **GPU sessions** use `venv_backend="none"` (run directly in container env) and `python -m pip/pytest` to avoid PATH mismatches - **uv** is set as the default venv backend (if available) for CPU sessions (faster installs) Also includes CI workflow simplifications: - **`_pr_gate.yml`** new reusable workflow centralizing file-change detection + linux-check wait logic (was duplicated across 3 workflow files) - **Collapsed pr/non-pr job pairs** into single jobs with conditional `runs-on` in `gpu_tests.yml`, `example_tests.yml`, `regression_tests.yml` - **Collapsed `multi-py` / `multi-torch` / `multi-transformers`** into a single `multi-version` matrix job in `unit_tests.yml` - **PR path filtering** for unit test secondary jobs (multi-version, launcher, partial-install) — skipped if no relevant files changed - **Fixed schedule/workflow_dispatch skipping** — jobs with `needs: [pr-gate]` were incorrectly skipped when all pr-gate internal jobs were skipped; fixed by making the gate job always run - **multi-version, launcher, partial-install** now also run on `schedule` / `workflow_dispatch` ### Usage ```bash python -m pip install nox uv # install nox and uv (once) nox -l # list all sessions nox -s gpu_megatron # run a GPU session (inside container) nox -s "unit-3.12(torch_211, tf_latest)" # run a specific unit test combination nox -s "unit-3.12(torch_211, tf_latest)" -R # force-recreate venv (e.g. after dep changes) COVERAGE_PROCESS_START=pyproject.toml nox -s "unit-3.12(torch_211, tf_latest)" # with coverage ``` ### Testing - Ran `nox -l` to verify all session names - Ran `gpu_megatron` session locally inside NeMo container — confirmed it uses `/opt/venv/bin/python` correctly - Manually triggered nightly-runs: - Unit: https://github.com/NVIDIA/Model-Optimizer/actions/runs/24608013657 - GPU: https://github.com/NVIDIA/Model-Optimizer/actions/runs/24608018763 - Examples: https://github.com/NVIDIA/Model-Optimizer/actions/runs/24608017322 ### 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 — CI infrastructure only - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: ✅ (added `nox` and `uv` to `dev-test`, both Apache-2.0) - 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 — no user-facing changes ### Additional Information Supersedes the tox-current-env workaround in the parent branch. --------- Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com> |
||
|
|
760c980727 |
Add ResNet50 support for torch_onnx quantization workflow (#1263)
## Summary - Add end-to-end ResNet50 support in the torch_onnx quantization → ONNX export → TRT engine pipeline - Fix multiple Conv2d-related export issues that blocked Conv2d-heavy models from working with FP8/INT8/MXFP8/NVFP4/auto quantization modes - Fix `configure_linear_module_onnx_quantizers` to handle all modules with block quantization (not just `nn.Linear`), fixing NVFP4/MXFP8 export for models with quantized non-Linear modules - Add `--trt_build` flag to `torch_quant_to_onnx.py` and simplify test infrastructure ### Files Changed - `modelopt/torch/_deploy/utils/torch_onnx.py` — Disable FP8 Conv2d weight quantizers and autocast during ONNX export - `modelopt/torch/quantization/export_onnx.py` — Fix `configure_linear_module_onnx_quantizers` for all module types with block quantization - `examples/torch_onnx/torch_quant_to_onnx.py` — Add `--trt_build` flag, calibration for FP8 override quantizers, Conv2d→FP8 override for auto mode, filter_func updates - `examples/torch_onnx/README.md` — Add ResNet50 to supported models table - `tests/examples/torch_onnx/test_torch_quant_to_onnx.py` — Add ResNet50 test entry, simplify using `--trt_build` - `tests/_test_utils/torch/vision_models.py` — Add ResNet50 to timm model registry ### Quantization modes passing - ✅ FP8, INT8, MXFP8, NVFP4, Auto (all 5 modes pass export + TRT build) - INT4_AWQ excluded (pre-existing limitation for all models) ## Test plan - [x] All 5 resnet50 test modes pass: `pytest tests/examples/torch_onnx/test_torch_quant_to_onnx.py -k resnet50` (5/5 passed) - [x] Full regression: 18 passed, 2 failed (pre-existing swinv2_tiny fp8/int8 failures) 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added ResNet50 to supported ONNX export vision models with FP8, INT8, MXFP8, and NVFP4 support. * Optional TensorRT engine build after export via a new CLI flag. * **Improvements** * Enhanced quantization calibration and export flows for FP8/INT8 models, including broader block-quantization support across module types and safer export handling. * Tests updated to include ResNet50 in the model matrix. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> Signed-off-by: ajrasane <arasane@nvidia.com> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Co-authored-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.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>
|
||
|
Keval MorabiaandClaude Sonnet 4.6
|
e4b054bf32 |
Fix and Speedup megatron_mmlu by >10x via prefill scoring and global batching (#1280)
### What does this PR do?
Type of change: new feature + bug fix
Two improvements to Megatron inference utilities:
**1. Pipeline Parallel (PP) correctness fixes**
PP inference was producing garbage output (MMLU ~0.24, random chance).
Two root causes:
- `megatron_generate` / `megatron_prefill` used
`get_forward_backward_func()` (the training pipeline scheduler), which
is not designed for inference. Rewrote both functions to use explicit
P2P communication via `recv_from_prev_pipeline_rank_` /
`send_to_next_pipeline_rank`, matching the `run_mcore_inference`
pattern.
- `import_mcore_gpt_from_hf` loads HF weights into stage 0's embedding
but never updates the output_layer on the last PP stage when
`share_embeddings_and_output_weights=True`. At model init,
`setup_embeddings_and_output_layer()` all-reduces from stage 0 to sync
the output layer; after importing HF weights that all-reduce is stale.
Fix: call `model.setup_embeddings_and_output_layer()` again after
import.
**2. `megatron_mmlu` speedup (~6x)**
Replaces the `megatron_mmlu` implementation with a significantly faster
approach that matches how `lm-evaluation-harness` scores multiple-choice
questions.
**Before:** autoregressive generation (`megatron_generate`, `osl=2`) per
example, 114 separate `load_dataset` calls, batch_size=1 — 260s for 5%
data.
**After:** single prefill forward pass + argmax over {A,B,C,D} logits, 2
`load_dataset` calls, configurable batch_size — 18s for 5% data (~6x
faster).
### Changes
**PP fixes:**
- `megatron_generate` / `megatron_prefill`: replace
`get_forward_backward_func` with explicit P2P
(`recv_from_prev_pipeline_rank_` / `send_to_next_pipeline_rank`)
- `import_mcore_gpt_from_hf`: call
`model.setup_embeddings_and_output_layer()` after HF weight import when
PP>1 and `share_embeddings_and_output_weights=True`
- `megatron_prefill`: add `skip_return_logits` param and VLM support
(needed for PP non-last stages)
**MMLU speedup:**
- **Log-likelihood scoring**: replace `megatron_generate` with
`megatron_prefill` — one forward pass per batch, no autoregressive
decode loop
- **Global batching**: collect all examples across all subjects, sort by
descending sequence length, run in `batch_size` chunks
- **2 dataset loads** instead of 114: use `load_dataset("cais/mmlu",
"all")` with per-subject grouping; skip dev load when `few_shots=0`
- **`percentage` → `fraction`** parameter rename for clarity
- **tqdm progress bar** (rank-0 only)
### Testing
- `test_megatron_generate_and_mmlu` parametrized over `tp` and `pp`.
Accuracy assertion: `0.36 < score < 0.39`. Manually checked generated
text is coherent.
- Re-ran M-Bridge Minitron MMLU based pruning for Nano v2 9B -> 7B and
all top 10 candidate's MMLU numbers are ballpark similar as before
### Before your PR is "*Ready for review*"
- Is this change backward compatible?: ❌ — `percentage` parameter
renamed to `fraction`; `enable_kv_cache` removed from `megatron_mmlu`
- 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?: ✅ — existing test updated and
parametrized for TP+PP
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
✅
🤖 Generated with [Claude Code](https://claude.ai/claude-code)
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit
* **Bug Fixes**
* Improved pipeline-parallel generation and MMLU evaluation reliability;
fixed output-layer synchronization in shared-embedding + pipeline
setups.
* **New Features**
* MMLU scoring now uses batched prefill logit scoring for faster,
batched evaluation.
* **Behavior Changes**
* Default MMLU sampling increased from 5% to 10%; calibration batch
sizing adjusted and related CLI/help text updated.
* **Tests**
* Distributed tests cover tensor- and pipeline-parallel modes and
tighten MMLU validation ranges.
* **Documentation**
* Updated pruning example and benchmark timing to reflect new sampling
and speedup.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---------
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
|
||
|
Keval MorabiaandClaude Sonnet 4.6
|
7e82a5cb03 |
[Serialization]: remove explicit weights_only default from safe_load to allow user to bypass if needed (#1279)
## Summary
- Remove the `kwargs.setdefault("weights_only", True)` call from
`safe_load`, deferring to torch's built-in default (which is `True` for
torch>=2.6)
- This allows users to override via the
`TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD=1` env var when they trust a
checkpoint but hit `pickle.UnpicklingError`
- Add a test that verifies the default fails on unsafe objects and the
env var bypass works
## Test plan
- [x] `python -m pytest tests/unit/torch/utils/test_serialization.py -v`
🤖 Generated with [Claude Code](https://claude.com/claude-code)
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit
* **Bug Fixes**
* Serialization utility now respects PyTorch's default behavior and
environment-variable configuration instead of forcibly enforcing
parameter overrides, providing greater configuration flexibility.
* **Tests**
* Added test coverage validating environment-variable override
functionality and default behavior in the serialization utility.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
|
||
|
yeyu-nvidiaandClaude Opus 4.6
|
07ae8e7128 |
Add LoRA co-training support for HF EAGLE speculative decoding (#1060)
### What does this PR do?
Type of change: New feature + bug fixes
Adds **LoRA co-training** support for HF EAGLE speculative decoding.
When `eagle_base_lora=True`, HF PEFT LoRA adapters are injected into the
base model and co-trained alongside the EAGLE draft module in a single
online training pass. A preservation loss (KL divergence between the
original frozen base model output and the LoRA-adapted output) prevents
base model drift. LoRA adapter weights are exported in standard peft
format alongside EAGLE draft artifacts.
### Key features
- **LoRA injection**: `peft.inject_adapter_in_model` applied in-place
(no wrapper), keeping the existing `HFEagleModel` structure intact.
- **Preservation loss**: Cross-entropy `H(ref, lora)` — equivalent
gradient to `KL(ref || lora)` since `H(ref)` is constant w.r.t. LoRA
params.
- **Warmup schedule**: `eagle_base_lora_warmup_steps` freezes LoRA for N
steps while the EAGLE head stabilizes, then enables co-training via a
`LoRAWarmupCallback`.
- **Logits detach regularization**: `eagle_base_lora_logits_detach_prob`
stochastically detaches base logits from the EAGLE loss path, preventing
LoRA from degenerating to maximize EAGLE accuracy at the cost of base
model quality.
- **Export**: Standard peft format (`adapter_model.safetensors` +
`adapter_config.json`) alongside EAGLE draft model.
- **Merge script**: `scripts/merge_lora.py` merges LoRA weights into the
base model and restores the original `config.json` (avoids transformers
5.x rewriting `rope_theta` → `rope_parameters` which breaks
vLLM/TRT-LLM).
- **Multinode fix**: `dp_shard_size` now uses `WORLD_SIZE` instead of
local GPU count.
### Config options
```python
mtsp.convert(model, mode=[("eagle", {
"eagle_base_lora": True, # enable LoRA co-training
"eagle_base_lora_rank": 64, # LoRA rank
"eagle_base_lora_alpha": 16.0, # LoRA scaling
"eagle_base_lora_target_modules": ["q_proj", "k_proj", "v_proj", "o_proj"],
"eagle_base_lora_preservation_loss_weight": 0.1, # preservation loss weight
"eagle_base_lora_warmup_steps": 0, # freeze LoRA for N steps
"eagle_base_lora_logits_detach_prob": 0.5, # detach prob (0=never, 1=always)
})])
```
### Experimental results (Qwen3-8B, checkpoint-60000)
Base model quality preserved across detach_prob sweep (lm_eval: IFEval,
ARC-C, Winogrande — results pending final collection).
**Acceptance rate** (mt_bench, draft_length=3, output_length=4096,
temperature=0):
| detach_prob | vLLM AR | TRT-LLM AR |
|---|---|---|
| baseline (no LoRA) | 2.14 | 2.15 |
| 0.5 | 1.45 | 1.44 |
| 0.8 | **3.06** | **3.01** |
| 0.85 | 2.90 | 2.90 |
| 0.9 | 2.76 | 2.77 |
| 0.95 | 2.51 | 2.58 |
| 0.99 | 2.37 | 2.37 |
| 0.999 | 2.30 | 2.27 |
| 0.9999 | 2.31 | 2.26 |
Best AR at `detach_prob=0.8`: ~40% improvement over baseline.
### Testing
`tests/unit/torch/speculative/plugins/test_hf_speculative_lora.py` (5
tests):
- `test_lora_layers_injected` — LoRA layers present after conversion
- `test_trainable_params` — only `lora_*` and `eagle_module` params are
trainable
- `test_forward_returns_loss` — forward returns non-zero scalar loss
- `test_eagle_offline_incompatible` — `eagle_base_lora=True` +
`eagle_offline=True` raises `ValueError`
- `test_export_lora_artifacts` — export produces standard peft adapter
files
### Bug fixes (included in this PR)
1. **`launch_train.sh` case pattern ordering**: glob
`--eagle_base_lora*` was before specific patterns
(`--eagle_base_lora_rank*`, etc.), silently swallowing LoRA args.
2. **LoRA optimizer exclusion during warmup**: warmup freezing excluded
LoRA from the optimizer entirely; fixed with `add_param_group` in the
callback.
3. **`merge_lora.py` config.json**: `save_pretrained()` with
transformers >=5.x rewrites `rope_theta` → `rope_parameters`, breaking
vLLM positional embeddings. Fixed by copying the original base model
config.
4. **Multinode `dp_shard_size`**: used local GPU count instead of
`WORLD_SIZE`.
### Checklist
- [x] Backward compatible (all new config fields have defaults)
- [x] Uses `peft` via lazy imports (no hard dependency)
- [x] Unit tests added
- [x] Online HF training only (`eagle_offline=True` blocked)
---------
Signed-off-by: Ye Yu <yeyu@nvidia.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.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> |
||
|
Shengliang Xu
|
c9b11559f1 |
Normalize .yml to .yaml in modelopt_recipes (#1260)
### What does this PR do?
Type of change: Chore
Standardize YAML file extensions in `modelopt_recipes/` to `.yaml` for
consistency. The existing recipes used a mix of `.yml` (PTQ recipes) and
`.yaml` (speculative decoding, model-specific recipes).
#### Changes
**Renamed files:**
- `modelopt_recipes/general/ptq/fp8_default-fp8_kv.yml` → `.yaml`
- `modelopt_recipes/general/ptq/nvfp4_default-fp8_kv.yml` → `.yaml`
- `modelopt_recipes/general/ptq/nvfp4_experts_only-fp8_kv.yml` → `.yaml`
- `modelopt_recipes/general/ptq/nvfp4_mlp_only-fp8_kv.yml` → `.yaml`
- `modelopt_recipes/general/ptq/nvfp4_omlp_only-fp8_kv.yml` → `.yaml`
**New pre-commit hook:** `normalize-yaml-ext`
- `tools/precommit/normalize_yaml_ext.py` — auto-renames `.yml` to
`.yaml`
for any staged file under `modelopt_recipes/`. Runs before recipe
validation so future contributions are caught automatically.
**Updated references:**
- `tests/unit/recipe/test_loader.py` — built-in recipe paths updated to
`.yaml`
Note: `load_recipe()` and `load_config()` probe both `.yml` and `.yaml`
suffixes, so callers using paths without extensions (e.g.,
`load_recipe("general/ptq/fp8_default-fp8_kv")`) are unaffected.
### Testing
Existing recipe loader tests pass with updated paths.
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit
* **Chores**
* Standardized recipe file extensions and added an automated pre-commit
normalization hook to enforce the convention.
* **Tests**
* Updated unit tests to reference the new recipe filename convention and
ensure consistency with configuration loading.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---------
Signed-off-by: Shengliang Xu <shengliangx@nvidia.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> |
||
|
Shengliang Xu
|
b6c6ec342c |
use typed quantize config instead of a raw dict (#1249)
### What does this PR do? But fix: Use typed QuantizeConfig instead using raw dict for formal typed ModelOpt configs. The dict typing was accidental. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Refactor** * Quantization recipe configuration is now implemented with a strongly-typed, structured schema that enforces type safety and provides enhanced validation with comprehensive error detection capabilities. * **Tests** * Updated recipe loading tests to correctly validate quantization configurations when recipes are loaded from directories, fully supporting the new structured object-based configuration format. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Shengliang Xu <shengliangx@nvidia.com> |
||
|
Ajinkya RasaneandClaude Opus 4.6
|
202c3d3894 |
Add SwinTransformer support for torch_onnx quantization workflow (#1235)
## Summary - Enable end-to-end quantize → ONNX export → TRT engine pipeline for SwinTransformer models (v1 and v2) across FP8, INT8, MXFP8, NVFP4, and auto precision modes - Add Conv2d quantization overrides for TRT compatibility (TRT only supports FP8/INT8 for convolutions) - Fix FP8 LayerNorm type mismatch in TRT stronglyTyped mode by adding `LayerNormalization` to `change_casts_to_fp16` - Fix `cast_initializer_to_dtype` crash when node has no initializer inputs - Simplify `download_example_onnx.py` to a single `--timm_model_name` (required) flag, removing redundant `--vit` and `--llama` flags - Add vision model support matrix to README (ViT, Swin, SwinV2) - Rewrite tests: parametrize over (ViT, Swin, SwinV2) × (fp8, int8, mxfp8, nvfp4, auto) with TRT engine build verification ## Test plan - [ ] `python -m pytest tests/examples/torch_onnx/test_torch_quant_to_onnx.py -v` — 15 tests (3 models × 5 modes), all pass - [ ] Verified Swin accuracy on ImageNet-1k across all precisions (FP8: 81.29%, INT8: 81.12%, MXFP8: 81.32%, NVFP4: 80.79%, Auto: 80.84% TRT top-1 vs 81.37% base) - [ ] INT4_AWQ deferred (TODO in test file) — requires INT4 exporter changes for non-MatMul/Gemm consumer patterns 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * ONNX export supports arbitrary timm vision models with auto device selection and new CLI options (--timm_model_name, --model_kwargs, --no_pretrained); batch-size/input sizing is now model-generic. * **Bug Fixes** * Expanded FP16/BF16 cast handling to additional ONNX ops. * Disabled inplace ReLU before auto-quantization to avoid incorrect transforms. * Conv2d quantization overrides added for improved TensorRT compatibility. * Safer handling when initializers are missing during dtype casting. * **Documentation** * README updated with supported models table, quantization mappings, and example CLI usage. * **Tests** * Tests expanded to multiple architectures/quant modes and now verify TensorRT engine build. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> |
||
|
h-guo18
|
6403389eb0 |
Feat: Configurable Eagle ROPE scaling during export (#1238)
### What does this PR do? JIRA ticket: https://jirasw.nvidia.com/browse/OMNIML-3469 Type of change: New feature Decouple EAGLE training rope configuration from export rope configuration, enabling separate YaRN rope scaling injection at export time for long-context inference. #### Changes **Configurable export rope scaling (`EagleConfig`)** - Add `eagle_export_rope_scaling` field to `EagleConfig` with default YaRN config (`factor=32.0`, `original_max_position_embeddings=2048`) - Set to `{}` to disable rope scaling injection at export **Simplified training defaults (`default_config.py`)** - Change default training rope from `llama3` (theta=500k) to `default` (theta=10k) — models now train with simple positional embeddings; rope scaling is applied only at export - Add `rope_theta` inside `rope_scaling` dict for transformers 5.x cross-version compatibility **Move config validation/rewriting into `EagleConfig` (`config.py`)** - `_derive_eagle_offline`: derives `eagle_offline` from `data_args.offline_data_path` via validation context, removing manual assignment in `main.py` - `_check_rope_scaling_consistency`: rejects configs where `eagle_export_rope_scaling` is set but training `rope_type` is not `"default"` - `_warn_rope_vs_training_seq_len`: warns when `original_max_position_embeddings` differs from `training_seq_len` **Export rope injection (`hf_spec_export.py`)** - Inject `eagle_export_rope_scaling` into the exported HF config when training rope_type is `"default"` - Fall back `rope_theta` from `rope_scaling` dict for transformers 5.x compatibility **Fix Megatron RotaryEmbedding crash (`megatron_eagle.py`)** - `dict_to_config()` set `rope_scaling=True` whenever the `rope_scaling` key existed, even without a `"factor"` — causing `RotaryEmbedding` to divide by `None` - Now only enables `rope_scaling` when the dict actually contains a `"factor"` key ### Usage Configure in YAML config (or use defaults from `eagle3.yaml`): ```yaml eagle: eagle_export_rope_scaling: rope_type: yarn factor: 32.0 original_max_position_embeddings: 2048 ``` Set to empty dict to disable export rope injection: ```yaml eagle: eagle_export_rope_scaling: {} ``` ### Testing - New unit tests: `tests/unit/torch/speculative/test_eagle_config.py` — rope consistency validator, seq_len warning, context-derived `eagle_offline` - New unit tests: `tests/unit/torch/export/test_hf_spec_rope_export.py` — export rope injection, fallback, and empty-config cases ### 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?: ✅ (new field has sensible default; existing configs work unchanged) - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: ✅ - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ❌ (should be added if merging as a feature) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Add export-time rope-scaling configuration for EAGLE models. * **Improvements** * Stronger validation and context-aware reconciliation between training and export configs. * Export now injects rope-scaling and rope-theta when appropriate. * Default rope-scaling values updated for EAGLE variants. * Model instances now expose export rope-scaling for downstream use. * **Tests** * Added unit tests covering rope-scaling export behavior and configuration validators. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> |
||
|
yeyu-nvidiaandClaude Sonnet 4.6
|
9050188034 |
Fix test_collect_hidden_states: use synthetic short conversations (#1234)
## Summary - `test_collect_hidden_states` was using real daring-anteater conversations (typically 1000+ tokens) but the tiny test model has `max_position_embeddings=32`. Both sampled conversations exceeded the default `--max-seq-len 3072` filter, producing zero `.pt` files and failing the assertion. - Added a `tiny_conversations_path` fixture with synthetic short single-turn conversations that tokenize within `max_position_embeddings=32`. - Changed `test_collect_hidden_states` to use this fixture with `--max-seq-len 32`. - Added a `None` guard for `tokenizer.chat_template.replace(...)` to avoid `AttributeError` when the tokenizer has no chat template. ## Test plan - [ ] `pytest tests/examples/speculative_decoding/test_eagle_offline_ptq.py::test_collect_hidden_states` passes - [ ] CI `speculative_decoding` job passes 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Bug Fixes** * Resolved compatibility issues when tokenizers do not have a chat template configuration by adding proper error handling. * Standardized tokenization input extraction logic across different transformer library versions for consistent behavior. * **Tests** * Enhanced test infrastructure with new conversation data fixtures and improved sequence length validation for speculative decoding examples. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Ye Yu <yeyu@nvidia.com> Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com> |
||
|
Shengliang Xu
|
da0e8ff132 |
Validate non-empty cfg when enabling quantizers in quant_cfg (#1192)
### What does this PR do?
An entry with enable=True and an empty or invalid cfg (e.g., cfg={},
cfg=[], or cfg=42) would pass normalize_quant_cfg_list but later crash
in set_quantizer_by_cfg when QuantizerAttributeConfig(enable=True) was
constructed with no actual quantizer attributes.
Add early validation in normalize_quant_cfg_list to reject such entries
at config time rather than surfacing a confusing Pydantic
ValidationError deep in the conversion pipeline.
### Testing
new unit tests added
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit
* **Bug Fixes**
* Stricter validation for quantizer entries: enabled quantizers must
have a non-empty dict or non-empty list of non-empty dicts; empty or
invalid cfg values now raise a clearer error. Normalization still
preserves explicit enable/disable flags.
* **Tests**
* Added unit tests covering rejection of empty/invalid configs when
enabled and acceptance when disabled, and ensuring normalized output
preserves disabled state.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---------
Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
|
||
|
Ajinkya RasaneandClaude Opus 4.6
|
3baa2da62e |
Upgrade ONNX from 1.19 to 1.21 (#1207)
### What does this PR do? Type of change: new feature Upgrade ONNX dependency from `~=1.19.0` to `~=1.21.0`. ONNX 1.20+ removed several deprecated helper functions (`float32_to_bfloat16`, `float32_to_float8e4m3`, `pack_float32_to_4bit`) that `onnx_graphsurgeon` 0.5.x still references at import time. This PR adds a compatibility shim (`modelopt/onnx/_onnx_compat.py`) that restores these functions using `ml_dtypes` before any `onnx_graphsurgeon` import occurs. This supersedes the partial inline fix from #1204 by also handling `float32_to_float8e4m3`. Changes: - Bump `onnx~=1.19.0` to `onnx~=1.21.0` in `pyproject.toml` - Add `modelopt/onnx/_onnx_compat.py` compatibility shim for removed ONNX APIs - Import shim in `modelopt/onnx/__init__.py` and `tests/unit/onnx/conftest.py` - Remove usage of removed `onnx.helper.pack_float32_to_4bit` in `test_quant_utils.py` - Update example requirements (`genai_llm`, `whisper`) to `onnx==1.21.0` **TensorRT Compatibility:** TRT 10.16-GA supports opsets 9–24. ModelOpt quantization modes use opsets 19–23, all within range. ONNX 1.21 does not force opset 26. ### Usage ```python # No API changes — the upgrade is transparent to users. # The compatibility shim is applied automatically on import. import modelopt.onnx ``` ### Testing - 469/470 ONNX unit tests pass inside `nvcr.io/nvidia/tensorrt:25.06-py3` (1 pre-existing ORT `CopyTensorAsync` EP issue, not ONNX-related) - 6/6 `torch_onnx` integration tests pass (fp8, int8, nvfp4, mxfp8, int4_awq, auto) - ViT FP8 quantization via `torch_onnx` → TRT engine build → ImageNet eval: **85.3% top-1, 97.8% top-5** - ViT FP8 quantization via `onnx_ptq` → TRT engine build succeeds - All pre-commit hooks pass (ruff, mypy, bandit, license headers) ### 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`: ✅ - Did you write any new necessary tests?: ✅ (updated existing tests, added conftest.py for compat shim) - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ❌ (dependency upgrade, no API change) ### Additional Information Related: #1204 (partial fix for `float32_to_bfloat16` only — this PR supersedes it with full coverage) 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Dependencies** * Removed unpinned ONNX from example requirement files and updated the ONNX optional dependency to ~=1.21.0. * **Refactor** * Centralized an ONNX compatibility shim to restore missing helper APIs when needed. * **Tests** * Added tests for the compatibility shim, adjusted quantization tests to remove reliance on removed ONNX helpers, and ensured shim runs before related tests. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> Co-authored-by: Claude Opus 4.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> |
||
|
Ajinkya RasaneandClaude Opus 4.6
|
b3feebfe02 |
Replace in-repo LLM ONNX export with TensorRT-Edge-LLM (#1210)
### What does this PR do? Type of change: Documentation, Cleanup This PR removes the in-repo LLM ONNX export pipeline (`llm_export.py` and `modelopt.onnx.llm_export_utils` package) and updates the `examples/torch_onnx/README.md` to direct users to [TensorRT-Edge-LLM](https://github.com/NVIDIA/TensorRT-Edge-LLM), which provides a more complete and actively maintained pipeline for quantizing LLMs/VLMs with ModelOpt and exporting them to optimized ONNX for edge deployment (Jetson, DRIVE). **Removed:** - `examples/torch_onnx/llm_export.py` — standalone LLM export script - `modelopt/onnx/llm_export_utils/` — supporting package (`export_utils.py`, `quantization_utils.py`, `surgeon_utils.py`) - `tests/examples/torch_onnx/test_llm_export.py` — associated tests **Updated:** - `examples/torch_onnx/README.md` — rewrote the LLM section with TensorRT-Edge-LLM installation, CLI tools, usage examples (LLM, VLM, EAGLE speculative decoding), supported model matrix, quantization methods, and troubleshooting guidance ### Usage Users should now use TensorRT-Edge-LLM CLI tools instead of the removed `llm_export.py`: ```bash # Quantize a model with ModelOpt tensorrt-edgellm-quantize-llm \ --model_dir Qwen/Qwen2.5-3B-Instruct \ --quantization fp8 \ --output_dir quantized/qwen2.5-3b-fp8 # Export to ONNX tensorrt-edgellm-export-llm \ --model_dir quantized/qwen2.5-3b-fp8 \ --output_dir onnx_models/qwen2.5-3b ``` ### Testing - No functional code changes — this is a removal of deprecated code and documentation update. - Vision model quantization/export (`torch_quant_to_onnx.py`) and mixed precision examples are unaffected. ### 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?: ❌ — Removes `llm_export.py` and `modelopt.onnx.llm_export_utils` package. Users should migrate to [TensorRT-Edge-LLM](https://github.com/NVIDIA/TensorRT-Edge-LLM). - 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 — This removes code and tests; no new functionality added. - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ❌ — Should be updated to note the removal of `llm_export_utils` and migration to TensorRT-Edge-LLM. ### Additional Information - Successor tool: [TensorRT-Edge-LLM](https://github.com/NVIDIA/TensorRT-Edge-LLM) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Documentation** * Restructured LLM ONNX export documentation to reflect new quantization → export → build → inference workflow with updated tooling guidance * Added comprehensive installation and setup instructions with system requirements verification and supported model information * Expanded examples to cover LLMs, VLMs, and speculative decoding techniques with step-by-step command guidance, parameters, and troubleshooting resources <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> |
||
|
Keval MorabiaandClaude Sonnet 4.6
|
9b4f43a9cc |
Improve Megatron Tokenization: streaming, reasoning_content support, HF in-memory tokenization, etc (#1221)
### What does this PR do?
Type of change: New feature
Improvements to `megatron_preprocess_data` for Nemotron v3 post-training
datasets and Megatron-Bridge distillation workflows:
- **`--reasoning_content`** flag (`strip` / `inline` / `native`) to
handle the `reasoning_content` field in Nemotron Post-Training v3
assistant messages
- **No intermediate JSONL** for HuggingFace datasets — load directly
from Arrow cache via `_iter_hf_as_json` + `process_hf_split`
- **Return output prefixes** (`list[str]`) from the Python API so
callers can build `--data_paths` without hardcoding paths; also printed
at end of run
- **Gzip input support** — `.jsonl.gz` files accepted directly;
`--input_dir` globs both `*.jsonl` and `*.jsonl.gz`
- **`--strip_newlines`** flag (opt-in) to replace newlines with spaces
in plain-text values; default preserves newlines (no breaking change for
code/structured-text datasets)
- **`--hf_streaming`** flag for very large datasets — only consumed rows
are downloaded; automatically falls back to non-streaming (with a
warning) if `--hf_max_samples_per_split` is not set, since streaming
without a cap is slower than cached non-streaming
- **Auto-shuffle** when `--hf_max_samples_per_split` is set — reservoir
sampling (buffer=10,000, seed=42) applied before capping to avoid biased
prefix sampling
- Remove `_document` suffix from output filenames (`_text.bin` instead
of `_text_document.bin`)
- Fix duplicate BOS token for chat-template data
(`add_special_tokens=False`)
- Fix `TypeError` in `process_hf_split` (`sum(list)` not `sum(int)`)
- Suppress duplicate prints across pool workers via
`_is_main_or_first_worker()`
- Raise `KeyError` instead of warning for missing JSON keys
- Default `hf_split=None` (all splits) instead of `"train"`
### Usage
```python
from modelopt.torch.utils.plugins.megatron_preprocess_data import megatron_preprocess_data
# Nemotron v3 with reasoning content preserved inline as <think>...</think>
prefixes = megatron_preprocess_data(
hf_dataset="nvidia/Nemotron-Post-Training-Dataset-v3",
json_keys=["messages"],
tokenizer_name_or_path="Qwen/Qwen3-0.6B",
output_dir="tokenized/",
workers=32,
reasoning_content="inline",
)
# prefixes == ["tokenized/nvidia--Nemotron-Post-Training-Dataset-v3_..._messages"]
data_paths = [x for p in prefixes for x in ("1.0", p)]
# Large pretraining dataset — stream + cap (auto-shuffled before capping)
prefixes = megatron_preprocess_data(
hf_dataset="nvidia/Nemotron-CC-v2.1",
hf_name="High-Quality",
hf_max_samples_per_split=5_000_000,
hf_streaming=True,
json_keys=["text"],
tokenizer_name_or_path="Qwen/Qwen3-0.6B",
output_dir="tokenized/",
workers=32,
append_eod=True,
strip_newlines=True,
)
```
### Testing
- New unit tests
- Tested tokenization on Nemotron Pretraining and Post-training v3
datasets
### 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?: ✅ (output filename changed:
`_text_document` → `_text`; existing callers need to re-tokenize or
rename files)
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: N/A
- Did you write any new necessary tests?: ✅
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
✅
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit
* **New Features**
* Preprocessing tool: configurable reasoning-content modes
(strip|inline|native), optional newline stripping, gzip (.jsonl.gz)
input support, HF streaming mode, auto-shuffle when per-split max
samples set, returns output-file prefixes, and processes all HF splits
by default without writing intermediate JSONL.
* **Documentation**
* Consolidated and reformatted dataset preparation and tokenization
guidance; updated install/auth instructions and examples.
* **Tests**
* Tests updated to validate returned prefixes, reasoning-content
behaviors, gzip input handling, HF streaming warnings, and HF output
assertions.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---------
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
|
||
|
Keval Morabia
|
14e9bea793 |
Fix code-quality
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> |
||
|
Keval Morabia
|
24b8dcf1e6 |
Disable test_quantize_cpu.py::test_save_restore on Windows as its flaky
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> |
||
|
Keval Morabia
|
04cd596d79 |
Add experimental support for transformers>=5.0 + min torch 2.8 (#975)
### What does this PR do? - Add experimental support for transformers >=5.0 and remove deprecated usages: https://github.com/huggingface/transformers/blob/main/MIGRATION_GUIDE_V5.md - ⚠️ For accelerate examples that used `--warmup-ratio: float` (deprecated in 5.x), we now change it to `--warmup-steps: float | int` which works as ratio if float but only for 5.x. For 4.x, it will error out if float and prompt user to change back to `--warmup-ratio` or pass an int absolute step count. - ⚠️ Unified Hugging Face checkpoint export for quantized checkpoints may not work for some models with transformers>=5.0 yet as it requires a lot of fixes (e.g. change in how MoE experts are organized) - ~Add Workaround for TRT-LLM's import of deprecated transformers functions so trt-llm based gpu unit tests work fine. Still deployment for models needs proper fixes directly in TRT-LLM hence llm/vlm ptq example tests still run with transformers 4.57~ - Everything except PTQ and Export (mainly MoE) should work fine with transformers>=5.0 - Bump min torch to 2.8 and enable 2.11 cicd testing - NOTE: Upcoming Nemo:26.04 container comes with transformers 5.3 ### Testing <!-- Mention how have you tested your change if applicable. --> - [x] CI/CD tests passing - [x] Manually tested unit tests, gpu tests with transformers 4.56 and 5.4 - [x] Manually tested example tests (except trt-llm container tests) with transformers 4.56 and 5.4 - [x] 2-gpu nightly CICD tests manually triggered and passing: [gpu tests](https://github.com/NVIDIA/Model-Optimizer/actions/runs/23867257540), [example tests](https://github.com/NVIDIA/Model-Optimizer/actions/runs/23867260643) ### 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)?: ✅ <!--- 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** * Make remote-code usage opt-in via a configurable --trust_remote_code flag across examples and tools. * **Bug Fixes** * Improve checkpoint/resume detection and related training guidance to avoid erroneous errors. * **Refactor** * Consolidate dtype/config naming, switch warmup settings from ratio → steps, and unify tokenizer invocation patterns. * **Documentation** * Simplify changelog title and add misc notes for release 0.44. * **Chores** * Remove scheduled PR-branch cleanup workflow and relax/remove several transformers version pins. * **Tests** * Adjust test gates, skips, and structures to align with updated deps and behaviors. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> |
||
|
yeyu-nvidiaandClaude Sonnet 4.6
|
cccfded8a9 |
Add support for offline speculative decoding model PTQ (#883)
## What does this PR do? **Type of change:** new feature **Overview:** This PR enables loading in a ModelOpt pretrained offline speculative decoding model (e.g., EAGLE3) and performs PTQ on it and export. ## Usage Follow the speculative_decoding examples to train an offline speculative decoding model first. Then follow the command below to quantize and export it: ```bash python hf_ptq.py --pyt_ckpt_path <dir_of_offline_specdec_model> --specdec_offline_dataset <dir_of_dataset> ``` ## 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** * Offline speculative decoding workflow: support loading a local dataset for calibration, generation, and export; new CLI option to specify the offline dataset. * **Improvements** * Export and quantization paths now accept and propagate offline speculative-decoding inputs. * Offline data loading honors a sample-size limit and enforces safe batch sizing for calibration. * **Bug Fixes** * Better handling of model/config mismatches and varied batch types in offline flows. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Ye Yu <yeyu@nvidia.com> Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com> |
||
|
Keval MorabiaandClaude Sonnet 4.6
|
8428f0610a |
fix: megatron export correctness for TP>1 GQA, single-file MTP, and Hub remote code (#1209)
### What does this PR do?
Type of change: Bug fix
Three correctness fixes for the Megatron Core GPT export pipeline:
**1. `_qkv_slicing`: reshape failure with TP>1 on GQA models**
When tensor parallelism is enabled, the `linear_qkv` weight tensor
arriving in `_qkv_slicing` is already TP-sharded, so `weight.shape[0]`
equals `per_rank_qkv_dim * head_size`, not `qkv_total_dim * head_size`.
All five reshape/`arange` operations were using the global
`qkv_total_dim`, causing a runtime shape mismatch for any GQA model with
TP > 1. The fix derives `per_rank_qkv_dim` and `num_query_groups_local`
from the actual tensor shape, making the logic correct for any TP degree
(a no-op for TP=1).
**2. `_get_mtp_state_dict`: `EntryNotFoundError` for non-sharded
models**
`hf_hub_download("model.safetensors.index.json")` raises
`EntryNotFoundError` for small models that ship a single
`model.safetensors` rather than a sharded index. The function now
catches this and falls back to downloading/reading `model.safetensors`
directly, scanning its keys with `safe_open`. The same two-path logic
applies to local directories.
**3. `copy_remote_code`: `ValueError` for Hub model IDs**
`copy_remote_code` only accepted local directory paths and raised
`ValueError` for HuggingFace Hub model IDs (e.g.
`"meta-llama/Llama-3.2-1B"`). The function now falls back to
`list_repo_files` + `hf_hub_download` to fetch and copy top-level `.py`
files (custom modeling code) when the path is not a local directory.
### Usage
```python
# TP>1 GQA export now works (previously raised RuntimeError on reshape)
export_mcore_gpt_to_hf(gqa_model, "meta-llama/Llama-3.2-1B", export_dir="./out", dtype=torch.bfloat16)
# Models with a single model.safetensors now have their MTP weights exported
export_mcore_gpt_to_hf(model, "./small_model_dir", export_dir="./out", dtype=torch.bfloat16)
# Hub model IDs no longer raise ValueError in copy_remote_code
export_mcore_gpt_to_hf(model, "org/custom-model-with-remote-code", export_dir="./out", dtype=torch.bfloat16)
```
### Testing
New tests added in `tests/gpu_megatron/torch/export/`:
- `test_unified_export_megatron.py::test_qkv_slicing_gqa_tp2` —
FP8-quantized GQA model export with TP=2 (`num_query_groups=2 <
num_attention_heads=8`), exercises both the weight reshape and
per-channel weight-scale reshape paths.
-
`test_unified_export_megatron.py::test_mtp_state_dict_single_safetensors`
— unit test verifying MTP weights are collected from a single
`model.safetensors` file.
- `test_unified_export_megatron.py::test_mtp_state_dict_index_file` —
unit test verifying MTP weights are collected from a sharded checkpoint.
- `test_unified_export_megatron.py::test_mtp_state_dict_no_mtp_keys` —
edge case: no MTP keys → empty dict, no side effects.
- `test_hf_checkpoint_utils.py` — four tests covering `copy_remote_code`
for local directories and Hub model IDs (with and without `.py` files).
### Before your PR is "*Ready for review*"
- Is this change backward compatible?: ✅
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: N/A
- Did you write any new necessary tests?: ✅
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
N/A
### Additional Information
Fixes reported against Megatron export when running quantization with
TP>1, small non-sharded HF models, and HuggingFace Hub model IDs passed
to `export_mcore_gpt_to_hf`.
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit
* **New Features**
* Export functionality now supports downloading code directly from
Hugging Face Hub model repositories in addition to local directories.
* **Bug Fixes**
* Improved safetensors loading with better error handling for missing
model entries and support for both single and sharded weight files.
* Enhanced tensor slicing behavior for multi-GPU distributed export
scenarios.
* **Tests**
* Added comprehensive test coverage for Hugging Face integration and
export functionality.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---------
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Co-authored-by: Claude Sonnet 4.6 <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 Morabia
|
ba4f42df1c |
Minor fix for example tests
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> |
||
|
Chenhan D. Yu
|
0246041b01 |
feat(speculative): add vLLM data synthesis pipeline and Nemotron dataset preparation scripts (#1176)
### What does this PR do? Type of change: New feature, new example, bug fix Adds a vLLM-based synthetic data generation pipeline for speculative decoding draft model training, along with dataset preparation scripts for NVIDIA's Nemotron Post-Training dataset collections. **Data synthesis pipeline** (`tools/launcher/common/vllm/query.sh` + `common/query.py`): - Launch a vLLM server and run multi-turn inference to synthesize training data from input conversation skeletons - Fork-safe OpenAI client: reinitializes HTTP connection pool after `datasets.map()` forks worker processes, preventing 400 errors from corrupted connections - Clear Docker `ENTRYPOINT` so vLLM containers (which default to `vllm serve`) work correctly under NeMo Run's executor - `--max-tokens` argument to bound generation length - Local file loading support (`--data /path/to/file.jsonl`) - Re-raise connection errors so `datasets.map()` halts the shard instead of silently producing empty rows - Map `developer` role to `system` (OpenAI format compatibility) **Multi-turn reasoning trace handling** (`common/query.py`): - Strip `<think>...</think>` blocks from intermediate assistant turns before re-feeding to the model; preserve the full trace only on the final turn **Nemotron dataset preparation** (`examples/dataset/`): - `make_nemotron_ptv2_dataset.py` — prepares [nvidia/Nemotron-Post-Training-Dataset-v2](https://huggingface.co/datasets/nvidia/Nemotron-Post-Training-Dataset-v2) (~3.3M rows generate, ~1.9M rows train) - `make_nemotron_ptv3_dataset.py` — prepares the [Nemotron PTv3 collection](https://huggingface.co/collections/nvidia/nemotron-post-training-v3) of 16 datasets (~3.4M rows generate, ~3.9M rows train) - Both support `generate` mode (strips assistant turns for synthesis input) and `train` mode (normalizes to clean OpenAI format for SFT) - `conversation_utils.py` — shared utilities: `strip_assistant_turns`, `normalize_messages`, `make_augment_fn`, `AugmentationSpec` - `augmentations.yaml` — 12 language-redirect variants + style/format hints, cycled across dataset rows - Scripts live in `examples/dataset/` (not under `speculative_decoding/`) to signal reusability beyond speculative decoding **Bug fixes**: - `strip_assistant_turns()`: return `{"messages": []}` when no user turns remain (system-only rows were previously passed through instead of being filtered) - `concatenate_datasets()`: guard against empty parts list - SSH tunnel user precedence: explicit `user` arg now correctly overrides `slurm_config.user` ### Usage ```bash # Prepare PTv3 input conversations for synthesis (~3.4M rows): python examples/dataset/make_nemotron_ptv3_dataset.py --output-dir /tmp/ptv3_gen # Launch vLLM server + synthesize responses: bash tools/launcher/common/vllm/query.sh \ --model /path/to/model \ --tensor-parallel-size 4 \ -- \ --data /tmp/ptv3_gen/default.jsonl \ --save /tmp/ptv3_responses \ --num-shards 10 --num-proc 4 --max-tokens 4096 # Prepare PTv2 for direct SFT training (~1.9M rows): python examples/dataset/make_nemotron_ptv2_dataset.py --mode train --output-dir /tmp/ptv2_train ``` ### Testing Tested end-to-end on an NVIDIA GB10 node (119 GiB GPU memory) with `vllm/vllm-openai:qwen3_5-cu130` container and `Qwen/Qwen3.5-4B`: - vLLM server starts correctly with cleared Docker entrypoint - `datasets.map(num_proc=4)` runs without connection errors (fork-safe client) - Multi-turn synthesis produces correct assistant responses with thinking traces handled ### 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 (data synthesis scripts; tested manually) - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added dataset generation and augmentation capabilities for Nemotron post-training datasets (v2 and v3) * Enhanced query functionality with thinking-block filtering and improved client management for robust parallel processing * Added support for local dataset file paths alongside HuggingFace Hub datasets * **Bug Fixes** * Fixed SLURM executor user resolution and Docker container entrypoint configuration * Improved error handling for connection failures during dataset synthesis * **Documentation** * Updated dataset preparation guide with new generation modes and augmentation configuration details <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: Chenhan D. Yu <5185878+ChenhanYu@users.noreply.github.com> Signed-off-by: Chenhan Yu <chenhany@nvidia.com> |
||
|
h-guo18
|
82d96a635f |
[Speculative Decoding] Refactor EAGLE3 training to YAML-based config and recipe system (#1134)
## What does this PR do?
Refactors EAGLE3 training to use a single base YAML config with
OmegaConf dotlist overrides.
**Type of change:** Refactor
## Changes
- Single base config
`modelopt_recipes/speculative_decoding/_base_eagle3.yaml` for all EAGLE3
training; removed per-model child YAMLs.
- `launch_train.sh` accepts `--config <yaml>` plus dotlist overrides
(e.g. `model.model_name_or_path=xxx`).
- Removed `__base__` YAML inheritance logic from `main.py`.
- `dp_shard_size` default changed from `0` sentinel to `None` for
clarity.
- Removed `eagle_config.json` and `fsdp_config.json`; architecture
config is now nested under `eagle.eagle_architecture_config` in YAML.
- `train_eagle3_and_export.sh` now uses base YAML + dotlist instead of
generating a temporary YAML.
- Updated README and tests accordingly.
## Usage
```bash
# Online training
./launch_train.sh \
--config ../../modelopt_recipes/speculative_decoding/_base_eagle3.yaml \
model.model_name_or_path=meta-llama/Llama-3.2-1B \
data.data_path=input_conversations/train.jsonl \
training.output_dir=ckpts/llama-3.2-1b-online
# Offline training
./launch_train.sh \
--config ../../modelopt_recipes/speculative_decoding/_base_eagle3.yaml \
model.model_name_or_path=meta-llama/Llama-3.2-1B \
data.offline_data_path=$HIDDEN_STATES_DIR \
training.output_dir=ckpts/llama-3.2-1b-offline
```
---------
Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
|
||
|
Gwena Cunha
|
bdc04f174d |
[6056809] Fix TRT dependency in ModelOpt ONNX quantization (#1189)
### What does this PR do? Type of change: Bug fix Regression bug introduced by the Autotune integration into ModelOpt ONNX quantization (https://github.com/NVIDIA/Model-Optimizer/pull/951), making ModelOpt dependent on TensorRT in all scenarios. This PR fixes this issue by requiring TensorRT only when `--autotune` is enabled. ### Usage ```python $ python -m modelopt.onnx.quantization --onnx_path=${MODEL_NAME}.onnx ``` ### Testing See bug 6056809. ### Before your PR is "*Ready for review*" Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors) (e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(..., weights_only=False)`, `pickle`, etc.). - Is this change backward compatible?: ✅ - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: ✅ - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Bug Fixes** * Autotune dependency failures now surface as clearer runtime errors instead of only logging warnings. * **Chores** * Centralized autotune presets and numeric defaults into a shared configuration. * Core autotune components are conditionally exposed so initialization succeeds when optional acceleration libraries are absent. * Deferred autotune imports to runtime to improve failure handling. * **Tests** * Added a test ensuring the quantization CLI/parser initializes correctly without optional acceleration libraries. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: gcunhase <4861122+gcunhase@users.noreply.github.com> |
||
|
Jenny ChenandClaude Opus 4.6
|
4255bc6a0c |
Fix MoE quantization tests and add sync_expert_weight_amax option (#1158)
Fix MOE quantization tests to ensure QuantSequentialMLP and QuantTEGroupedMLP produce identical outputs when SequentialMLP has weight & input quantizer sync both enabled. By default SequentialMLP does not sync weight scales (but does sync input quantizer scales) so the outputs of QuantSequentialMLP and QuantTEGroupedMLP are not the same. However when weight scales are synced the amax values and outputs of both modules should be identical ### What does this PR do? Type of change: Bug fix - Add sync_expert_weight_amax flag to MaxCalibConfig to optionally sync weight quantizer amax across SequentialMLP experts (matches TEGroupedMLP) - Update layer_sync_moe_local_experts_amax to accept sync_weight_amax param ## Summary by CodeRabbit * **New Features** * Added a calibration option (default: off) to optionally synchronize weight-quantizer amax values across local experts in MoE models, aligning sequential-expert calibration with grouped-expert behavior. * **Documentation** * Calibration and layer synchronization docs updated to describe the new option and its semantics. * **Tests** * Updated and expanded quantization tests to cover grouped vs sequential MoE calibration, the new expert-weight synchronization option, and adjusted parametrizations and skip logic. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Jennifer Chen <jennifchen@nvidia.com> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.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>
|
||
|
kaix-nv
|
df80a0f7b2 |
[5/n] Add VSA for Video Diffusion (#1053)
### 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. Adds Video Sparse Attention (VSA) as a new sparse attention
method in ModelOpt. VSA implements a two-branch architecture
(compression + sparse) using 3D block tiling for video diffusion model.
VSA integrates with HuggingFace models by registering as
attn_implementation="modelopt_vsa" in HF's ALL_ATTENTION_FUNCTIONS,
which is the same pattern used by the existing Triton FA backend. After
`sparsify()`, HF dispatches Q, K, V directly to the VSA kernel with no
monkey-patching needed.
### Usage
```python
# Load any HuggingFace model
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.1-8B")
# Define VSA config
vsa_config = {
"sparse_cfg": {
"*attn*": {
"method": "vsa",
"block_size_3d": (4, 4, 4), # 3D tile dimensions (T, H, W)
"top_k_ratio": 0.5, # keep top 50% of blocks
"video_shape": (8, 16, 16), # video dims after patchification
"enable": True,
},
"default": {"enable": False},
},
}
# Apply — registers modelopt_vsa with HF automatically
model = mtsa.sparsify(model, vsa_config)
```
### Testing
<!-- Mention how have you tested your change if applicable. -->
`pytest tests/unit/torch/sparsity/attention_sparsity/test_vsa.py`
### Before your PR is "*Ready for review*"
Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)
and your commits are signed (`git commit -s -S`).
Make sure you read and follow the [Security Best
Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors)
(e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(...,
weights_only=False)`, `pickle`, etc.).
- Is this change backward compatible?: ✅ / ❌ / 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 Video Sparse Attention (VSA) as a new sparse attention method
for attention optimization.
* Introduced VSA configuration support with customizable block sizes and
sparsity ratios.
* Integrated VSA with HuggingFace transformers for enhanced model
compatibility.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---------
Signed-off-by: Kai Xu <kaix@nvidia.com>
|
||
|
|
00c002ffd1 |
Exclude small-channel Conv nodes from FP8 quantization (#1083)
### What does this PR do? Small-channel convolutions (IC or OC <= 16) do not benefit from FP8 quantization and cause performance regressions due to poor kernel selection. Hardcode the threshold directly in graph_utils rather than exposing it as a CLI flag. <!-- 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 * **Bug Fixes** * Improved FP8 quantization filtering to exclude convolutions with very small input or output channel counts, reducing incorrect FP8 quantization and improving model accuracy and stability. * **Tests** * Added unit tests validating FP8 small-channel exclusion across channel combinations and confirming no change in int8 mode. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: samcheng <samcheng@nvidia.com> Signed-off-by: nv-samcheng <130026689+nv-samcheng@users.noreply.github.com> Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> Co-authored-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> |
||
|
Rohan Joshi
|
d397caea5a |
Remove --skip-softmax option from hf_sa (only allow calibration option) (#1123)
### What does this PR do? The `hf_sa.py` script applies an attention sparsity method to a model and exports a "sparsified" checkpoint. The SKIP_SOFTMAX_DEFAULT sparse attention config doesn't run calibration or make any change to the checkpoint. If we use this method with `hf_sa.py` the exported checkpoint is identical to the original checkpoint. This could cause confusion so it's best to remove the option altogether. This PR removes the option and switches to SKIP_SOFTMAX_CALIB as the default config. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Chores** * Updated sparse-attention options: removed the legacy "skip_softmax" choice and set the default to "skip_softmax_calib", so the old configuration path is no longer available. * **Tests** * Aligned tests with the new default sparse-attention setting to reflect the calibration-enabled choice. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: Rohan Joshi <rohjoshi@nvidia.com> |
||
|
Ajinkya Rasane
|
7c1f33c2f8 |
[BugFix][5271237][ONNX] Add Q/DQ placement for Conv->LayerNorm patterns (#1117)
### What does this PR do?
Type of change: BugFix
- ModelOpt was not placing Q/DQ nodes between Conv and
LayerNormalization, causing TensorRT to select slower i8f16 kernels
instead of faster i8i8 kernels for Conv layers whose output feeds into
LayerNorm (e.g., ConvNext models).
### Changes:
- Register LayerNormalization in ORT's QDQ registry (ort_utils.py)
- Add find_conv_to_layernorm_nodes() to detect Conv→(copy ops)→LayerNorm
patterns (graph_utils.py)
- Add detected LayerNorm nodes to quantizable nodes list so Q/DQ pairs
are inserted on Conv output (int8.py)
### Usage
```python
# No API change — existing quantize() call now automatically handles Conv->LayerNorm patterns
import modelopt.onnx.quantization as moq
moq.quantize("convnext.onnx", quantize_mode="int8")
# Output model will now have: Conv -> Transpose -> Q -> DQ -> LayerNorm
````
### Testing
- Added unit test test_conv_layernorm_quantization with a ConvNext-like
test model (Conv→Transpose→LayerNorm→Transpose→Conv)
- All 237 unit tests pass (236 existing + 1 new)
- Validated on ConvNext-tiny (opset 17): 23/23 LayerNorm nodes get Q/DQ
on activation input
- Built TRT engine in nvcr.io/nvidia/tensorrt:26.02-py3 — Conv layers
after LayerNorm select i8i8 kernels
### 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?: N/A
### Additional Information
- NVBug: 5271237
- JIRA: OMNIML-2380
- Note: The residual Add output quantization discussed in comments #4-9
of the bug is not addressed here and will be handled in a separate PR.
- Accuracy regression observed for the ConvNext model. Please have a
look at this
[comment](https://github.com/NVIDIA/Model-Optimizer/pull/1117#issuecomment-4156915554)
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit
* **New Features**
* Expanded quantization to include LayerNormalization in Q/DQ flows and
detect Conv→LayerNormalization patterns for quantization.
* **Refactor**
* Improved internal graph-analysis helpers to make detection logic
reusable and more reliable.
* **Tests**
* Added a model builder for Conv→LayerNorm graphs and a unit test
validating end-to-end Conv→LayerNormalization quantization.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---------
Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.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> |
||
|
h-guo18
|
7f5fd65003 |
[Feat]FakeBaseModel for offline eagle; Kimi-K2.5 fixes; (#1052)
### What does this PR do? Adds `FakeBaseModel` for offline EAGLE training and several Kimi-K2.5 compatibility fixes. - **New**: `FakeBaseModel` — lightweight model that loads only `lm_head` and `embed_tokens` from a local checkpoint, avoiding full model weight loading during offline training. Configured via `FakeBaseArguments` and integrated into `load_vlm_or_llm`. - **Fix**: `_find_base_model_parts` — support Kimi-K2.5 VLM layout (`language_model.model` path) - **Fix**: offline mode lm_head access and CompressedTensors ignore path - **Fix**: Kimi-K2.5 decoder `past_key_value`/`past_key_values` argument mismatch - **Fix**: `rglob` for `.pt` discovery in nested offline data dirs; single-node GPU count respects `CUDA_VISIBLE_DEVICES` Type of change: Bug fix, new feature ### Testing Tested offline EAGLE training for Kimi-K2.5 end-to-end. ### 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** * Lightweight fake-base model support for offline speculative-decoding training * **Improvements** * Added CLI flags: --use_fake_base_for_offline, --trust_remote_code, and --fsdp * Expanded offline .pt discovery to include nested subdirectories * Better GPU detection with explicit single-node logging; FSDP enabled only when requested * Model loading and launch tooling now honor offline and trust-remote-code flags * **Bug Fixes** * Improved compatibility with legacy transformer / Kimi-K2 call signatures * **Tests** * Added tests covering fake-base loading and offline training workflows <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: h-guo18 <67671475+h-guo18@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> |
||
|
Keval Morabia
|
b61fb4e16d |
Allow searcher ckpt dir for per-rank ckpt files (#1091)
### What does this PR do? Type of change: minor improvement <!-- Use one of the following: Bug fix, new feature, new example, new tests, documentation. --> Currently we modify user's `path.pth` to `path<rank>.pth` which may be confusing. Instead we now add alternative to provide `path/` and store `path/rank<rank>.pth` per-rank searcher state files ### 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?: ✅ <!--- 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 * **Documentation** * Updated pruning examples and CLI help to reference directory-based checkpoint paths for intermediate pruning scores and clarified default wording. * **Improvements** * Checkpoint handling now accepts directory paths for intermediate score storage, produces per-rank files automatically, and resolves to a directory default when unset. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> |
||
|
dthienan-nv
|
aad14d159c |
[5963347] Expose iterator interface for calibration data (#1075)
### What does this PR do?
Expose the internal calibration_data_reader parameter to the public
quantize function to allow providing an external CalibrationDataReader
object. This enables an iterator style interface for large calibration
sets.
<!-- Details about the change. -->
### Usage
```python
input_tensor = torch.randn(2, 16, 16)
# Calibration data comes from a custom data reader, enabling iterator based reading functionality
class ExampleCalibrationDataReader(CalibrationDataReader):
def __init__(self, input_data):
self.data_list = [{"input": input_data.numpy()}]
self.iter = iter(self.data_list)
def get_next(self):
return next(self.iter, None)
def get_first(self):
return self.data_list[0]
def rewind(self):
self.iter = iter(self.data_list)
calibration_reader = ExampleCalibrationDataReader(input_tensor)
moq.quantize(
onnx_path,
quantize_mode="int8",
high_precision_dtype=high_precision_dtype,
calibration_data_reader=calibration_reader,
)
```
### Testing
Unit test added to exercise this interface.
### Before your PR is "*Ready for review*"
- Is this change backward compatible?: ✅
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: N/A
- Did you write any new necessary tests?: ✅
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
✅
### Additional Information
Resolves bug: 5963347
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit
* **New Features**
* ONNX quantization accepts an optional calibration-data reader,
allowing users to supply custom calibration workflows while preserving
existing behavior when none is provided.
* **Documentation**
* Changelog updated to document the new quantization workflow option.
* **Tests**
* Added tests validating both custom calibration-reader and precomputed
calibration-data workflows.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---------
Signed-off-by: dmoodie <dmoodie@nvidia.com>
|
||
|
Keval Morabia
|
f22f4f5522 |
Add Full TE Spec support for Megatron Pruning DynamicModules + MoE bug fixes (#1024)
### What does this PR do? Type of change: Improvement + Bug Fix <!-- Use one of the following: Bug fix, new feature, new example, new tests, documentation. --> Quantization recently added support for Full TE spec. Adding same for Pruning as well so we can retire ModelOpt spec and just use standard TE spec. **NOTE: We still dont support TEGroupedGemm and instead use TE SequentialMLP for now (but this can be configured in standard TE Spec so we dont need modelopt spec)** Note that this does not affect the usage of the pruning workflow but makes pruning slightly faster and may result in slightly different pruned model because of different kernel and numerics. [Bug fix]: Previously NAS-based pruning for MoE models would hang when evaluating MMLU for pruned candidate models because of a bug. Fixed in this PR as well [Bug fix]: Previously hidden size importance hooks were not applied to pre_mlp_layernorm for MoE layers. Fixed in this PR as well resulting in a significant improvement in MMLU for Qwen3-30B-A3B ### Testing <!-- Mention how have you tested your change if applicable. --> - [x] Unit tests updated and passing - [x] Compare pruning results for Qwen3-8B -> 6B. ⚠️ Difference in MMLU scores resulting in a different best picked model. But scores more or less in similar range - difference may be because of different kernel for TE layers ``` Least important 6 layers: ModelOpt Spec: 27, 28, 29, 31, 32, 33 TE Spec: 27, 28, 30, 31, 32, 33 Top 10 pruned candidates: | num_layers | hidden_size | ffn_hidden_size | Params (B) | MMLU (ModelOpt Spec) | MMLU (TE Spec) | |------------|-------------|-----------------|------------|----------------------|----------------| | 34 | 3328 | 11264 | 5.99 | 0.390 | 0.397 | | 30 | 3584 | 11776 | 5.99 | 0.572 [BEST] | 0.575 | | 36 | 3840 | 8192 | 5.98 | 0.511 | 0.511 | | 36 | 3584 | 9216 | 5.98 | 0.477 | 0.497 | | 36 | 3072 | 11776 | 5.97 | 0.278 | 0.252 | | 32 | 3584 | 10752 | 5.96 | 0.542 | 0.541 | | 36 | 3328 | 10240 | 5.92 | 0.365 | 0.412 | | 34 | 3840 | 8704 | 5.91 | 0.537 | 0.539 | | 30 | 4096 | 9216 | 5.90 | 0.566 | 0.591 [BEST] | | 34 | 3584 | 9728 | 5.89 | 0.499 | 0.510 | ``` - [x] Compare pruning results for Nemotron-Nano-9B-v2 -> 7B. MMLU scores slight difference but best pruned model selection same ``` Least important 8 layers (Before and After): [43, 44, 45, 46, 47, 48, 50, 52] Top 10 pruned candidates: | num_layers | hidden_size | mamba_num_heads | mamba_head_dim | ffn_hidden_size | Params (B) | MMLU (ModelOpt Spec) | MMLU (TE Spec) | |------------|-------------|------------------|---------------|-----------------|------------|----------------------|----------------| | 50 | 4480 | 128 | 56 | 15680 | 7.00 | 0.211 | 0.202 | | 56 | 4096 | 96 | 80 | 14336 | 7.00 | 0.438 | 0.436 | | 48 | 4352 | 120 | 80 | 13824 | 7.00 | 0.679 [BEST] | 0.679 [BEST] | | 56 | 4352 | 112 | 80 | 10240 | 7.00 | 0.516 | 0.520 | | 54 | 4480 | 104 | 80 | 11264 | 7.00 | 0.263 | 0.262 | | 46 | 4480 | 128 | 72 | 14848 | 7.00 | 0.610 | 0.617 | | 50 | 4480 | 112 | 64 | 15680 | 7.00 | 0.426 | 0.421 | | 54 | 4096 | 112 | 80 | 13312 | 7.00 | 0.579 | 0.589 | | 56 | 4352 | 120 | 72 | 10752 | 7.00 | 0.466 | 0.469 | | 52 | 4352 | 120 | 72 | 12800 | 7.00 | 0.561 | 0.560 | ``` - [x] Compare pruning results for Qwen3-30B-A3B -> 24B. Previously there was a bug in hooks added so now we see a big improvement ``` Top 10 pruned candidates (~1 hour per candidate MMLU computation so skipped after 3): | num_layers | hidden_size | num_attention_heads | num_moe_experts | Params (B)| MMLU (ModelOpt Spec) | MMLU (TE Spec) | |------------|-------------|---------------------|-----------------|-----------|----------------------|----------------| | 46 | 2048 | 28 | 104 | 23.98B | 0.663 | 0.698 | | 40 | 2048 | 28 | 120 | 23.95B | 0.577 | 0.668 | | 46 | 1792 | 24 | 120 | 23.94B | 0.435 | 0.500 | | 46 | 2048 | 24 | 104 | 23.88B | | | | 40 | 2048 | 24 | 120 | 23.87B | | | | 46 | 1792 | 20 | 120 | 23.85B | | | | 40 | 2048 | 20 | 120 | 23.78B | | | | 46 | 2048 | 20 | 104 | 23.78B | | | | 42 | 2048 | 32 | 112 | 23.62B | | | | 48 | 1792 | 32 | 112 | 23.54B | | | ``` - [x] Run pruning experiments for gptoss-20b (21B actually) -> 18B with TESpec. Seems like GPTOSS MMLU is dropping steeply even in 15% pruning ``` "Only considering atmost 40% for width and 20% for depth pruning hparams Skipping hparams_to_skip=['num_attention_heads'] during search space generation... Search space for num_layers: [20, 22, 24] Search space for hidden_size: [2048, 2304, 2560, 2816, 2880] Search space for num_moe_experts: [24, 32] Search space for moe_ffn_hidden_size: [2048, 2304, 2560, 2816, 2880] Total search space in consideration: 150 Top 10 candidates with scores: {'num_layers': 20, 'hidden_size': 2880, 'num_moe_experts': 32, 'moe_ffn_hidden_size': 2880} -> 17.62B params, 0.3780 score {'num_layers': 22, 'hidden_size': 2880, 'num_moe_experts': 32, 'moe_ffn_hidden_size': 2560} -> 17.32B params, 0.4160 score [BEST SUBNET] {'num_layers': 20, 'hidden_size': 2880, 'num_moe_experts': 32, 'moe_ffn_hidden_size': 2816} -> 17.27B params, 0.3523 score {'num_layers': 20, 'hidden_size': 2816, 'num_moe_experts': 32, 'moe_ffn_hidden_size': 2880} -> 17.23B params, 0.3848 score {'num_layers': 22, 'hidden_size': 2560, 'num_moe_experts': 32, 'moe_ffn_hidden_size': 2880} -> 17.13B params, 0.3062 score {'num_layers': 24, 'hidden_size': 2880, 'num_moe_experts': 32, 'moe_ffn_hidden_size': 2304} -> 17.09B params, 0.3984 score {'num_layers': 22, 'hidden_size': 2816, 'num_moe_experts': 32, 'moe_ffn_hidden_size': 2560} -> 16.94B params, 0.3957 score {'num_layers': 20, 'hidden_size': 2816, 'num_moe_experts': 32, 'moe_ffn_hidden_size': 2816} -> 16.88B params, 0.3835 score {'num_layers': 22, 'hidden_size': 2560, 'num_moe_experts': 32, 'moe_ffn_hidden_size': 2816} -> 16.78B params, 0.2154 score {'num_layers': 24, 'hidden_size': 2304, 'num_moe_experts': 32, 'moe_ffn_hidden_size': 2880} -> 16.73B params, 0.0014 score" ``` - [x] Run pruning experiments for Nemotron-3-Nano-30B-A3B (31.5B actually) -> 24B with TESpec ``` Top 10 candidates with scores: {'num_layers': 46, 'hidden_size': 2688, 'mamba_num_heads': 64, 'num_moe_experts': 96, 'moe_ffn_hidden_size': 1856, 'moe_shared_expert_intermediate_size': 3712} -> 24.00B params, 0.0000 score {'num_layers': 52, 'hidden_size': 2048, 'mamba_num_heads': 64, 'num_moe_experts': 128, 'moe_ffn_hidden_size': 1856, 'moe_shared_expert_intermediate_size': 3072} -> 24.00B params, 0.2764 score {'num_layers': 48, 'hidden_size': 2688, 'mamba_num_heads': 56, 'num_moe_experts': 96, 'moe_ffn_hidden_size': 1856, 'moe_shared_expert_intermediate_size': 3712} -> 24.00B params, 0.6098 score {'num_layers': 52, 'hidden_size': 2560, 'mamba_num_heads': 64, 'num_moe_experts': 104, 'moe_ffn_hidden_size': 1792, 'moe_shared_expert_intermediate_size': 3328} -> 24.00B params, 0.6233 score {'num_layers': 48, 'hidden_size': 2688, 'mamba_num_heads': 64, 'num_moe_experts': 96, 'moe_ffn_hidden_size': 1856, 'moe_shared_expert_intermediate_size': 3072} -> 24.00B params, 0.6301 score [BEST SUBNET] {'num_layers': 52, 'hidden_size': 2688, 'mamba_num_heads': 56, 'num_moe_experts': 96, 'moe_ffn_hidden_size': 1856, 'moe_shared_expert_intermediate_size': 2560} -> 23.99B params, 0.6125 score {'num_layers': 52, 'hidden_size': 2688, 'mamba_num_heads': 48, 'num_moe_experts': 96, 'moe_ffn_hidden_size': 1856, 'moe_shared_expert_intermediate_size': 3328} -> 23.99B params, 0.4255 score {'num_layers': 50, 'hidden_size': 2048, 'mamba_num_heads': 64, 'num_moe_experts': 128, 'moe_ffn_hidden_size': 1856, 'moe_shared_expert_intermediate_size': 3584} -> 23.99B params, 0.2859 score {'num_layers': 50, 'hidden_size': 2688, 'mamba_num_heads': 56, 'num_moe_experts': 96, 'moe_ffn_hidden_size': 1856, 'moe_shared_expert_intermediate_size': 3072} -> 23.99B params, 0.6125 score {'num_layers': 42, 'hidden_size': 2304, 'mamba_num_heads': 40, 'num_moe_experts': 120, 'moe_ffn_hidden_size': 1856, 'moe_shared_expert_intermediate_size': 3584} -> 23.99B params, 0.0366 score ``` ### 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?: ⚠️ TE has different kernels so pruned model may be slightly different because of different numerics <!--- 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)?: ✅ <!--- Only for new features, API changes, critical bug fixes or backward incompatible changes. --> ### Additional Information OMNIML-3504 <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Full Transformer Engine support for Minitron pruning; no custom model spec required. * **Bug Fixes** * Resolved pruning hang on MoE models by correcting importance-hook behavior. * **Documentation** * Updated changelog and example README; bumped recommended container tag and expanded Docker run/mount guidance; adjusted release dates. * **Improvements** * Per-rank local activation handling, broader candidate caching, runtime router scoring mitigation, clearer pruning status messages. * **Tests** * Updated tests to validate Transformer Engine backend and adjusted dynamic-module expectations. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> |
||
|
Chenjie Luo
|
1dc890d971 |
Remove _moe_count_expert_calib_tokens flag; tie token counting to moe_calib_experts_ratio (#1062)
Cherry-pick for 0.43.0 ## Summary - **Remove `moe_count_expert_calib_tokens`** config field and the `_moe_count_expert_calib_tokens` internal flag. Token counting is now implicitly enabled when `moe_calib_experts_ratio` is set, removing a redundant knob. - **Change `--moe_calib_experts_ratio` default to `None`** in `hf_ptq.py` (was `1.0`). Previously all experts were force-calibrated by default; now the feature is opt-in and non-MoE models are unaffected without any flag. - **Disable `layer_sync_moe_local_experts_amax`** when `moe_calib_experts_ratio` is set, since each expert is calibrated independently with sufficient token coverage in that mode. - **Simplify `_QuantSparseMoe.forward`**: remove redundant truthy checks on `_moe_calib_experts_ratio` inside the branch that already assumes it is set. ## Changed files | File | Change | |------|--------| | `modelopt/torch/quantization/config.py` | Remove `moe_count_expert_calib_tokens` field; update `moe_calib_experts_ratio` description to document amax sync behavior | | `modelopt/torch/quantization/mode.py` | Remove `moe_count_expert_calib_tokens` propagation in `wrapped_calib_func` | | `modelopt/torch/quantization/plugins/huggingface.py` | Remove `_moe_count_expert_calib_tokens` from `_QuantSparseMoe`; simplify `forward`; skip `layer_sync_moe_local_experts_amax` when ratio is set | | `examples/llm_ptq/hf_ptq.py` | Default `--moe_calib_experts_ratio` to `None`; guard validation | | `tests/unit/.../test_sparse_moe.py` | Update tests to use `_moe_calib_experts_ratio` instead of removed flag | ## Test plan - [x] Verify `hf_ptq.py` works without `--moe_calib_experts_ratio` (non-MoE model, default `None`) 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Configuration Changes** * moe_calib_experts_ratio now defaults to None (disabled) instead of 1.0; validation only occurs when a value is provided. * **Refactor** * Simplified MoE calibration flow and token-counting behavior; removed a deprecated expert-calibration configuration field. * **Documentation** * Changelog and docstrings updated to reflect the new default and calibration behavior. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Chenjie Luo <chenjiel@nvidia.com> |