From df1b76ba70e5c396e824941d8ea78151b8d01467 Mon Sep 17 00:00:00 2001 From: Ritesh Patel Date: Tue, 11 Aug 2026 13:36:29 -0700 Subject: [PATCH 1/2] [PyTorch] Fix MXFP8 master weight cast when a rank owns no shard. Signed-off-by: Ritesh Patel --- .../test_mxfp8_master_weight_empty_shard.py | 66 +++++++++++++++++++ transformer_engine/pytorch/tensor/utils.py | 4 ++ 2 files changed, 70 insertions(+) create mode 100644 tests/pytorch/mxfp8/test_mxfp8_master_weight_empty_shard.py diff --git a/tests/pytorch/mxfp8/test_mxfp8_master_weight_empty_shard.py b/tests/pytorch/mxfp8/test_mxfp8_master_weight_empty_shard.py new file mode 100644 index 0000000000..748433f14e --- /dev/null +++ b/tests/pytorch/mxfp8/test_mxfp8_master_weight_empty_shard.py @@ -0,0 +1,66 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +import pathlib +import tempfile + +import pytest +import torch + +import transformer_engine.pytorch as te +from transformer_engine.pytorch import MXFP8Quantizer +from transformer_engine.pytorch.tensor.utils import quantize_master_weights + +recipe_available, reason_for_no_recipe = te.is_mxfp8_available(return_reason=True) + + +def _single_rank_group(): + if not torch.distributed.is_initialized(): + torch.cuda.set_device(0) + with tempfile.NamedTemporaryFile(delete=False) as f: + rendezvous_file = pathlib.Path(f.name) + torch.distributed.init_process_group( + backend="nccl", + init_method=rendezvous_file.resolve().as_uri(), + rank=0, + world_size=1, + ) + return torch.distributed.GroupMember.WORLD + + +def _make_weight(dtype): + quantizer = MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3, rowwise=True, columnwise=True) + weight = quantizer.make_empty((128, 128), dtype=dtype, device="cuda") + quantizer.update_quantized(torch.randn(128, 128, dtype=dtype, device="cuda"), weight) + return weight + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +def test_empty_master_shard_agrees_with_populated_rank(monkeypatch, dtype): + """A rank owning no shard must reduce the same amax dtype as one that owns data. + + Wide FSDP sharding pads the parameter bucket, so the tail ranks can end up with an + empty shard of every weight. Those ranks still join the amax all-reduce. + """ + group = _single_rank_group() + + amax_dtypes = [] + real_all_reduce = torch.distributed.all_reduce + + def spy(tensor, *args, **kwargs): + amax_dtypes.append(tensor.dtype) + return real_all_reduce(tensor, *args, **kwargs) + + monkeypatch.setattr(torch.distributed, "all_reduce", spy) + + populated = _make_weight(dtype) + master = torch.randn(populated.numel(), dtype=torch.float32, device="cuda") + quantize_master_weights([populated], [master], [0], group=group) + + # Used to raise UnboundLocalError instead of reaching the all-reduce. + quantize_master_weights([_make_weight(dtype)], [None], [None], group=group) + + assert len(amax_dtypes) == 2 + assert amax_dtypes[0] == amax_dtypes[1] diff --git a/transformer_engine/pytorch/tensor/utils.py b/transformer_engine/pytorch/tensor/utils.py index e35d57b363..cef45c0223 100644 --- a/transformer_engine/pytorch/tensor/utils.py +++ b/transformer_engine/pytorch/tensor/utils.py @@ -1028,6 +1028,10 @@ def _cast_master_weights_to_fp8_mxfp8_scaling( # Parameter attributes device = params[0][0].device + # Every shard can be empty on a rank. Master weights are cast to the model dtype in + # quantize_master_weights, so use that as the fallback: the amax buffer below is + # all-reduced and its dtype has to agree with the ranks that do own a shard. + master_weight_dtype = params[0][0].dtype for _, master_weight, _, _ in params: if master_weight is not None: master_weight_dtype = master_weight.dtype From c1a9445ac9b91fd968682033511c75917e433c06 Mon Sep 17 00:00:00 2001 From: Ritesh Patel Date: Tue, 11 Aug 2026 14:14:32 -0700 Subject: [PATCH 2/2] [PyTorch] Destroy the process group created by the MXFP8 empty-shard test. --- .../test_mxfp8_master_weight_empty_shard.py | 31 +++++++++---------- 1 file changed, 14 insertions(+), 17 deletions(-) diff --git a/tests/pytorch/mxfp8/test_mxfp8_master_weight_empty_shard.py b/tests/pytorch/mxfp8/test_mxfp8_master_weight_empty_shard.py index 748433f14e..f49231e2fb 100644 --- a/tests/pytorch/mxfp8/test_mxfp8_master_weight_empty_shard.py +++ b/tests/pytorch/mxfp8/test_mxfp8_master_weight_empty_shard.py @@ -2,9 +2,6 @@ # # See LICENSE for license information. -import pathlib -import tempfile - import pytest import torch @@ -15,18 +12,20 @@ recipe_available, reason_for_no_recipe = te.is_mxfp8_available(return_reason=True) -def _single_rank_group(): - if not torch.distributed.is_initialized(): +@pytest.fixture +def single_rank_group(): + # Only tear down a group this fixture owns; another test may have set one up. + created = not torch.distributed.is_initialized() + if created: torch.cuda.set_device(0) - with tempfile.NamedTemporaryFile(delete=False) as f: - rendezvous_file = pathlib.Path(f.name) torch.distributed.init_process_group( - backend="nccl", - init_method=rendezvous_file.resolve().as_uri(), - rank=0, - world_size=1, + backend="nccl", store=torch.distributed.HashStore(), rank=0, world_size=1 ) - return torch.distributed.GroupMember.WORLD + try: + yield torch.distributed.GroupMember.WORLD + finally: + if created: + torch.distributed.destroy_process_group() def _make_weight(dtype): @@ -38,14 +37,12 @@ def _make_weight(dtype): @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) -def test_empty_master_shard_agrees_with_populated_rank(monkeypatch, dtype): +def test_empty_master_shard_agrees_with_populated_rank(monkeypatch, single_rank_group, dtype): """A rank owning no shard must reduce the same amax dtype as one that owns data. Wide FSDP sharding pads the parameter bucket, so the tail ranks can end up with an empty shard of every weight. Those ranks still join the amax all-reduce. """ - group = _single_rank_group() - amax_dtypes = [] real_all_reduce = torch.distributed.all_reduce @@ -57,10 +54,10 @@ def spy(tensor, *args, **kwargs): populated = _make_weight(dtype) master = torch.randn(populated.numel(), dtype=torch.float32, device="cuda") - quantize_master_weights([populated], [master], [0], group=group) + quantize_master_weights([populated], [master], [0], group=single_rank_group) # Used to raise UnboundLocalError instead of reaching the all-reduce. - quantize_master_weights([_make_weight(dtype)], [None], [None], group=group) + quantize_master_weights([_make_weight(dtype)], [None], [None], group=single_rank_group) assert len(amax_dtypes) == 2 assert amax_dtypes[0] == amax_dtypes[1]