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 new file mode 100644 index 0000000000..8e7cfde82d --- /dev/null +++ b/tests/pytorch/mxfp8/test_mxfp8_master_weight_empty_shard.py @@ -0,0 +1,71 @@ +# 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() + + +# 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=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) +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 + + 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() + 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()], [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