Fail fast on non-finite AutoQuantize output gradients

[meenchen]

What does this PR do?

Type of change: Bug fix

Fail fast when AutoQuantize receives non-finite output gradients, before accumulating sensitivity scores. The error names the affected module and suggests checking the model, data, and loss; it also identifies cuDNN SDPA backward on fully masked rows as one possible cause and gives an explicit retry workaround.

Unlike the earlier revision, this does not disable cuDNN or change any attention backend settings. Invalid gradients are not zeroed or ignored, and there is no automatic retry.

Usage

No API or recipe changes. For the reproduced cuDNN failure, the caller can explicitly set torch.backends.cuda.enable_cudnn_sdp(False) before a fresh AutoQuantize run.

Testing

• AutoQuantize unit suite: 110 passed. Coverage includes NaN and positive/negative infinity, module diagnostics, preventing invalid score accumulation, model-state cleanup, and unchanged SDPA backend settings.
• Real-model E2E on four GB300 GPUs, Qwen/Qwen3.6-35B-A3B, main 8025a3dc5481129aa21fef99cb13a879e1b5847e plus this patch, batch size 8, 512 calibration samples, and w4a16_nvfp4_fp8_at_6p0bits-active_moe.yaml:
  • Default backend: the new diagnostic fired at model.language_model.layers.39.self_attn.q_proj; cuDNN remained enabled and no quantized model was exported. Expected-error check passed.
  • Explicit cuDNN-disabled fresh run: both 64-batch calibration passes, all 16 scoring batches, optimization at 5.99 effective bits, and checkpoint export completed with exit code 0. Verified all three indexed safetensors shards and quantization configuration; the index contains 93,563 tensors.
• All applicable pre-commit checks and git diff --check passed.

Before your PR is "Ready for review"

Contributor guidelines and security guidance followed; commits are signed and signed off.

• Is this change backward compatible?: Yes; finite-gradient behavior and backend settings are unchanged.
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: N/A; neither added.
• Did you write any new necessary tests?: Yes.
• Did you update Changelog?: Yes, 0.48.0 bug fixes.
• Did you get Claude approval on this PR?: No; awaiting review.

Additional Information

This improves error reporting rather than fixing the upstream cuDNN kernel. Blackwell-specificity is not established. Checkpoint deployment/reload was not tested.

A calibration-only checkpoint-resume attempt completed scoring but encountered a separate candidate_stats KeyError. That issue is outside this patch; the successful export validation above used a fresh run without search-state resume.

Summary by CodeRabbit

• Bug Fixes
  • AutoQuantize now fails fast when output gradients contain non-finite values, with an actionable error identifying the affected module and troubleshooting guidance.
  • Attention backend settings are preserved and restored after successful runs and failures.
  • Model state is restored when sensitivity scoring encounters an error during setup or execution, including temporary configuration changes.

[copy-pr-bot]

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.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 41dff836-c670-4005-9172-2a82756df513

📥 Commits

Reviewing files that changed from the base of the PR and between 13b0bd65abdf4b7d418fb5dce375dc5ee9914ac3 and e6eca8e.

📒 Files selected for processing (3)

• CHANGELOG.rst
• modelopt/torch/quantization/algorithms.py
• tests/unit/torch/quantization/test_autoquant.py

🚧 Files skipped from review as they are similar to previous changes (1)

• CHANGELOG.rst

Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review.

---

📝 Walkthrough

Walkthrough

AutoQuantize now rejects non-finite output gradients before sensitivity accumulation. Tests verify error handling and restoration of model, scoring, and attention backend state.

Changes

AutoQuantize gradient validation

Layer / File(s)	Summary
Fail-fast gradient validation modelopt/torch/quantization/algorithms.py	_accumulate_scores validates output gradients and raises a module-specific RuntimeError when values are non-finite.
Scoring cleanup coverage tests/unit/torch/quantization/test_autoquant.py, CHANGELOG.rst	Tests verify state restoration across successful and failed scoring paths, including CUDA SDPA settings. The changelog records the fix.

Priority: ⬇️ Low

Estimated code review effort: 2 (Simple) | ~10 minutes

Change: Bug fix

Merge Risk: ⚪ Minimal · up to e6eca

The change has no identified merge-blocking risk.

🚥 Pre-merge checks | ✅ 5 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 16.67% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 6 functions across 2 files. (1 skipped: 1…	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and concisely describes the main change: AutoQuantize now fails fast when output gradients are non-finite.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Security Anti-Patterns	✅ Passed	PASS. The reviewed range changes one package Python file, modelopt/torch/quantization/algorithms.py, and adds only a torch.isfinite check plus diagnostic error text. The added code does not use `t…

Full details: Docstring Coverage

Explanation

Docstring coverage is 16.67% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 6 functions across 2 files. (1 skipped: 1 unsupported.)

• Fix all pre-merge checks with AI

✨ Finishing Touches 💡 1

📝 Generate docstrings 💡

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Commit to this branch
• Create a new PR

---

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

[github-actions]

PR Preview Action v1.8.1
Preview removed because the pull request was closed.
2026-09-17 17:53 UTC

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 78.97%. Comparing base (b9cfdce) to head (e6eca8e).

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2432      +/-   ##
==========================================
+ Coverage   71.50%   78.97%   +7.47%     
==========================================
  Files         590      590              
  Lines       64749    64750       +1     
==========================================
+ Hits        46297    51136    +4839     
+ Misses      18452    13614    -4838     

Flag	Coverage Δ
examples-diffusers	20.88% <0.00%> (-0.01%)	⬇️
examples-gpt-oss	13.39% <0.00%> (-0.01%)	⬇️
examples-hf_ptq	22.49% <0.00%> (-0.04%)	⬇️
examples-llm_distill	13.46% <0.00%> (-0.01%)	⬇️
examples-llm_eval	17.37% <0.00%> (-0.01%)	⬇️
examples-llm_qat	17.67% <0.00%> (-0.01%)	⬇️
examples-llm_sparsity	15.93% <0.00%> (-0.01%)	⬇️
examples-megatron_bridge	26.27% <0.00%> (-0.12%)	⬇️
examples-specdec_bench	13.15% <0.00%> (-0.01%)	⬇️
examples-speculative_decoding	17.79% <0.00%> (-0.07%)	⬇️
examples-torch_onnx	21.89% <0.00%> (-0.01%)	⬇️
examples-torch_trt	15.22% <0.00%> (-0.01%)	⬇️
gpu	58.35% <0.00%> (+25.91%)	⬆️
regression	15.15% <0.00%> (+0.29%)	⬆️
unit	57.86% <100.00%> (+<0.01%)	⬆️

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.

[cjluo-nv]

Bot review (gpt-5.6-sol) — DM the bot to share feedback.

LGTM: the fail-fast check is scoped correctly, restores session state, preserves backend settings, and has focused non-finite-gradient coverage.

Needs action:

• Provide human sign-off for the justified edit to test_backward_scoring_session_restores_partial_setup, which adds coverage that failure cleanup leaves cuDNN SDPA unchanged.

[cjluo-nv]

Bot review (gpt-5.6-sol) — DM the bot to share feedback.

Approved: the fail-fast check is scoped correctly, preserves scoring cleanup, and has focused NaN and infinity coverage.
No action needed:

• ✔️ The prior sign-off request is no longer applicable: test_backward_scoring_session_restores_partial_setup remains unchanged; this PR only inserts new tests before it.