Skip to content

Support combined expert and data parallelism in TrainerRank #917

Description

@bradhilton

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.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions