Skip to content

Make group-boundary AutoQuant scoring the default#1988

Draft
meenchen wants to merge 3 commits into
mainfrom
weimingc/autoquant-group-scoring
Draft

Make group-boundary AutoQuant scoring the default#1988
meenchen wants to merge 3 commits into
mainfrom
weimingc/autoquant-group-scoring

Conversation

@meenchen

@meenchen meenchen commented Jul 17, 2026

Copy link
Copy Markdown
Contributor

What does this PR do?

Type of change: New feature

AutoQuant gradient scoring currently measures attention projections at their leaf outputs. That can understate perturbations amplified by the surrounding attention computation. Expert projections already use their parent MLP/mixer output as the scoring boundary.

This PR:

  • makes group-boundary scoring the default for gradient AutoQuant;
  • scores quantizable descendants of attention modules at the parent attention output, covering common Hugging Face and Megatron layouts such as self_attn.qkv_proj, self_attention.linear_qkv/linear_proj, and linear_attn.*;
  • resolves named ModuleList/ModuleDict attention containers to the callable child on the projection path, avoiding hooks on uncalled containers;
  • warns instead of silently falling back when a likely attention projection cannot be mapped;
  • preserves the established parent MLP/mixer boundary for unfused routed experts, shared experts, and fused expert containers;
  • keeps recipe decisions independent, so this does not impose an attention-family quantization mask or force projections to share a format;
  • preserves leaf-output attention scoring through score_boundary: local;
  • keeps KL-divergence on local scoring;
  • transactionally disables and restores configuration-level KV caching while parent-boundary gradient scores are collected;
  • rejects explicit mutable cache/recurrent-state inputs, positional or keyword, before replay while permitting use_cache=False;
  • guarantees forward, backward, parameter-hook, and cache-setting cleanup if cache disabling or scoring raises;
  • validates the method/boundary combination before quantization mode conversion, leaving rejected models unchanged and retryable;
  • exposes the switch through the recipe/config system only;
  • records the resolved boundary in AutoQuant checkpoints and rejects incompatible replay;
  • documents the behavior and adds focused coverage.

This branch intentionally contains only group-boundary scoring. Per-element score normalization remains a separate change.

Usage

Gradient AutoQuant recipes use group-boundary scoring by default:

auto_quantize:
  auto_quantize_method: gradient

To restore leaf-output attention scoring:

auto_quantize:
  auto_quantize_method: gradient
  score_boundary: local

KL-divergence remains local because its threshold search does not support parent/group-boundary scoring.

Testing

  • Focused unit/recipe/example suite:
    336 passed, 1 deselected
    
    from test_autoquant.py, test_loader.py, and test_hf_ptq_args.py.
  • New regressions cover:
    • an executed attention child nested under an uncalled ModuleList;
    • positional and keyword state/cache replay inputs;
    • explicit use_cache=False acceptance and use_cache=True rejection;
    • transactional restoration and hook cleanup after a partially failing cache setter.
  • Full pre-commit suite for every changed file: passed.
  • GPU end-to-end PTQ/export on aws-cmh:
    • Slurm job: 2663167
    • Tested commit: 64abbf4bf940036845a19abeac1e6ae0c7b927b9
    • Model: Qwen3.5-9B
    • Recipe: general/auto_quantize/nvfp4_fp8_at_5p4bits
    • The recipe omits score_boundary; the saved state records the resolved value as group
    • Result: mixed FP8/NVFP4 recipe selected at 5.40 effective bits
    • Unified HF checkpoint export and artifact validation: passed
    • Job status: COMPLETED, exit 0:0

Before your PR is "Ready for review"

  • Is this change backward compatible?: ❌ The default gradient attention scoring boundary changes; score_boundary: local restores the previous behavior.
  • 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?: ✅
  • Did you get Claude approval on this PR?: N/A

Additional Information

The scoring boundary is distinct from module grouping: this PR changes where sensitivity is measured, not which modules share one AutoQuant assignment.

@copy-pr-bot

copy-pr-bot Bot commented Jul 17, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@coderabbitai

coderabbitai Bot commented Jul 17, 2026

Copy link
Copy Markdown
Contributor

Important

Review skipped

Draft detected.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 8b967899-59dd-4638-ab71-cd5c6a346f64

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

  • 🔍 Trigger review
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch weimingc/autoquant-group-scoring

Comment @coderabbitai help to get the list of available commands.

@github-actions

github-actions Bot commented Jul 17, 2026

Copy link
Copy Markdown
Contributor
PR Preview Action v1.8.1

QR code for preview link

🚀 View preview at
https://NVIDIA.github.io/Model-Optimizer/pr-preview/pr-1988/

Built to branch gh-pages at 2026-07-23 22:14 UTC.
Preview will be ready when the GitHub Pages deployment is complete.

@codecov

codecov Bot commented Jul 17, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 93.54839% with 10 lines in your changes missing coverage. Please review.
✅ Project coverage is 75.83%. Comparing base (01c708e) to head (64abbf4).
⚠️ Report is 3 commits behind head on main.

Files with missing lines Patch % Lines
modelopt/torch/quantization/algorithms.py 93.10% 10 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main    #1988      +/-   ##
==========================================
+ Coverage   75.81%   75.83%   +0.02%     
==========================================
  Files         518      518              
  Lines       58574    58737     +163     
==========================================
+ Hits        44406    44544     +138     
- Misses      14168    14193      +25     
Flag Coverage Δ
unit 55.14% <93.54%> (+0.11%) ⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

Score gradient perturbations at parent attention and MLP outputs by default, with an explicit local-boundary opt-out and local-only KL-divergence behavior.

Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
@meenchen
meenchen force-pushed the weimingc/autoquant-group-scoring branch from 82b7171 to 3fcb35f Compare July 23, 2026 17:00
meenchen added 2 commits July 23, 2026 12:40
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant