[DistMuon] Support tensor-parallel storage layouts

tests/unit_tests/gpu/flex_shard/test_dist_muon.py

@@ -403,5 +403,130 @@ def capture_compute(_compute_layout, compute):
         self.assertTrue(compute_ready_layout.storage_is_compute_ready)
 
 
+@unittest.skipUnless(torch.cuda.device_count() >= 4, "requires four CUDA devices")
+class TestDistMuonTensorParallel(DTensorTestBase):
+    @property
+    def world_size(self):
+        return 4
+
+    @property
+    def device_type(self):
+        return "cuda"
+
+    @with_comms
+    def test_dp_tp_storage_matches_unsharded_update(self):
+        mesh = init_device_mesh(
+            self.device_type,
+            (2, 2),
+            mesh_dim_names=("dp_shard", "tp"),
+        )
+        device = torch.device(self.device_type, self.rank)
+        lr = 0.03
+        weight_decay = 0.2
+        matrix_rows = 2
+
+        values = {
+            "layers.0.per_head": (
+                torch.arange(24, device=device).reshape(8, 3).float().div_(13)
+            ),
+            "layers.0.whole": (
+                torch.arange(16, device=device).reshape(4, 4).float().div_(11)
+            ),
+            "layers.0.experts": (
+                torch.arange(24, device=device).reshape(4, 3, 2).float().div_(7)
+            ),
+        }
+        placements = {
+            "layers.0.per_head": (
+                _StridedShard(0, split_factor=mesh["tp"].size()),
+                Shard(0),
+            ),
+            "layers.0.whole": (Shard(0), Shard(1)),
+            "layers.0.experts": (Shard(0), Shard(1)),
+        }
+        parameters = {
+            fqn: torch.nn.Parameter(
+                distribute_tensor(value.clone(), mesh, placements[fqn])
+            )
+            for fqn, value in values.items()
+        }
+        gradients = {
+            fqn: value.clone().mul_(0.37).add_(0.2).sin_()
+            for fqn, value in values.items()
+        }
+        for fqn, parameter in parameters.items():
+            parameter.grad = distribute_tensor(
+                gradients[fqn].clone(),
+                mesh,
+                placements[fqn],
+            )
+
+        optimizer = build_dist_muon(
+            [
+                {
+                    "params": list(parameters.values()),
+                    "param_names": list(parameters),
+                }
+            ],
+            compute_sharding_by_fqn={
+                "layers.0.per_head": ComputeLayout(
+                    shardings_by_mesh_axis={
+                        "dp_shard": BlockShard(0, matrix_rows),
+                        "tp": BlockShard(0, matrix_rows),
+                    }
+                ),
+                "layers.0.whole": ComputeLayout(
+                    shardings_by_mesh_axis={
+                        "dp_shard": Owned(),
+                        "tp": Owned(),
+                    }
+                ),
+                "layers.0.experts": ComputeLayout(
+                    shardings_by_mesh_axis={
+                        "dp_shard": Shard(0),
+                        "tp": Shard(0),
+                    }
+                ),
+            },
+            bucket_configs=[
+                BucketConfig(patterns=("layers.0.per_head", "layers.0.whole")),
+                BucketConfig(patterns=("layers.0.experts",)),
+            ],
+            lr=lr,
+            weight_decay=weight_decay,
+            momentum=0.0,
+            nesterov=False,
+            ns_steps=2,
+        )
+
+        def make_direction(_compute_layout, compute):
+            compute.mul_(0.5).add_(0.25)
+
+        with mock.patch.object(
+            optimizer,
+            "_compute_update",
+            side_effect=make_direction,
+        ):
+            optimizer.step()
+
+        compute_shapes = {
+            "layers.0.per_head": (matrix_rows, values["layers.0.per_head"].shape[1]),
+            "layers.0.whole": values["layers.0.whole"].shape,
+            "layers.0.experts": values["layers.0.experts"].shape[1:],
+        }
+        for fqn, parameter in parameters.items():
+            expected = values[fqn].mul(1 - lr * weight_decay)
+            expected.add_(
+                gradients[fqn].mul(0.5).add(0.25),
+                alpha=-_adjust_muon_learning_rate(lr, None, compute_shapes[fqn]),
+            )
+            torch.testing.assert_close(
+                parameter.full_tensor(),
+                expected,
+                rtol=0,
+                atol=0,
+            )
+
+
 if __name__ == "__main__":
     unittest.main()

tests/unit_tests/gpu/test_qk_clip.py

@@ -16,8 +16,9 @@
 import torch
 import torch.nn as nn
 from torch.distributed.device_mesh import DeviceMesh, init_device_mesh
-from torch.distributed.tensor import distribute_tensor, Shard
+from torch.distributed.tensor import distribute_tensor, DTensor, Shard
 from torch.distributed.tensor.debug import CommDebugMode
+from torch.distributed.tensor.placement_types import _StridedShard
 from torch.nn.attention.flex_attention import create_block_mask
 from torch.testing._internal.distributed._tensor.common_dtensor import (
     DTensorTestBase,
@@ -285,3 +286,111 @@ def weight_module(num_rows: int) -> nn.Module:
         # One packed MAX all-reduce covers every layer and head; the 8 weights
         # this model clips must add nothing on top of it.
         self.assertEqual(total, 1, f"expected one collective, got {collectives}")
+
+
+@pytest.mark.multi_gpu
+@unittest.skipUnless(torch.cuda.device_count() >= 4, "requires four CUDA devices")
+class QKClipTensorParallelTest(DTensorTestBase):
+    @property
+    def world_size(self) -> int:
+        return 4
+
+    @property
+    def device_type(self) -> str:
+        return "cuda"
+
+    @with_comms
+    def test_dp_tp_head_scales_match_flex_shard_storage(self) -> None:
+        mesh = init_device_mesh(
+            self.device_type,
+            (2, 2),
+            mesh_dim_names=("dp_shard", "tp"),
+        )
+        device = torch.device(self.device_type, self.rank)
+        num_heads = 4
+        qk_nope_head_dim = 2
+        qk_rope_head_dim = 1
+        v_head_dim = 2
+        in_features = 2
+        storage_placements = (
+            _StridedShard(0, split_factor=mesh["tp"].size()),
+            Shard(0),
+        )
+
+        def weight_module(num_rows: int) -> nn.Module:
+            module = nn.Module()
+            module.register_parameter(
+                "weight",
+                nn.Parameter(
+                    distribute_tensor(
+                        torch.ones(num_rows, in_features, device=device),
+                        mesh,
+                        storage_placements,
+                    )
+                ),
+            )
+            return module
+
+        attention = Attention.__new__(Attention)
+        nn.Module.__init__(attention)
+        attention.q_lora_rank = 1
+        attention.qk_nope_head_dim = qk_nope_head_dim
+        attention.qk_rope_head_dim = qk_rope_head_dim
+        attention.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
+        attention.v_head_dim = v_head_dim
+        attention.wq_b = weight_module(num_heads * attention.qk_head_dim)
+        attention.wkv_b = weight_module(num_heads * (qk_nope_head_dim + v_head_dim))
+        attention.inner_attention = QKClipFlexAttention.Config().build()
+
+        coordinate = mesh.get_coordinate()
+        assert coordinate is not None
+        dp_rank, tp_rank = coordinate
+        maxima_by_coordinate = (
+            ((50.0, 400.0), (25.0, 800.0)),
+            ((200.0, 50.0), (500.0, 40.0)),
+        )
+        attention.inner_attention.max_attention_logits_N.append(
+            torch.tensor(maxima_by_coordinate[dp_rank][tp_rank], device=device)
+        )
+        model = nn.Module()
+        model.add_module("attention", attention)
+
+        with CommDebugMode() as comm_mode:
+            qk_clip([model], reduction_mesh=mesh["dp_shard"])
+
+        collectives = {
+            str(op): count for op, count in comm_mode.get_comm_counts().items() if count
+        }
+        self.assertEqual(
+            sum(collectives.values()),
+            1,
+            f"expected one collective, got {collectives}",
+        )
+
+        scales_N = torch.tensor((0.5, 0.25, 0.2, 0.125), device=device)
+        q_expected_NDI = torch.ones(
+            num_heads,
+            attention.qk_head_dim,
+            in_features,
+            device=device,
+        )
+        q_expected_NDI[:, :qk_nope_head_dim].mul_(scales_N.view(-1, 1, 1).pow(0.5))
+        q_expected_NDI[:, qk_nope_head_dim:].mul_(scales_N.view(-1, 1, 1))
+        q_weight = attention.wq_b.weight
+        assert isinstance(q_weight, DTensor)
+        torch.testing.assert_close(q_weight.full_tensor(), q_expected_NDI.flatten(0, 1))
+
+        kv_expected_NDI = torch.ones(
+            num_heads,
+            qk_nope_head_dim + v_head_dim,
+            in_features,
+            device=device,
+        )
+        kv_expected_NDI[:, :qk_nope_head_dim].mul_(scales_N.view(-1, 1, 1).pow(0.5))
+        kv_weight = attention.wkv_b.weight
+        assert isinstance(kv_weight, DTensor)
+        torch.testing.assert_close(
+            kv_weight.full_tensor(),
+            kv_expected_NDI.flatten(0, 1),
+        )
+        self.assertFalse(attention.inner_attention.max_attention_logits_N)

torchtitan/distributed/flex_shard/README.md

@@ -33,9 +33,13 @@ The public API is exported from `torchtitan.distributed.flex_shard`:
   transitions.
 
 Storage placements describe persistent ownership only; they do not define
-Muon matrix boundaries. Flat matrix-batch compute supports `BlockShard` on at
-most one non-unit mesh axis. Storage on that axis may use exact `Shard(0)` or
-`Replicate`; every other non-unit storage mesh axis must be replicated.
+Muon matrix boundaries. Flat matrix-batch compute may use `BlockShard` on one
+or more non-unit mesh axes. FlexShard combines those axes into one optimizer
+transport group, routes complete matrices across the group, and restores the
+original DTensor storage layout after the update. This supports combined DP
+and TP storage, including the strided shards produced when TP and FSDP shard
+the same tensor dimension. Non-transport mesh axes retain their storage
+placement.
 
 Several mesh axes may shard the same tensor dimension. By default they apply
 in storage-mesh order; `shard_order_by_tensor_dim` states a different order,