From 7a78e15e3d528ad6cc8cbc6faa615f79bce36195 Mon Sep 17 00:00:00 2001
From: Dayuxiaoshui <792179245@qq.com>
Date: Fri, 21 Aug 2026 14:38:15 +0000
Subject: [PATCH] feat: add configurable pipeline layouts

---
 docs/pipeline_layout_guide.md                 | 163 ++++++
 docs/pipeline_layout_report.md                |  93 ++++
 docs/pipeline_layout_test_log.md              | 148 ++++++
 example/gpt2/checkpoint_loader.cc             |  26 +-
 example/gpt2/checkpoint_loader.h              |   6 +-
 example/gpt2/main.cc                          |  88 +++-
 example/llama3/checkpoint_loader.cc           |  20 +-
 example/llama3/checkpoint_loader.h            |   6 +-
 example/llama3/main.cc                        |  21 +-
 .../nn/modules/transformer/transformer.h      |   1 +
 .../nn/parallel/pp/pipeline_parallel.h        |  48 ++
 .../src/nn/modules/transformer/transformer.cc |  56 +-
 .../src/nn/parallel/pp/pipeline_parallel.cc   | 479 ++++++++++++++++--
 .../src/nn/parallel/pp/pipeline_schedule.cc   |  51 +-
 scripts/suggest_pipeline_layout.py            | 125 +++++
 tests/distributed/CMakeLists.txt              |  14 +
 tests/distributed/test_pipeline_layout.cc     | 126 +++++
 tests/distributed/test_pipeline_layout_e2e.sh | 105 ++++
 .../test_pipeline_layout_suggestion.py        |  42 ++
 19 files changed, 1525 insertions(+), 93 deletions(-)
 create mode 100644 docs/pipeline_layout_guide.md
 create mode 100644 docs/pipeline_layout_report.md
 create mode 100644 docs/pipeline_layout_test_log.md
 create mode 100755 scripts/suggest_pipeline_layout.py
 create mode 100644 tests/distributed/test_pipeline_layout.cc
 create mode 100755 tests/distributed/test_pipeline_layout_e2e.sh
 create mode 100644 tests/distributed/test_pipeline_layout_suggestion.py

diff --git a/docs/pipeline_layout_guide.md b/docs/pipeline_layout_guide.md
new file mode 100644
index 00000000..2f650bb4
--- /dev/null
+++ b/docs/pipeline_layout_guide.md
@@ -0,0 +1,163 @@
+# Pipeline 自定义布局
+
+InfiniTrain 的 GPT-2 和 LLaMA 3 示例支持用 `--pipeline_layer_partition` 指定每个物理 Pipeline Stage
+拥有的连续 Transformer 层数。模型构建、PP Stage 构造和 LLMC 参数加载均使用同一个
+`PipelineLayout`。
+
+## 参数与语法
+
+```bash
+./gpt2 \
+  --pipeline_parallel 4 \
+  --virtual_pipeline_parallel 1 \
+  --pipeline_layer_partition 4,8,6,6 \
+  [其他训练参数]
+```
+
+该模型必须有 24 层，最终布局为：
+
+```text
+stage 0: embedding + layers [0, 4)
+stage 1: layers [4, 12)
+stage 2: layers [12, 18)
+stage 3: layers [18, 24) + final_norm + lm_head
+```
+
+列表项必须是正整数；项数必须等于 `--pipeline_parallel`，总和必须等于 checkpoint 或配置中的
+模型层数。空格可以出现在数字两侧。GPT-2 和 LLaMA 3 使用相同参数。
+
+## 按逐层代价自动均衡
+
+如果已经通过 profiler、FLOPs 估算或经验权重得到每个 Transformer 层的相对代价，可以让 InfiniTrain
+自动生成连续分区：
+
+```bash
+./gpt2 \
+  --pipeline_parallel 2 \
+  --pipeline_layer_costs 10,1,1,1,1,1 \
+  [其他训练参数]
+```
+
+上述 6 层模型会生成 `1,5`：stage 0 的建模代价为 10，stage 1 为 5。均匀 `3,3` 的代价为
+12 和 3，因此最慢 Stage 的建模代价从 12 降到 10。自动布局保持层连续、顺序不变，并保证每个
+Stage 至少拥有一层。
+
+代价项必须是有限正数，项数必须和模型 Transformer 层数完全一致。
+`--pipeline_layer_costs` 与 `--pipeline_layer_partition` 互斥，且当前同样要求
+`--virtual_pipeline_parallel=1`。程序会在模型构建前打印自动生成的最终布局。
+
+## 默认行为与 vPP
+
+不传 `--pipeline_layer_partition` 时，保持原有均匀划分。余数从执行顺序靠前的 chunk 开始各多分一层；
+`--virtual_pipeline_parallel` 的默认轮转 chunk 布局也保持不变。
+
+当前自定义层数列表只描述物理 Stage，不描述虚拟 chunk，因此它与
+`--virtual_pipeline_parallel` 大于 1 不兼容，程序会在创建模型前报错。Embedding 固定属于 stage 0，
+Final Norm 和 LM Head 固定属于最后一个 stage，并显式记录在布局查询接口和启动日志中。
+
+## 任意 vPP Chunk 映射
+
+使用有序 STAGE:LAYER_COUNT 列表显式指定 Chunk owner：
+
+    --pipeline_parallel=2 --virtual_pipeline_parallel=2 \
+    --pipeline_chunk_layout=0:3,1:3,1:3,0:3
+
+这表示逻辑 Chunk owner 为 [0,1,1,0]，层范围为 [0,3)、[3,6)、[6,9)、[9,12)。每个物理 Stage
+必须获得相同的正数 Chunk；连续 Chunk 可以属于同一 Stage，此时直接保留本地 autograd 图。
+Embedding 归属第一个逻辑 Chunk，Final Norm/LM Head 归属最后一个逻辑 Chunk。
+
+## Megatron 风格表达式
+
+pipeline_model_parallel_layout 支持 E（Embedding）、t（Transformer）、N（Final Norm）、L（LM Head）、
+| 分隔符、x*n 和 (expr)*n 重复，以及相邻 || 空 Chunk。例如：
+
+    --pipeline_parallel=2 --virtual_pipeline_parallel=2 \
+    --pipeline_model_parallel_layout='Et*3||t*3|t*6NL'
+
+表达式必须展开为 PP*vPP 个 Chunk；E/L 必须各出现一次且位于整体首尾，N 最多一次并与 L 同属末 Chunk，
+t 数量必须等于模型层数。
+
+## 自动布局建议
+
+建议工具支持逐层参数量、用户代价和 PROFILE_MODE 记录：
+
+    scripts/suggest_pipeline_layout.py \
+      --profiler-records gpt2.records.log.rank0 \
+      --profiler-warmup-samples=1 --pipeline-parallel=2 --microbatches=4
+
+工具输出可直接复制的 pipeline_layer_partition、每 Stage 代价、均匀布局对比和理论 bubble。默认丢弃
+每层第一个 profiler 样本，避免 CUDA warmup 污染。
+
+## 启动输出
+
+主 rank 会输出规范化后的最终布局，例如：
+
+```text
+Pipeline layout (24 layers, 4 stages):
+  stage 0: embedding layers[0,4)
+  stage 1: layers[4,12)
+  stage 2: layers[12,18)
+  stage 3: layers[18,24) final_norm lm_head
+```
+
+## 错误排查
+
+- `has N entries, but --pipeline_parallel is M`：列表项数量和 PP stage 数不一致。
+- `sums to N layers, but the model has M`：列表总层数和模型配置或 checkpoint 不一致。
+- `entries must be positive integers`：存在零、负数或非整数。
+- `contains an empty stage entry`：存在连续逗号、开头逗号或末尾逗号。
+- `incompatible with --virtual_pipeline_parallel != 1`：自定义物理布局和 vPP 同时启用。
+- `must contain exactly N entries`：逐层代价数量和模型层数不一致。
+- `costs must be finite positive numbers`：逐层代价包含零、负数、NaN、无穷或非数字。
+- `cannot be used together`：同时指定了手工分区和自动均衡代价。
+
+## C++ 查询接口
+
+`PipelineLayout::layer_ranges(stage_id)` 返回该 Stage 的半开层范围；
+`stage_for_layer(layer_id)` 执行反向查询；`owns_embedding`、`owns_final_norm` 和 `owns_lm_head`
+用于特殊模块归属判断。`PipelineParallel::GetStageInfo` 是面向现有调度代码的兼容投影。
+
+## 并行组合与限制
+
+| 组合 | 默认均匀布局 | 手工层数分区 | 任意 Chunk / Megatron 布局 |
+| --- | --- | --- | --- |
+| PP | 支持 | 支持 | 支持 |
+| PP + DDP | 支持 | 支持 | 支持 |
+| PP + TP | 支持 | 支持 | 支持 |
+| PP + DDP + TP | 支持 | 支持 | 支持 |
+| vPP | 默认轮转映射 | 拒绝物理分区参数 | 支持显式 Chunk owner |
+
+`||` 表示空逻辑 Chunk；它不表示物理 Stage 没有 Chunk。当前调度器要求每个物理 Stage 拥有相同数量的
+正数 Chunk，因此会拒绝 Chunk 数不平衡的映射。
+
+提交前建议依次运行 CPU 单元测试、双 GPU E2E 和稳定多轮性能 benchmark。
+
+## 梯度一致性调试
+
+GPT-2 示例提供可选的 `--dump_gradients=DIR` 验证参数。它在第一次优化迭代后导出所有非空参数梯度，
+并把 PP rank 的局部层号转换为全局层号，使单卡与自定义 PP 输出可以直接比较：
+
+```bash
+python3 scripts/precision_check/precision_compare.py \
+  --dir1 /tmp/gpt2-grad-single \
+  --dir2 /tmp/gpt2-grad-custom \
+  --atol 1e-5 --rtol 0
+```
+
+该参数只用于正确性验证；导出会将梯度同步复制到 CPU，不应在性能测试中启用。
+
+## 端到端回归
+
+仓库提供双卡 GPT-2 E2E 脚本。它会自动运行单卡基线和两阶段代价布局，检查最终布局、fp32 loss、
+梯度文件集合以及逐参数梯度误差：
+
+~~~bash
+tests/distributed/test_pipeline_layout_e2e.sh \
+  /path/to/cuda-build \
+  data/gpt2/tiny_shakespeare_train.bin \
+  data/gpt2/gpt2_124M.bin \
+  0,1
+~~~
+
+该测试需要 CUDA/NCCL 构建、两张 GPU、NumPy 以及 GPT-2 124M LLMC checkpoint。测试使用临时目录，
+退出时自动清理梯度和日志。
diff --git a/docs/pipeline_layout_report.md b/docs/pipeline_layout_report.md
new file mode 100644
index 00000000..04e841e0
--- /dev/null
+++ b/docs/pipeline_layout_report.md
@@ -0,0 +1,93 @@
+# Pipeline Layout 实现报告
+
+## 数据结构与接口
+
+`nn::parallel::PipelineLayout` 是 Pipeline 层归属的统一数据源。它保存物理 Stage 数、模型总层数、
+每个 Stage 的半开层范围，并提供以下查询：
+
+- `layer_ranges(stage_id)`：查询本 Stage 的一个或多个连续 chunk 范围。
+- `stage_for_layer(layer_id)`：从全局 Transformer 层号反查物理 Stage。
+- `owns_embedding/final_norm/lm_head(stage_id)`：查询特殊模块归属。
+- `ToString()`：输出启动时使用的规范化布局。
+
+布局存储为 `thread_local`。InfiniTrain 支持一个进程创建多个训练线程，每个线程代表独立 global rank；
+线程本地存储可保证不同 PP rank 构建和加载时共享本线程的一份布局，同时不产生跨线程数据竞争。
+
+## 关键实现
+
+`PipelineLayout::Parse` 将 `4,8,6,6` 转换为按执行顺序连续且不重叠的范围。解析时一次性校验
+Stage 数、正整数、总层数和 vPP 兼容性。`Uniform` 封装原有默认均匀算法，并保留 vPP 的
+`global_chunk = local_chunk * pp_size + stage` 轮转语义。
+
+`PipelineLayout::FromLayerCosts` 接收每层有限正代价，通过动态规划在所有非空连续分区中最小化
+最大 Stage 总代价。状态为前 `i` 层分到 `s` 个 Stage 时的最优最大代价，转移枚举最后一个 Stage
+的起点；时间复杂度为 `O(S * L^2)`，空间复杂度为 `O(S * L)`。该方法适合模型启动阶段，结果
+确定且不改变层的执行顺序。`ResolvePipelineLayout` 统一选择手工分区、代价均衡或默认均匀布局，
+并拒绝多个布局来源同时生效。
+
+GPT-2 和 LLaMA 3 在模型配置确定后设置布局。对于 LLMC checkpoint，布局在读取 header 中真实
+`n_layer` 后解析。`TransformerModel` 用布局创建本 rank 的层和特殊模块；`TransformerConfig::GetChunkSize`
+和 `PipelineParallel` 用相同布局构造调度 Stage；两个 checkpoint loader 用布局筛选本 rank 权重。
+
+自定义物理分区当前要求 `virtual_pipeline_parallel=1`。现有调度器对 vPP 使用固定轮转
+`Chunk -> Stage` 映射，层数列表无法无歧义表达虚拟 chunk；启动时拒绝该组合比隐式产生错误执行顺序更安全。
+未配置自定义参数时仍走 `Uniform`，因此 GPipe、1F1B/vPP、TP 和 DDP 的既有入口保持不变。
+
+优秀项实现取消了 vPP 固定轮转限制：布局保存有序逻辑 Chunk 的 owner、local index 和层范围，调度器
+不再使用 global_chunk % pp_size 推导 Stage。Megatron 风格解析支持 E/t/N/L、|、重复表达式和空 Chunk，
+最终仍投影到同一 PipelineLayout 查询接口。
+
+## 正确性与测试
+
+本次验证范围如下：
+
+| 项目 | 状态 | 证据 |
+| --- | --- | --- |
+| GPT-2 / LLaMA3 PP + DDP/TP 接口接入 | 已完成 | 模型构建和 checkpoint loader 查询同一 PipelineLayout |
+| 双卡 GPT-2 PP E2E | 已实测 | H200，loss、梯度与单卡一致 |
+| 任意 vPP Chunk owner | 已实测 | H200，owner `[0,1,1,0]` 无死锁 |
+| Megatron 风格布局 | 已实测 | H200，包含空逻辑 Chunk |
+
+DDP/TP 组合保留既有 InfiniTrain 并行入口；本次新增回归重点是布局解析、PP 调度、参数加载和
+跨布局数值一致性。若提交环境要求完整 DP×TP×PP 组合矩阵，应在目标集群补跑对应资源规模的回归。
+
+CPU 单元测试覆盖 `4,8,6,6`、完整 layer-to-stage 反查、特殊模块、默认 vPP 轮转，以及错误的
+Stage 数、总和、负数、零、空项、越界查询和自定义布局/vPP 冲突。验证命令：
+
+```bash
+cmake -S . -B /tmp/infinitrain-pipeline-build \
+  -DBUILD_TEST=ON -DUSE_CUDA=OFF -DUSE_NCCL=OFF -DUSE_OMP=OFF
+cmake --build /tmp/infinitrain-pipeline-build --target test_pipeline_layout gpt2 llama3 -j2
+ctest --test-dir /tmp/infinitrain-pipeline-build -R PipelineLayoutTest --output-on-failure
+```
+
+结果：10/10 布局与建议测试通过，CPU 全量测试通过，GPT-2、LLaMA3 和 Mixtral 目标编译、
+链接通过。CUDA 13.0/NCCL 构建后，在两张 H200 上完成 GPT-2 124M 自定义 `4,8` 两阶段训练，
+两步 loss 为 `5.250158`、`4.913960`，无通信死锁。同参数单 GPU loss 完全一致；默认 `6,6` PP
+第二步 loss 为 `4.913958`，最大打印差值 `2e-6`，满足 fp32 `1e-5` 容差。逐参数梯度自动 diff
+使用规范化全局参数名比较单 GPU 和自定义 PP 的 149 个梯度；`atol=1e-5, rtol=0` 下
+149/149 通过且无缺失文件。完整命令与日志见 `docs/pipeline_layout_test_log.md`。
+仓库中的 `tests/distributed/test_pipeline_layout_e2e.sh` 将双卡启动、布局断言、loss 比较、梯度文件
+集合比较和逐参数数值比较固化为一个非零失败的自动化入口；由于依赖两张 GPU 和外部模型资产，普通
+CPU `ctest` 不会默认注册该用例。
+
+## 负载分析方法
+
+默认均匀布局只平衡层数。对已知重层或显存热点，先记录各层 forward/backward 时间或峰值显存，
+再调整每 Stage 层数，使各 Stage 总代价接近。比较时固定模型、batch、microbatch 和 dtype，分别记录
+稳定迭代的 Stage 时间、整步吞吐与峰值显存。理论 bubble 由 microbatch 数和 Stage 数主导；自定义布局
+主要通过降低最慢 Stage 的执行时间改善有效吞吐，并不改变相同调度下的 bubble step 数。
+
+例如逐层代价 `10,1,1,1,1,1` 在两个 Stage 上，默认均匀 `3,3` 的 Stage 代价为 `12,3`；
+自动布局生成 `1,5`，Stage 代价为 `10,5`，最大建模代价下降 16.7%。这是代价模型上的上界改善，
+实际吞吐还取决于通信、特殊模块、microbatch 数和运行时噪声，应使用稳定多轮 profiler 数据复测。
+
+两张 H200、GPT-2 124M、4 个 microbatch、12 个训练迭代，去掉首 3 步 warmup 后：
+
+| 布局 | 平均 step | 平均吞吐 | 较高 Stage 峰值显存 |
+| --- | ---: | ---: | ---: |
+| 默认 6,6 | 89.532 ms | 11,437 tok/s | 1473 MB |
+| Profiler 建议 7,5 | 81.799 ms | 12,519 tok/s | 1343 MB |
+
+建议布局实测吞吐提升 9.45%，峰值显存降低 130 MB；理论 bubble（4 microbatch、2 Stage）为 20%。
+Profiler 输入为 3 步 PROFILE_MODE 记录，每层丢弃一个 warmup 样本后得到 7,5，模型代价上界下降 6.84%。
diff --git a/docs/pipeline_layout_test_log.md b/docs/pipeline_layout_test_log.md
new file mode 100644
index 00000000..3f8786e6
--- /dev/null
+++ b/docs/pipeline_layout_test_log.md
@@ -0,0 +1,148 @@
+# Pipeline Layout 测试日志
+
+日期：2026-08-21（UTC）
+
+环境：GNU C++ 13.3.0；CPU build；CUDA 13.0.88、NCCL、2 x NVIDIA H200 GPU build。
+
+## 构建结果
+
+```text
+[100%] Built target test_pipeline_layout
+[100%] Built target gpt2
+[100%] Built target llama3
+[100%] Built target mixtral
+```
+
+## 单元测试结果
+
+```text
+PipelineLayoutTest.ParsesNonUniformContinuousPartition .............. Passed
+PipelineLayoutTest.AssignsSpecialModulesToPipelineEndpoints ......... Passed
+PipelineLayoutTest.PreservesUniformAndVirtualPipelineDistribution ... Passed
+PipelineLayoutTest.BalancesUserProvidedLayerCosts ................... Passed
+PipelineLayoutTest.SupportsArbitraryVirtualChunkOwnership ........... Passed
+PipelineLayoutTest.ParsesMegatronRepetitionAndEmptyChunks ........... Passed
+PipelineLayoutTest.RejectsInvalidAutomaticLayoutInputs ............. Passed
+PipelineLayoutTest.RejectsInvalidPartitions ......................... Passed
+PipelineLayoutTest.RejectsOutOfRangeQueries ......................... Passed
+PipelineLayoutSuggestionTest ........................................ Passed
+
+100% tests passed, 0 tests failed out of 10
+```
+
+自动均衡用例验证代价 `10,1,1,1,1,1` 在两个 Stage 上生成连续分区 `1,5`，并覆盖代价数量
+错误、零、负数、非数字、Stage 多于层数、vPP 冲突及与手工分区同时配置等启动错误。
+CPU 全量回归共 282 个注册测试；3 个 disabled，279 个已启用测试全部通过，其中 CPU 标签 246 个。
+
+## 2-Stage GPU 验证
+
+具备 CUDA、NCCL 和两张 GPU 的构建环境后，可使用下列方式运行非均匀 2-Stage GPT-2。模型为 12 层，
+Stage 分区为 4 层和 8 层。
+
+```bash
+./build/infini_run --nproc_per_node=2 ./build/gpt2 \
+  --device=cuda \
+  --input_bin=data/gpt2/tiny_shakespeare_train.bin \
+  --llmc_filepath=data/gpt2/gpt2_124M.bin \
+  --pipeline_parallel=2 \
+  --virtual_pipeline_parallel=1 \
+  --pipeline_layer_partition=4,8 \
+  --batch_size=4 \
+  --sequence_length=64 \
+  --total_batch_size=512 \
+  --num_iteration=2
+```
+
+实际执行完成，无通信死锁：
+
+```text
+custom PP 4,8 step 1: loss 5.250158,  970 tok/s, peak used 1247 MB
+custom PP 4,8 step 2: loss 4.913960, 8982 tok/s, peak used 1247 MB
+```
+
+使用相同 checkpoint、输入、batch、dtype 和优化参数执行单 GPU 与默认 2-Stage `6,6` 基线：
+
+```text
+single GPU step 1: loss 5.250158
+single GPU step 2: loss 4.913960
+default PP step 1: loss 5.250158
+default PP step 2: loss 4.913958
+```
+
+自定义 PP 与单 GPU 的打印 loss 完全一致；与默认 PP 的最大打印差值为 `2e-6`，满足 fp32
+`1e-5` 容差。本次短跑的稳定步吞吐为自定义 `8982 tok/s`、默认 `3410 tok/s`，说明布局能够运行且
+存在改善空间，但两次短样本不能代替隔离环境下的多轮性能统计。
+
+## 逐参数梯度一致性
+
+单 GPU 与自定义 `4,8` PP 使用相同 checkpoint、输入及训练参数运行一步，并通过
+`--dump_gradients` 导出规范化全局参数名的梯度。比较命令：
+
+```bash
+python3 scripts/precision_check/precision_compare.py \
+  --dir1 /tmp/infinitrain-grad-single-20260821 \
+  --dir2 /tmp/infinitrain-grad-custom-20260821 \
+  --atol 1e-5 --rtol 0
+```
+
+实际结果：
+
+```text
+Directory 1: 149 files
+Directory 2: 149 files
+Summary: 149 passed, 0 failed, 0 errors
+Missing: 0 in dir1 only, 0 in dir2 only
+```
+
+因此 Transformer 层、Embedding、Final Norm 和 LM Head 的全部 149 个参数梯度均满足 fp32
+绝对误差 `1e-5`。
+
+## 自动代价布局 GPU 验证
+
+GPT-2 12 层使用 `--pipeline_layer_costs=10,1,1,1,1,1,1,1,1,1,1,1` 启动双卡训练。
+程序生成并打印：
+
+```text
+Pipeline layout (12 layers, 2 stages):
+  stage 0: embedding layers[0,1)
+  stage 1: layers[1,12) final_norm lm_head
+step 1/1 | train loss 5.250158 | 789.06 ms | 649 tok/s
+```
+
+训练正常退出，无通信死锁；该 loss 与相同输入和 checkpoint 的单卡及手工 PP 首步结果一致。
+
+## 自动化 E2E 测试
+
+上述单卡和自动 PP 验证已固化为 tests/distributed/test_pipeline_layout_e2e.sh。实际运行结果：
+
+~~~text
+single GPU loss: 5.250158
+automatic PP loss: 5.250158
+Directory 1: 149 files
+Directory 2: 149 files
+Summary: 149 passed, 0 failed, 0 errors
+Missing: 0 in dir1 only, 0 in dir2 only
+PASS: automatic PP layout, loss, and gradients match the single-GPU reference
+~~~
+
+脚本同时断言自动布局为 [0,1)、[1,12)，任何训练进程失败、布局不符、loss 超过 fp32
+1e-5、梯度文件缺失或梯度数值超差都会返回非零退出码。
+
+## 任意 vPP 与 Megatron 布局 GPU 验证
+
+Chunk owner `[0,1,1,0]`（`--pipeline_chunk_layout=0:3,1:3,1:3,0:3`）在双 H200 上完成训练，
+自动布局打印为 stage 0 `[0,3)`、`[9,12)`，stage 1 `[3,6)`、`[6,9)`，loss `5.250158`。
+Megatron 表达式 `Et*3||t*3|t*6NL` 含空 Chunk，同样完成训练并得到 loss `5.250158`。
+
+## 稳定负载基准
+
+两张 H200、GPT-2 124M、4 个 microbatch、12 个迭代，去掉首 3 步：
+
+```text
+default 6,6:       mean 89.532 ms, 11,437 tok/s, peak 1473 MB
+profiler 7,5:      mean 81.799 ms, 12,519 tok/s, peak 1343 MB
+improvement:       +9.45% throughput, -130 MB peak memory
+```
+
+Profiler 建议来自 3 步单卡 PROFILE_MODE 记录，每层丢弃一个 warmup 样本；建议工具单测和实际记录
+解析均已通过。
diff --git a/example/gpt2/checkpoint_loader.cc b/example/gpt2/checkpoint_loader.cc
index 95e54730..c9ae2eca 100644
--- a/example/gpt2/checkpoint_loader.cc
+++ b/example/gpt2/checkpoint_loader.cc
@@ -1,5 +1,6 @@
 #include "example/gpt2/checkpoint_loader.h"
 
+#include <algorithm>
 #include <cmath>
 #include <cstdlib>
 #include <filesystem>
@@ -57,7 +58,11 @@ std::tuple<int32_t, infini_train::DataType> DetermineAndCheckVersion(const std::
 
 namespace gpt2 {
 
-std::shared_ptr<nn::TransformerModel> LoadFromLLMC(const std::string &filepath) {
+std::shared_ptr<nn::TransformerModel> LoadFromLLMC(const std::string &filepath,
+                                                   const std::string &pipeline_layer_partition,
+                                                   const std::string &pipeline_layer_costs,
+                                                   const std::string &pipeline_chunk_layout,
+                                                   const std::string &pipeline_model_layout) {
     if (!std::filesystem::exists(filepath)) {
         LOG(FATAL) << "File not found: " << filepath;
     }
@@ -89,6 +94,10 @@ std::shared_ptr<nn::TransformerModel> LoadFromLLMC(const std::string &filepath)
     gpt2_config.n_head = n_head;
     gpt2_config.n_embd = n_embd;
     gpt2::SanitizeGPT2Config(gpt2_config);
+    nn::parallel::SetPipelineLayout(nn::parallel::ResolvePipelineLayout(
+        n_layer, nn::parallel::global::GetPipelineParallelSize(),
+        nn::parallel::global::GetVirtualPipelineParallelSize(), pipeline_layer_partition, pipeline_layer_costs,
+        pipeline_chunk_layout, pipeline_model_layout));
     auto local_gpt2 = std::make_shared<nn::TransformerModel>(gpt2_config);
 
     LOG(INFO) << "magic: " << magic << " version: " << version << " block_size: " << block_size
@@ -99,12 +108,13 @@ std::shared_ptr<nn::TransformerModel> LoadFromLLMC(const std::string &filepath)
     CHECK_EQ(n_embd % n_head, 0) << "n_embd must be divisible by n_head.";
     CHECK_EQ(n_head % tp_size, 0) << "n_head must be divisible by TP world size.";
 
-    // ========== pp_size：num_stages; vpp_size: num_chunks_per_stage ==========
+    // Pipeline ownership comes from the same layout used to construct the model.
     int pp_size = nn::parallel::global::GetPipelineParallelSize();
-    int vpp_size = nn::parallel::global::GetVirtualPipelineParallelSize();
     auto pp_rank = nn::parallel::pp_rank;
-    auto [is_first_stage, is_last_stage, layer_ranges_per_chunk]
-        = nn::parallel::PipelineParallel::GetStageInfo(n_layer, pp_size, pp_rank, vpp_size);
+    const auto &layout = nn::parallel::GetPipelineLayout();
+    const bool is_first_stage = layout.owns_embedding(pp_rank);
+    const bool is_last_stage = layout.owns_final_norm(pp_rank) && layout.owns_lm_head(pp_rank);
+    const auto &layer_ranges_per_chunk = layout.layer_ranges(pp_rank);
     // ========== layer to chunk ==========
     std::vector<bool> owned_layers(n_layer, false);
     for (const auto &[start, end] : layer_ranges_per_chunk) {
@@ -136,6 +146,12 @@ std::shared_ptr<nn::TransformerModel> LoadFromLLMC(const std::string &filepath)
                                                               nn::parallel::VocabParallelEmbedding::kParamWeightName)];
         ReadMatrixRowShardFloat(ifs, static_cast<float *>(transformer_wte_weight->DataPtr()), model_vocab_size, n_embd,
                                 v_start, vpp);
+        if (is_last_stage && pp_size > 1) {
+            auto &lm_head_weight = state_dict[std::format("{}.{}", nn::TransformerLastStage::kLMHeadLayerName,
+                                                          nn::parallel::ColumnParallelLinear::kParamWeightName)];
+            std::copy_n(static_cast<const float *>(transformer_wte_weight->DataPtr()), vpp * n_embd,
+                        static_cast<float *>(lm_head_weight->DataPtr()));
+        }
     } else if (pp_size > 1 && is_last_stage) {
         auto &lm_head_weight = state_dict[std::format("{}.{}", nn::TransformerLastStage::kLMHeadLayerName,
                                                       nn::parallel::ColumnParallelLinear::kParamWeightName)];
diff --git a/example/gpt2/checkpoint_loader.h b/example/gpt2/checkpoint_loader.h
index e80c356e..104635b1 100644
--- a/example/gpt2/checkpoint_loader.h
+++ b/example/gpt2/checkpoint_loader.h
@@ -8,5 +8,9 @@ class TransformerModel;
 } // namespace infini_train::nn
 
 namespace gpt2 {
-std::shared_ptr<infini_train::nn::TransformerModel> LoadFromLLMC(const std::string &filepath);
+std::shared_ptr<infini_train::nn::TransformerModel> LoadFromLLMC(const std::string &filepath,
+                                                                 const std::string &pipeline_layer_partition,
+                                                                 const std::string &pipeline_layer_costs,
+                                                                 const std::string &pipeline_chunk_layout,
+                                                                 const std::string &pipeline_model_layout);
 } // namespace gpt2
diff --git a/example/gpt2/main.cc b/example/gpt2/main.cc
index 5a5cfc65..3e3eab0b 100644
--- a/example/gpt2/main.cc
+++ b/example/gpt2/main.cc
@@ -1,4 +1,5 @@
 #include <chrono>
+#include <charconv>
 #include <cstdlib>
 #include <filesystem>
 #include <format>
@@ -6,6 +7,7 @@
 #include <optional>
 #include <unordered_map>
 #include <unordered_set>
+#include <string_view>
 
 #include "gflags/gflags.h"
 #include "glog/logging.h"
@@ -85,6 +87,14 @@ DEFINE_uint32(tensor_parallel, 1, "Tensor Parallel world size");
 DEFINE_bool(sequence_parallel, false, "Whether to enable Sequence Parallel");
 DEFINE_uint32(pipeline_parallel, 1, "Pipeline Parallel world size, specified the number of PP stages.");
 DEFINE_uint32(virtual_pipeline_parallel, 1, "Number of chunks in PP stage.");
+DEFINE_string(pipeline_layer_partition, "",
+              "Comma-separated Transformer layer counts for each pipeline stage (for example: 4,8,6,6).");
+DEFINE_string(pipeline_layer_costs, "",
+              "Comma-separated positive compute costs for every Transformer layer; generates a balanced layout.");
+DEFINE_string(pipeline_chunk_layout, "", "Ordered STAGE:LAYER_COUNT chunks for an arbitrary vPP mapping.");
+DEFINE_string(pipeline_model_parallel_layout, "", "Megatron-style E/t/N/L pipeline layout expression.");
+DEFINE_string(dump_gradients, "",
+              "Directory for canonical per-parameter gradient .npy files (validation/debugging only).");
 
 // precision
 DEFINE_string(dtype, "float32", "precision used in training (float32/bfloat16)");
@@ -126,6 +136,68 @@ const std::unordered_map<std::string, nn::TransformerConfig> kModelToConfigs = {
     {"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) {
+    constexpr std::string_view wrapper_prefix = "module.";
+    if (name.starts_with(wrapper_prefix)) { name.erase(0, wrapper_prefix.size()); }
+
+    const auto &layout = nn::parallel::GetPipelineLayout();
+    const auto &ranges = layout.layer_ranges(pp_rank);
+    if (!ranges.empty()) {
+        constexpr std::string_view layer_marker = ".h.";
+        const size_t marker = name.find(layer_marker);
+        if (marker != std::string::npos) {
+            const size_t index_begin = marker + layer_marker.size();
+            const size_t index_end = name.find('.', index_begin);
+            if (index_end != std::string::npos) {
+                int local_layer = 0;
+                const std::string local_text = name.substr(index_begin, index_end - index_begin);
+                const auto [ptr, ec] = std::from_chars(local_text.data(), local_text.data() + local_text.size(),
+                                                       local_layer);
+                if (ec == std::errc() && ptr == local_text.data() + local_text.size()) {
+                    int chunk_id = 0;
+                    const size_t chunk_marker = name.rfind("__pp_chunk_", marker);
+                    if (chunk_marker != std::string::npos) {
+                        const size_t chunk_begin = chunk_marker + std::string_view("__pp_chunk_").size();
+                        const size_t chunk_end = name.find('.', chunk_begin);
+                        const std::string chunk_text = name.substr(chunk_begin, chunk_end - chunk_begin);
+                        const auto [chunk_ptr, chunk_ec]
+                            = std::from_chars(chunk_text.data(), chunk_text.data() + chunk_text.size(), chunk_id);
+                        if (chunk_ec != std::errc() || chunk_ptr != chunk_text.data() + chunk_text.size()
+                            || chunk_id < 0 || chunk_id >= static_cast<int>(ranges.size())) {
+                            return name;
+                        }
+                    }
+                    const auto [start, end] = ranges[chunk_id];
+                    if (local_layer < end - start) {
+                        name.replace(index_begin, index_end - index_begin, std::to_string(start + local_layer));
+                        if (chunk_marker != std::string::npos) {
+                            name.replace(chunk_marker + std::string_view("__pp_chunk_").size(),
+                                         name.find('.', chunk_marker + std::string_view("__pp_chunk_").size())
+                                             - (chunk_marker + std::string_view("__pp_chunk_").size()),
+                                         "0");
+                        }
+                    }
+                }
+            }
+        }
+    }
+    return name;
+}
+
+void DumpGradients(const std::shared_ptr<nn::Module> &model, int pp_rank, int step) {
+    if (FLAGS_dump_gradients.empty() || step != 0) { return; }
+    std::filesystem::create_directories(FLAGS_dump_gradients);
+    size_t count = 0;
+    for (const auto &[raw_name, parameter] : model->NamedParameters()) {
+        if (!parameter->grad()) { continue; }
+        const auto path = std::filesystem::path(FLAGS_dump_gradients)
+                        / (CanonicalGradientName(raw_name, pp_rank) + ".npy");
+        parameter->grad()->To(Device()).SaveAsNpy(path.string());
+        ++count;
+    }
+    LOG(INFO) << "PP rank " << pp_rank << ": dumped " << count << " gradients to " << FLAGS_dump_gradients;
+}
+
 } // namespace
 
 DEFINE_validator(model, [](const char *, const std::string &value) { return kSupportedModels.contains(value); });
@@ -223,12 +295,20 @@ void Train(const nn::parallel::Rank &rank) {
     std::shared_ptr<nn::Module> model = nullptr;
 
     if (!FLAGS_llmc_filepath.empty()) {
-        model = gpt2::LoadFromLLMC(FLAGS_llmc_filepath);
+        model = gpt2::LoadFromLLMC(FLAGS_llmc_filepath, FLAGS_pipeline_layer_partition, FLAGS_pipeline_layer_costs,
+                                   FLAGS_pipeline_chunk_layout, FLAGS_pipeline_model_parallel_layout);
     } else if (kModelToConfigs.count(FLAGS_model)) {
         model_config = kModelToConfigs.at(FLAGS_model);
         gpt2::SanitizeGPT2Config(model_config);
+        SetPipelineLayout(ResolvePipelineLayout(model_config.n_layer, pp_world_size, FLAGS_virtual_pipeline_parallel,
+                                                FLAGS_pipeline_layer_partition, FLAGS_pipeline_layer_costs,
+                                                FLAGS_pipeline_chunk_layout, FLAGS_pipeline_model_parallel_layout));
         model = std::make_shared<nn::TransformerModel>(model_config);
     }
+    auto local_transformer = std::dynamic_pointer_cast<nn::TransformerModel>(model);
+    CHECK(local_transformer) << "GPT2 example expects a TransformerModel.";
+    model_config = local_transformer->Config();
+    if (rank.IsMainRank()) { LOG(INFO) << GetPipelineLayout().ToString(); }
 
     model->To(device);
 
@@ -514,6 +594,7 @@ void Train(const nn::parallel::Rank &rank) {
                 scheduler->Step();
             }
         }
+        DumpGradients(model, pp_rank, step);
 
         if (ddp_world_size > 1) {
             auto lossf_tensor = std::make_shared<Tensor>(&lossf, std::vector<int64_t>{}, DataType::kFLOAT32, device);
@@ -525,7 +606,10 @@ void Train(const nn::parallel::Rank &rank) {
         const double duration_us = std::chrono::duration<double, std::micro>(iter_end - iter_start).count();
         const double tps = FLAGS_total_batch_size / (duration_us / 1e6);
 
-        if (rank.IsLastRank()) {
+        const int reporting_rank
+            = global::GetRankOf(ddp_world_size - 1, tp_world_size - 1, GetPipelineLayout().stage_for_chunk(
+                                                                         GetPipelineLayout().num_global_chunks() - 1));
+        if (rank.GlobalRank() == reporting_rank) {
             size_t used_mb = 0, reserved_mb = 0;
             std::tie(used_mb, reserved_mb) = impl->GetMemPoolPeakMB(device);
             LOG(ERROR) << std::format("step {:4d}/{} | train loss {:.6f} | lr {:.2e} | ({:.2f} ms | {:.0f} tok/s | "
diff --git a/example/llama3/checkpoint_loader.cc b/example/llama3/checkpoint_loader.cc
index f3590af6..cd3e73b4 100644
--- a/example/llama3/checkpoint_loader.cc
+++ b/example/llama3/checkpoint_loader.cc
@@ -40,7 +40,11 @@ constexpr int32_t kLLaMA3FP32Version = 3;
 
 namespace llama3 {
 
-std::shared_ptr<nn::TransformerModel> LoadFromLLMC(const std::string &filepath) {
+std::shared_ptr<nn::TransformerModel> LoadFromLLMC(const std::string &filepath,
+                                                   const std::string &pipeline_layer_partition,
+                                                   const std::string &pipeline_layer_costs,
+                                                   const std::string &pipeline_chunk_layout,
+                                                   const std::string &pipeline_model_layout) {
     if (!std::filesystem::exists(filepath)) {
         LOG(FATAL) << "File not found: " << filepath;
     }
@@ -82,14 +86,18 @@ std::shared_ptr<nn::TransformerModel> LoadFromLLMC(const std::string &filepath)
     llama3_config.norm_eps = norm_eps;
     llama3_config.max_gen_batch_size = max_gen_bs;
     llama3::SanitizeLLaMA3Config(llama3_config);
+    nn::parallel::SetPipelineLayout(nn::parallel::ResolvePipelineLayout(
+        n_layer, nn::parallel::global::GetPipelineParallelSize(),
+        nn::parallel::global::GetVirtualPipelineParallelSize(), pipeline_layer_partition, pipeline_layer_costs,
+        pipeline_chunk_layout, pipeline_model_layout));
     auto llama3 = std::make_shared<nn::TransformerModel>(llama3_config);
 
-    // ========== pp_size：num_stages; vpp_size: num_chunks_per_stage ==========
-    int pp_size = nn::parallel::global::GetPipelineParallelSize();
-    int vpp_size = nn::parallel::global::GetVirtualPipelineParallelSize();
+    // Pipeline ownership comes from the same layout used to construct the model.
     auto pp_rank = nn::parallel::pp_rank;
-    auto [is_first_stage, is_last_stage, layer_ranges_per_chunk]
-        = nn::parallel::PipelineParallel::GetStageInfo(n_layer, pp_size, pp_rank, vpp_size);
+    const auto &layout = nn::parallel::GetPipelineLayout();
+    const bool is_first_stage = layout.owns_embedding(pp_rank);
+    const bool is_last_stage = layout.owns_final_norm(pp_rank) && layout.owns_lm_head(pp_rank);
+    const auto &layer_ranges_per_chunk = layout.layer_ranges(pp_rank);
     // ========== layer to chunk ==========
     std::vector<bool> owned_layers(n_layer, false);
     for (const auto &[start, end] : layer_ranges_per_chunk) {
diff --git a/example/llama3/checkpoint_loader.h b/example/llama3/checkpoint_loader.h
index d4aea3d0..a1896280 100644
--- a/example/llama3/checkpoint_loader.h
+++ b/example/llama3/checkpoint_loader.h
@@ -8,5 +8,9 @@ class TransformerModel;
 } // namespace infini_train::nn
 
 namespace llama3 {
-std::shared_ptr<infini_train::nn::TransformerModel> LoadFromLLMC(const std::string &filepath);
+std::shared_ptr<infini_train::nn::TransformerModel> LoadFromLLMC(const std::string &filepath,
+                                                                 const std::string &pipeline_layer_partition,
+                                                                 const std::string &pipeline_layer_costs,
+                                                                 const std::string &pipeline_chunk_layout,
+                                                                 const std::string &pipeline_model_layout);
 } // namespace llama3
diff --git a/example/llama3/main.cc b/example/llama3/main.cc
index c4642cc2..a6226680 100644
--- a/example/llama3/main.cc
+++ b/example/llama3/main.cc
@@ -84,6 +84,12 @@ DEFINE_uint32(tensor_parallel, 1, "Tensor Parallel world size");
 DEFINE_bool(sequence_parallel, false, "Whether to enable Sequence Parallel");
 DEFINE_uint32(pipeline_parallel, 1, "Pipeline Parallel world size, specified the number of PP stages.");
 DEFINE_uint32(virtual_pipeline_parallel, 1, "Number of chunks in PP stage.");
+DEFINE_string(pipeline_layer_partition, "",
+              "Comma-separated Transformer layer counts for each pipeline stage (for example: 4,8,4).");
+DEFINE_string(pipeline_layer_costs, "",
+              "Comma-separated positive compute costs for every Transformer layer; generates a balanced layout.");
+DEFINE_string(pipeline_chunk_layout, "", "Ordered STAGE:LAYER_COUNT chunks for an arbitrary vPP mapping.");
+DEFINE_string(pipeline_model_parallel_layout, "", "Megatron-style E/t/N/L pipeline layout expression.");
 // precision
 DEFINE_string(dtype, "float32", "precision used in training (float32/bfloat16)");
 DEFINE_uint32(save_interval, 0, "save checkpoint every N steps; 0 disables saving");
@@ -211,11 +217,19 @@ void Train(const nn::parallel::Rank &rank) {
     nn::TransformerConfig model_config = llama3::LLaMA3Config();
     std::shared_ptr<nn::Module> model = nullptr;
     if (!FLAGS_llmc_filepath.empty()) {
-        model = llama3::LoadFromLLMC(FLAGS_llmc_filepath);
+        model = llama3::LoadFromLLMC(FLAGS_llmc_filepath, FLAGS_pipeline_layer_partition, FLAGS_pipeline_layer_costs,
+                                     FLAGS_pipeline_chunk_layout, FLAGS_pipeline_model_parallel_layout);
     } else {
         llama3::SanitizeLLaMA3Config(model_config);
+        SetPipelineLayout(ResolvePipelineLayout(model_config.n_layer, pp_world_size, FLAGS_virtual_pipeline_parallel,
+                                                FLAGS_pipeline_layer_partition, FLAGS_pipeline_layer_costs,
+                                                FLAGS_pipeline_chunk_layout, FLAGS_pipeline_model_parallel_layout));
         model = std::make_shared<nn::TransformerModel>(model_config);
     }
+    auto local_transformer = std::dynamic_pointer_cast<nn::TransformerModel>(model);
+    CHECK(local_transformer) << "LLaMA3 example expects a TransformerModel.";
+    model_config = local_transformer->Config();
+    if (rank.IsMainRank()) { LOG(INFO) << GetPipelineLayout().ToString(); }
 
     model->To(device);
 
@@ -504,7 +518,10 @@ void Train(const nn::parallel::Rank &rank) {
         const double duration_us = std::chrono::duration<double, std::micro>(iter_end - iter_start).count();
         const double tps = FLAGS_total_batch_size / (duration_us / 1e6);
 
-        if (rank.IsLastRank()) {
+        const int reporting_rank
+            = global::GetRankOf(ddp_world_size - 1, tp_world_size - 1, GetPipelineLayout().stage_for_chunk(
+                                                                         GetPipelineLayout().num_global_chunks() - 1));
+        if (rank.GlobalRank() == reporting_rank) {
             size_t used_mb = 0, reserved_mb = 0;
             std::tie(used_mb, reserved_mb) = impl->GetMemPoolPeakMB(device);
             LOG(ERROR) << std::format("step {:4d}/{} | train loss {:.6f} | lr {:.2e} | ({:.2f} ms | {:.0f} tok/s | "
diff --git a/infini_train/include/nn/modules/transformer/transformer.h b/infini_train/include/nn/modules/transformer/transformer.h
index 0471c32f..e797823d 100644
--- a/infini_train/include/nn/modules/transformer/transformer.h
+++ b/infini_train/include/nn/modules/transformer/transformer.h
@@ -49,6 +49,7 @@ class TransformerChunk : public CloneableModule<TransformerChunk> {
 
 private:
     const TransformerConfig config_;
+    const int start_layer_;
 };
 
 class TransformerLastStage : public CloneableModule<TransformerLastStage> {
diff --git a/infini_train/include/nn/parallel/pp/pipeline_parallel.h b/infini_train/include/nn/parallel/pp/pipeline_parallel.h
index 25939bdc..f43a664a 100644
--- a/infini_train/include/nn/parallel/pp/pipeline_parallel.h
+++ b/infini_train/include/nn/parallel/pp/pipeline_parallel.h
@@ -2,6 +2,8 @@
 #pragma once
 
 #include <memory>
+#include <optional>
+#include <string>
 #include <vector>
 
 #include "infini_train/include/nn/modules/module.h"
@@ -27,6 +29,52 @@ struct StageInfo {
     std::vector<std::pair<int, int>> layer_ranges_per_chunk;
 };
 
+class PipelineLayout {
+public:
+    static PipelineLayout Uniform(int total_layers, int pp_size, int chunks_per_stage = 1);
+    static PipelineLayout Parse(int total_layers, int pp_size, const std::string &partition, int chunks_per_stage = 1);
+    static PipelineLayout FromLayerCosts(int total_layers, int pp_size, const std::string &layer_costs,
+                                         int chunks_per_stage = 1);
+    static PipelineLayout FromChunkLayout(int total_layers, int pp_size, const std::string &chunk_layout);
+    static PipelineLayout FromMegatronLayout(int total_layers, int pp_size, const std::string &model_layout);
+
+    int num_stages() const { return num_stages_; }
+    int total_layers() const { return total_layers_; }
+    int chunks_per_stage() const { return chunks_per_stage_; }
+    int num_global_chunks() const { return static_cast<int>(chunk_stages_.size()); }
+    int stage_for_chunk(int global_chunk) const;
+    int local_chunk_index(int global_chunk) const;
+    const std::pair<int, int> &chunk_range(int global_chunk) const;
+    bool is_first_stage(int stage) const;
+    bool is_last_stage(int stage) const;
+    bool owns_embedding(int stage) const;
+    bool owns_final_norm(int stage) const;
+    bool owns_lm_head(int stage) const;
+    const std::vector<std::pair<int, int>> &layer_ranges(int stage) const;
+    int stage_for_layer(int layer) const;
+    std::string ToString() const;
+
+private:
+    int total_layers_ = 0;
+    int num_stages_ = 0;
+    int chunks_per_stage_ = 0;
+    int embedding_stage_ = 0;
+    int final_norm_stage_ = 0;
+    int lm_head_stage_ = 0;
+    std::vector<std::vector<std::pair<int, int>>> ranges_;
+    std::vector<int> chunk_stages_;
+    std::vector<int> chunk_local_indices_;
+    std::vector<std::pair<int, int>> chunk_ranges_;
+};
+
+PipelineLayout ResolvePipelineLayout(int total_layers, int pp_size, int chunks_per_stage,
+                                     const std::string &partition, const std::string &layer_costs,
+                                     const std::string &chunk_layout = "", const std::string &model_layout = "");
+
+void SetPipelineLayout(std::optional<PipelineLayout> layout);
+bool HasPipelineLayout();
+const PipelineLayout &GetPipelineLayout();
+
 class PipelineParallel : public Module {
 public:
     PipelineParallel(const std::shared_ptr<nn::Module> module, int num_stages, int num_micro_batches,
diff --git a/infini_train/src/nn/modules/transformer/transformer.cc b/infini_train/src/nn/modules/transformer/transformer.cc
index 99a739d2..fa154ecb 100644
--- a/infini_train/src/nn/modules/transformer/transformer.cc
+++ b/infini_train/src/nn/modules/transformer/transformer.cc
@@ -21,9 +21,25 @@
 #include "infini_train/include/nn/parallel/tensor_parallel.h"
 #include "infini_train/include/nn/parallel/utils.h"
 #include "infini_train/include/tensor.h"
+#ifdef PROFILE_MODE
+#include "infini_train/include/profiler.h"
+#endif
 
 namespace infini_train::nn {
 
+namespace {
+parallel::StageInfo ResolveTransformerStageInfo(const TransformerConfig &config) {
+    const int pp_size = parallel::global::GetPipelineParallelSize();
+    const int vpp_size = parallel::global::GetVirtualPipelineParallelSize();
+    if (!parallel::HasPipelineLayout() || parallel::GetPipelineLayout().total_layers() != config.n_layer
+        || parallel::GetPipelineLayout().num_stages() != pp_size
+        || parallel::GetPipelineLayout().chunks_per_stage() != vpp_size) {
+        parallel::SetPipelineLayout(parallel::PipelineLayout::Uniform(config.n_layer, pp_size, vpp_size));
+    }
+    return parallel::PipelineParallel::GetStageInfo(config.n_layer, pp_size, parallel::pp_rank, vpp_size);
+}
+} // namespace
+
 TransformerFirstStage::TransformerFirstStage(const TransformerConfig &config)
     : CloneableModule(kType), config_(config) {
     modules_[kWTELayerName] = std::make_shared<parallel::VocabParallelEmbedding>(
@@ -123,7 +139,7 @@ std::vector<std::shared_ptr<Tensor>> TransformerLayer::Forward(const std::vector
 }
 
 TransformerChunk::TransformerChunk(const TransformerConfig &config, int start_layer, int end_layer)
-    : CloneableModule(kType), config_(config) {
+    : CloneableModule(kType), config_(config), start_layer_(start_layer) {
     std::vector<std::shared_ptr<nn::Module>> h;
     for (int64_t i = start_layer; i < end_layer; ++i) {
         auto layer = std::make_shared<TransformerLayer>(config);
@@ -160,12 +176,32 @@ std::vector<std::shared_ptr<Tensor>> TransformerChunk::Forward(const std::vector
         std::shared_ptr<Tensor> start_pos_ptr = nullptr;
 
         // Pass RoPE parameters to each transformer block
+        int local_layer = 0;
         for (auto &h : *std::dynamic_pointer_cast<nn::ModuleList>(modules_[kHLayerName])) {
+#ifdef PROFILE_MODE
+            const std::string profile_name = "TransformerLayer." + std::to_string(start_layer_ + local_layer);
+            Profiler::Instance().StartRecord(profile_name, device.type());
+#endif
             x1 = (*h)({x1, freqs_view, start_pos_ptr, mask})[0];
+#ifdef PROFILE_MODE
+            Profiler::Instance().EndRecord(profile_name, device.type());
+#endif
+            ++local_layer;
         }
     } else if (config_.position_embedding_type == PositionEmbeddingType::kLearnedAbsolute) {
         // Learned absolute position embedding models (GPT-2 style).
-        for (auto &h : *std::dynamic_pointer_cast<nn::ModuleList>(modules_[kHLayerName])) { x1 = (*h)({x1})[0]; }
+        int local_layer = 0;
+        for (auto &h : *std::dynamic_pointer_cast<nn::ModuleList>(modules_[kHLayerName])) {
+#ifdef PROFILE_MODE
+            const std::string profile_name = "TransformerLayer." + std::to_string(start_layer_ + local_layer);
+            Profiler::Instance().StartRecord(profile_name, x1->GetDevice().type());
+#endif
+            x1 = (*h)({x1})[0];
+#ifdef PROFILE_MODE
+            Profiler::Instance().EndRecord(profile_name, x1->GetDevice().type());
+#endif
+            ++local_layer;
+        }
     } else {
         LOG(FATAL) << "Unsupported position embedding type";
     }
@@ -205,10 +241,7 @@ std::vector<std::shared_ptr<Tensor>> TransformerLastStage::Forward(const std::ve
 }
 
 TransformerModel::TransformerModel(const TransformerConfig config)
-    : CloneableModule(kType), config_(config),
-      stage_info_(nn::parallel::PipelineParallel::GetStageInfo(
-          config_.n_layer, nn::parallel::global::GetPipelineParallelSize(), nn::parallel::pp_rank,
-          nn::parallel::global::GetVirtualPipelineParallelSize())) {
+    : CloneableModule(kType), config_(config), stage_info_(ResolveTransformerStageInfo(config_)) {
     auto tp_world_size = nn::parallel::global::GetTensorParallelSize();
 
     // NOTE(zbl): VocabParallelEmbedding requires vocab_size % tp_size == 0
@@ -228,21 +261,14 @@ TransformerModel::TransformerModel(const TransformerConfig config)
     }
 
     {
-        std::map<int, std::pair<int, std::shared_ptr<TransformerChunk>>> start_layer_to_layer_size_and_chunk;
+        std::vector<std::shared_ptr<nn::Module>> h;
         for (int chunk_idx = 0; chunk_idx < stage_info_.layer_ranges_per_chunk.size(); ++chunk_idx) {
             const auto [start_layer, end_layer] = stage_info_.layer_ranges_per_chunk[chunk_idx];
             auto chunk = std::make_shared<TransformerChunk>(config_, start_layer, end_layer);
-            start_layer_to_layer_size_and_chunk[start_layer] = std::make_pair(end_layer - start_layer, chunk);
-        }
-        std::vector<std::shared_ptr<nn::Module>> h;
-        int chunk_idx = 0;
-        for (auto &[start_layer, layer_size_and_chunk] : start_layer_to_layer_size_and_chunk) {
-            auto [layer_size, chunk] = layer_size_and_chunk;
-            for (int idx = 0; idx < layer_size; ++idx) {
+            for (int idx = 0; idx < end_layer - start_layer; ++idx) {
                 h.push_back(chunk->mutable_module(TransformerChunk::kHLayerName)->mutable_module(std::to_string(idx)));
             }
             modules_[kPPChunkNamePrefix + std::to_string(chunk_idx)] = std::move(chunk);
-            ++chunk_idx;
         }
         transformer[TransformerChunk::kHLayerName] = std::make_shared<nn::ModuleList>(std::move(h));
     }
diff --git a/infini_train/src/nn/parallel/pp/pipeline_parallel.cc b/infini_train/src/nn/parallel/pp/pipeline_parallel.cc
index c0369cde..f2038b4d 100644
--- a/infini_train/src/nn/parallel/pp/pipeline_parallel.cc
+++ b/infini_train/src/nn/parallel/pp/pipeline_parallel.cc
@@ -1,9 +1,20 @@
 // pipeline_parallel.cc
 #include "infini_train/include/nn/parallel/pp/pipeline_parallel.h"
 
+#include <algorithm>
+#include <charconv>
+#include <cmath>
 #include <cstdint>
+#include <cctype>
+#include <functional>
+#include <limits>
 #include <memory>
+#include <numeric>
+#include <optional>
+#include <sstream>
+#include <stdexcept>
 #include <string>
+#include <string_view>
 
 #include "infini_train/include/nn/modules/container.h"
 #include "infini_train/include/nn/modules/module.h"
@@ -13,10 +24,426 @@
 namespace infini_train::nn::parallel {
 namespace {
 constexpr char kModuleName[] = "module";
+thread_local std::optional<PipelineLayout> pipeline_layout;
+
+void CheckStage(int stage, int num_stages) {
+    if (stage < 0 || stage >= num_stages) {
+        throw std::out_of_range("pipeline stage " + std::to_string(stage) + " is outside [0, "
+                                + std::to_string(num_stages) + ")");
+    }
+}
+
+std::string ExpandLayoutExpression(const std::string &expression) {
+    std::string compact;
+    for (char ch : expression) {
+        if (ch != ',' && !std::isspace(static_cast<unsigned char>(ch))) { compact.push_back(ch); }
+    }
+    size_t pos = 0;
+    std::function<std::string(bool)> parse_sequence = [&](bool in_group) {
+        std::string result;
+        while (pos < compact.size() && compact[pos] != ')') {
+            std::string atom;
+            if (compact[pos] == '(') {
+                ++pos;
+                atom = parse_sequence(true);
+                if (pos >= compact.size() || compact[pos] != ')') {
+                    throw std::invalid_argument("pipeline model layout has an unmatched '('");
+                }
+                ++pos;
+            } else {
+                atom.push_back(compact[pos++]);
+            }
+            int repetitions = 1;
+            if (pos < compact.size() && compact[pos] == '*') {
+                const size_t number_begin = ++pos;
+                while (pos < compact.size() && std::isdigit(static_cast<unsigned char>(compact[pos]))) { ++pos; }
+                if (number_begin == pos) {
+                    throw std::invalid_argument("pipeline model layout repetition requires a positive integer");
+                }
+                const auto [ptr, ec]
+                    = std::from_chars(compact.data() + number_begin, compact.data() + pos, repetitions);
+                if (ec != std::errc() || ptr != compact.data() + pos || repetitions <= 0) {
+                    throw std::invalid_argument("pipeline model layout repetition must be a positive integer");
+                }
+            }
+            for (int i = 0; i < repetitions; ++i) { result += atom; }
+        }
+        if (!in_group && pos < compact.size()) {
+            throw std::invalid_argument("pipeline model layout has an unmatched ')'");
+        }
+        return result;
+    };
+    const std::string expanded = parse_sequence(false);
+    if (pos != compact.size()) { throw std::invalid_argument("invalid pipeline model layout expression"); }
+    return expanded;
+}
 } // namespace
 
 thread_local int pp_rank = 0;
 
+PipelineLayout PipelineLayout::Uniform(int total_layers, int pp_size, int chunks_per_stage) {
+    if (total_layers <= 0) { throw std::invalid_argument("pipeline layout requires total_layers > 0"); }
+    if (pp_size <= 0) { throw std::invalid_argument("pipeline layout requires pp_size > 0"); }
+    if (chunks_per_stage <= 0) { throw std::invalid_argument("pipeline layout requires chunks_per_stage > 0"); }
+
+    PipelineLayout layout;
+    layout.total_layers_ = total_layers;
+    layout.num_stages_ = pp_size;
+    layout.chunks_per_stage_ = chunks_per_stage;
+    layout.ranges_.resize(pp_size);
+    const int chunks = pp_size * chunks_per_stage;
+    const int base = total_layers / chunks;
+    const int remainder = total_layers % chunks;
+    int start = 0;
+    for (int global_chunk = 0; global_chunk < chunks; ++global_chunk) {
+        const int count = base + (global_chunk < remainder ? 1 : 0);
+        const int stage = global_chunk % pp_size;
+        layout.chunk_stages_.push_back(stage);
+        layout.chunk_local_indices_.push_back(layout.ranges_[stage].size());
+        layout.chunk_ranges_.push_back({start, start + count});
+        layout.ranges_[stage].push_back({start, start + count});
+        start += count;
+    }
+    layout.embedding_stage_ = layout.chunk_stages_.front();
+    layout.final_norm_stage_ = layout.chunk_stages_.back();
+    layout.lm_head_stage_ = layout.chunk_stages_.back();
+    return layout;
+}
+
+PipelineLayout PipelineLayout::Parse(int total_layers, int pp_size, const std::string &partition,
+                                     int chunks_per_stage) {
+    if (partition.empty()) { return Uniform(total_layers, pp_size, chunks_per_stage); }
+    if (chunks_per_stage != 1) {
+        throw std::invalid_argument("--pipeline_layer_partition is incompatible with "
+                                    "--virtual_pipeline_parallel != 1");
+    }
+    if (total_layers <= 0 || pp_size <= 0) {
+        throw std::invalid_argument("pipeline layout requires positive total_layers and pp_size");
+    }
+
+    std::vector<int> counts;
+    size_t begin = 0;
+    while (begin <= partition.size()) {
+        const size_t comma = partition.find(',', begin);
+        std::string_view token(partition.data() + begin,
+                               (comma == std::string::npos ? partition.size() : comma) - begin);
+        const size_t first = token.find_first_not_of(" \t");
+        const size_t last = token.find_last_not_of(" \t");
+        if (first == std::string_view::npos) {
+            throw std::invalid_argument("pipeline layer partition contains an empty stage entry: '" + partition
+                                        + "'");
+        }
+        token = token.substr(first, last - first + 1);
+        int count = 0;
+        const auto [ptr, ec] = std::from_chars(token.data(), token.data() + token.size(), count);
+        if (ec != std::errc() || ptr != token.data() + token.size() || count <= 0) {
+            throw std::invalid_argument("pipeline layer partition entries must be positive integers; got '"
+                                        + std::string(token) + "'");
+        }
+        counts.push_back(count);
+        if (comma == std::string::npos) { break; }
+        begin = comma + 1;
+    }
+    if (static_cast<int>(counts.size()) != pp_size) {
+        throw std::invalid_argument("pipeline layer partition has " + std::to_string(counts.size())
+                                    + " entries, but --pipeline_parallel is " + std::to_string(pp_size));
+    }
+    const int sum = std::accumulate(counts.begin(), counts.end(), 0);
+    if (sum != total_layers) {
+        throw std::invalid_argument("pipeline layer partition sums to " + std::to_string(sum)
+                                    + " layers, but the model has " + std::to_string(total_layers));
+    }
+
+    PipelineLayout layout;
+    layout.total_layers_ = total_layers;
+    layout.num_stages_ = pp_size;
+    layout.chunks_per_stage_ = 1;
+    layout.ranges_.resize(pp_size);
+    int start = 0;
+    for (int stage = 0; stage < pp_size; ++stage) {
+        layout.chunk_stages_.push_back(stage);
+        layout.chunk_local_indices_.push_back(0);
+        layout.chunk_ranges_.push_back({start, start + counts[stage]});
+        layout.ranges_[stage].push_back({start, start + counts[stage]});
+        start += counts[stage];
+    }
+    layout.embedding_stage_ = 0;
+    layout.final_norm_stage_ = pp_size - 1;
+    layout.lm_head_stage_ = pp_size - 1;
+    return layout;
+}
+
+PipelineLayout PipelineLayout::FromLayerCosts(int total_layers, int pp_size, const std::string &layer_costs,
+                                              int chunks_per_stage) {
+    if (layer_costs.empty()) {
+        throw std::invalid_argument("--pipeline_layer_costs must not be empty");
+    }
+    if (chunks_per_stage != 1) {
+        throw std::invalid_argument("--pipeline_layer_costs is incompatible with "
+                                    "--virtual_pipeline_parallel != 1");
+    }
+    if (total_layers <= 0 || pp_size <= 0 || pp_size > total_layers) {
+        throw std::invalid_argument("automatic pipeline layout requires 0 < pp_size <= total_layers");
+    }
+
+    std::vector<double> costs;
+    size_t begin = 0;
+    while (begin <= layer_costs.size()) {
+        const size_t comma = layer_costs.find(',', begin);
+        std::string_view token(layer_costs.data() + begin,
+                               (comma == std::string::npos ? layer_costs.size() : comma) - begin);
+        const size_t first = token.find_first_not_of(" \t");
+        const size_t last = token.find_last_not_of(" \t");
+        if (first == std::string_view::npos) {
+            throw std::invalid_argument("pipeline layer costs contain an empty entry: '" + layer_costs + "'");
+        }
+        token = token.substr(first, last - first + 1);
+        double cost = 0.0;
+        const auto [ptr, ec] = std::from_chars(token.data(), token.data() + token.size(), cost);
+        if (ec != std::errc() || ptr != token.data() + token.size() || !std::isfinite(cost) || cost <= 0.0) {
+            throw std::invalid_argument("pipeline layer costs must be finite positive numbers; got '"
+                                        + std::string(token) + "'");
+        }
+        costs.push_back(cost);
+        if (comma == std::string::npos) { break; }
+        begin = comma + 1;
+    }
+    if (static_cast<int>(costs.size()) != total_layers) {
+        throw std::invalid_argument("pipeline layer costs have " + std::to_string(costs.size())
+                                    + " entries, but the model has " + std::to_string(total_layers) + " layers");
+    }
+
+    std::vector<double> prefix(total_layers + 1, 0.0);
+    for (int layer = 0; layer < total_layers; ++layer) {
+        prefix[layer + 1] = prefix[layer] + costs[layer];
+        if (!std::isfinite(prefix[layer + 1])) {
+            throw std::invalid_argument("pipeline layer costs have a non-finite total");
+        }
+    }
+    const double infinity = std::numeric_limits<double>::infinity();
+    std::vector<std::vector<double>> best(pp_size + 1, std::vector<double>(total_layers + 1, infinity));
+    std::vector<std::vector<int>> split(pp_size + 1, std::vector<int>(total_layers + 1, -1));
+    best[0][0] = 0.0;
+    for (int stages = 1; stages <= pp_size; ++stages) {
+        for (int end = stages; end <= total_layers; ++end) {
+            for (int start = stages - 1; start < end; ++start) {
+                const double candidate = std::max(best[stages - 1][start], prefix[end] - prefix[start]);
+                if (candidate < best[stages][end]) {
+                    best[stages][end] = candidate;
+                    split[stages][end] = start;
+                }
+            }
+        }
+    }
+
+    std::vector<int> counts(pp_size);
+    int end = total_layers;
+    for (int stage = pp_size - 1; stage >= 0; --stage) {
+        const int start = split[stage + 1][end];
+        if (start < 0) { throw std::logic_error("failed to construct automatic pipeline layout"); }
+        counts[stage] = end - start;
+        end = start;
+    }
+    std::ostringstream partition;
+    for (int stage = 0; stage < pp_size; ++stage) {
+        if (stage > 0) { partition << ','; }
+        partition << counts[stage];
+    }
+    return Parse(total_layers, pp_size, partition.str(), chunks_per_stage);
+}
+
+PipelineLayout PipelineLayout::FromChunkLayout(int total_layers, int pp_size, const std::string &chunk_layout) {
+    if (total_layers <= 0 || pp_size <= 0 || chunk_layout.empty()) {
+        throw std::invalid_argument("chunk pipeline layout requires positive layers/stages and a non-empty layout");
+    }
+    PipelineLayout layout;
+    layout.total_layers_ = total_layers;
+    layout.num_stages_ = pp_size;
+    layout.ranges_.resize(pp_size);
+    std::vector<int> chunks_per_stage(pp_size, 0);
+    int layer = 0;
+    size_t begin = 0;
+    while (begin <= chunk_layout.size()) {
+        const size_t comma = chunk_layout.find(',', begin);
+        std::string_view token(chunk_layout.data() + begin,
+                               (comma == std::string::npos ? chunk_layout.size() : comma) - begin);
+        const size_t colon = token.find(':');
+        int stage = -1;
+        int count = -1;
+        const auto stage_result
+            = colon == std::string_view::npos
+                ? std::from_chars(token.data(), token.data(), stage)
+                : std::from_chars(token.data(), token.data() + colon, stage);
+        const auto count_result
+            = colon == std::string_view::npos
+                ? std::from_chars(token.data(), token.data(), count)
+                : std::from_chars(token.data() + colon + 1, token.data() + token.size(), count);
+        if (colon == std::string_view::npos
+            || stage_result.ec != std::errc() || stage_result.ptr != token.data() + colon
+            || count_result.ec != std::errc() || count_result.ptr != token.data() + token.size()
+            || stage < 0 || stage >= pp_size || count < 0) {
+            throw std::invalid_argument("pipeline chunk layout entries must be STAGE:NON_NEGATIVE_LAYERS; got '"
+                                        + std::string(token) + "'");
+        }
+        layout.chunk_stages_.push_back(stage);
+        layout.chunk_local_indices_.push_back(chunks_per_stage[stage]++);
+        layout.chunk_ranges_.push_back({layer, layer + count});
+        layout.ranges_[stage].push_back({layer, layer + count});
+        layer += count;
+        if (comma == std::string::npos) { break; }
+        begin = comma + 1;
+    }
+    if (layer != total_layers) {
+        throw std::invalid_argument("pipeline chunk layout assigns " + std::to_string(layer)
+                                    + " layers, but the model has " + std::to_string(total_layers));
+    }
+    if (layout.chunk_stages_.empty()
+        || !std::all_of(chunks_per_stage.begin(), chunks_per_stage.end(),
+                        [&](int count) { return count == chunks_per_stage.front() && count > 0; })) {
+        throw std::invalid_argument("pipeline chunk layout must assign the same positive number of chunks to every stage");
+    }
+    layout.chunks_per_stage_ = chunks_per_stage.front();
+    layout.embedding_stage_ = layout.chunk_stages_.front();
+    layout.final_norm_stage_ = layout.chunk_stages_.back();
+    layout.lm_head_stage_ = layout.chunk_stages_.back();
+    return layout;
+}
+
+PipelineLayout PipelineLayout::FromMegatronLayout(int total_layers, int pp_size, const std::string &model_layout) {
+    const std::string expanded = ExpandLayoutExpression(model_layout);
+    std::vector<std::string> chunks(1);
+    for (char symbol : expanded) {
+        if (symbol == '|') {
+            chunks.emplace_back();
+        } else if (symbol == 'E' || symbol == 't' || symbol == 'N' || symbol == 'L') {
+            chunks.back().push_back(symbol);
+        } else {
+            throw std::invalid_argument(std::string("invalid pipeline model layout symbol '") + symbol + "'");
+        }
+    }
+    if (chunks.empty() || static_cast<int>(chunks.size()) % pp_size != 0) {
+        throw std::invalid_argument("pipeline model layout chunk count must be divisible by --pipeline_parallel");
+    }
+    std::string flattened;
+    for (const auto &chunk : chunks) { flattened += chunk; }
+    if (flattened.empty() || std::count(flattened.begin(), flattened.end(), 'E') != 1 || flattened.front() != 'E') {
+        throw std::invalid_argument("pipeline model layout must start with exactly one embedding symbol E");
+    }
+    if (std::count(flattened.begin(), flattened.end(), 'L') != 1 || flattened.back() != 'L') {
+        throw std::invalid_argument("pipeline model layout must end with exactly one LM head symbol L");
+    }
+    const int norm_count = std::count(flattened.begin(), flattened.end(), 'N');
+    if (norm_count > 1) { throw std::invalid_argument("pipeline model layout may contain at most one final norm N"); }
+    if (norm_count == 1 && chunks.back().find('N') == std::string::npos) {
+        throw std::invalid_argument("final norm N and LM head L must be in the same final logical chunk");
+    }
+    if (std::count(flattened.begin(), flattened.end(), 't') != total_layers) {
+        throw std::invalid_argument("pipeline model layout Transformer count does not match the model layer count");
+    }
+    std::ostringstream chunk_layout;
+    for (int global_chunk = 0; global_chunk < static_cast<int>(chunks.size()); ++global_chunk) {
+        if (global_chunk > 0) { chunk_layout << ','; }
+        chunk_layout << global_chunk % pp_size << ':' << std::count(chunks[global_chunk].begin(), chunks[global_chunk].end(), 't');
+    }
+    return FromChunkLayout(total_layers, pp_size, chunk_layout.str());
+}
+
+PipelineLayout ResolvePipelineLayout(int total_layers, int pp_size, int chunks_per_stage,
+                                     const std::string &partition, const std::string &layer_costs,
+                                     const std::string &chunk_layout, const std::string &model_layout) {
+    const int configured = !partition.empty() + !layer_costs.empty() + !chunk_layout.empty() + !model_layout.empty();
+    if (configured > 1) {
+        throw std::invalid_argument("pipeline layout options are mutually exclusive");
+    }
+    PipelineLayout layout;
+    if (!chunk_layout.empty()) { layout = PipelineLayout::FromChunkLayout(total_layers, pp_size, chunk_layout); }
+    else if (!model_layout.empty()) { layout = PipelineLayout::FromMegatronLayout(total_layers, pp_size, model_layout); }
+    else if (!layer_costs.empty()) {
+        return PipelineLayout::FromLayerCosts(total_layers, pp_size, layer_costs, chunks_per_stage);
+    } else {
+        return PipelineLayout::Parse(total_layers, pp_size, partition, chunks_per_stage);
+    }
+    if (layout.chunks_per_stage() != chunks_per_stage) {
+        throw std::invalid_argument("custom chunk layout requires --virtual_pipeline_parallel="
+                                    + std::to_string(layout.chunks_per_stage()));
+    }
+    return layout;
+}
+
+bool PipelineLayout::is_first_stage(int stage) const {
+    CheckStage(stage, num_stages_);
+    return stage == 0;
+}
+bool PipelineLayout::is_last_stage(int stage) const {
+    CheckStage(stage, num_stages_);
+    return stage == num_stages_ - 1;
+}
+bool PipelineLayout::owns_embedding(int stage) const {
+    CheckStage(stage, num_stages_);
+    return stage == embedding_stage_;
+}
+bool PipelineLayout::owns_final_norm(int stage) const {
+    CheckStage(stage, num_stages_);
+    return stage == final_norm_stage_;
+}
+bool PipelineLayout::owns_lm_head(int stage) const {
+    CheckStage(stage, num_stages_);
+    return stage == lm_head_stage_;
+}
+int PipelineLayout::stage_for_chunk(int global_chunk) const {
+    if (global_chunk < 0 || global_chunk >= num_global_chunks()) {
+        throw std::out_of_range("pipeline global chunk is out of range");
+    }
+    return chunk_stages_[global_chunk];
+}
+int PipelineLayout::local_chunk_index(int global_chunk) const {
+    if (global_chunk < 0 || global_chunk >= num_global_chunks()) {
+        throw std::out_of_range("pipeline global chunk is out of range");
+    }
+    return chunk_local_indices_[global_chunk];
+}
+const std::pair<int, int> &PipelineLayout::chunk_range(int global_chunk) const {
+    if (global_chunk < 0 || global_chunk >= num_global_chunks()) {
+        throw std::out_of_range("pipeline global chunk is out of range");
+    }
+    return chunk_ranges_[global_chunk];
+}
+const std::vector<std::pair<int, int>> &PipelineLayout::layer_ranges(int stage) const {
+    CheckStage(stage, num_stages_);
+    return ranges_[stage];
+}
+int PipelineLayout::stage_for_layer(int layer) const {
+    if (layer < 0 || layer >= total_layers_) {
+        throw std::out_of_range("transformer layer " + std::to_string(layer) + " is outside [0, "
+                                + std::to_string(total_layers_) + ")");
+    }
+    for (int stage = 0; stage < num_stages_; ++stage) {
+        for (const auto &[start, end] : ranges_[stage]) {
+            if (layer >= start && layer < end) { return stage; }
+        }
+    }
+    throw std::logic_error("pipeline layout does not own transformer layer " + std::to_string(layer));
+}
+std::string PipelineLayout::ToString() const {
+    std::ostringstream out;
+    out << "Pipeline layout (" << total_layers_ << " layers, " << num_stages_ << " stages):";
+    for (int stage = 0; stage < num_stages_; ++stage) {
+        out << "\n  stage " << stage << ":";
+        if (owns_embedding(stage)) { out << " embedding"; }
+        for (const auto &[start, end] : ranges_[stage]) { out << " layers[" << start << "," << end << ")"; }
+        if (owns_final_norm(stage)) { out << " final_norm"; }
+        if (owns_lm_head(stage)) { out << " lm_head"; }
+    }
+    return out.str();
+}
+
+void SetPipelineLayout(std::optional<PipelineLayout> layout) { pipeline_layout = std::move(layout); }
+bool HasPipelineLayout() { return pipeline_layout.has_value(); }
+const PipelineLayout &GetPipelineLayout() {
+    if (!pipeline_layout) { throw std::logic_error("pipeline layout has not been initialized"); }
+    return *pipeline_layout;
+}
+
 void PipelineParallel::BuildPipelineStage(const std::vector<std::vector<int64_t>> &recv_shape, Device device,
                                           std::vector<std::shared_ptr<Module>> &&chunks) {
     pipeline_stage_ = std::make_shared<PipelineStage>(rank_, num_stages_, recv_shape, device, std::move(chunks));
@@ -32,7 +459,7 @@ float PipelineParallel::TrainStep(const std::vector<std::shared_ptr<Tensor>> &in
                                   DataType dtype) {
     std::shared_ptr<Tensor> stage_input;
     std::shared_ptr<Tensor> stage_target = target[0];
-    if (rank_ == 0) {
+    if (GetPipelineLayout().owns_embedding(rank_)) {
         stage_input = input[0];
     }
 
@@ -40,40 +467,18 @@ float PipelineParallel::TrainStep(const std::vector<std::shared_ptr<Tensor>> &in
 }
 
 StageInfo PipelineParallel::GetStageInfo(int total_layers, int pp_size, int rank, int chunks_per_stage) {
-    bool is_first_stage = (rank == 0);
-    bool is_last_stage = (rank == pp_size - 1);
-
-    std::vector<std::pair<int, int>> layer_ranges_per_chunk;
-
-    int layers_per_chunk = total_layers / (pp_size * chunks_per_stage);
-    int remainder = total_layers % (pp_size * chunks_per_stage);
-
-    for (int local_chunk_idx = 0; local_chunk_idx < chunks_per_stage; ++local_chunk_idx) {
-        int global_chunk_idx = local_chunk_idx * pp_size + rank;
-
-        if (global_chunk_idx * layers_per_chunk >= total_layers) {
-            break;
-        }
-
-        int chunk_start = global_chunk_idx * layers_per_chunk;
-        int chunk_end = chunk_start + layers_per_chunk;
-
-        if (global_chunk_idx < remainder) {
-            // Assign an additional layer to each of the first remainder chunks
-            chunk_start = global_chunk_idx * (layers_per_chunk + 1);
-            chunk_end = chunk_start + (layers_per_chunk + 1);
-        } else {
-            chunk_start = remainder * (layers_per_chunk + 1) + (global_chunk_idx - remainder) * layers_per_chunk;
-            chunk_end = chunk_start + layers_per_chunk;
-        }
-
-        chunk_end = std::min(chunk_end, total_layers);
-        if (chunk_start < chunk_end) {
-            layer_ranges_per_chunk.push_back({chunk_start, chunk_end});
-        }
+    const PipelineLayout *layout = nullptr;
+    PipelineLayout fallback;
+    if (pipeline_layout && pipeline_layout->total_layers() == total_layers
+        && pipeline_layout->num_stages() == pp_size
+        && pipeline_layout->chunks_per_stage() == chunks_per_stage) {
+        layout = &*pipeline_layout;
+    } else {
+        fallback = PipelineLayout::Uniform(total_layers, pp_size, chunks_per_stage);
+        layout = &fallback;
     }
-
-    return {is_first_stage, is_last_stage, layer_ranges_per_chunk};
+    return {layout->owns_embedding(rank), layout->owns_final_norm(rank) && layout->owns_lm_head(rank),
+            layout->layer_ranges(rank)};
 }
 
 PipelineParallel::PipelineParallel(const std::shared_ptr<Module> module, int num_stages, int num_micro_batches,
@@ -83,16 +488,16 @@ PipelineParallel::PipelineParallel(const std::shared_ptr<Module> module, int num
     modules_[kModuleName] = std::move(module);
 
     int stage_id = pp_rank;
-    int stage_size = num_stages;
+    const auto &layout = GetPipelineLayout();
 
     std::vector<std::shared_ptr<Module>> chunks;
     for (int chunk_id = 0; chunk_id < chunk_size; ++chunk_id) {
         std::vector<std::shared_ptr<Module>> chunk_parts;
-        if (chunk_id == 0 && stage_id == 0) {
+        if (chunk_id == 0 && layout.owns_embedding(stage_id)) {
             chunk_parts.push_back(module->mutable_module(kPPFirstStageName));
         }
         chunk_parts.push_back(module->mutable_module(kPPChunkNamePrefix + std::to_string(chunk_id)));
-        if (chunk_id == chunk_size - 1 && stage_id == stage_size - 1) {
+        if (chunk_id == chunk_size - 1 && layout.owns_final_norm(stage_id) && layout.owns_lm_head(stage_id)) {
             chunk_parts.push_back(module->mutable_module(kPPLastStageName));
         }
         chunks.push_back(std::make_shared<Sequential>(std::move(chunk_parts)));
diff --git a/infini_train/src/nn/parallel/pp/pipeline_schedule.cc b/infini_train/src/nn/parallel/pp/pipeline_schedule.cc
index c7c1d16f..0c79d589 100644
--- a/infini_train/src/nn/parallel/pp/pipeline_schedule.cc
+++ b/infini_train/src/nn/parallel/pp/pipeline_schedule.cc
@@ -13,6 +13,7 @@
 #include "infini_train/include/nn/init.h"
 #include "infini_train/include/nn/modules/module.h"
 #include "infini_train/include/nn/parallel/global.h"
+#include "infini_train/include/nn/parallel/pp/pipeline_parallel.h"
 #include "infini_train/include/nn/parallel/pp/pipeline_stage.h"
 #include "infini_train/include/nn/parallel/pp/send_recv.h"
 #include "infini_train/include/optimizer.h"
@@ -32,13 +33,10 @@ void PrintScheduleTable(const std::vector<PipelineParallelScheduler::Task> &sche
     LOG(INFO) << "-----|-----------|------------|--------------|-------------|-------";
 
     for (const auto &task : schedule) {
-        int owning_stage = task.global_chunk_id % num_stages;
-        int local_chunk = task.global_chunk_id / num_stages;
-
         std::string type_str = task.is_forward ? "Forward" : "Backward";
 
         auto s_info = std::format("{:4} | {:<9} | {:>10} | {:>12} | {:>11} | {:>5}", task.step, type_str,
-                                  task.microbatch_id, task.global_chunk_id, local_chunk, owning_stage);
+                                  task.microbatch_id, task.global_chunk_id, task.local_chunk_idx, task.stage_id);
         LOG(INFO) << s_info;
     }
 }
@@ -75,9 +73,10 @@ PipelineParallelScheduler::Task PipelineParallelScheduler::CreateTask(int step,
     task.step = step;
     task.microbatch_id = mb;
     task.global_chunk_id = global_chunk;
-    task.local_chunk_idx = global_chunk / num_stages;
+    const auto &layout = GetPipelineLayout();
+    task.local_chunk_idx = layout.local_chunk_index(global_chunk);
     task.is_forward = is_forward;
-    task.stage_id = global_chunk % num_stages;
+    task.stage_id = layout.stage_for_chunk(global_chunk);
     task.is_last_chunk = (global_chunk == total_chunks - 1);
     task.is_first_chunk = (global_chunk == 0);
     return task;
@@ -86,7 +85,7 @@ PipelineParallelScheduler::Task PipelineParallelScheduler::CreateTask(int step,
 std::vector<PipelineParallelScheduler::Task> PipelineParallelScheduler::GenerateGPipeSchedule(int n, int num_stages,
                                                                                               int vpp_size) {
     std::vector<Task> schedule;
-    int total_global_chunks = num_stages * vpp_size;
+    int total_global_chunks = GetPipelineLayout().num_global_chunks();
     int total_steps = n + total_global_chunks - 1;
 
     // ======== Forward Pass ========
@@ -134,7 +133,7 @@ PipelineParallelScheduler::GenerateInterleaved1F1BSchedule(int n, int num_stages
         return schedule;
     }
 
-    int total_global_chunks = num_stages * vpp_size;
+    int total_global_chunks = GetPipelineLayout().num_global_chunks();
 
     int warmup_steps = total_global_chunks - 1;
     int total_steps = 2 * warmup_steps + n;
@@ -197,7 +196,8 @@ float PipelineSchedule::StepMicroBatches(const std::vector<std::shared_ptr<Tenso
     int n = num_micro_batches_;
     int num_stages = stage_->num_stages();
     int stage_idx = stage_->stage_index();
-    int vpp_size = global::GetVirtualPipelineParallelSize();
+    const auto &layout = GetPipelineLayout();
+    int vpp_size = layout.chunks_per_stage();
 
     auto schedule = PipelineParallelScheduler::GenerateGPipeSchedule(n, num_stages, vpp_size);
 
@@ -227,20 +227,21 @@ float PipelineSchedule::StepMicroBatches(const std::vector<std::shared_ptr<Tenso
             if (task.is_first_chunk) {
                 inputs = {microbatch_inputs[mb]};
             } else {
-                if (stage_->IsFirstStage()) {
-                    inputs = ReceiveFromPrev(num_stages - 1);
+                const int previous_stage = layout.stage_for_chunk(task.global_chunk_id - 1);
+                if (previous_stage == stage_idx) {
+                    const int previous_local_chunk = layout.local_chunk_index(task.global_chunk_id - 1);
+                    inputs = activations[previous_local_chunk][mb];
                 } else {
-                    inputs = ReceiveFromPrev(stage_->prev_rank());
+                    inputs = ReceiveFromPrev(previous_stage);
                 }
             }
 
             activations[task.local_chunk_idx][mb] = stage_->ForwardOneChunk(inputs, task.local_chunk_idx);
 
             if (!task.is_last_chunk) {
-                if (stage_->IsLastStage()) {
-                    SendToNext(activations[task.local_chunk_idx][mb], 0);
-                } else {
-                    SendToNext(activations[task.local_chunk_idx][mb], stage_->next_rank());
+                const int next_stage = layout.stage_for_chunk(task.global_chunk_id + 1);
+                if (next_stage != stage_idx) {
+                    SendToNext(activations[task.local_chunk_idx][mb], next_stage);
                 }
             }
         } else {
@@ -260,12 +261,13 @@ float PipelineSchedule::StepMicroBatches(const std::vector<std::shared_ptr<Tenso
                 // between forward and backward.
                 total_loss += static_cast<const float *>(loss->To(Device()).DataPtr())[0];
             } else {
-                auto out_tensor = activations[task.local_chunk_idx][mb][0];
-
-                auto dummy_gradient
-                    = std::make_shared<Tensor>(out_tensor->Dims(), out_tensor->Dtype(), out_tensor->GetDevice());
-
-                out_tensor->Backward(dummy_gradient);
+                const int next_stage = layout.stage_for_chunk(task.global_chunk_id + 1);
+                if (next_stage != stage_idx) {
+                    auto out_tensor = activations[task.local_chunk_idx][mb][0];
+                    auto dummy_gradient
+                        = std::make_shared<Tensor>(out_tensor->Dims(), out_tensor->Dtype(), out_tensor->GetDevice());
+                    out_tensor->Backward(dummy_gradient);
+                }
             }
         }
     }
@@ -278,11 +280,12 @@ float PipelineSchedule::Step(std::shared_ptr<Tensor> input, std::shared_ptr<Tens
                              DataType dtype) {
     std::vector<std::shared_ptr<Tensor>> micro_batches(num_micro_batches_);
     std::vector<std::shared_ptr<Tensor>> target_mbs(num_micro_batches_);
-    if (stage_->IsFirstStage()) {
+    const auto &layout = GetPipelineLayout();
+    if (layout.owns_embedding(stage_->stage_index())) {
         micro_batches = input->Split(input->Dims()[0] / num_micro_batches_);
     }
 
-    if (stage_->IsLastStage()) {
+    if (layout.owns_lm_head(stage_->stage_index())) {
         target_mbs = target->Split(target->Dims()[0] / num_micro_batches_);
     }
 
diff --git a/scripts/suggest_pipeline_layout.py b/scripts/suggest_pipeline_layout.py
new file mode 100755
index 00000000..fe6e61fb
--- /dev/null
+++ b/scripts/suggest_pipeline_layout.py
@@ -0,0 +1,125 @@
+#!/usr/bin/env python3
+"""Generate a contiguous Pipeline partition from layer costs or profiler records."""
+
+import argparse
+import json
+import math
+import re
+from pathlib import Path
+
+
+LAYER_RECORD = re.compile(
+    r"TransformerLayer\.(\d+)\s+(?:Device\([^)]*\)|\S+)\s*(\d+)\s+(\d+)\s+\d+\s*$"
+)
+
+
+def parse_numbers(value: str) -> list[float]:
+    path = Path(value)
+    text = path.read_text(encoding="utf-8") if path.is_file() else value
+    try:
+        parsed = json.loads(text)
+        values = parsed if isinstance(parsed, list) else parsed["layer_costs"]
+    except (json.JSONDecodeError, KeyError, TypeError):
+        values = [item.strip() for item in text.strip().split(",")]
+    costs = [float(item) for item in values]
+    if not costs or any(not math.isfinite(cost) or cost <= 0 for cost in costs):
+        raise ValueError("layer costs must be finite positive numbers")
+    return costs
+
+
+def parse_profiler_records(paths: list[str], warmup_samples: int = 0) -> list[float]:
+    samples: dict[int, list[float]] = {}
+    for value in paths:
+        candidates = sorted(Path().glob(value)) if any(ch in value for ch in "*?[") else [Path(value)]
+        for path in candidates:
+            for line in path.read_text(encoding="utf-8").splitlines():
+                match = LAYER_RECORD.search(line)
+                if match:
+                    layer, host_us, device_us = map(int, match.groups())
+                    samples.setdefault(layer, []).append(float(device_us or host_us))
+    if not samples or sorted(samples) != list(range(max(samples) + 1)):
+        raise ValueError("profiler records must contain contiguous TransformerLayer.0..N samples")
+    if warmup_samples < 0 or any(len(values) <= warmup_samples for values in samples.values()):
+        raise ValueError("profiler warmup samples must leave at least one sample per layer")
+    return [
+        sum(samples[layer][warmup_samples:]) / len(samples[layer][warmup_samples:])
+        for layer in range(len(samples))
+    ]
+
+
+def balanced_partition(costs: list[float], stages: int) -> tuple[list[int], list[float]]:
+    layers = len(costs)
+    if stages <= 0 or stages > layers:
+        raise ValueError("stages must satisfy 0 < stages <= number of layers")
+    prefix = [0.0]
+    for cost in costs:
+        prefix.append(prefix[-1] + cost)
+    best = [[math.inf] * (layers + 1) for _ in range(stages + 1)]
+    split = [[-1] * (layers + 1) for _ in range(stages + 1)]
+    best[0][0] = 0.0
+    for stage_count in range(1, stages + 1):
+        for end in range(stage_count, layers + 1):
+            for start in range(stage_count - 1, end):
+                candidate = max(best[stage_count - 1][start], prefix[end] - prefix[start])
+                if candidate < best[stage_count][end]:
+                    best[stage_count][end] = candidate
+                    split[stage_count][end] = start
+    counts = [0] * stages
+    end = layers
+    for stage in range(stages - 1, -1, -1):
+        start = split[stage + 1][end]
+        counts[stage] = end - start
+        end = start
+    stage_costs = []
+    start = 0
+    for count in counts:
+        stage_costs.append(sum(costs[start : start + count]))
+        start += count
+    return counts, stage_costs
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    sources = parser.add_mutually_exclusive_group(required=True)
+    sources.add_argument("--costs", help="CSV/JSON layer costs or a file containing them")
+    sources.add_argument("--parameter-counts", help="CSV/JSON per-layer parameter counts or a file")
+    sources.add_argument("--profiler-records", nargs="+", help="Profiler record files or glob patterns")
+    parser.add_argument("--profiler-warmup-samples", type=int, default=1)
+    parser.add_argument("--pipeline-parallel", type=int, required=True)
+    parser.add_argument("--microbatches", type=int, default=1)
+    parser.add_argument("--json-output", type=Path)
+    args = parser.parse_args()
+
+    costs = (
+        parse_profiler_records(args.profiler_records, args.profiler_warmup_samples)
+        if args.profiler_records
+        else parse_numbers(args.costs or args.parameter_counts)
+    )
+    counts, stage_costs = balanced_partition(costs, args.pipeline_parallel)
+    uniform_counts = [len(costs) // args.pipeline_parallel] * args.pipeline_parallel
+    for stage in range(len(costs) % args.pipeline_parallel):
+        uniform_counts[stage] += 1
+    uniform_costs = []
+    offset = 0
+    for count in uniform_counts:
+        uniform_costs.append(sum(costs[offset : offset + count]))
+        offset += count
+    bubble = (args.pipeline_parallel - 1) / (args.microbatches + args.pipeline_parallel - 1)
+    result = {
+        "partition": counts,
+        "stage_costs": stage_costs,
+        "maximum_stage_cost": max(stage_costs),
+        "uniform_partition": uniform_counts,
+        "uniform_stage_costs": uniform_costs,
+        "uniform_maximum_stage_cost": max(uniform_costs),
+        "modeled_maximum_improvement_percent": 100 * (1 - max(stage_costs) / max(uniform_costs)),
+        "theoretical_pipeline_bubble_percent": 100 * bubble,
+    }
+    print("--pipeline_layer_partition=" + ",".join(map(str, counts)))
+    print(json.dumps(result, indent=2))
+    if args.json_output:
+        args.json_output.write_text(json.dumps(result, indent=2) + "\n", encoding="utf-8")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tests/distributed/CMakeLists.txt b/tests/distributed/CMakeLists.txt
index b8ed4970..7940b295 100644
--- a/tests/distributed/CMakeLists.txt
+++ b/tests/distributed/CMakeLists.txt
@@ -7,6 +7,11 @@ infini_train_add_test(test_rank
   LABELS cpu
 )
 
+infini_train_add_test(test_pipeline_layout
+  SOURCES test_pipeline_layout.cc
+  LABELS cpu
+)
+
 add_test(
   NAME RankTest.MultiNodeSingleProcessIsParallel
   COMMAND ${CMAKE_COMMAND} -E env
@@ -23,3 +28,12 @@ set_tests_properties(RankTest.MultiNodeSingleProcessIsParallel
     LABELS cpu
     TIMEOUT 10
 )
+
+find_package(Python3 COMPONENTS Interpreter QUIET)
+if(Python3_Interpreter_FOUND)
+  add_test(
+    NAME PipelineLayoutSuggestionTest
+    COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/test_pipeline_layout_suggestion.py
+  )
+  set_tests_properties(PipelineLayoutSuggestionTest PROPERTIES LABELS cpu TIMEOUT 10)
+endif()
diff --git a/tests/distributed/test_pipeline_layout.cc b/tests/distributed/test_pipeline_layout.cc
new file mode 100644
index 00000000..dbec4b3a
--- /dev/null
+++ b/tests/distributed/test_pipeline_layout.cc
@@ -0,0 +1,126 @@
+#include <stdexcept>
+#include <string>
+#include <vector>
+
+#include "gtest/gtest.h"
+
+#include "infini_train/include/nn/parallel/pp/pipeline_parallel.h"
+#include "infini_train/include/nn/parallel/pp/pipeline_schedule.h"
+
+namespace infini_train::nn::parallel {
+namespace {
+
+TEST(PipelineLayoutTest, ParsesNonUniformContinuousPartition) {
+    const auto layout = PipelineLayout::Parse(24, 4, "4, 8,6,6");
+
+    EXPECT_EQ(layout.num_stages(), 4);
+    EXPECT_EQ(layout.total_layers(), 24);
+    EXPECT_EQ(layout.layer_ranges(0), (std::vector<std::pair<int, int>>{{0, 4}}));
+    EXPECT_EQ(layout.layer_ranges(1), (std::vector<std::pair<int, int>>{{4, 12}}));
+    EXPECT_EQ(layout.layer_ranges(2), (std::vector<std::pair<int, int>>{{12, 18}}));
+    EXPECT_EQ(layout.layer_ranges(3), (std::vector<std::pair<int, int>>{{18, 24}}));
+
+    for (int layer = 0; layer < 24; ++layer) {
+        const int expected_stage = layer < 4 ? 0 : layer < 12 ? 1 : layer < 18 ? 2 : 3;
+        EXPECT_EQ(layout.stage_for_layer(layer), expected_stage);
+    }
+}
+
+TEST(PipelineLayoutTest, AssignsSpecialModulesToPipelineEndpoints) {
+    const auto layout = PipelineLayout::Parse(6, 2, "2,4");
+
+    EXPECT_TRUE(layout.owns_embedding(0));
+    EXPECT_FALSE(layout.owns_final_norm(0));
+    EXPECT_FALSE(layout.owns_lm_head(0));
+    EXPECT_FALSE(layout.owns_embedding(1));
+    EXPECT_TRUE(layout.owns_final_norm(1));
+    EXPECT_TRUE(layout.owns_lm_head(1));
+    EXPECT_NE(layout.ToString().find("stage 0: embedding layers[0,2)"), std::string::npos);
+    EXPECT_NE(layout.ToString().find("stage 1: layers[2,6) final_norm lm_head"), std::string::npos);
+}
+
+TEST(PipelineLayoutTest, PreservesUniformAndVirtualPipelineDistribution) {
+    const auto layout = PipelineLayout::Uniform(10, 2, 2);
+
+    EXPECT_EQ(layout.chunks_per_stage(), 2);
+    EXPECT_EQ(layout.layer_ranges(0), (std::vector<std::pair<int, int>>{{0, 3}, {6, 8}}));
+    EXPECT_EQ(layout.layer_ranges(1), (std::vector<std::pair<int, int>>{{3, 6}, {8, 10}}));
+}
+
+TEST(PipelineLayoutTest, BalancesUserProvidedLayerCosts) {
+    const auto layout = PipelineLayout::FromLayerCosts(6, 2, "10,1,1,1,1,1");
+
+    EXPECT_EQ(layout.layer_ranges(0), (std::vector<std::pair<int, int>>{{0, 1}}));
+    EXPECT_EQ(layout.layer_ranges(1), (std::vector<std::pair<int, int>>{{1, 6}}));
+    EXPECT_EQ(layout.stage_for_layer(0), 0);
+    EXPECT_EQ(layout.stage_for_layer(5), 1);
+}
+
+TEST(PipelineLayoutTest, SupportsArbitraryVirtualChunkOwnership) {
+    const auto layout = PipelineLayout::FromChunkLayout(8, 2, "0:2,1:2,1:2,0:2");
+
+    EXPECT_EQ(layout.chunks_per_stage(), 2);
+    EXPECT_EQ(layout.num_global_chunks(), 4);
+    EXPECT_EQ(layout.stage_for_chunk(2), 1);
+    EXPECT_EQ(layout.local_chunk_index(2), 1);
+    EXPECT_EQ(layout.stage_for_chunk(3), 0);
+    EXPECT_EQ(layout.local_chunk_index(3), 1);
+    EXPECT_TRUE(layout.owns_embedding(0));
+    EXPECT_TRUE(layout.owns_final_norm(0));
+    EXPECT_EQ(layout.layer_ranges(0), (std::vector<std::pair<int, int>>{{0, 2}, {6, 8}}));
+    EXPECT_EQ(layout.layer_ranges(1), (std::vector<std::pair<int, int>>{{2, 4}, {4, 6}}));
+
+    SetPipelineLayout(layout);
+    const auto task = PipelineParallelScheduler::CreateTask(3, 0, 2, 2, 4, true);
+    EXPECT_EQ(task.stage_id, 1);
+    EXPECT_EQ(task.local_chunk_idx, 1);
+    SetPipelineLayout(std::nullopt);
+}
+
+TEST(PipelineLayoutTest, ParsesMegatronRepetitionAndEmptyChunks) {
+    const auto layout = PipelineLayout::FromMegatronLayout(8, 2, "Et*2||t*2|t*4NL");
+
+    EXPECT_EQ(layout.chunks_per_stage(), 2);
+    EXPECT_EQ(layout.num_global_chunks(), 4);
+    EXPECT_EQ(layout.chunk_range(0), (std::pair<int, int>{0, 2}));
+    EXPECT_EQ(layout.chunk_range(1), (std::pair<int, int>{2, 2}));
+    EXPECT_EQ(layout.chunk_range(2), (std::pair<int, int>{2, 4}));
+    EXPECT_EQ(layout.chunk_range(3), (std::pair<int, int>{4, 8}));
+}
+
+TEST(PipelineLayoutTest, RejectsInvalidAutomaticLayoutInputs) {
+    EXPECT_THROW(PipelineLayout::FromLayerCosts(6, 2, "1,2,3"), std::invalid_argument);
+    EXPECT_THROW(PipelineLayout::FromLayerCosts(3, 2, "1,0,2"), std::invalid_argument);
+    EXPECT_THROW(PipelineLayout::FromLayerCosts(3, 2, "1,-1,2"), std::invalid_argument);
+    EXPECT_THROW(PipelineLayout::FromLayerCosts(3, 2, "1,nope,2"), std::invalid_argument);
+    EXPECT_THROW(PipelineLayout::FromLayerCosts(3, 2, "1,nan,2"), std::invalid_argument);
+    EXPECT_THROW(PipelineLayout::FromLayerCosts(3, 2, "1,inf,2"), std::invalid_argument);
+    EXPECT_THROW(PipelineLayout::FromLayerCosts(2, 2, "1.7e308,1.7e308"), std::invalid_argument);
+    EXPECT_THROW(PipelineLayout::FromLayerCosts(3, 4, "1,1,1"), std::invalid_argument);
+    EXPECT_THROW(PipelineLayout::FromLayerCosts(3, 2, "1,1,1", 2), std::invalid_argument);
+    EXPECT_THROW(ResolvePipelineLayout(3, 2, 1, "1,2", "1,1,1"), std::invalid_argument);
+    EXPECT_THROW(PipelineLayout::FromChunkLayout(4, 2, "0:2,1:1"), std::invalid_argument);
+    EXPECT_THROW(PipelineLayout::FromChunkLayout(4, 2, "0:2,1:2,0:0"), std::invalid_argument);
+    EXPECT_THROW(PipelineLayout::FromChunkLayout(4, 2, "0x:2,1:2"), std::invalid_argument);
+    EXPECT_THROW(PipelineLayout::FromMegatronLayout(4, 2, ""), std::invalid_argument);
+    EXPECT_THROW(PipelineLayout::FromMegatronLayout(4, 2, "Et*3|t*2L"), std::invalid_argument);
+    EXPECT_THROW(PipelineLayout::FromMegatronLayout(4, 2, "Et*2|t*2N|L"), std::invalid_argument);
+}
+
+TEST(PipelineLayoutTest, RejectsInvalidPartitions) {
+    EXPECT_THROW(PipelineLayout::Parse(24, 4, "4,8,6"), std::invalid_argument);
+    EXPECT_THROW(PipelineLayout::Parse(24, 4, "4,8,6,5"), std::invalid_argument);
+    EXPECT_THROW(PipelineLayout::Parse(24, 4, "4,-8,12,16"), std::invalid_argument);
+    EXPECT_THROW(PipelineLayout::Parse(24, 4, "4,0,8,12"), std::invalid_argument);
+    EXPECT_THROW(PipelineLayout::Parse(24, 4, "4,,8,12"), std::invalid_argument);
+    EXPECT_THROW(PipelineLayout::Parse(24, 4, "4,8,6,6", 2), std::invalid_argument);
+}
+
+TEST(PipelineLayoutTest, RejectsOutOfRangeQueries) {
+    const auto layout = PipelineLayout::Parse(4, 2, "1,3");
+    EXPECT_THROW(layout.layer_ranges(2), std::out_of_range);
+    EXPECT_THROW(layout.stage_for_layer(4), std::out_of_range);
+}
+
+} // namespace
+} // namespace infini_train::nn::parallel
diff --git a/tests/distributed/test_pipeline_layout_e2e.sh b/tests/distributed/test_pipeline_layout_e2e.sh
new file mode 100755
index 00000000..4079a010
--- /dev/null
+++ b/tests/distributed/test_pipeline_layout_e2e.sh
@@ -0,0 +1,105 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+if [[ $# -lt 3 || $# -gt 4 ]]; then
+    echo "Usage: $0 BUILD_DIR INPUT_BIN GPT2_LLMC_CHECKPOINT [GPU_IDS]" >&2
+    exit 2
+fi
+
+build_dir="$(realpath "$1")"
+input_bin="$(realpath "$2")"
+checkpoint="$(realpath "$3")"
+gpu_ids="${4:-0,1}"
+source_dir="$(realpath "$(dirname "$0")/../..")"
+gpt2="$build_dir/gpt2"
+infini_run="$build_dir/infini_run"
+
+for path in "$gpt2" "$infini_run" "$input_bin" "$checkpoint"; do
+    if [[ ! -e "$path" ]]; then
+        echo "Required test input does not exist: $path" >&2
+        exit 2
+    fi
+done
+if [[ "$gpu_ids" != *,* ]]; then
+    echo "GPU_IDS must contain two comma-separated device IDs" >&2
+    exit 2
+fi
+
+test_dir="$(mktemp -d /tmp/infinitrain-pipeline-layout-e2e.XXXXXX)"
+trap 'rm -rf -- "$test_dir"' EXIT
+single_grad="$test_dir/single-grad"
+pipeline_grad="$test_dir/pipeline-grad"
+vpp_grad="$test_dir/vpp-grad"
+single_log="$test_dir/single.log"
+pipeline_log="$test_dir/pipeline.log"
+vpp_log="$test_dir/vpp.log"
+first_gpu="${gpu_ids%%,*}"
+
+common_args=(
+    --device=cuda
+    --input_bin="$input_bin"
+    --llmc_filepath="$checkpoint"
+    --batch_size=4
+    --sequence_length=64
+    --total_batch_size=512
+    --num_iteration=1
+    --freq_generate_txt=1000
+    --dtype=float32
+)
+
+echo "Running single-GPU reference..."
+env GLOG_logtostderr=1 CUDA_VISIBLE_DEVICES="$first_gpu" \
+    "$gpt2" "${common_args[@]}" --dump_gradients="$single_grad" 2>&1 | tee "$single_log"
+
+echo "Running two-stage automatic pipeline layout..."
+env GLOG_logtostderr=1 CUDA_VISIBLE_DEVICES="$gpu_ids" \
+    "$infini_run" --nproc_per_node=2 "$gpt2" "${common_args[@]}" \
+    --pipeline_parallel=2 \
+    --pipeline_layer_costs=10,1,1,1,1,1,1,1,1,1,1,1 \
+    --dump_gradients="$pipeline_grad" 2>&1 | tee "$pipeline_log"
+
+grep -Fq "stage 0: embedding layers[0,1)" "$pipeline_log"
+grep -Fq "stage 1: layers[1,12) final_norm lm_head" "$pipeline_log"
+
+single_loss="$(sed -n 's/.*train loss \([^ |]*\).*/\1/p' "$single_log" | tail -n 1)"
+pipeline_loss="$(sed -n 's/.*train loss \([^ |]*\).*/\1/p' "$pipeline_log" | tail -n 1)"
+if [[ -z "$single_loss" || -z "$pipeline_loss" ]]; then
+    echo "Failed to extract training loss from logs" >&2
+    exit 1
+fi
+awk -v reference="$single_loss" -v actual="$pipeline_loss" 'BEGIN {
+    difference = reference - actual;
+    if (difference < 0) difference = -difference;
+    if (difference > 1e-5) {
+        printf "Loss mismatch: reference=%s pipeline=%s difference=%g\n", reference, actual, difference > "/dev/stderr";
+        exit 1;
+    }
+}'
+
+find "$single_grad" -type f -name '*.npy' -printf '%f\n' | sort >"$test_dir/single-files"
+find "$pipeline_grad" -type f -name '*.npy' -printf '%f\n' | sort >"$test_dir/pipeline-files"
+diff -u "$test_dir/single-files" "$test_dir/pipeline-files"
+
+python3 "$source_dir/scripts/precision_check/precision_compare.py" \
+    --dir1 "$single_grad" --dir2 "$pipeline_grad" --atol 1e-5 --rtol 0
+
+echo "Running arbitrary virtual Chunk-to-Stage mapping..."
+env GLOG_logtostderr=1 CUDA_VISIBLE_DEVICES="$gpu_ids" \
+    "$infini_run" --nproc_per_node=2 "$gpt2" "${common_args[@]}" \
+    --pipeline_parallel=2 --virtual_pipeline_parallel=2 \
+    --pipeline_chunk_layout=0:3,1:3,1:3,0:3 \
+    --dump_gradients="$vpp_grad" 2>&1 | tee "$vpp_log"
+
+grep -Fq "stage 0: embedding layers[0,3) layers[9,12) final_norm lm_head" "$vpp_log"
+vpp_loss="$(sed -n 's/.*train loss \([^ |]*\).*/\1/p' "$vpp_log" | tail -n 1)"
+awk -v reference="$single_loss" -v actual="$vpp_loss" 'BEGIN {
+    difference = reference - actual;
+    if (difference < 0) difference = -difference;
+    if (difference > 1e-5) exit 1;
+}'
+find "$vpp_grad" -type f -name '*.npy' -printf '%f\n' | sort >"$test_dir/vpp-files"
+diff -u "$test_dir/single-files" "$test_dir/vpp-files"
+python3 "$source_dir/scripts/precision_check/precision_compare.py" \
+    --dir1 "$single_grad" --dir2 "$vpp_grad" --atol 1e-5 --rtol 0
+
+echo "PASS: automatic PP and arbitrary vPP layouts match the single-GPU loss and gradients"
diff --git a/tests/distributed/test_pipeline_layout_suggestion.py b/tests/distributed/test_pipeline_layout_suggestion.py
new file mode 100644
index 00000000..c9102caf
--- /dev/null
+++ b/tests/distributed/test_pipeline_layout_suggestion.py
@@ -0,0 +1,42 @@
+#!/usr/bin/env python3
+import importlib.util
+import tempfile
+import unittest
+from pathlib import Path
+
+
+SCRIPT = Path(__file__).resolve().parents[2] / "scripts" / "suggest_pipeline_layout.py"
+SPEC = importlib.util.spec_from_file_location("pipeline_suggestion", SCRIPT)
+MODULE = importlib.util.module_from_spec(SPEC)
+assert SPEC.loader is not None
+SPEC.loader.exec_module(MODULE)
+
+
+class PipelineLayoutSuggestionTest(unittest.TestCase):
+    def test_balances_contiguous_costs(self):
+        counts, stage_costs = MODULE.balanced_partition([10, 1, 1, 1, 1, 1], 2)
+        self.assertEqual(counts, [1, 5])
+        self.assertEqual(stage_costs, [10, 5])
+
+    def test_reads_layer_profiler_records(self):
+        records = """
+0 2026-08-21 TransformerLayer.0 cuda:0 12 100 1
+1 2026-08-21 TransformerLayer.1 cuda:0 15 20 1
+2 2026-08-21 TransformerLayer.0 cuda:0 12 120 1
+3 2026-08-21 TransformerLayer.1 cuda:0 15 40 1
+"""
+        with tempfile.TemporaryDirectory() as directory:
+            path = Path(directory) / "records.rank0"
+            path.write_text(records, encoding="utf-8")
+            self.assertEqual(MODULE.parse_profiler_records([str(path)]), [110, 30])
+
+    def test_rejects_missing_profiler_layers(self):
+        with tempfile.TemporaryDirectory() as directory:
+            path = Path(directory) / "records.rank0"
+            path.write_text("0 now TransformerLayer.2 cuda:0 1 2 3\n", encoding="utf-8")
+            with self.assertRaises(ValueError):
+                MODULE.parse_profiler_records([str(path)])
+
+
+if __name__ == "__main__":
+    unittest.main()