fix: distribute routing replay padding across experts

slime/backends/megatron_utils/cp_utils.py

@@ -236,19 +236,11 @@ def slice_log_prob_with_cp(
         return torch.cat([chunk_1, chunk_2], dim=0)
 
 
-def _pad_routed_experts(experts: torch.Tensor, pad: int, num_experts: int) -> torch.Tensor:
+def _pad_routed_experts(experts: torch.Tensor, pad: int) -> torch.Tensor:
+    """Mark padding until the final context-parallel layout is known."""
     if pad == 0:
         return experts
-    _, num_layers, topk = experts.shape
-    pad_experts = (
-        torch.arange(
-            pad * num_layers * topk,
-            device=experts.device,
-            dtype=experts.dtype,
-        ).reshape((pad, num_layers, topk))
-        % num_experts
-    )
-    return torch.cat([experts, pad_experts], dim=0)
+    return F.pad(experts, (0, 0, 0, 0, 0, pad), value=-1)
 
 
 def prepare_routed_experts_for_routing_replay(
@@ -265,7 +257,7 @@ def prepare_routed_experts_for_routing_replay(
     for experts, token_ids in zip(rollout_routed_experts, tokens, strict=False):
         assert experts.shape[0] == token_ids.shape[0] - 1, f"{experts.shape}, {token_ids.shape}"
 
-    padded_experts = [_pad_routed_experts(experts, 1, num_experts) for experts in rollout_routed_experts]
+    padded_experts = [_pad_routed_experts(experts, 1) for experts in rollout_routed_experts]
     pad_size = mpu.get_tensor_model_parallel_world_size() * data_pad_size_multiplier
 
     if allgather_cp:
@@ -274,16 +266,35 @@ def prepare_routed_experts_for_routing_replay(
         cp_rank = mpu.get_context_parallel_rank()
         global_pad_size = cp_size * pad_size
         pad = (global_pad_size - routed_experts.size(0) % global_pad_size) % global_pad_size
-        routed_experts = _pad_routed_experts(routed_experts, pad, num_experts)
+        routed_experts = _pad_routed_experts(routed_experts, pad)
         routed_experts = routed_experts.chunk(cp_size, dim=0)[cp_rank]
     else:
-        routed_experts = [
-            slice_with_cp(experts, lambda x, pad: _pad_routed_experts(x, pad, num_experts))
-            for experts in padded_experts
-        ]
+        routed_experts = [slice_with_cp(experts, _pad_routed_experts) for experts in padded_experts]
         routed_experts = torch.cat(routed_experts, dim=0)
         pad = (pad_size - routed_experts.size(0) % pad_size) % pad_size
-        routed_experts = _pad_routed_experts(routed_experts, pad, num_experts)
+        routed_experts = _pad_routed_experts(routed_experts, pad)
+
+    # Number only the pad rows that survive CP slicing; SP partitions these
+    # assignments without changing them. Padding marks the whole row, so one
+    # element is sufficient to identify it.
+    padding_rows = (routed_experts[:, 0, 0] == -1).nonzero().flatten()
+    if padding_rows.numel():
+        _, num_layers, topk = routed_experts.shape
+        ep_size = mpu.get_expert_model_parallel_world_size()
+        assert num_experts % ep_size == 0, f"{num_experts=}, {ep_size=}"
+        assert topk <= num_experts, f"{topk=}, {num_experts=}"
+        experts_per_rank = num_experts // ep_size
+
+        pad_ordinal = torch.arange(padding_rows.numel(), device=routed_experts.device, dtype=torch.int64).reshape(
+            (-1, 1, 1)
+        )
+        layer = torch.arange(num_layers, device=routed_experts.device, dtype=torch.int64).reshape((1, num_layers, 1))
+        column = torch.arange(topk, device=routed_experts.device, dtype=torch.int64).reshape((1, 1, topk))
+        assignment = (pad_ordinal + layer * ep_size) * topk + column
+        expert_rank = assignment.remainder(ep_size)
+        local_expert = assignment.div(ep_size, rounding_mode="floor").remainder(experts_per_rank)
+        pad_experts = expert_rank * experts_per_rank + local_expert
+        routed_experts[padding_rows] = pad_experts.to(routed_experts.dtype)
 
     if sequence_parallel:
         tp_rank = mpu.get_tensor_model_parallel_rank()

tests/test_cp_utils.py

@@ -1,10 +1,13 @@
-"""CPU unit tests for ``slime.backends.megatron_utils.cp_utils.get_sum_of_sample_mean``.
+"""CPU unit tests for ``slime.backends.megatron_utils.cp_utils``.
 
-Pins the per-rollout reducer contract: a rollout split into N training
+The reducer tests pin the per-rollout contract: a rollout split into N training
 samples (compact / subagent) must contribute exactly one token-weighted
 mean to the sum, even when first-fit packing puts those siblings into
 different micro-batches at training time.
 
+The routing-replay tests pin padding load distribution across samples,
+context-parallel layouts, and sequence-parallel slices.
+
 The CPU-only CI image does not ship megatron — ``_cp_dist_helpers``
 stubs ``megatron.core.mpu`` at import time so the subsequent
 ``cp_utils`` import binds against the stub.
@@ -26,6 +29,7 @@
 from slime.backends.megatron_utils.cp_utils import (  # noqa: E402
     get_logits_and_tokens_offset_with_cp,
     get_sum_of_sample_mean,
+    prepare_routed_experts_for_routing_replay,
 )
 
 
@@ -51,6 +55,114 @@ def _denoms(*values: int) -> torch.Tensor:
     return torch.tensor(values, dtype=torch.float32)
 
 
+def _set_parallel_ranks(monkeypatch, *, cp_size=1, cp_rank=0, tp_size=1, tp_rank=0, ep_size=1):
+    from megatron.core import mpu
+
+    monkeypatch.setattr(mpu, "get_context_parallel_world_size", lambda: cp_size)
+    monkeypatch.setattr(mpu, "get_context_parallel_rank", lambda: cp_rank)
+    monkeypatch.setattr(mpu, "get_tensor_model_parallel_world_size", lambda: tp_size, raising=False)
+    monkeypatch.setattr(mpu, "get_tensor_model_parallel_rank", lambda: tp_rank, raising=False)
+    monkeypatch.setattr(mpu, "get_expert_model_parallel_world_size", lambda: ep_size, raising=False)
+
+
+def _make_routing_replay_samples(num_samples=4, *, num_layers=4, topk=2):
+    routed_experts = [torch.full((2, num_layers, topk), 7, dtype=torch.int32) for _ in range(num_samples)]
+    tokens = [torch.zeros(3, dtype=torch.int64) for _ in range(num_samples)]
+    return routed_experts, tokens
+
+
+@pytest.mark.unit
+def test_routing_replay_small_padding_balances_expert_parallel_ranks(monkeypatch):
+    topk = 8
+    for num_experts, ep_size in ((128, 8), (128, 32), (256, 8), (256, 32)):
+        experts_per_rank = num_experts // ep_size
+        for num_padding_rows in range(1, 6):
+            _set_parallel_ranks(monkeypatch, ep_size=ep_size)
+            routed_experts, tokens = _make_routing_replay_samples(
+                num_padding_rows,
+                num_layers=4,
+                topk=topk,
+            )
+
+            actual = prepare_routed_experts_for_routing_replay(
+                routed_experts,
+                tokens,
+                num_experts=num_experts,
+                data_pad_size_multiplier=1,
+                sequence_parallel=False,
+                allgather_cp=False,
+            )
+
+            for sample_idx, expected in enumerate(routed_experts):
+                torch.testing.assert_close(actual[sample_idx * 3 : sample_idx * 3 + 2], expected)
+
+            padding = actual[torch.arange(2, num_padding_rows * 3, 3)]
+            assert padding.dtype == torch.int32
+            assert torch.all((padding >= 0) & (padding < num_experts))
+            sorted_padding = padding.sort(dim=2).values
+            assert torch.all(sorted_padding[:, :, 1:] != sorted_padding[:, :, :-1])
+
+            ep_ranks = padding // experts_per_rank
+            expected_row_fanout = min(ep_size, topk)
+            for row_layer_ranks in ep_ranks.reshape((-1, topk)):
+                assert torch.unique(row_layer_ranks).numel() == expected_row_fanout
+
+            expected_active_ranks = min(ep_size, num_padding_rows * topk)
+            for layer_experts, layer_ranks in zip(
+                padding.transpose(0, 1),
+                ep_ranks.transpose(0, 1),
+                strict=True,
+            ):
+                assert torch.unique(layer_experts).numel() == num_padding_rows * topk
+                histogram = torch.bincount(layer_ranks.flatten(), minlength=ep_size)
+                assert torch.count_nonzero(histogram).item() == expected_active_ranks
+                assert histogram.max().item() - histogram.min().item() <= 1
+
+
+@pytest.mark.unit
+def test_routing_replay_batch_padding_advances_per_row(monkeypatch):
+    _set_parallel_ranks(monkeypatch)
+    routed_experts, tokens = _make_routing_replay_samples(2)
+
+    actual = prepare_routed_experts_for_routing_replay(
+        routed_experts,
+        tokens,
+        num_experts=8,
+        data_pad_size_multiplier=8,
+        sequence_parallel=False,
+        allgather_cp=False,
+    )
+
+    expected = torch.tensor([[4, 5], [6, 7]], dtype=torch.int32)
+    torch.testing.assert_close(actual[-2:, 0], expected)
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize(("allgather_cp", "expected_padding_rows"), [(False, (3, 3)), (True, (2, 4))])
+def test_routing_replay_padding_uses_post_cp_order_before_sp(monkeypatch, allgather_cp, expected_padding_rows):
+    routed_experts, tokens = _make_routing_replay_samples(3)
+
+    for cp_rank in range(2):
+        tp_outputs = []
+        for tp_rank in range(2):
+            _set_parallel_ranks(monkeypatch, cp_size=2, cp_rank=cp_rank, tp_size=2, tp_rank=tp_rank)
+            tp_outputs.append(
+                prepare_routed_experts_for_routing_replay(
+                    routed_experts,
+                    tokens,
+                    num_experts=8,
+                    data_pad_size_multiplier=1,
+                    sequence_parallel=True,
+                    allgather_cp=allgather_cp,
+                )
+            )
+
+        padding_rows = torch.cat(tp_outputs)
+        padding_rows = padding_rows[~(padding_rows == 7).all(dim=(1, 2))]
+        expected = torch.arange(expected_padding_rows[cp_rank] * 2, dtype=torch.int32).reshape((-1, 2))
+        torch.testing.assert_close(padding_rows[:, 0], expected)
+
+
 @pytest.mark.unit
 def test_default_reduces_to_per_sample_mean():
     """``sample_denoms=None`` reproduces the legacy per-sample-mean."""