mirror of
https://github.com/NVIDIA/Model-Optimizer.git
synced 2026-10-02 03:14:52 +08:00
Wei-Ming Chen
f377b77116
### 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](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: 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. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Bug Fixes** * AutoQuantize now fails fast when output gradients contain non-finite values, with an actionable error identifying the affected module. * 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. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>