DeepSpeed

35d1c7d3 - Derive AutoEP rank splits from the per-expert count exchange (#8190)

Derive AutoEP rank splits from the per-expert count exchange (#8190) ## Redundancy `compute_split_plan` (`deepspeed/module_inject/auto_ep_layer.py`) runs two all-to-all exchanges over overlapping metadata: ```python # exchange 1: [ep_size] per-rank totals -> output_splits local_counts_tensor = count_matrix.sum(dim=1).clone() dist.all_to_all_single(remote_counts_tensor, local_counts_tensor, group=ep_group) output_splits = remote_counts_tensor.cpu().tolist() # exchange 2: [ep_size, E_local] per-expert matrix -> local_counts dist.all_to_all_single(received_counts_flat, local_expert_counts_flat, group=ep_group) received_counts = received_counts_flat.view(ep_size, num_local_experts) ``` The second exchange already contains the first. `received_counts[s, e]` is the number of tokens source rank `s` sends to local expert `e`, so ``` received_counts.sum(dim=1)[s] == (tokens rank s sends to this rank) == output_splits[s] ``` element-wise and exactly — the first exchange transmits a strict projection of what the second one transmits. The planner also issued two separate `.cpu()` calls, each of which is a device-to-host sync. ## Fix Drop the rank-total exchange and derive `output_splits` from the detailed one: ```python received_counts = received_counts_flat.view(ep_size, num_local_experts) input_splits, output_splits = torch.stack(( count_matrix.sum(dim=1), received_counts.sum(dim=1), )).cpu().tolist() ``` Per split-plan call this changes: | | before | after | |---|---|---| | all-to-all collectives | 2 | 1 | | device-to-host syncs | 2 | 1 | `compute_split_plan_from_expert_indices` (folded-TP path) now shares the same helper instead of duplicating the logic. Only metadata is affected. The dispatch/combine payload `_AllToAllV` calls are untouched. ## Measurements EP=8 (2 nodes x 8 H100), 256 experts, 1024 tokens/rank, top_k=8, 50 trials, p50 of split-plan time per layer: | token distribution | before | after | saved | |---|---|---|---| | balanced | 0.366 ms | 0.190 ms | 0.176 ms (-48%) | | mild skew | 0.352 ms | 0.182 ms | 0.171 ms (-49%) | | severe skew | 0.360 ms | 0.183 ms | 0.176 ms (-49%) | Collective time alone goes 0.120 ms -> 0.058 ms, consistent with halving the number of exchanges. The saving is flat across skew because what is removed is fixed per-call overhead, not payload-dependent work. For a 20-layer MoE model this is ~3.5 ms per forward pass. ## Tests `tests/unit/v1/moe/test_autoep_unit.py::TestSplitPlan` adds: - an oracle check that the derived `output_splits` equal the per-rank totals, for every rank and three token distributions; - a differential test against a local reimplementation of the old two-collective planner, asserting the old path issues 2 collectives, the new one issues 1, and all four `SplitPlan` fields are identical; - a test that `compute_split_plan_from_expert_indices` produces an identical plan for the same routing assignment. Signed-off-by: yh0903 <helloyu0903@gmail.com> Co-authored-by: Hongwei Chen <33092912+hwchen2017@users.noreply.github.com>