Add MoE uneven shard test with mixed expert and non-expert params [skip-build]

Test parallel Muon with uneven dims (33, 19) mixing 2D DTensor params
(parallel pipeline) and 3D expert plain tensors (batched NS path).

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

test/test_muon_moe.py

@@ -7,7 +7,7 @@ import pytest
 import torch
 import torch.distributed as dist
 from optimizer.muon import Muon, get_default_muon_param_groups
-from torch.distributed.tensor import DTensor, Replicate
+from torch.distributed.tensor import DTensor, Replicate, Shard, distribute_tensor
 from torch.profiler import ProfilerActivity, profile
 
 from .utils import ParallelDims, assert_params_equal, parallelize_llama4
@@ -287,3 +287,117 @@ def test_parallel_muon_moe_few_experts(
     else:
         assert_params_equal(parallelized_model,
                             sequential_moe_result_few_experts)
+
+
+# ---------------------------------------------------------------------------
+# Uneven shard test: mixed expert (3D plain) + non-expert (2D DTensor)
+# with dimensions not evenly divisible by shard count.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("uneven_dim", [
+    pytest.param(33, id="33"),
+    pytest.param(19, id="19"),
+])
+def test_parallel_muon_moe_uneven_shard(init_dist, uneven_dim):
+    """Test MoE parallel Muon with uneven shard dimensions.
+
+    Mixes non-expert 2D DTensor params (uneven FSDP sharding, parallel
+    pipeline path) with expert 3D plain-tensor params (batched NS path).
+    Verifies the combination produces correct results vs sequential baseline.
+    """
+    from optimizer.newton_schulz import set_ns_compile
+
+    rank = dist.get_rank()
+    world_size = dist.get_world_size()
+    mesh = dist.init_device_mesh("cuda", (world_size, ),
+                                 mesh_dim_names=("dp", ))
+
+    set_ns_compile(False)
+    torch.manual_seed(42)
+
+    other_dim = 64
+    num_experts = 4
+
+    muon_params = []
+    muon_names = []
+    full_params = []
+    full_grads = []
+
+    # 2D non-expert params with uneven dims → parallel pipeline
+    for i in range(2):
+        full = torch.randn(uneven_dim, other_dim, device="cuda")
+        full_params.append(full.clone())
+        dt = distribute_tensor(full, mesh, [Shard(0)])
+        p = torch.nn.Parameter(dt)
+        g = torch.randn(uneven_dim, other_dim, device="cuda")
+        full_grads.append(g.clone())
+        p.grad = distribute_tensor(g, mesh, [Shard(0)])
+        muon_params.append(p)
+        muon_names.append(f"layers.{i}.weight")
+
+    # 3D expert params (plain tensors) → batched NS path
+    full = torch.randn(num_experts, uneven_dim, other_dim, device="cuda")
+    full_params.append(full.clone())
+    p = torch.nn.Parameter(full)
+    g = torch.randn(num_experts, uneven_dim, other_dim, device="cuda")
+    full_grads.append(g.clone())
+    p.grad = g
+    muon_params.append(p)
+    muon_names.append("layers.2.experts.w1.weight")
+
+    # --- Parallel path ---
+    param_groups_par = [{
+        "params": muon_params,
+        "names": muon_names,
+        "use_muon": True,
+        "lr": 0.02,
+        "weight_decay": 0.01,
+        "momentum": 0.95,
+        "nesterov": True,
+        "ns_steps": 5,
+        "none_grad": False,
+    }]
+    optim_par = Muon(params=param_groups_par,
+                     chunk_size=1,
+                     warmup_step=0,
+                     expert_keys=["experts"])
+    optim_par.step()
+
+    # --- Sequential baseline ---
+    seq_params = []
+    for fp in full_params:
+        p = torch.nn.Parameter(fp.clone())
+        seq_params.append(p)
+
+    for p, g in zip(seq_params, full_grads):
+        p.grad = g.clone()
+
+    param_groups_seq = [{
+        "params": seq_params,
+        "names": list(muon_names),
+        "use_muon": True,
+        "lr": 0.02,
+        "weight_decay": 0.01,
+        "momentum": 0.95,
+        "nesterov": True,
+        "ns_steps": 5,
+        "none_grad": False,
+    }]
+    optim_seq = Muon(params=param_groups_seq, expert_keys=["experts"])
+    optim_seq.step()
+
+    # --- Compare ---
+    for i in range(len(muon_params)):
+        par_data = muon_params[i].data
+        if isinstance(par_data, DTensor):
+            par_data = par_data.full_tensor()
+        torch.testing.assert_close(par_data,
+                                   seq_params[i].data,
+                                   atol=0,
+                                   rtol=0)
+
+    set_ns_compile(True)
+    logger.info(
+        "test_parallel_muon_moe_uneven_shard (dim=%d) PASSED (rank %d)",
+        uneven_dim, rank)