docs: add AutoQuantize mixed-precision search blog#1979
Conversation
Signed-off-by: realAsma <akuriparambi@nvidia.com>
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
Signed-off-by: realAsma <akuriparambi@nvidia.com>
| AutoQuantize gradient is fast! | ||
| ============================== | ||
|
|
||
| **Speed.** The direct way to measure sensitivity — quantize one layer at a time and measure a downstream evaluation such as loss, accuracy — runs the whole model per measurement. The AutoQuantize gradient scores all layer × format combinations in one forward + backward sweep. On Qwen3.6-35B-A3B, with everything else identical, that's a ~51× difference (Table 1). |
There was a problem hiding this comment.
Should we add a brief line explaining what the KL divergence sensitivity score represents?
Since we already define the AutoQuantize Gradient-based sensitivity earlier, this would give readers the full context for the comparison.
There was a problem hiding this comment.
I added this, could you take a look at the latest?
Signed-off-by: realAsma <akuriparambi@nvidia.com>
|
|
||
| A mixed-precision assignment is only useful if the runtime can execute it. So AutoQuantize performs a deployment-aware search — runtime coupling constraints are folded into the search rather than patched up afterwards, meaning the searched model is deployable out of the box in vLLM, SGLang, TensorRT-LLM, and similar inference runtimes. Any restriction of the form "this group of operators takes one joint format decision" becomes a merged knapsack item with aggregated sensitivity and cost. | ||
|
|
||
| **1) Joint quantization for fused linear layers.** Inference runtimes often fuse linear operators, which imposes a shared quantization format across the fused group. This constraint is applied within each layer: that layer's Q, K, and V projections are fused and must share one format, so the fused QKV projection becomes a single decision variable. The naive score would just sum the three per-projection sensitivities — but that treats their Hessians as independent, when the three outputs actually interact through the attention operation. Instead, AutoQuantize quantizes all three projections jointly with format :math:`f` and measures the sensitivity at the attention output, so the metric naturally captures how the projections' quantization errors combine through attention: |
There was a problem hiding this comment.
Looks like there's no score rule to redirect QKV to an attention output, we only have that for MoEs: https://github.com/NVIDIA/Model-Optimizer/blob/main/modelopt/torch/quantization/algorithms.py#L1027
In my understanding the QKV will actually get a summed score
There was a problem hiding this comment.
Yes, @meenchen had recently implemented that and found that it was better. So I am keeping that here instead of making it confusing.
There was a problem hiding this comment.
I see, is it this PR:#1878? We should consider merging it if it works better :)
Signed-off-by: realAsma <akuriparambi@nvidia.com>
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
| Results | ||
| ******* | ||
|
|
||
| .. image:: assets/autoquantize-qwen3-mmlu-effective-bits.png |
There was a problem hiding this comment.
BB: Can you get BF16 baselines for the models in the figure, regenerate this figure with BF16 baseline (plotted as a horizontal line) and send the figure to me in slack? Do not upload it here before my approval.
There was a problem hiding this comment.
🤖 Bot comment.
I’ve queued the BF16 baseline runs and a regenerated figure with horizontal BF16 reference lines. I’ll send the candidate privately in Slack and will not upload it to this PR before your approval.
meenchen
left a comment
meenchen
left a comment
There was a problem hiding this comment.
High-priority documentation correctness findings are inline.
| Results | ||
| ******* | ||
|
|
||
| .. image:: assets/autoquantize-qwen3-mmlu-effective-bits.png |
There was a problem hiding this comment.
This plot starts at 4.0, but current NVFP4 accounting is 4.5 bits including scale overhead.
| Joint quantization for fused linear layers | ||
| =========================================== | ||
|
|
||
| TensorRT-LLM fuses each layer's Q, K, and V projections, so AutoQuantize constrains the three projections to share one format. Sensitivity remains measured at each projection's individual module output, and the three scores are aggregated for the shared-format decision: |
There was a problem hiding this comment.
nit: this is not specific to TensorRT-LLM but applies to most of the inference libraries
| Performance cost | ||
| ================ | ||
|
|
||
| ModelOpt uses *effective bits* as the cost: the total compressed model size, counting each weight's bit-width plus scale-factor overhead. Sweeping the budget over effective-bits targets directly controls model size and makes the accuracy-vs-compression frontier easy to interpret. |
There was a problem hiding this comment.
[P2] “Total compressed model size” overstates what this constraint measures. The LP denominator includes only AutoQuantize-eligible quantizable weights, so embeddings, norms, and other excluded parameters are not counted. Format overhead is included when a format supplies an explicit effective_bits value; otherwise the implementation falls back to a num_bits heuristic. Could we describe this as a modeled average bit cost over eligible weights, with format-provided overhead where available?
| AutoQuantize gradient is fast! | ||
| ============================== | ||
|
|
||
| Direct sensitivity measurement runs the whole model once per layer per format. KL-divergence scoring is one such direct measurement: quantize one layer, run a forward pass, and compare the output distributions of the quantized and unquantized models. The gradient score instead scores all layers at once, so its cost scales with the number of formats, not the number of layers. On Qwen3.6-35B-A3B that is a ~51× difference (Table 1). |
There was a problem hiding this comment.
[P2] “Cost scales with the number of formats, not the number of layers” contradicts both Table 1 and the implementation. For each scoring batch, the patched forward visits every score module and locally replays every candidate format, then performs one backward; the table’s O(N_layers × N_formats) is the correct total-work scaling. The saving is that AutoQuantize avoids a full-model evaluation per layer and format, not that layer scaling disappears. Please rephrase accordingly and change line 53 to “one backward pass per scoring batch.”
| How AutoQuantize works | ||
| ********************** | ||
|
|
||
| AutoQuantize is a neural architecture search (NAS) inspired method that works in three steps: score how sensitive each operation is to quantization, model the performance cost of each available format, and solve for the best layer-wise assignment under the cost budget with a knapsack-style optimization. The sensitivity score is a second-order Taylor approximation in the spirit of Optimal Brain Surgeon [1]_, as introduced in LLM-MQ [2]_. |
There was a problem hiding this comment.
[P2] LLM-MQ does not introduce this second-order diagonal-Fisher score. Its §2.2, Eq. 1–2 uses the first-order weight-space score |g_iᵀ(W_i − Q_b(W_i))| and explicitly focuses on first-order information instead of the Hessian: https://nicsefc.ee.tsinghua.edu.cn/nics_file/pdf/5c805adc-b555-499f-9882-5ca35ce674b5.pdf. Please credit LLM-MQ for the mixed-precision allocation inspiration and cite the actual source for the output-side Fisher score.
|
|
||
| The figure shows the accuracy-vs-compression frontier AutoQuantize traces on Qwen3 models: sweep the bit budget, solve at each point, evaluate on MMLU. Accuracy rises with the budget. | ||
|
|
||
| Two things to notice. First, the rise is essentially monotonic, which means the sensitivity score is ranking layers correctly — a noisy proxy would give a jagged frontier. Second, adding FP8 to the format menu helps: at every budget, searching over NVFP4, FP8, and BF16 matches or beats NVFP4 and BF16 alone. A sensitive layer doesn't need to fall back all the way to BF16 — FP8 is a good middle ground, protecting moderately sensitive layers at a fraction of the cost. |
There was a problem hiding this comment.
[P2] These results do not form an empirical Pareto frontier or establish that the sensitivity ranking is correct. For example, Qwen3-1.7B at 6 bits (~55.0) dominates its 6.5–8-bit points, and Qwen3-14B drops from 7.5 to 8 bits. A solver optimum for the additive proxy is not necessarily Pareto-optimal in measured MMLU; without repeated runs or error bars, the dips could also be evaluation noise. Please call this a budget sweep or proxy-optimal assignment and soften the text to an overall upward trend, or add direct ranking validation.
Signed-off-by: realAsma <akuriparambi@nvidia.com>
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Claude review passed — no blocking issues found. LGTM
Scope: Documentation-only PR (3 files: new docs/source/announcements/autoquantize.rst, one PNG asset, and a landing-page card/filter in docs/source/index.rst). Reviewed the RST blog and index changes; the binary PNG was not opened.
What I verified:
- The
mtq.auto_quantize(...)code example matches the real public API: signature/argument names (constraints,quantization_formats,data_loader,forward_step,loss_func,num_calib_steps,num_score_steps), themodel, search_state = ...return-value unpacking, and theNVFP4_DEFAULT_CFG/FP8_DEFAULT_CFGconfig constants all exist inmodelopt/torch/quantization. - The sensitivity-score derivation (Taylor → OBS → diagonal-Fisher/squared-gradient approximation) is internally consistent and correctly presented.
- The landing-page card in
index.rstis consistent with the RST (title, date, tags, author order all match); newest-first ordering is preserved.
Findings (1 SUGGESTION, non-blocking):
- The benchmark model name
Qwen3.6-35B-A3B(Table 1 timing paragraph) does not correspond to a real Qwen3 release — the MoE model with 3B active params isQwen3-30B-A3B. Since this is customer-facing published content, please confirm the exact model used and correct the identifier before publishing. (Inline comment posted.)
Risk: Low — docs-only, no code/API/state changes.
4 participants
What does this PR do?
Type of change: documentation.
Adds the AutoQuantize technical blog to the announcements system introduced by #1971.
This is a stacked PR targeting the branch for #1971. After #1971 merges, this PR can be retargeted to main.
Usage
N/A; documentation only.
Testing
Before your PR is "Ready for review"
Additional Information
Depends on #1971.