Fix distributed amax error handling

modelopt/torch/quantization/nn/modules/tensor_quantizer.py

@@ -1352,16 +1352,16 @@ def set_from_modelopt_state(self, modelopt_state, properties_only: bool = False)
     def sync_amax_across_distributed_group(self, parallel_group: DistributedProcessGroup):
         """Synchronize the amax across all ranks in the given group."""
         if parallel_group.is_initialized() and getattr(self, "_amax", None) is not None:
-            try:
-                dist.all_reduce(self._amax, op=dist.ReduceOp.MAX, group=parallel_group.group)
-            except RuntimeError as e:
-                # This error happens if the distributed backend is using GPU and
-                # the tensor is not on GPU (or vice versa).
+            if (
+                self._amax.device.type == "cpu"
+                and dist.get_backend(parallel_group.group) == dist.Backend.NCCL
+            ):
                 warnings.warn(
-                    f"Failed to synchronize amax: {e}, probably because the tensor is on a device which is not"
-                    "supported by the current distributed backend. This warning can be ignored"
-                    "if happening during modelopt restore."
+                    "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."
                 )
+                return
+            dist.all_reduce(self._amax, op=dist.ReduceOp.MAX, group=parallel_group.group)
 
     @contextlib.contextmanager
     def disable_pre_quant_scale(self):

tests/gpu/torch/quantization/test_tensor_quantizer_cuda.py

@@ -19,6 +19,7 @@
 
 import pytest
 import torch
+from _test_utils.torch.distributed.utils import spawn_multiprocess_job
 from _test_utils.torch.quantization.tensor_quantizer_common import (
     BlockQuantTester,
     TensorQuantizerTester,
@@ -29,6 +30,22 @@
 from modelopt.torch.quantization.config import QuantizerAttributeConfig
 from modelopt.torch.quantization.extensions import get_cuda_ext_mx
 from modelopt.torch.quantization.nn.modules import tensor_quantizer
+from modelopt.torch.utils.distributed import DistributedProcessGroup
+
+
+def _test_sync_cpu_amax_with_nccl(rank, size):
+    quantizer = tensor_quantizer.TensorQuantizer()
+    quantizer._amax = torch.tensor(float(rank))
+
+    with pytest.warns(UserWarning, match="tensor is on CPU, but NCCL only supports CUDA tensors"):
+        quantizer.sync_amax_across_distributed_group(DistributedProcessGroup())
+
+    assert quantizer.amax.device.type == "cpu"
+    assert quantizer.amax == rank
+
+
+def test_sync_cpu_amax_with_nccl(need_2_gpus):
+    spawn_multiprocess_job(2, _test_sync_cpu_amax_with_nccl, backend="nccl")
 
 
 class TestTensorQuantizerCuda(TensorQuantizerTester):