TrainerRank currently rejects expert parallelism (EP > 1) combined with data parallelism (DP > 1). Support this combination internally so MoE experiments can distribute distinct batches across DP workers without managing expert-routing communication.
Current behavior
At ART main 7496cc09252c52ec7a63ab74d11caabce173f5b0, _validate_hybridep_topology raises NotImplementedError for active expert parallelism with DP>1:
TrainerRank does not support combining data parallelism with expert parallelism because uneven DP inputs can desynchronize HybridEP collectives.
The current error recommends DP=1 with CP and EP set to the world size. This issue records a source-confirmed limitation; no new GPU reproduction was run for this report.
Desired behavior
- Support EP>1 with DP>1, including different token counts and uneven or empty work on DP ranks, through
forward_micro_batches and dp_rank_forward.
- Internally coordinate expert communication groups, collective ordering, and HybridEP buffer capacity/lifetime across participating ranks. Determine the required grouping/scheduling changes before lifting the guard.
- Preserve full source-order outputs, registered custom heads, checkpoint selection, accumulation, and correct model/expert/custom-head gradient reductions.
dp_reduce should continue combining distinct data batches without counting model-parallel replicas again.
- Users should not need to pad their DP batches or orchestrate expert collectives to prevent hangs. Keep explicit errors for combinations that remain unsupported.
Acceptance
- A native EP=2, DP=2 MoE canary completes forward, backward, and optimizer update on unequal DP batches, including an empty-work case.
- Compare outputs and appropriately normalized model/expert/custom-head gradients and updates against a matched supported reference within stated tolerances.
- Exercise CP combinations, multiple microbatches, and live-graph/buffer-capacity transitions without collective mismatches or hangs; document the supported topology matrix.
Related: #851 concerns a separate EP=1, CP>1 grouped-GEMM crash. #911 / #912 concern the public full-output CP contract; neither tracks this EP/DP restriction.
TrainerRank currently rejects expert parallelism (
EP > 1) combined with data parallelism (DP > 1). Support this combination internally so MoE experiments can distribute distinct batches across DP workers without managing expert-routing communication.Current behavior
At ART main
7496cc09252c52ec7a63ab74d11caabce173f5b0,_validate_hybridep_topologyraisesNotImplementedErrorfor active expert parallelism with DP>1:The current error recommends DP=1 with CP and EP set to the world size. This issue records a source-confirmed limitation; no new GPU reproduction was run for this report.
Desired behavior
forward_micro_batchesanddp_rank_forward.dp_reduceshould continue combining distinct data batches without counting model-parallel replicas again.Acceptance
Related: #851 concerns a separate EP=1, CP>1 grouped-GEMM crash. #911 / #912 concern the public full-output CP contract; neither tracks this EP/DP restriction.