Fix distributed amax error handling

[realAsma]

Description

TensorQuantizer.sync_amax_across_distributed_group currently catches every RuntimeError from dist.all_reduce and converts it to a warning. This can mask failures such as NCCL OOMs, allowing ranks to diverge and hang in later collectives.

This PR handles only the intended unsupported-device case before entering the collective:

• skip distributed amax synchronization when the amax tensor is on CPU and the process-group backend is NCCL
• retain the existing restore-oriented warning for that case
• let all errors from supported collectives propagate normally

Tests

• pytest_pwd tests/unit/torch/quantization/test_dist.py -q
• pre-commit run --files modelopt/torch/quantization/nn/modules/tensor_quantizer.py tests/unit/torch/quantization/test_dist.py

Summary by CodeRabbit

• Bug Fixes

  • Improved distributed quantization synchronization for NCCL runs when the quantization range tensor is on CPU.
  • Added an explicit warning and an early exit to prevent unsupported all-reduce attempts.

• Tests

  • Added a multi-process CUDA/NCCL test covering CPU-based range tensors, verifying the warning is emitted and the value stays on CPU as expected.

[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]

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

📝 Walkthrough

Walkthrough

Adds a CPU-and-NCCL guard in TensorQuantizer.sync_amax_across_distributed_group that warns and exits before dist.all_reduce, and adds a distributed CUDA test covering the CPU amax case.

Changes

CPU/NCCL amax sync guard

Layer / File(s)	Summary
Guard implementation and distributed test modelopt/torch/quantization/nn/modules/tensor_quantizer.py, tests/gpu/torch/quantization/test_tensor_quantizer_cuda.py	Replaces try/except-based dist.all_reduce handling with an explicit CPU+NCCL check; adds a 2-process NCCL test that verifies the warning and preserves the CPU amax value.

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

🚥 Pre-merge checks | ✅ 6

✅ Passed checks (6 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly summarizes the main change to distributed amax error handling.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
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	Reviewed the PR-mentioned Python changes; they only add a CPU/NCCL guard and test, with no unsafe load/eval/exec/trust_remote_code/nosec patterns or new deps.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch asma/distributed-amax-error-handling

---

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

[github-actions]

PR Preview Action v1.8.1
🚀 View preview at https://NVIDIA.github.io/Model-Optimizer/pr-preview/pr-1944/
Built to branch gh-pages at 2026-07-09 15:55 UTC. Preview will be ready when the GitHub Pages deployment is complete.

[realAsma]

BB: Approved; make this a regular PR

[realAsma]

🤖 Bot comment.

🐝 Approved—I'll mark this PR ready for review now.

[realAsma]

🤖 Bot comment.

Done — PR #1944 is now a regular, ready-for-review PR. GitHub confirms it remains open and mergeable at 73e88e476d; no code or branch content changed.

[coderabbitai]

🧹 Nitpick comments (1)

modelopt/torch/quantization/nn/modules/tensor_quantizer.py (1)

1350-1362: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Guard correctly narrows error handling to the CPU+NCCL case; consider rank-aware warning.

The explicit CPU+NCCL check replacing the blanket try/except RuntimeError is a solid fix — genuine dist.all_reduce errors for other backends/cases now propagate instead of being silently swallowed as warnings, matching the PR's intent and the accompanying test.

One nit: since this runs identically on every rank in the group, warnings.warn will fire redundantly on all ranks. Coding guidelines call for rank-aware logging in distributed code paths to avoid noisy duplicate output.

♻️ Suggested tweak (if a rank-aware warn helper exists in the codebase)

-                warnings.warn(
+                warn_rank_0(
                     "Failed to synchronize amax because the tensor is on CPU, but NCCL only supports CUDA "
                     "tensors. This warning can be ignored if happening during modelopt restore."
                 )

As per coding guidelines, "Develop with distributed processing in mind: use print_rank_0 or warn_rank_0 when possible and guard shared side effects against race conditions between ranks." Please confirm whether such a helper exists in modelopt.torch.utils before applying.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@modelopt/torch/quantization/nn/modules/tensor_quantizer.py` around lines 1350
- 1362, The CPU+NCCL warning in sync_amax_across_distributed_group is emitted on
every rank, causing duplicate distributed logs. Update the warning path to use a
rank-aware helper such as warn_rank_0 or print_rank_0 if it exists in
modelopt.torch.utils, and keep the warning only for the CPU+NCCL branch while
preserving the current all_reduce behavior for valid cases.

Source: Coding guidelines

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@modelopt/torch/quantization/nn/modules/tensor_quantizer.py`:
- Around line 1350-1362: The CPU+NCCL warning in
sync_amax_across_distributed_group is emitted on every rank, causing duplicate
distributed logs. Update the warning path to use a rank-aware helper such as
warn_rank_0 or print_rank_0 if it exists in modelopt.torch.utils, and keep the
warning only for the CPU+NCCL branch while preserving the current all_reduce
behavior for valid cases.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 89f4e79c-a26b-43cd-85aa-7c7dc8f16f7e

📥 Commits

Reviewing files that changed from the base of the PR and between b012e4c and 73e88e4.

📒 Files selected for processing (2)

• modelopt/torch/quantization/nn/modules/tensor_quantizer.py
• tests/gpu/torch/quantization/test_tensor_quantizer_cuda.py

[codecov]

Codecov Report

❌ Patch coverage is 66.66667% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 77.16%. Comparing base (aeb1ab4) to head (c4db72d).

Files with missing lines	Patch %	Lines
.../torch/quantization/nn/modules/tensor_quantizer.py	66.66%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1944      +/-   ##
==========================================
- Coverage   77.79%   77.16%   -0.64%     
==========================================
  Files         519      519              
  Lines       57928    57928              
==========================================
- Hits        45067    44699     -368     
- Misses      12861    13229     +368     

Flag	Coverage Δ
examples	43.28% <66.66%> (-0.11%)	⬇️
gpu	57.81% <66.66%> (-0.64%)	⬇️
regression	14.93% <0.00%> (+0.06%)	⬆️
unit	55.22% <66.66%> (ø)

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.

[realAsma]

/claude review

[claude]

Claude review passed — no blocking issues found. LGTM

Scope: Reviewed the 2 files actually in this PR (modelopt/torch/quantization/nn/modules/tensor_quantizer.py, tests/gpu/torch/quantization/test_tensor_quantizer_cuda.py). Note: a stale shallow local base surfaced unrelated transformer_engine.py hunks in the raw diff; those are base drift (real base 2d4e472), not part of this PR, and were excluded.

Assessment (low risk):

• The core change narrows a blanket try/except RuntimeError around dist.all_reduce to an explicit CPU-amax + NCCL-backend pre-check that warns and returns early, letting genuine collective failures (e.g. NCCL OOM) propagate instead of being silently downgraded to warnings. This matches the stated intent.
• Correctness checks pass: dist.get_backend(...) == dist.Backend.NCCL is valid (Backend is a str subclass), get_backend is safely gated behind is_initialized(), and the early-return leaves _amax unsynchronized — consistent with the prior behavior and the new test's assertions.
• No mode/state, export, or backward-compatibility surface is touched; the method signature and modelopt_state schema are unchanged.
• The new 2-GPU NCCL test is well-placed and correctly validates the warning, the CPU device, and the per-rank amax value.

Findings — CRITICAL: 0, IMPORTANT: 0, SUGGESTION: 0 (new). CodeRabbit already raised the rank-aware warn_rank_0 nit (non-blocking); not duplicating it here.