diff --git a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py index 98d9e0dcb1e..acaf83dc4fc 100644 --- a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py +++ b/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): diff --git a/tests/gpu/torch/quantization/test_tensor_quantizer_cuda.py b/tests/gpu/torch/quantization/test_tensor_quantizer_cuda.py index af884c878b2..44eac4d0034 100644 --- a/tests/gpu/torch/quantization/test_tensor_quantizer_cuda.py +++ b/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):