Add PP + dp_replicate deadlock regression tests [skip-build]

Add tests simulating Pipeline Parallelism with dp_replicate > 1 to verify
construct_shard_mesh doesn't deadlock when different PP stages call
dist.new_group for different parameters independently.

- test_muon.py: Dense model test (PP=2, dp_replicate=2, dp_shard=2)
- test_muon_moe.py: MoE test with asymmetric expert/non-expert params
- test_pp.yaml: K8s job spec for running on GPU pod

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

test/test_muon.py

@@ -391,3 +391,124 @@ def test_parallel_muon_uneven_shard(init_dist, uneven_dim):
     set_ns_compile(True)
     logger.info("test_parallel_muon_uneven_shard (dim=%d) PASSED (rank %d)",
                 uneven_dim, rank)
+
+
+def test_pp_dp_replicate_no_deadlock(init_dist):
+    """Regression: PP-like setup where different rank subsets call
+    construct_shard_mesh for different parameters must not deadlock.
+
+    Simulates PP=2 with dp_replicate=2, dp_shard=2. Each PP stage has
+    4 ranks with a (2,2) mesh and [Replicate, Shard(0)] placements.
+    Stages create different numbers of parameters, forcing
+    construct_shard_mesh to be called independently per stage.
+    Without use_local_synchronization=True in dist.new_group(),
+    this would deadlock.
+    """
+    from optimizer.distributed.utils import _ranks_to_dist_cache
+    from optimizer.newton_schulz import set_ns_compile
+
+    rank = dist.get_rank()
+    world_size = dist.get_world_size()
+    assert world_size == 8
+
+    set_ns_compile(False)
+
+    # Clear cache to ensure dist.new_group is actually called
+    _ranks_to_dist_cache.clear()
+
+    # Create full mesh: PP=2, dp_replicate=2, dp_shard=2
+    full_mesh = dist.init_device_mesh(
+        "cuda",
+        (2, 2, 2),
+        mesh_dim_names=("pp", "dp_replicate", "dp_shard"),
+    )
+
+    # Per-stage submesh (shape (2,2), 4 ranks each)
+    stage_mesh = full_mesh["dp_replicate", "dp_shard"]
+    pp_rank = full_mesh.get_local_rank("pp")
+
+    torch.manual_seed(42 + pp_rank)
+
+    # Asymmetric param counts: stage 0 gets 3, stage 1 gets 5
+    num_params = 3 if pp_rank == 0 else 5
+    placements = [Replicate(), Shard(0)]
+
+    muon_params = []
+    muon_names = []
+    full_params_snapshot = []
+    full_grads = []
+
+    for i in range(num_params):
+        full = torch.randn(32, 64, device="cuda")
+        full_params_snapshot.append(full.clone())
+        dt = distribute_tensor(full, stage_mesh, placements)
+        p = torch.nn.Parameter(dt)
+        grad_full = torch.randn(32, 64, device="cuda")
+        full_grads.append(grad_full.clone())
+        p.grad = distribute_tensor(grad_full, stage_mesh, placements)
+        muon_params.append(p)
+        muon_names.append(f"stage{pp_rank}.layer.{i}.weight")
+
+    param_groups = [{
+        "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,
+    }]
+
+    # Must not deadlock
+    optim = Muon(params=param_groups, chunk_size=1, warmup_step=0)
+    optim.step()
+
+    # Second step to verify cached path
+    for p in muon_params:
+        grad_full = torch.randn(32, 64, device="cuda")
+        p.grad = distribute_tensor(grad_full, stage_mesh, placements)
+    optim.step()
+
+    # Correctness: compare against sequential baseline
+    seq_params = []
+    for fp, fg in zip(full_params_snapshot, full_grads):
+        p = torch.nn.Parameter(fp.clone())
+        p.grad = fg.clone()
+        seq_params.append(p)
+
+    param_groups_seq = [{
+        "params":
+        seq_params,
+        "names":
+        [f"stage{pp_rank}.layer.{i}.weight" for i in range(num_params)],
+        "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)
+    optim_seq.step()
+
+    for i in range(num_params):
+        par_full = muon_params[i].data.full_tensor()
+        torch.testing.assert_close(par_full,
+                                   seq_params[i].data,
+                                   atol=0,
+                                   rtol=0)
+
+    set_ns_compile(True)
+    logger.info(
+        "test_pp_dp_replicate_no_deadlock PASSED (rank %d, pp_rank %d)", rank,
+        pp_rank)

test/test_muon_moe.py

@@ -402,3 +402,133 @@ def test_parallel_muon_moe_uneven_shard(init_dist, uneven_dim):
     logger.info(
         "test_parallel_muon_moe_uneven_shard (dim=%d) PASSED (rank %d)",
         uneven_dim, rank)
+
+
+def test_pp_dp_replicate_moe_no_deadlock(init_dist):
+    """Regression: PP-like MoE setup where different stages have different
+    parameter types must not deadlock in construct_shard_mesh.
+
+    Simulates PP=2 with dp_replicate=2, dp_shard=2. Stage 0 has only
+    non-expert 2D DTensor params; stage 1 has non-expert 2D DTensor params
+    plus 3D expert plain-tensor params. This mirrors real PP+MoE where
+    expert layers exist only in certain stages.
+    """
+    from optimizer.distributed.utils import _ranks_to_dist_cache
+    from optimizer.newton_schulz import set_ns_compile
+
+    rank = dist.get_rank()
+    world_size = dist.get_world_size()
+    assert world_size == 8
+
+    set_ns_compile(False)
+
+    # Clear cache to ensure dist.new_group is actually called
+    _ranks_to_dist_cache.clear()
+
+    # Create full mesh: PP=2, dp_replicate=2, dp_shard=2
+    full_mesh = dist.init_device_mesh(
+        "cuda",
+        (2, 2, 2),
+        mesh_dim_names=("pp", "dp_replicate", "dp_shard"),
+    )
+
+    stage_mesh = full_mesh["dp_replicate", "dp_shard"]
+    pp_rank = full_mesh.get_local_rank("pp")
+
+    torch.manual_seed(42 + pp_rank)
+
+    placements = [Replicate(), Shard(0)]
+    num_experts = 4
+
+    muon_params = []
+    muon_names = []
+    full_params_snapshot = []
+    full_grads = []
+
+    # Non-expert 2D DTensor params (both stages, different counts)
+    num_dense = 2 if pp_rank == 0 else 3
+    for i in range(num_dense):
+        full = torch.randn(32, 64, device="cuda")
+        full_params_snapshot.append(full.clone())
+        dt = distribute_tensor(full, stage_mesh, placements)
+        p = torch.nn.Parameter(dt)
+        g = torch.randn(32, 64, device="cuda")
+        full_grads.append(g.clone())
+        p.grad = distribute_tensor(g, stage_mesh, placements)
+        muon_params.append(p)
+        muon_names.append(f"stage{pp_rank}.layers.{i}.weight")
+
+    # Stage 1 only: 3D expert plain-tensor params
+    if pp_rank == 1:
+        full = torch.randn(num_experts, 32, 64, device="cuda")
+        full_params_snapshot.append(full.clone())
+        p = torch.nn.Parameter(full)
+        g = torch.randn(num_experts, 32, 64, device="cuda")
+        full_grads.append(g.clone())
+        p.grad = g
+        muon_params.append(p)
+        muon_names.append(
+            f"stage{pp_rank}.layers.{num_dense}.experts.w1.weight")
+
+    param_groups = [{
+        "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,
+    }]
+
+    # Must not deadlock
+    optim = Muon(params=param_groups,
+                 chunk_size=1,
+                 warmup_step=0,
+                 expert_keys=["experts"])
+    optim.step()
+
+    # Second step to verify cached path
+    for p in muon_params:
+        if isinstance(p.data, DTensor):
+            g = torch.randn(32, 64, device="cuda")
+            p.grad = distribute_tensor(g, stage_mesh, placements)
+        else:
+            p.grad = torch.randn_like(p.data)
+    optim.step()
+
+    # Correctness: compare against sequential baseline
+    seq_params = []
+    for fp in full_params_snapshot:
+        seq_params.append(torch.nn.Parameter(fp.clone()))
+    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()
+
+    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_pp_dp_replicate_moe_no_deadlock PASSED (rank %d, pp_rank %d)",
+        rank, pp_rank)