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

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")

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]

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