【RFC】SAPP-PPB for Pipeline Parallelism Balancing

Background

In large-model training, pipeline parallelism is a common approach to scaling training across more devices. However, training efficiency is highly sensitive to how layers are partitioned across pipeline stages.

In practice, suboptimal pipeline partitioning often leads to several problems:

• stage-wise compute imbalance, which increases pipeline bubble and idle waiting time
• stage-wise memory imbalance, which makes some devices become memory bottlenecks
• higher tuning cost when interleaving and recomputation are involved
• difficulty in finding a good stage-layer allocation manually within a reasonable time

Large-model performance tuning usually needs to jointly consider memory constraints and pipeline-related performance factors. Engineers often have to try many combinations on a real cluster before reaching an acceptable strategy, which can take weeks and consume substantial compute resources. SAPP-related tooling is designed to reduce this cost by constructing an optimization problem from model memory/time information and pipeline-related hyperparameters, then solving it globally to generate stage-layer allocation and layer-wise recomputation decisions.

To address this, we propose to introduce SAPP-PPB into hyper_parallel/auto_parallel as a pipeline parallelism balancing module.

Proposal

The main logic of PPB is to use a symbolic cost model to simulate memory cost and compute-related cost under pipeline-parallel settings, and then generate a better pipeline strategy through optimization.

PPB takes model-level symbolic information and pipeline-related configuration as input, builds a cost model for candidate strategies, and formulates the balancing process as an optimization problem. The optimization target is to improve stage balance, reduce idle waiting time, and control memory overhead under pipeline-related constraints such as stage count, micro-batch settings, interleaving configuration, and memory budget.

Inputs / Outputs

Inputs

• model description, such as layer/block structure, repeated layers, symbolic memory information, and symbolic compute/time information
• pipeline-related configuration, such as number of stages, number of micro-batches, interleaving configuration, pipeline scheduling assumptions, and memory budget
• optional strategy-related hints, such as recomputation impact, solver time limit, and output options

Outputs

• pipeline balancing strategy, such as stage-layer allocation and recommended partition strategy
• memory / performance related decisions, such as recomputation decisions and per-stage balance result
• optional analysis artifacts, such as visualization, simulator view, and optimization summary

Expected Benefits

After introducing PPB, we expect:

• lower manual tuning cost
• better stage-wise balance
• reduced idle waiting time
• better compute and memory utilization
• improved end-to-end training performance

schema_version: 1
source: gitcode
gitcode_repo: mindspore/hyper-parallel
gitcode_issue: 78
source_url: https://gitcode.com/mindspore/hyper-parallel/issues/78

The proposal targets hyper_parallel/auto_parallel; start by reviewing that module's existing structure and interfaces. Define how the listed model and pipeline configuration inputs should produce stage allocation, recomputation decisions, balance results, and optional analysis artifacts. Done means the proposed PPB module and its optimization behavior are specified and validated against those outputs.