From df1b76ba70e5c396e824941d8ea78151b8d01467 Mon Sep 17 00:00:00 2001 From: Ritesh Patel Date: Tue, 11 Aug 2026 13:36:29 -0700 Subject: [PATCH 1/3] [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 89dd3cb9a51ed9fef58b2cf926b19f0981adfc9d Mon Sep 17 00:00:00 2001 From: Ritesh Patel Date: Tue, 11 Aug 2026 14:14:32 -0700 Subject: [PATCH 2/3] [PyTorch] Destroy the process group created by the MXFP8 empty-shard test. Signed-off-by: Ritesh Patel --- .../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] From 554ca0a780fe5f62918c71e30238500f356ef2c3 Mon Sep 17 00:00:00 2001 From: Ritesh Patel Date: Mon, 17 Aug 2026 10:32:11 -0700 Subject: [PATCH 3/3] [PyTorch] Add a 2-rank MXFP8 empty-shard test where one rank owns the master shard. Signed-off-by: Ritesh Patel --- .../test_cast_master_weights_to_fp8.py | 47 +++++++++++++++++++ .../test_mxfp8_master_weight_empty_shard.py | 22 ++++++--- 2 files changed, 62 insertions(+), 7 deletions(-) diff --git a/tests/pytorch/distributed/test_cast_master_weights_to_fp8.py b/tests/pytorch/distributed/test_cast_master_weights_to_fp8.py index 16c766e495..ab047dc32f 100644 --- a/tests/pytorch/distributed/test_cast_master_weights_to_fp8.py +++ b/tests/pytorch/distributed/test_cast_master_weights_to_fp8.py @@ -37,6 +37,7 @@ Float8BlockwiseQTensor, NVFP4Tensor, is_mxfp8_available, + MXFP8Quantizer, MXFP8Tensor, ) from transformer_engine.pytorch.tensor.utils import ( @@ -914,6 +915,49 @@ def _test_cast_master_weights_to_nvfp4(dp_group, manual_post_all_gather_processi torch.testing.assert_close(loss_nvfp4, loss, atol=0, rtol=0) +def _test_mxfp8_empty_master_shard(dp_group): + """One rank owns the whole master shard, every other rank passes None. + + Wide FSDP sharding pads the parameter bucket, so a tail rank can own an empty shard of + every weight in it. Those ranks still join the amax all-reduce, so the packed amax + buffer they allocate has to match the dtype used by the ranks that do own data. This + used to raise UnboundLocalError before reaching the all-reduce. + """ + rank = dist.get_rank(dp_group) + world_size = dist.get_world_size(dp_group) + + for owner_rank in range(world_size): + # Same seed on every rank, so the model weight is identical to start with. + torch.manual_seed(1234) + torch.cuda.manual_seed(1234) + high_precision = torch.randn(128, 128, dtype=torch.bfloat16, device="cuda") + quantizer = MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3, rowwise=True, columnwise=True) + model_weight = quantizer.make_empty((128, 128), dtype=torch.bfloat16, device="cuda") + quantizer.update_quantized(high_precision, model_weight) + + if rank == owner_rank: + # Optimizers hold fp32 master weights; the cast to the model dtype happens inside. + master_weight = high_precision.to(torch.float32).reshape(-1) + start_offset = 0 + else: + master_weight = None + start_offset = None + + quantize_master_weights([model_weight], [master_weight], [start_offset], dp_group) + + # The amax is reduced with MAX over the group and every rank computes its scales from + # the result, so the scales must agree even on the ranks that contributed nothing. A + # dtype disagreement in that collective shows up here (or hangs the reduction). + for scale_inv in (model_weight._rowwise_scale_inv, model_weight._columnwise_scale_inv): + gathered = [torch.empty_like(scale_inv) for _ in range(world_size)] + dist.all_gather(gathered, scale_inv, group=dp_group) + for other_rank, other in enumerate(gathered): + assert torch.equal(gathered[owner_rank], other), ( + f"MXFP8 scale_inv mismatch between rank {owner_rank} (owns the shard) and " + f"rank {other_rank} (empty shard)" + ) + + def run_parallel_tests() -> None: """Run parallel tests""" @@ -953,6 +997,9 @@ def run_parallel_tests() -> None: for post_ag_processing in manual_post_all_gather_processings: _test_cast_master_weights_to_fp8(quantization, dp_group, post_ag_processing) _test_fsdp_cast_master_weights_to_fp8(quantization, dp_group, post_ag_processing) + if is_mxfp8_available(): + print("starting mxfp8 empty master shard test") + _test_mxfp8_empty_master_shard(dp_group) nvfp4_available, _ = is_nvfp4_available(return_reason=True) if nvfp4_available: print("starting cast master weights to nvfp4 test") 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 f49231e2fb..8e7cfde82d 100644 --- a/tests/pytorch/mxfp8/test_mxfp8_master_weight_empty_shard.py +++ b/tests/pytorch/mxfp8/test_mxfp8_master_weight_empty_shard.py @@ -28,20 +28,28 @@ def single_rank_group(): torch.distributed.destroy_process_group() -def _make_weight(dtype): +# multi_tensor_compute_scale_inv_e8m0 requires a bf16 amax, and the amax buffer takes the +# model weight dtype, so bf16 is the only model dtype this path supports. +MODEL_DTYPE = torch.bfloat16 + + +def _make_weight(): 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) + weight = quantizer.make_empty((128, 128), dtype=MODEL_DTYPE, device="cuda") + quantizer.update_quantized(torch.randn(128, 128, dtype=MODEL_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, single_rank_group, dtype): +def test_empty_master_shard_agrees_with_populated_rank(monkeypatch, single_rank_group): """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. + + Single-GPU counterpart of the 2-rank case in + tests/pytorch/distributed/test_cast_master_weights_to_fp8.py, which checks the same + agreement over a real collective. """ amax_dtypes = [] real_all_reduce = torch.distributed.all_reduce @@ -52,12 +60,12 @@ def spy(tensor, *args, **kwargs): monkeypatch.setattr(torch.distributed, "all_reduce", spy) - populated = _make_weight(dtype) + populated = _make_weight() master = torch.randn(populated.numel(), dtype=torch.float32, device="cuda") 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=single_rank_group) + quantize_master_weights([_make_weight()], [None], [None], group=single_rank_group) assert len(amax_dtypes) == 2 assert amax_dtypes[0] == amax_dtypes[1]