[PyTorch] Fix MXFP8 master weight cast when a rank owns no shard.

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]

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