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..f49231e2fb --- /dev/null +++ b/tests/pytorch/mxfp8/test_mxfp8_master_weight_empty_shard.py @@ -0,0 +1,63 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +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) + + +@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) + torch.distributed.init_process_group( + backend="nccl", store=torch.distributed.HashStore(), rank=0, world_size=1 + ) + try: + yield torch.distributed.GroupMember.WORLD + finally: + if created: + torch.distributed.destroy_process_group() + + +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, 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. + """ + 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=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) + + 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