Price the shared-expert return under context and expert parallelism

[bradhilton]

Part of #949. The trainer-rank planner prices a MoE layer's shared-expert return, the output held beside the routed experts' work. It drops that return whenever context or expert parallelism is on. The per-token return does not depend on either: CP shards rows, and shared experts are not expert-parallel. So at CP>1 and EP>1 the MoE workspace was missing 8 KB per local token in checkpoint mode.

Evidence

The data comes from Qwen3.6-35B-A3B allocator traces on H200 (8 layers, 194,753 tokens), replayed on CPU, plus ART's own planner reports from the same runs.

• The measured return is identical in every topology. At the recompute peak, the gated shared-expert pair (lora.py:1328 output and shared_experts.py:127 gated output) is 4,096 + 4,096 B per local token at CP1, CP2/EP1 and EP2/CP2. That is exactly what the planner charges in checkpoint mode for a gated shared expert (shared + shared).
• The planner drops it outside CP1/EP1. The MoE coefficient is 192,512 B/token at CP1 (routed 188,416 plus the 4,096 shared return), 188,416 at CP2 and 282,624 at EP2 (1.5 × 188,416, the routed allowance). The shared term is 0 at CP2 and EP2.
• Cause. _shared_expert_output_bytes_per_token required TP, CP, PP, EP and ETP all to be 1, which Account for grouped checkpoint, head, and enclosing forward buffers #898 introduced as blanket support gating.
• The routed-expert terms are already right. At the CP2 peak, the routed tensors total 11,792 elements per routed row, against 11,776 priced (FC2 input, FC2 base/adapter/sum, both dispatched inputs, the FC1 gate/up sum).

Change

• CP and EP no longer disable the shared-expert return. TP, PP and ETP must still be 1; they are unmeasured.
• At CP>1 or EP>1, the checkpoint MoE coefficient gains 8,192 B per local token for a gated shared expert (4,096 in forward). Nothing changes at CP1/EP1.

Tests

tests/unit/test_trainer_rank_shared_memory.py:

• the shared return is 4,096 B per local token at CP2 and at EP2;
• a CP2 layer's MoE coefficient includes it (192,512);
• at CP2/EP2 with the HybridEP dispatcher, the shared return escapes Price HybridEP expert work and buffer growth at EP>1 #954's 1.5× routed-row allowance: 12 × 11,776 × 2 + 4,096 forward, + 8,192 in checkpoint mode;
• the unsupported-topology mutations now use PP=2 and CP=0.

All five new test cases fail on main. Locally, on a heavily loaded host, 1,583 trainer-rank unit and acceptance tests pass and three fail:

• test_planning_failures_and_empty_ranks_use_aligned_status and test_optim_step_live_graph_error_is_collective fail the same way on main;
• test_forward_micro_batches_yield_modes_collectively[4] hung once and passed on two reruns.

Not covered

The traces also show about 6.4 KB per local token of per-token MoE state that the planner has never modeled, at any CP:

• router scores (E × fp32);
• the routing map (E bytes);
• TE's row-id map ((2E+1) × int32);
• the shared expert's FC1 gate/up sum and GLU output.

At EP2 the router and dispatch metadata differ (HybridEP), so these are left for a follow-up with its own validation rather than folded into this fix. Today the planner's input-gradient allowance, which is not live at the peak, absorbs both this residual and the shared term fixed here. Neither should rely on that.

This change is independent of #963 (the recomputed-mixer floor); the two touch different code.

🤖 Generated with Claude Code