[OMNIML-2244] Add support for auto quantizing a model - #571
Merged
Conversation
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #571 +/- ##
==========================================
+ Coverage 74.43% 74.45% +0.02%
==========================================
Files 182 182
Lines 18234 18250 +16
==========================================
+ Hits 13572 13588 +16
Misses 4662 4662 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
cjluo-nv
reviewed
Nov 17, 2025
cjluo-nv
reviewed
Nov 17, 2025
cjluo-nv
reviewed
Nov 17, 2025
cjluo-nv
approved these changes
Nov 18, 2025
ajrasane
force-pushed
the
ajrasane/onnx_mixed_precision
branch
from
November 19, 2025 21:07
3362da0 to
f12e08c
Compare
Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> Return before exporting an auto quantized model Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com>
ajrasane
force-pushed
the
ajrasane/onnx_mixed_precision
branch
from
November 19, 2025 21:19
f12e08c to
33bd81b
Compare
Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com>
gcunhase
reviewed
Nov 20, 2025
| choices=["fp8", "mxfp8", "int8", "nvfp4", "int4_awq", "auto"], | ||
| default="mxfp8", | ||
| help="Type of quantization to apply (mxfp8, nvfp4, int4_awq)", | ||
| help="Type of quantization to apply (fp8, mxfp8, int8, nvfp4, int4_awq, auto)", |
Contributor
There was a problem hiding this comment.
Small suggestion: remove (fp8, mxfp8, int8, nvfp4, int4_awq, auto), as argparse should automatically include the choices in the help output.
Removing that from the help string should result in:
--quantize_mode {fp8,mxfp8,int8,nvfp4,int4_awq,auto}
Type of quantization to apply (default: mxfp8)
gcunhase
approved these changes
Nov 20, 2025
Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com>
ajrasane
enabled auto-merge (squash)
November 20, 2025 18:46
Merged
kevalmorabia97
added a commit
that referenced
this pull request
Jun 27, 2026
### What does this PR do? Type of change: New feature Adds **Domino** speculative decoding: the parallel DFlash draft backbone plus a lightweight **GRU causal correction head**. The backbone produces *base* logits for a full draft block in one forward; a GRU over the block's teacher-forced tokens produces a causal state that is fused with the backbone hidden state and projected to a vocab-sized logit correction on the block suffix — injecting the intra-block causal dependency the parallel backbone lacks. Trained with a dual loss `(1-λ)*final + λ*base`, where `λ_base` decays linearly 1→0 (curriculum: learn the parallel backbone first, then the correction). Reuses the DFlash mode/config/recipe; selected via `dflash_architecture_config.projector_type=domino` and routed to its own registry so `HFDominoModel` does not shadow `HFDFlashModel`. Exports in the z-lab/SpecForge drafter format (`prefix_gru.*` / `embed_proj.*`). > Note: the inference side (vLLM / AR evaluation) is intentionally **not** wired up yet — the correction head is not applied in serving. To be added once the inference path lands. ### Usage ```bash # Online training (recipe: projector_type=domino) uv run launch.py --yaml examples/Qwen/Qwen3-8B/hf_online_domino.yaml --yes ``` ### Testing CPU unit tests in `tests/unit/torch/speculative/plugins/test_hf_domino.py` cover conversion routing, the training forward (dual loss + grads), the λ schedule, and the export format. Online Qwen3-8B training validated end-to-end (loss curve below). <img width="1803" height="809" alt="image" src="https://github.com/user-attachments/assets/7c9d2001-bd80-4dec-919b-443e61089cca" /> ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ (opt-in via `projector_type=domino`; DFlash path unchanged) - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A (no new dependency) - Did you write any new necessary tests?: ✅ - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ - Did you get Claude approval on this PR?: ❌ ### Additional Information Reference: SpecForge PR #571 (z-lab); drafter format `huggingface.co/Huang2020/Qwen3-8B-Domino-b16`. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit ## Release Notes * **New Features** * Added Domino speculative-decoding training with a decaying base/final dual-loss curriculum and Domino-specific lambda scheduling (training-only; inference wiring not yet included). * Added Domino draft-head export support for training checkpoints. * **Documentation & Configuration** * Added a Domino speculative-decoding training recipe and an HF Online Domino launcher configuration for Qwen3-8B. * **Refactor** * Updated speculative model conversion/export to route to Domino variants based on the configured projector type. * **Tests** * Added unit tests for Domino conversion, training loss/metrics, lambda decay behavior, and exporter output layout. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Type of change:
Example update
Overview:
mtq.auto_quantize()Usage
Testing
Able to auto quantize ViT model
Accuracy comparison for the ViT model
Before your PR is "Ready for review"