feat: add configurable pipeline layouts

[Dayuxiaoshui]

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