Skip to content

Support pipeline parallelism in TrainerRank without exposing stage scheduling to callers #916

Description

@bradhilton

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.

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