【训练营】Pipeline 并行自定义布局 by Dayuxiaoshui · Pull Request #209 · InfiniTensor/InfiniTrain

Dayuxiaoshui · 2026-08-24T13:39:12Z

PR: Configurable Pipeline Layouts

Commit covered

7a78e15 feat: add configurable pipeline layouts

19 files changed, 1525 insertions(+), 93 deletions(-).

Summary

GPT-2 and LLaMA 3 can now assign a custom, non-uniform number of Transformer layers to each pipeline-parallel stage, instead of only the built-in uniform split. A new nn::parallel::PipelineLayout is the single source of truth for "which Transformer layers (and the embedding / final-norm / lm-head) live on which physical PP stage," and it's consumed identically by model construction, PipelineParallel stage setup, and both LLMC checkpoint loaders — so layer assignment can't drift between subsystems.

Layout can be specified four ways

Uniform (default, unchanged) — existing even split, remainder layers assigned to the earliest stages. Existing GPipe/1F1B/vPP/TP/DDP entry points are untouched when no new flag is passed.
Manual partition — --pipeline_layer_partition=4,8,6,6: one entry per PP stage, entries must be positive integers and sum to the model's real layer count. Requires --virtual_pipeline_parallel=1.
Cost-based auto-balance — --pipeline_layer_costs=10,1,1,1,1,1: one relative cost per layer (e.g. from a profiler); a DP over contiguous partitions (O(stages * layers^2)) picks the split that minimizes the slowest stage's total cost, keeping layers contiguous and in order. Mutually exclusive with manual partition.
Arbitrary vPP chunk / Megatron-style layout — --pipeline_chunk_layout=0:3,1:3,1:3,0:3 or --pipeline_model_parallel_layout='Et*3||t*3|t*6NL': explicit ordered chunk-to-stage ownership, removing the previous fixed global_chunk = local_chunk * pp_size + stage round-robin restriction for virtual_pipeline_parallel > 1.

Other pieces

PipelineLayout::layer_ranges/stage_for_layer/owns_embedding/owns_final_norm/owns_lm_head — query API used by model construction and checkpoint loaders; stored thread_local since one process can host multiple training threads, each representing an independent global rank.
PipelineParallel::GetStageInfo — compatibility projection for existing scheduler code.
scripts/suggest_pipeline_layout.py — offline tool that turns PROFILE_MODE records (or user-supplied costs) into a ready-to-paste --pipeline_layer_partition value, dropping each layer's first profiler sample by default to avoid CUDA-warmup skew.
GPT-2 / LLaMA3 checkpoint_loader.{h,cc} and main.cc updated to resolve and validate the layout (after reading the real layer count from the LLMC checkpoint header / config) before model construction, and to print the normalized final layout on the main rank at startup.
Validation fails fast, before model construction, on: wrong entry count vs. --pipeline_parallel, wrong layer-count sum, non-positive/empty entries, custom physical partition combined with --virtual_pipeline_parallel != 1, and manual partition + auto-balance cost both set. Full error-message list in docs/pipeline_layout_guide.md.
New optional --dump_gradients=DIR on the GPT-2 example exports all non-empty parameter gradients after the first optimizer step, remapping each PP rank's local layer number to a global one so single-GPU and custom-PP gradient dumps can be diffed directly via scripts/precision_check/precision_compare.py.

Docs added

docs/pipeline_layout_guide.md — user-facing flag reference and syntax.
docs/pipeline_layout_report.md — design/implementation notes, the DP algorithm, and test results.
docs/pipeline_layout_test_log.md — raw logs backing the results below.

Tests

tests/distributed/test_pipeline_layout.cc (CPU, ctest-registered) — manual partition (4,8,6,6), full layer-to-stage reverse lookup, special-module ownership, default vPP rotation, and the error cases above.
tests/distributed/test_pipeline_layout_suggestion.py — unit tests for the cost-based auto-balance DP and profiler-record parsing.
tests/distributed/test_pipeline_layout_e2e.sh — 2-GPU CUDA/NCCL script (needs GPT-2 124M LLMC checkpoint; not registered in default ctest since it needs 2 GPUs + external model assets): runs a single-GPU baseline vs. a custom two-stage layout, then diffs the final layout string, fp32 loss, gradient-file set, and per-parameter gradients.

Reported results (see docs/pipeline_layout_report.md for full detail): 10/10 CPU layout/suggestion tests pass; CPU build of GPT-2/LLaMA3/Mixtral compiles and links cleanly. On 2×H200, GPT-2 124M with a custom 4,8 two-stage layout matches single-GPU loss to within 2e-6 (well inside the fp32 1e-5 tolerance used), and all 149/149 per-parameter gradients match at atol=1e-5, rtol=0 with no missing files. A profiler-suggested 7,5 layout (vs. default 6,6) measured 9.45% higher throughput and 130MB lower peak stage memory on the same 2×H200 setup.

Risk / compatibility

No behavior change for existing runs: uniform layout stays the default when none of the new flags are passed.
Manual/auto-balanced physical partitions explicitly reject --virtual_pipeline_parallel != 1 rather than silently producing wrong execution order — the fixed vPP round-robin can't unambiguously absorb a physical layer-count list.
New GPU E2E coverage is 2-GPU only; the report notes a full DP×TP×PP combination matrix should still be regression-tested on target cluster hardware before wider rollout.
--dump_gradients copies gradients to CPU synchronously and is intended for correctness verification only — should not be enabled during performance benchmarking.

Test plan

cmake -S . -B build -DBUILD_TEST=ON -DUSE_CUDA=OFF -DUSE_NCCL=OFF -DUSE_OMP=OFF
cmake --build build --target test_pipeline_layout gpt2 llama3 -j$(nproc)
ctest --test-dir build -R PipelineLayoutTest --output-on-failure

For the CUDA/NCCL 2-GPU pipeline-layout regression:

tests/distributed/test_pipeline_layout_e2e.sh \
  <cuda-build-dir> data/gpt2/tiny_shakespeare_train.bin data/gpt2/gpt2_124M.bin 0,1

Balancing only the Transformer layers misses the two modules that are pinned to the ends of the pipeline: the embedding always sits on stage 0 and the final norm plus lm head on the last stage. On a large-vocabulary model the lm head, not the layers, is the slowest part of the last stage. --pipeline_layer_costs now accepts optional "E:<cost>" and "L:<cost>" entries alongside the per-layer costs. They take no layer slot; they only add a constant to the first and last segment in the balancing DP. The trailing constant applies only to the state that closes the last stage, so the optimal substructure still holds. suggest_pipeline_layout.py grows the matching --embedding-cost / --lm-head-cost so an offline suggestion and the runtime solver agree. Also add the regression that was missing for the parallelism claims: a layout x PP/DDP/TP matrix that checks every combination against one single-GPU reference, and a synthetic LLMC asset generator so the matrix can run where the llm.c starter pack cannot be downloaded. Measured on 2 x H20, 8 layers / 384 hidden / vocab 50304 / seq 512: uniform 4,4 runs at 86,628 tok/s, while "1x8,L:4.5" picks 6,2 and runs at 102,904 tok/s (+18.8%). The raw parameter-count ratio (L:10.9) overshoots to 7,1 and only gains 6.8%, so the cost still has to be measured rather than derived from parameter counts. Verified: 11/11 layout unit tests, and 14/14 matrix cases on 8 x H20 matching the single-GPU reference loss within 1e-6, with 101/101 per-parameter gradients matching at atol=1e-5 on the 11 non-TP cases.

JYMiracle305 · 2026-09-30T02:44:50Z

 namespace gpt2 {

-std::shared_ptr<nn::TransformerModel> LoadFromLLMC(const std::string &filepath) {
+std::shared_ptr<nn::TransformerModel> LoadFromLLMC(const std::string &filepath,


不建议在这里以传参形式传原始字符串，可以把处理好的布局放在 global，提供 api 供这里获取

JYMiracle305 · 2026-09-30T02:47:21Z

    {"d48", {.block_size = 1024, .vocab_size = 50257, .n_layer = 48, .n_head = 25, .n_embd = 1600}},
 };

+std::string CanonicalGradientName(std::string name, int pp_rank) {


不需要用业务代码来验证，最终打印 loss 和原测例对比即可

JYMiracle305 · 2026-09-30T03:07:13Z

@@ -0,0 +1,126 @@
+#!/usr/bin/env python3


不需要虚拟生成 checkpoint 和 token 数据

JYMiracle305 · 2026-09-30T03:10:28Z

@@ -0,0 +1,105 @@
+#!/usr/bin/env bash


新增的不均匀布局测例配置可以加在 test_config.json，不需要新增测试入口。利用测试入口，可以同一套配置测试GPT2 和 LLaMA3

JYMiracle305 · 2026-09-30T03:13:38Z

@@ -0,0 +1,215 @@
+#!/usr/bin/env bash


这个针对不同布局的测试也可以加入 test_config.json，单独新增一个标签包含这些测例

JYMiracle305 · 2026-09-30T03:17:35Z

    std::vector<std::pair<int, int>> layer_ranges_per_chunk;
 };

+class PipelineLayout {


这个布局结构体考虑单独划分一个文件，否则 pipeline_parallel.h/cc 太重了，而且责任不明确

JYMiracle305 · 2026-09-30T03:21:27Z

@@ -0,0 +1,61 @@
+#!/usr/bin/env python3


这个脚本就是为了验证脚本 suggest_pipeline_layout.py 吗，还有哪些必要性

Dayuxiaoshui added 2 commits August 21, 2026 14:38

feat: add configurable pipeline layouts

7a78e15

Dayuxiaoshui changed the title ~~feat: add configurable pipeline layouts~~ 【训练营】Pipeline 并行自定义布局 Sep 19, 2026

JYMiracle305 self-requested a review September 20, 2026 01:48

JYMiracle305 reviewed Sep 30, 2026

View reviewed changes

Navigation Menu

Repository navigation

【训练营】Pipeline 并行自定义布局 - #209

【训练营】Pipeline 并行自定义布局#209
Dayuxiaoshui wants to merge 2 commits into
InfiniTensor:masterfrom
Dayuxiaoshui:master

Dayuxiaoshui commented Aug 24, 2026

Uh oh!

JYMiracle305 Sep 30, 2026

Uh oh!

JYMiracle305 Sep 30, 2026

Uh oh!

JYMiracle305 Sep 30, 2026

Uh oh!

JYMiracle305 Sep 30, 2026

Uh oh!

JYMiracle305 Sep 30, 2026

Uh oh!

JYMiracle305 Sep 30, 2026

Uh oh!

JYMiracle305 Sep 30, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

Dayuxiaoshui commented Aug 24, 2026

PR: Configurable Pipeline Layouts

Commit covered

Summary

Layout can be specified four ways

Other pieces

Docs added

Tests

Risk / compatibility

Test plan

Uh oh!

JYMiracle305 Sep 30, 2026

Choose a reason for hiding this comment

Uh oh!

JYMiracle305 Sep 30, 2026

Choose a reason for hiding this comment

Uh oh!

JYMiracle305 Sep 30, 2026

Choose a reason for hiding this comment

Uh oh!

JYMiracle305 Sep 30, 2026

Choose a reason for hiding this comment

Uh oh!

JYMiracle305 Sep 30, 2026

Choose a reason for hiding this comment

Uh oh!

JYMiracle305 Sep 30, 2026

Choose a reason for hiding this comment

Uh oh!

JYMiracle305 Sep 30, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants