[mxfp8] Handle unevenly sharded expert weights in FSDP compute weights

tests/unit_tests/gpu/test_mxfp8_linear.py

@@ -413,3 +413,78 @@ def test_mxfp8_input_activation_formats_for_backward_match():
         rtol=0,
         atol=0,
     )
+
+
+class _StubMesh:
+    """Stands in for a DeviceMesh: fsdp_pre_all_gather only reads ``size()``."""
+
+    def __init__(self, size: int) -> None:
+        self._size = size
+
+    def size(self) -> int:
+        return self._size
+
+
+class _StubMixedPrecisionPolicy:
+    param_dtype = torch.bfloat16
+
+
+def test_fsdp_pre_all_gather_pads_an_uneven_shard():
+    """A dim-0 size that does not divide the mesh leaves the last rank short.
+
+    All-gather needs every rank to contribute the same number of elements, so
+    FSDP's contract is that the hook returns the *padded* shard and passes the
+    logical size through as metadata. Driven directly rather than through FSDP
+    because a dense MXFP8 weight cannot be unevenly sharded on two ranks: the
+    kernels require out_features divisible by 32, which always divides 2.
+    """
+    # Logical (96, 128) over five ranks: ceil(96 / 5) == 20 rows each, so the
+    # last rank holds only 16 and pads up to 20.
+    shard_NK = torch.randn(16, 128, device="cuda", dtype=torch.bfloat16)
+    sharded_weight = _LinearShardedTensorWithMXFP8Compute(shard_NK)
+
+    (comm_NK,), metadata = sharded_weight.fsdp_pre_all_gather(
+        _StubMesh(5),
+        torch.Size([96, 128]),
+        None,
+        None,
+        _StubMixedPrecisionPolicy(),
+    )
+
+    assert comm_NK.shape == (20, 128)
+    assert torch.equal(comm_NK[:16], shard_NK)
+    assert torch.count_nonzero(comm_NK[16:]) == 0
+    # The logical size rides along so post-all-gather can drop the padding.
+    assert tuple(metadata) == (96, 128)
+
+
+def test_fsdp_post_all_gather_drops_the_padding():
+    """The gathered buffer includes padding; quantization must not see it."""
+    sharded_weight = _LinearShardedTensorWithMXFP8Compute(
+        torch.randn(16, 128, device="cuda", dtype=torch.bfloat16)
+    )
+    # Five ranks contributing 20 padded rows each.
+    gathered_NK = torch.randn(100, 128, device="cuda", dtype=torch.bfloat16)
+
+    unsharded_tensor, managed_tensors = sharded_weight.fsdp_post_all_gather(
+        (gathered_NK,), torch.Size([96, 128]), torch.bfloat16
+    )
+
+    assert isinstance(unsharded_tensor, _UnshardedFSDPTensor)
+    assert unsharded_tensor.shape == (96, 128)
+    assert len(managed_tensors) == 3
+
+
+def test_fsdp_pre_all_gather_rejects_a_non_zero_shard_dim():
+    """Only dim 0 is supported; the all-gather concatenates along it."""
+    sharded_weight = _LinearShardedTensorWithMXFP8Compute(
+        torch.randn(96, 64, device="cuda", dtype=torch.bfloat16)
+    )
+    with pytest.raises(NotImplementedError, match="sharding dimension 0 only"):
+        sharded_weight.fsdp_pre_all_gather(
+            _StubMesh(2),
+            torch.Size([96, 128]),
+            None,
+            None,
+            _StubMixedPrecisionPolicy(),
+        )

torchtitan/components/quantization/_fsdp_tensor.py

@@ -133,6 +133,7 @@
 
 from __future__ import annotations
 
+import math
 from dataclasses import fields, is_dataclass
 from typing import Any
 
@@ -326,22 +327,68 @@ def fsdp_should_release_all_gather_outputs_after_post_all_gather(self) -> bool:
         return True
 
     def fsdp_pre_all_gather(self, mesh, outer_size, outer_stride, module, mp_policy):
-        """Return the high-precision communication tensor."""
+        """Return the high-precision communication tensor and the logical size.
+
+        All-gather needs every rank to contribute the same number of elements,
+        so an expert count that does not divide the mesh size leaves the last
+        rank short. FSDP's contract is that this returns the *padded* shard;
+        the logical size travels in the metadata so ``fsdp_post_all_gather``
+        can drop the padding before quantizing it as if it were weight.
+        """
         del outer_stride, module
-        if outer_size[0] % mesh.size() != 0:
-            raise ValueError(
-                "FSDP unsharded tensors require dimension 0 to be evenly divisible "
-                "by the FSDP shard mesh size"
+        # FSDP hands the hook no shard dimension, but the local shard differs
+        # from the logical size exactly along it. The default, non-extension
+        # path has it directly as ``fsdp_placement.dim`` and rejects the same
+        # case there:
+        # https://github.com/pytorch/pytorch/blob/c7da99c173f2b67905ee798576a644b6b32cbfee/torch/distributed/fsdp/_fully_shard/_fsdp_param.py#L323-L331
+        sharded_dims = [
+            dim
+            for dim, (local, logical) in enumerate(
+                zip(self._tensor.shape, outer_size, strict=True)
+            )
+            if local != logical
+        ]
+        if len(sharded_dims) > 1:
+            raise RuntimeError(
+                f"FSDP sharded more than one dimension: local "
+                f"{tuple(self._tensor.shape)} against logical {tuple(outer_size)}"
+            )
+        if sharded_dims and sharded_dims[0] != 0:
+            raise NotImplementedError(
+                "FSDP unsharded tensors support sharding dimension 0 only, but "
+                f"this parameter of shape {tuple(outer_size)} is sharded on "
+                f"dimension {sharded_dims[0]}. TorchTitan selects Shard(1) for "
+                "grouped experts when the FSDP degree exceeds the expert "
+                "count, so either lower the degree or raise the expert count."
             )
         dtype = mp_policy.param_dtype or self._tensor.dtype
-        return (self._tensor.to(dtype),), None
+        source = self._tensor.to(dtype)
+        # Pad to what FSDP calls ``padded_sharded_param_size``. The default
+        # path pre-pads to ``chunks[0].size()``, and torch.chunk puts the
+        # remainder in the earlier chunks, so that equals ceil(dim0 / world):
+        # https://github.com/pytorch/pytorch/blob/c7da99c173f2b67905ee798576a644b6b32cbfee/torch/distributed/fsdp/_fully_shard/_fsdp_param.py#L332-L345
+        # An extension must return exactly that size; only the short ranks
+        # would trip the check, so the rest hang in the all-gather instead:
+        # https://github.com/pytorch/pytorch/blob/c7da99c173f2b67905ee798576a644b6b32cbfee/torch/distributed/fsdp/_fully_shard/_fsdp_param.py#L1143-L1158
+        padded_rows = math.ceil(outer_size[0] / mesh.size())
+        if source.size(0) != padded_rows:
+            padded = source.new_zeros((padded_rows, *source.shape[1:]))
+            padded[: source.size(0)] = source
+            source = padded
+        return (source,), outer_size
 
     def fsdp_post_all_gather(
         self, all_gather_outputs, metadata, param_dtype, *, out=None
     ):
         """Create or refill the unsharded tensor operands after all-gather."""
-        del metadata, param_dtype
+        del param_dtype
         (gathered_weight,) = all_gather_outputs
+        # ``metadata`` is the logical size returned by fsdp_pre_all_gather. An
+        # unevenly sharded parameter gathers padding rows past it, which must
+        # not reach the quantizer: they would occupy real scale tiles and, for
+        # a grouped weight, appear as extra experts.
+        if metadata is not None and gathered_weight.size(0) != metadata[0]:
+            gathered_weight = gathered_weight.narrow(0, 0, metadata[0])
 
         # On the first unshard, FSDP has no unsharded-tensor container or managed
         # tensors yet. Build both and return them to FSDP. With RAF=False, FSDP