fix(autoquant): score grouped QKV at attention output#1993
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>
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #1993 +/- ##
=======================================
Coverage 77.79% 77.80%
=======================================
Files 518 518
Lines 58204 58205 +1
=======================================
+ Hits 45279 45285 +6
+ Misses 12925 12920 -5
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
/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
Summary
Findings — CRITICAL: 0, IMPORTANT: 0, SUGGESTION: 1
This is a tightly-scoped, well-tested correctness fix (2 files, +87/-4).
What was reviewed
modelopt/torch/quantization/algorithms.py— the substantive change.tests/unit/torch/quantization/test_autoquant.py— new focused regression coverage.
Algorithm correctness (OK)
The change adds _ATTN_QKV_RULE to AutoQuantizeGradientSearcher.score_module_rules, so grouped q_proj/k_proj/v_proj candidates are scored via joint perturbation at the enclosing self-attention output instead of the previous sum of three independently-measured projection-output scores. QKV already share one QuantRecipeHparam via grouping; the three identical score_modules dedupe to the single attn module (list(set(...))), and the full-backward-hook on attn accumulates one importance entry keyed by grad_output w.r.t. the attention output. Toggling hparam.active swaps quantizers on all three projections together, so output_diff correctly captures the joint QKV perturbation. Parameter/cost accounting still aggregates all three projections (verified by the test's cost assertion).
Backward compatibility (OK)
The forward-wrapper signature change from (module, input, *args, **kwargs) to (module, *args, **kwargs) is strictly more general: identical for positional calls, and it additionally supports keyword-bound inputs (x=x) which the new attention score module needs and which the parametrized test exercises. No public API, config schema, or modelopt_state change.
Most impactful note
One non-blocking SUGGESTION (inline): scoring now re-invokes the full attention forward multiple times per step rather than a pure Linear. This is safe under the usual use_cache=False calibration path, but that idempotency assumption is now load-bearing — worth documenting or guarding for score modules that mutate state (KV cache).
Risk assessment
Low. Internal scoring correction, no API change, fully covered by new tests, coverage green.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
1 participant
What does this PR do?
Type of change: Bug fix.
Changes gradient-based AutoQuantize sensitivity scoring for grouped
q_proj,k_proj, andv_projcandidates to measure their joint perturbation at the enclosing self-attention module output. QKV projections still share one format decision, and parameter/cost accounting still aggregates all three projections.This follows the output-side sensitivity formulation described in #1979 while avoiding the previous sum of three independently measured projection-output scores.
Usage
No API changes.
Testing
pre-commit run --files modelopt/torch/quantization/algorithms.py tests/unit/torch/quantization/test_autoquant.pypytest_pwd -q -x tests/unit/torch/quantization/test_autoquant.py(71 passed)Before your PR is "Ready for review"
CONTRIBUTING.md: N/AAdditional Information
Related technical background: #1979.