From 28db895ef755870ec5b6132169c25e357b5b0218 Mon Sep 17 00:00:00 2001 From: realAsma Date: Wed, 8 Jul 2026 17:44:00 +0000 Subject: [PATCH 1/3] Fix distributed amax error handling Signed-off-by: realAsma --- .../nn/modules/tensor_quantizer.py | 16 ++++----- tests/unit/torch/quantization/test_dist.py | 33 +++++++++++++++++++ 2 files changed, 41 insertions(+), 8 deletions(-) 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/unit/torch/quantization/test_dist.py b/tests/unit/torch/quantization/test_dist.py index 4130b8ce82a..13742b1bcc1 100644 --- a/tests/unit/torch/quantization/test_dist.py +++ b/tests/unit/torch/quantization/test_dist.py @@ -13,6 +13,9 @@ # See the License for the specific language governing permissions and # limitations under the License. +from unittest.mock import Mock, patch + +import pytest import torch import torch.distributed as dist from _test_utils.torch.distributed.utils import spawn_multiprocess_job @@ -45,3 +48,33 @@ def forward_loop(model): def test_data_parallel(skip_on_windows): spawn_multiprocess_job(2, _test_data_parallel_helper, backend="gloo") + + +def test_sync_amax_skips_cpu_tensor_with_nccl(): + quantizer = TensorQuantizer() + quantizer._amax = torch.tensor(1.0) + parallel_group = Mock() + parallel_group.is_initialized.return_value = True + + with ( + patch.object(dist, "get_backend", return_value=dist.Backend.NCCL), + patch.object(dist, "all_reduce") as all_reduce, + pytest.warns(UserWarning, match="tensor is on CPU, but NCCL only supports CUDA tensors"), + ): + quantizer.sync_amax_across_distributed_group(parallel_group) + + all_reduce.assert_not_called() + + +def test_sync_amax_propagates_collective_errors(): + quantizer = TensorQuantizer() + quantizer._amax = torch.tensor(1.0) + parallel_group = Mock() + parallel_group.is_initialized.return_value = True + + with ( + patch.object(dist, "get_backend", return_value=dist.Backend.GLOO), + patch.object(dist, "all_reduce", side_effect=RuntimeError("out of memory")), + pytest.raises(RuntimeError, match="out of memory"), + ): + quantizer.sync_amax_across_distributed_group(parallel_group) From c5433ad927b61a02565fc58d892130f46a423d65 Mon Sep 17 00:00:00 2001 From: realAsma Date: Wed, 8 Jul 2026 17:56:45 +0000 Subject: [PATCH 2/3] Move CPU NCCL amax test to GPU suite Signed-off-by: realAsma --- .../quantization/test_tensor_quantizer_cuda.py | 17 +++++++++++++++++ tests/unit/torch/quantization/test_dist.py | 16 ---------------- 2 files changed, 17 insertions(+), 16 deletions(-) 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): diff --git a/tests/unit/torch/quantization/test_dist.py b/tests/unit/torch/quantization/test_dist.py index 13742b1bcc1..a98d239f22b 100644 --- a/tests/unit/torch/quantization/test_dist.py +++ b/tests/unit/torch/quantization/test_dist.py @@ -50,22 +50,6 @@ def test_data_parallel(skip_on_windows): spawn_multiprocess_job(2, _test_data_parallel_helper, backend="gloo") -def test_sync_amax_skips_cpu_tensor_with_nccl(): - quantizer = TensorQuantizer() - quantizer._amax = torch.tensor(1.0) - parallel_group = Mock() - parallel_group.is_initialized.return_value = True - - with ( - patch.object(dist, "get_backend", return_value=dist.Backend.NCCL), - patch.object(dist, "all_reduce") as all_reduce, - pytest.warns(UserWarning, match="tensor is on CPU, but NCCL only supports CUDA tensors"), - ): - quantizer.sync_amax_across_distributed_group(parallel_group) - - all_reduce.assert_not_called() - - def test_sync_amax_propagates_collective_errors(): quantizer = TensorQuantizer() quantizer._amax = torch.tensor(1.0) From c4db72d7ed33440386d7109edcc04f52d5e1bc4a Mon Sep 17 00:00:00 2001 From: realAsma Date: Wed, 8 Jul 2026 18:07:20 +0000 Subject: [PATCH 3/3] Remove redundant distributed amax unit test Signed-off-by: realAsma --- tests/unit/torch/quantization/test_dist.py | 17 ----------------- 1 file changed, 17 deletions(-) diff --git a/tests/unit/torch/quantization/test_dist.py b/tests/unit/torch/quantization/test_dist.py index a98d239f22b..4130b8ce82a 100644 --- a/tests/unit/torch/quantization/test_dist.py +++ b/tests/unit/torch/quantization/test_dist.py @@ -13,9 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from unittest.mock import Mock, patch - -import pytest import torch import torch.distributed as dist from _test_utils.torch.distributed.utils import spawn_multiprocess_job @@ -48,17 +45,3 @@ def forward_loop(model): def test_data_parallel(skip_on_windows): spawn_multiprocess_job(2, _test_data_parallel_helper, backend="gloo") - - -def test_sync_amax_propagates_collective_errors(): - quantizer = TensorQuantizer() - quantizer._amax = torch.tensor(1.0) - parallel_group = Mock() - parallel_group.is_initialized.return_value = True - - with ( - patch.object(dist, "get_backend", return_value=dist.Backend.GLOO), - patch.object(dist, "all_reduce", side_effect=RuntimeError("out of memory")), - pytest.raises(RuntimeError, match="out of memory"), - ): - quantizer.sync_amax_across_distributed_group(parallel_group)