OpenPipe / OpenPipe/ART

Support pipeline parallelism in TrainerRank without exposing stage scheduling to callers

Open
#916 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 pipeline parallelism (PP > 1) and multiple local model chunks. Support pipeline-parallel execution internally so experiment code can keep using the same forwards, full outputs, custom losses, backward, and optimizer API without managing pipeline stages.

Current behavior

At ART main 7496cc09252c52ec7a63ab74d11caabce173f5b0, the constructor raises TrainerRankRuntimeSupportError when runtime.provider.pipeline_model_parallel_size > 1 or len(runtime.model) > 1. The error explains that TrainerRank does not use the MCore forward/backward schedule and requires PP=1 with exactly one local model chunk.

This is a source-confirmed unsupported configuration, not a newly observed GPU failure. Removing the guard alone would not provide the missing scheduling and communication.

Desired behavior
  • Internally schedule stage forwards/backwards and activation/gradient transfers while preserving forward_micro_batches and dp_rank_forward semantics.
  • Preserve full source-order outputs, caller-defined losses and registered custom heads, checkpoint selection, gradient accumulation, and optimizer behavior. Public DP reductions must not count pipeline stages as independent data batches.
  • Keep pipeline stage ownership and scheduling out of experiment code. Explicitly define supported combinations with TP/CP and any initial limitations, including multiple local chunks.
Acceptance
  • A native PP=2 canary completes forward, caller-side loss, backward, and optimizer update through the public API.
  • Compare outputs, model/custom-head gradients, and parameter updates with a matched PP=1 reference within stated numerical tolerances; cover no-grad execution and accumulation across multiple forwards.
  • Exercise uneven microbatch workloads and checkpoint selection without mismatched communication or hangs. Keep explicit errors for configurations that remain unsupported.

Related: #911 / #912 establish the full-output contract for context parallelism; this issue tracks the separate pipeline execution limitation.

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 the constructor guard referenced in the issue, then trace the public forward, backward, and optimizer paths. Define and implement internal pipeline scheduling and communication while preserving the listed API semantics. Done means the PP=2 canary and PP=1 comparisons pass, including no-grad, accumulation, uneven microbatches, checkpoint selection, and explicit errors for unsupported combinations.

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.