#917·ART

Support combined expert and data parallelism in TrainerRank

Author: bradhiltonCreated Sep 16, 2026Updated Sep 16, 2026

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.