OpenPipe / OpenPipe/ART

Support combined expert and data parallelism in TrainerRank

Open
#917 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
10.8k
Forks
989
Avg merge
6h 29m
Merged PRs (30d)
85

Description

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.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start in src/art/trainer_rank/_impl.py at _validate_hybridep_topology, then trace forward_micro_batches and dp_rank_forward to understand the current EP/DP restriction. The work is done when EP=2 and DP=2 complete forward, backward, and optimizer updates on unequal and empty DP batches, with CP and multiple-microbatch combinations preserving outputs, gradients, and collective ordering; document the supported topology matrix.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
distributed-systems, machine-learning
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.