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 1/2] 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 000000000..2f650bb44
--- /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 000000000..04e841e00
--- /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 000000000..3f8786e6d
--- /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 95e54730b..c9ae2eca5 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 e80c356e3..104635b17 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 5a5cfc656..3e3eab0b2 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 f3590af6e..cd3e73b43 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 d4aea3d04..a1896280a 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 c4642cc2b..a62266806 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 0471c32fe..e797823df 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 25939bdc2..f43a664ad 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 99a739d2d..fa154ecbd 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 c0369cdeb..f2038b4d1 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 c7c1d16fa..0c79d5890 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 000000000..fe6e61fb6
--- /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 b8ed49700..7940b2956 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 000000000..dbec4b3a1
--- /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 000000000..4079a0105
--- /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 000000000..c9102caf3
--- /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()

From f141d92202a2dcf59273456a07d32426143e2347 Mon Sep 17 00:00:00 2001
From: Dayuxiaoshui <792179245@qq.com>
Date: Sat, 19 Sep 2026 14:03:52 +0000
Subject: [PATCH 2/2] feat: weigh embedding and lm head in pipeline layout
 balancing

Balancing only the Transformer layers misses the two modules that are
pinned to the ends of the pipeline: the embedding always sits on stage 0
and the final norm plus lm head on the last stage. On a large-vocabulary
model the lm head, not the layers, is the slowest part of the last stage.

--pipeline_layer_costs now accepts optional "E:<cost>" and "L:<cost>"
entries alongside the per-layer costs. They take no layer slot; they only
add a constant to the first and last segment in the balancing DP. The
trailing constant applies only to the state that closes the last stage,
so the optimal substructure still holds. suggest_pipeline_layout.py grows
the matching --embedding-cost / --lm-head-cost so an offline suggestion
and the runtime solver agree.

Also add the regression that was missing for the parallelism claims: a
layout x PP/DDP/TP matrix that checks every combination against one
single-GPU reference, and a synthetic LLMC asset generator so the matrix
can run where the llm.c starter pack cannot be downloaded.

Measured on 2 x H20, 8 layers / 384 hidden / vocab 50304 / seq 512:
uniform 4,4 runs at 86,628 tok/s, while "1x8,L:4.5" picks 6,2 and runs at
102,904 tok/s (+18.8%). The raw parameter-count ratio (L:10.9) overshoots
to 7,1 and only gains 6.8%, so the cost still has to be measured rather
than derived from parameter counts.

Verified: 11/11 layout unit tests, and 14/14 matrix cases on 8 x H20
matching the single-GPU reference loss within 1e-6, with 101/101
per-parameter gradients matching at atol=1e-5 on the 11 non-TP cases.
---
 docs/pipeline_layout_guide.md                 |  72 +++++-
 docs/pipeline_layout_report.md                |  51 ++++-
 docs/pipeline_layout_test_log.md              |  86 +++++++
 example/gpt2/main.cc                          |   3 +-
 .../nn/parallel/pp/pipeline_parallel.h        |   3 +
 .../src/nn/parallel/pp/pipeline_parallel.cc   |  54 ++++-
 scripts/assets/make_synthetic_gpt2_assets.py  | 126 ++++++++++
 scripts/suggest_pipeline_layout.py            |  39 +++-
 tests/distributed/test_pipeline_layout.cc     |  35 +++
 .../test_pipeline_layout_parallel_matrix.sh   | 215 ++++++++++++++++++
 .../test_pipeline_layout_suggestion.py        |  19 ++
 11 files changed, 678 insertions(+), 25 deletions(-)
 create mode 100644 scripts/assets/make_synthetic_gpt2_assets.py
 create mode 100755 tests/distributed/test_pipeline_layout_parallel_matrix.sh

diff --git a/docs/pipeline_layout_guide.md b/docs/pipeline_layout_guide.md
index 2f650bb44..8a49611ef 100644
--- a/docs/pipeline_layout_guide.md
+++ b/docs/pipeline_layout_guide.md
@@ -46,6 +46,27 @@ Stage 至少拥有一层。
 `--pipeline_layer_costs` 与 `--pipeline_layer_partition` 互斥，且当前同样要求
 `--virtual_pipeline_parallel=1`。程序会在模型构建前打印自动生成的最终布局。
 
+### 把 Embedding / LM Head 的代价算进去
+
+只按 Transformer 层均衡会漏掉两端：Embedding 恒在 stage 0，Final Norm 和 LM Head 恒在最后
+一个 Stage。大词表模型上 LM Head 常常才是最慢的部分。代价串支持两个可选标签，单位与逐层代价
+相同，不占层数：
+
+```bash
+./gpt2 \
+  --pipeline_parallel 2 \
+  --pipeline_layer_costs '1,1,1,1,1,1,1,1,L:4.5' \
+  [其他训练参数]
+```
+
+8 层等代价模型默认会切成 `4,4`；加上 `L:4.5` 后末级 Stage 的代价变成 `2+4.5`，求解器改为
+`6,2`。同理 `E:<cost>` 会把层从 stage 0 推走。标签大小写不敏感，各自最多出现一次，位置任意，
+代价允许为 `0` 表示可忽略；未标注项仍必须是正数。
+
+标签代价要实测标定，不要直接用参数量比值。实测中 `V·d / 12·d²` 给出的 `L:10.9` 会过度倾斜，
+比标定值 `L:4.5` 少拿一半收益，详见 `docs/pipeline_layout_test_log.md`。`scripts/suggest_pipeline_layout.py`
+的 `--embedding-cost` / `--lm-head-cost` 使用同一套求解，可先线下看分区和代价分布再上机。
+
 ## 默认行为与 vPP
 
 不传 `--pipeline_layer_partition` 时，保持原有均匀划分。余数从执行顺序靠前的 chunk 开始各多分一层；
@@ -109,6 +130,8 @@ Pipeline layout (24 layers, 4 stages):
 - `incompatible with --virtual_pipeline_parallel != 1`：自定义物理布局和 vPP 同时启用。
 - `must contain exactly N entries`：逐层代价数量和模型层数不一致。
 - `costs must be finite positive numbers`：逐层代价包含零、负数、NaN、无穷或非数字。
+- `only accept the 'E:' ... and 'L:' ... tags`：代价串里出现了 `E:`/`L:` 之外的标签。
+- `repeat the 'E:' entry`：同一个标签出现了多次。
 - `cannot be used together`：同时指定了手工分区和自动均衡代价。
 
 ## C++ 查询接口
@@ -119,13 +142,19 @@ Pipeline layout (24 layers, 4 stages):
 
 ## 并行组合与限制
 
+下表中「已实测」指 `tests/distributed/test_pipeline_layout_parallel_matrix.sh` 在 8×H20 上
+跑过且与单卡结果一致（见 `docs/pipeline_layout_test_log.md`）；「支持」指代码路径打通但该格
+没有单独跑过回归。
+
 | 组合 | 默认均匀布局 | 手工层数分区 | 任意 Chunk / Megatron 布局 |
 | --- | --- | --- | --- |
-| PP | 支持 | 支持 | 支持 |
-| PP + DDP | 支持 | 支持 | 支持 |
-| PP + TP | 支持 | 支持 | 支持 |
-| PP + DDP + TP | 支持 | 支持 | 支持 |
-| vPP | 默认轮转映射 | 拒绝物理分区参数 | 支持显式 Chunk owner |
+| PP | 支持 | 已实测 PP2 / PP4 | 已实测 |
+| PP + DDP | 支持 | 已实测 DP2 / DP4 | 已实测 DP2 |
+| PP + TP | 支持 | 已实测 TP2 | 支持 |
+| PP + DDP + TP | 支持 | 已实测 TP2 × DP2 | 支持 |
+| vPP | 默认轮转映射 | 拒绝物理分区参数 | 已实测显式 Chunk owner |
+
+逐层代价布局（含 `E:`/`L:` 标签）生成的是普通物理分区，已实测 PP2 和 PP2 × DP2。
 
 `||` 表示空逻辑 Chunk；它不表示物理 Stage 没有 Chunk。当前调度器要求每个物理 Stage 拥有相同数量的
 正数 Chunk，因此会拒绝 Chunk 数不平衡的映射。
@@ -161,3 +190,36 @@ tests/distributed/test_pipeline_layout_e2e.sh \
 
 该测试需要 CUDA/NCCL 构建、两张 GPU、NumPy 以及 GPT-2 124M LLMC checkpoint。测试使用临时目录，
 退出时自动清理梯度和日志。
+
+## 并行组合矩阵回归
+
+`test_pipeline_layout_parallel_matrix.sh` 把手工分区、逐层代价、特殊模块代价、任意 Chunk 映射
+和 Megatron 表达式五类布局与 PP / DDP / TP 的组合跑成一个矩阵，
+每个用例都和同一份单卡参考比较 loss；参数未被 TP 切分的用例还会逐参数比较梯度：
+
+~~~bash
+tests/distributed/test_pipeline_layout_parallel_matrix.sh \
+  /path/to/cuda-build \
+  data/gpt2/tiny_shakespeare_train.bin \
+  data/gpt2/gpt2_124M.bin \
+  8
+~~~
+
+最后一个参数是可用 GPU 数（默认 8）。分区由 checkpoint header 里的真实层数推导，因此换模型
+不用改脚本；需要的 GPU 多于实际数量的用例会标记为 SKIP，所以在 2 卡或 4 卡机器上同样可用。
+任何用例失败都会返回非零退出码。
+
+TP 用例只比较 loss：TP 切分参数后，各 rank 导出的是梯度分片，无法与单卡逐参数对齐。DP 用例
+中各 DP rank 会把同名梯度写进同一目录（内容相同）。
+
+## 离线生成测试资产
+
+如果机器无法下载 llm.c starter pack，可以本地合成格式兼容的 checkpoint 和 token 文件：
+
+~~~bash
+python3 scripts/assets/make_synthetic_gpt2_assets.py --out-dir data/gpt2-synthetic
+~~~
+
+输出按 `--seed` 逐字节可复现，因此不同布局读到的是同一份权重和同一批 token——这正是布局回归
+需要的性质。权重是随机初始化的，loss 接近 `ln(vocab)`，不能用来判断收敛质量。
+注意 `--n-head` 必须保持 12：LLMC loader 不覆盖 `n_kv_head`，而配置校验要求两者相等。
diff --git a/docs/pipeline_layout_report.md b/docs/pipeline_layout_report.md
index 04e841e00..b5ca8591d 100644
--- a/docs/pipeline_layout_report.md
+++ b/docs/pipeline_layout_report.md
@@ -25,6 +25,13 @@ Stage 数、正整数、总层数和 vPP 兼容性。`Uniform` 封装原有默
 确定且不改变层的执行顺序。`ResolvePipelineLayout` 统一选择手工分区、代价均衡或默认均匀布局，
 并拒绝多个布局来源同时生效。
 
+只按 Transformer 层均衡会漏掉两端的特殊模块：Embedding 恒在 stage 0，Final Norm 和 LM Head
+恒在最后一个 Stage，大词表模型上 LM Head 往往是真正的瓶颈。因此代价串额外接受 `E:<cost>` 和
+`L:<cost>` 两个可选标签，它们不占层数，只在 DP 转移里作为常数加到首段和末段的代价上（末段常数
+只在 `s == S && i == L` 的状态生效，中间状态不受影响，因此最优子结构保持成立）。标签代价允许为 0
+表示「可忽略」，未标注项仍必须是正数并且个数等于模型层数。`scripts/suggest_pipeline_layout.py`
+用 `--embedding-cost` / `--lm-head-cost` 实现同一套转移，保证线下建议和运行时求解结果一致。
+
 GPT-2 和 LLaMA 3 在模型配置确定后设置布局。对于 LLMC checkpoint，布局在读取 header 中真实
 `n_layer` 后解析。`TransformerModel` 用布局创建本 rank 的层和特殊模块；`TransformerConfig::GetChunkSize`
 和 `PipelineParallel` 用相同布局构造调度 Stage；两个 checkpoint loader 用布局筛选本 rank 权重。
@@ -47,12 +54,28 @@ GPT-2 和 LLaMA 3 在模型配置确定后设置布局。对于 LLMC checkpoint
 | 双卡 GPT-2 PP E2E | 已实测 | H200，loss、梯度与单卡一致 |
 | 任意 vPP Chunk owner | 已实测 | H200，owner `[0,1,1,0]` 无死锁 |
 | Megatron 风格布局 | 已实测 | H200，包含空逻辑 Chunk |
+| PP × DP × TP 组合矩阵 | 已实测 | 8×H20，14 个组合全部与单卡一致 |
+| 特殊模块代价均衡 | 已实测 | 8×H20，`L:` 代价自动选中最优分区，吞吐 +18.8% |
+
+组合矩阵由 `tests/distributed/test_pipeline_layout_parallel_matrix.sh` 固化，在 8 张 H20 上
+覆盖手工分区、逐层代价、特殊模块代价、任意 Chunk 映射和 Megatron 表达式五类布局与
+PP / PP×DP2 / PP×DP4 / PP×TP2 / PP×TP2×DP2 的组合：
+
+| 布局来源 | PP | PP + DDP | PP + TP | PP + DDP + TP |
+| --- | --- | --- | --- | --- |
+| 手工层数分区 | PASS (PP2/PP4) | PASS (DP2/DP4) | PASS (PP2/PP4 × TP2) | PASS (PP2×TP2×DP2) |
+| 逐层代价 | PASS (PP2) | — | — | — |
+| 逐层代价 + `L:` 标签 | PASS (PP2) | PASS (×DP2) | — | — |
+| 任意 Chunk 映射 | PASS (PP2×vPP2) | PASS (×DP2) | — | — |
+| Megatron 表达式 | PASS (PP4) | PASS (×DP2) | — | — |
 
-DDP/TP 组合保留既有 InfiniTrain 并行入口；本次新增回归重点是布局解析、PP 调度、参数加载和
-跨布局数值一致性。若提交环境要求完整 DP×TP×PP 组合矩阵，应在目标集群补跑对应资源规模的回归。
+单卡参考 loss `7.016952`；全部 14 个用例 loss 最大偏差 `1e-6`，非 TP 用例的 101 个逐参数
+梯度在 `atol=1e-5, rtol=0` 下全部一致。TP 用例只比较 loss，因为 TP 切分参数后各 rank 导出的
+是梯度分片，无法与单卡逐参数对齐。表中 `—` 表示未跑该格，不表示不支持。
 
-CPU 单元测试覆盖 `4,8,6,6`、完整 layer-to-stage 反查、特殊模块、默认 vPP 轮转，以及错误的
-Stage 数、总和、负数、零、空项、越界查询和自定义布局/vPP 冲突。验证命令：
+CPU 单元测试覆盖 `4,8,6,6`、完整 layer-to-stage 反查、特殊模块、默认 vPP 轮转、`E:`/`L:`
+代价标签，以及错误的 Stage 数、总和、负数、零、空项、未知标签、重复标签、越界查询和
+自定义布局/vPP 冲突。验证命令：
 
 ```bash
 cmake -S . -B /tmp/infinitrain-pipeline-build \
@@ -61,8 +84,9 @@ cmake --build /tmp/infinitrain-pipeline-build --target test_pipeline_layout gpt2
 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` 两阶段训练，
+结果：11/11 布局与建议测试通过（10 个 GTest 用例 + 1 个建议脚本用例集），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` 下
@@ -91,3 +115,18 @@ CPU `ctest` 不会默认注册该用例。
 
 建议布局实测吞吐提升 9.45%，峰值显存降低 130 MB；理论 bubble（4 microbatch、2 Stage）为 20%。
 Profiler 输入为 3 步 PROFILE_MODE 记录，每层丢弃一个 warmup 样本后得到 7,5，模型代价上界下降 6.84%。
+
+特殊模块代价在两张 H20 上单独做了验证。模型为 8 层 / 384 hidden / vocab 50304 / seq 512，
+LM Head 的参数量约等于 10.9 层，是明确的瓶颈端：
+
+| 布局 | 中位 step | 吞吐 | 末级 Stage 峰值显存 |
+| --- | ---: | ---: | ---: |
+| 默认 4,4 | 189.13 ms | 86,628 tok/s | 8308 MB |
+| `1×8,L:4.5` → 6,2 | 159.22 ms | 102,904 tok/s | 6020 MB |
+| `1×8,L:10.9` → 7,1 | 177.09 ms | 92,516 tok/s | 4876 MB |
+
+标定后的 `L:4.5` 自动选中的 `6,2` 与手工扫描 `5,3`/`6,2`/`7,1` 得到的最优点一致，吞吐提升
+18.8%。直接用参数量比值 `V·d / 12·d²` 得到的 `L:10.9` 会过度倾斜，只提升 6.8%：大 GEMM 的
+实际效率高于参数量假设。因此特殊模块代价和逐层代价适用同一条结论——参数量只能当上界，真实
+代价要用 profiler 或一次扫描标定。当 LM Head 代价小于一层时（如 768 hidden 的同款模型），
+求解器会保持默认均匀分区不变，这是期望行为。
diff --git a/docs/pipeline_layout_test_log.md b/docs/pipeline_layout_test_log.md
index 3f8786e6d..7a6ca69f1 100644
--- a/docs/pipeline_layout_test_log.md
+++ b/docs/pipeline_layout_test_log.md
@@ -146,3 +146,89 @@ improvement:       +9.45% throughput, -130 MB peak memory
 
 Profiler 建议来自 3 步单卡 PROFILE_MODE 记录，每层丢弃一个 warmup 样本；建议工具单测和实际记录
 解析均已通过。
+
+## 8×H20 并行组合矩阵（2026-09-19）
+
+### 环境
+
+| 项目 | 值 |
+| --- | --- |
+| GPU | 8 × NVIDIA H20 96 GB |
+| 构建 | CUDA 12.8 + NCCL，`-DUSE_CUDA=ON -DUSE_NCCL=ON -DBUILD_TEST=ON` |
+| 模型 | 合成 LLMC fp32 checkpoint，8 层 / 12 头 / 384 hidden / vocab 1024 |
+| 训练参数 | `--batch_size=4 --sequence_length=64 --total_batch_size=2048 --dtype=float32` |
+
+该机器的出口代理屏蔽了 HuggingFace CDN，无法下载 GPT-2 124M starter pack，因此改用
+`scripts/assets/make_synthetic_gpt2_assets.py` 生成格式兼容、按 seed 逐字节可复现的
+checkpoint 和 token 文件。权重是随机初始化的，首步 loss 接近 `ln(vocab)`；本节比较的是
+「同一份输入在不同布局/并行组合下是否得到同一个结果」，与权重是否预训练无关。
+
+### 单元测试
+
+```text
+ctest -R PipelineLayout --output-on-failure
+100% tests passed, 0 tests failed out of 11
+```
+
+### 组合矩阵
+
+```bash
+tests/distributed/test_pipeline_layout_parallel_matrix.sh \
+  build data/gpt2-synthetic/tokens_train.bin data/gpt2-synthetic/gpt2_synthetic.bin 8
+```
+
+单卡参考：loss `7.016952`，101 个梯度。全部 14 个用例通过：
+
+| 用例 | GPU | 布局来源 | 组合 | 比较 | 结果 |
+| --- | ---: | --- | --- | --- | --- |
+| pp2_partition | 2 | 手工 `2,6` | PP2 | loss + 101 梯度 | PASS |
+| pp4_partition | 4 | 手工 `1,5,1,1` | PP4 | loss + 101 梯度 | PASS |
+| pp2_ddp2 | 4 | 手工 `2,6` | PP2 × DP2 | loss + 101 梯度 | PASS |
+| pp2_ddp4 | 8 | 手工 `2,6` | PP2 × DP4 | loss + 101 梯度 | PASS |
+| pp2_tp2 | 4 | 手工 `2,6` | PP2 × TP2 | loss | PASS |
+| pp2_tp2_ddp2 | 8 | 手工 `2,6` | PP2 × TP2 × DP2 | loss | PASS |
+| pp4_tp2 | 8 | 手工 `1,5,1,1` | PP4 × TP2 | loss | PASS |
+| pp2_costs | 2 | 逐层代价 `10,1,…` | PP2 | loss + 101 梯度 | PASS |
+| pp2_costs_lm_head | 2 | 代价 `1×8,L:2` → `5,3` | PP2 | loss + 101 梯度 | PASS |
+| pp2_costs_lm_head_ddp2 | 4 | 代价 `1×8,L:2` → `5,3` | PP2 × DP2 | loss + 101 梯度 | PASS |
+| pp2_vpp2_chunk | 2 | `0:2,1:2,1:2,0:2` | PP2 × vPP2 | loss + 101 梯度 | PASS |
+| pp2_vpp2_chunk_ddp2 | 4 | `0:2,1:2,1:2,0:2` | PP2 × vPP2 × DP2 | loss + 101 梯度 | PASS |
+| pp4_megatron | 4 | `Et*2\|t*2\|t*2\|t*2NL` | PP4 | loss + 101 梯度 | PASS |
+| pp4_megatron_ddp2 | 8 | `Et*2\|t*2\|t*2\|t*2NL` | PP4 × DP2 | loss + 101 梯度 | PASS |
+
+所有用例的打印 loss 落在 `7.016951`~`7.016952`，最大差值 `1e-6`，满足 fp32 `1e-5` 容差；
+非 TP 用例（11 个）的 101 个梯度在 `atol=1e-5, rtol=0` 下全部通过且文件集合一致。
+
+两点限制需要声明：
+
+- TP 用例只比较 loss。TP 会切分参数，各 rank 导出的梯度是分片，无法与单卡逐参数直接对齐。
+- DP 用例中每个 DP rank 会把同名梯度写进同一个目录。比较全部通过，但这是重复写同一份内容，
+  严格来说存在并发写的理论风险；如需逐 DP rank 分别校验，应给 `--dump_gradients` 加 rank 后缀。
+
+## 特殊模块代价均衡实测（8×H20 中的 2 卡）
+
+`--pipeline_layer_costs` 的 `E:`/`L:` 标签把 Embedding 和 LM Head 的代价纳入均衡。用一个
+LM Head 明显占主导的形状验证：8 层 / 384 hidden / vocab 50304 / seq 512，PP=2，
+8 个 microbatch，12 迭代取后 8 步中位数。
+
+```text
+layout                     median step      throughput    last-stage peak
+uniform 4,4 (默认)           189.13 ms     86,628 tok/s          8308 MB
+手工 5,3                     169.05 ms     96,916 tok/s          7164 MB
+手工 6,2                     160.77 ms    101,934 tok/s          6020 MB
+手工 7,1                     177.75 ms     92,172 tok/s          4876 MB
+--pipeline_layer_costs=1×8,L:4.5   159.22 ms    102,904 tok/s    6020 MB   <- 自动选中 6,2
+--pipeline_layer_costs=1×8,L:10.9  177.09 ms     92,516 tok/s    4876 MB   <- 自动选中 7,1
+```
+
+结论有两层：
+
+1. 代价标签确实生效。`L:4.5` 自动生成 `6,2`，与手工扫描出的最优点一致，相对默认均匀布局
+   吞吐提升 **18.8%**、单步耗时下降 **15.8%**、末级 Stage 峰值显存下降 2288 MB。
+2. 代价必须实测，不能直接用参数量比值。按 `V·d / 12·d²` 估算 LM Head 相当于 10.9 层，
+   据此得到的 `7,1` 反而比默认布局只快 6.8%——大 GEMM 的实际效率高于参数量比值的假设。
+   这与逐层代价的既有建议一致：用 profiler 或扫描标定代价，参数量只能作为上界参考。
+
+作为对照，把同一实验换成 768 hidden / vocab 50304（LM Head 约等于 5.5 层）时，默认 `4,4`
+已经是 8 层粒度下的最优点（`4,4` 380.93 ms vs `5,3` 390.72 ms vs `3,5` 425.50 ms），
+`L:` 代价在小于一层时不会改变布局——这是期望行为，不是失效。
diff --git a/example/gpt2/main.cc b/example/gpt2/main.cc
index 3e3eab0b2..767e8b5fd 100644
--- a/example/gpt2/main.cc
+++ b/example/gpt2/main.cc
@@ -90,7 +90,8 @@ 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.");
+              "Comma-separated positive compute costs for every Transformer layer; generates a balanced layout. "
+              "Optional 'E:<cost>' and 'L:<cost>' entries add the embedding and lm head cost to the balancing.");
 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, "",
diff --git a/infini_train/include/nn/parallel/pp/pipeline_parallel.h b/infini_train/include/nn/parallel/pp/pipeline_parallel.h
index f43a664ad..d56dc07dc 100644
--- a/infini_train/include/nn/parallel/pp/pipeline_parallel.h
+++ b/infini_train/include/nn/parallel/pp/pipeline_parallel.h
@@ -33,6 +33,9 @@ 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);
+    // `layer_costs` is one positive cost per Transformer layer, plus optional "E:<cost>"
+    // and "L:<cost>" entries for the embedding and lm head that the first and last stage
+    // carry on top of their layers.
     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);
diff --git a/infini_train/src/nn/parallel/pp/pipeline_parallel.cc b/infini_train/src/nn/parallel/pp/pipeline_parallel.cc
index f2038b4d1..e4c19bcf9 100644
--- a/infini_train/src/nn/parallel/pp/pipeline_parallel.cc
+++ b/infini_train/src/nn/parallel/pp/pipeline_parallel.cc
@@ -187,6 +187,13 @@ PipelineLayout PipelineLayout::FromLayerCosts(int total_layers, int pp_size, con
     }
 
     std::vector<double> costs;
+    // Embedding and lm head run on the first and last stage no matter how the layers are
+    // split, so their cost has to enter the balancing too. They are optional entries tagged
+    // "E:<cost>" and "L:<cost>"; every untagged entry is a Transformer layer, in layer order.
+    double embedding_cost = 0.0;
+    double lm_head_cost = 0.0;
+    bool has_embedding_cost = false;
+    bool has_lm_head_cost = false;
     size_t begin = 0;
     while (begin <= layer_costs.size()) {
         const size_t comma = layer_costs.find(',', begin);
@@ -198,13 +205,44 @@ PipelineLayout PipelineLayout::FromLayerCosts(int total_layers, int pp_size, con
             throw std::invalid_argument("pipeline layer costs contain an empty entry: '" + layer_costs + "'");
         }
         token = token.substr(first, last - first + 1);
+        const std::string_view entry = token;
+
+        double *special = nullptr;
+        bool *already_seen = nullptr;
+        if (token.size() > 1 && token[1] == ':') {
+            if (token[0] == 'E' || token[0] == 'e') {
+                special = &embedding_cost;
+                already_seen = &has_embedding_cost;
+            } else if (token[0] == 'L' || token[0] == 'l') {
+                special = &lm_head_cost;
+                already_seen = &has_lm_head_cost;
+            } else {
+                throw std::invalid_argument("pipeline layer costs only accept the 'E:' (embedding) and "
+                                            "'L:' (lm head) tags; got '"
+                                            + std::string(entry) + "'");
+            }
+            if (*already_seen) {
+                throw std::invalid_argument("pipeline layer costs repeat the '" + std::string(1, entry[0])
+                                            + ":' entry: '" + layer_costs + "'");
+            }
+            *already_seen = true;
+            token = token.substr(2);
+        }
+
         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) {
+        // A zero cost is a meaningful "this module is negligible" for the tagged entries,
+        // but an unweighted Transformer layer is not.
+        const bool out_of_range = special == nullptr ? cost <= 0.0 : cost < 0.0;
+        if (ec != std::errc() || ptr != token.data() + token.size() || !std::isfinite(cost) || out_of_range) {
             throw std::invalid_argument("pipeline layer costs must be finite positive numbers; got '"
-                                        + std::string(token) + "'");
+                                        + std::string(entry) + "'");
+        }
+        if (special != nullptr) {
+            *special = cost;
+        } else {
+            costs.push_back(cost);
         }
-        costs.push_back(cost);
         if (comma == std::string::npos) { break; }
         begin = comma + 1;
     }
@@ -220,14 +258,22 @@ PipelineLayout PipelineLayout::FromLayerCosts(int total_layers, int pp_size, con
             throw std::invalid_argument("pipeline layer costs have a non-finite total");
         }
     }
+    if (!std::isfinite(prefix[total_layers] + embedding_cost + lm_head_cost)) {
+        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) {
+        // Stage 0 always owns the embedding; the lm head always lands on the last stage,
+        // which is the one closed off by best[pp_size][total_layers].
+        const double leading = stages == 1 ? embedding_cost : 0.0;
         for (int end = stages; end <= total_layers; ++end) {
+            const double trailing = (stages == pp_size && end == total_layers) ? lm_head_cost : 0.0;
             for (int start = stages - 1; start < end; ++start) {
-                const double candidate = std::max(best[stages - 1][start], prefix[end] - prefix[start]);
+                const double segment = prefix[end] - prefix[start] + leading + trailing;
+                const double candidate = std::max(best[stages - 1][start], segment);
                 if (candidate < best[stages][end]) {
                     best[stages][end] = candidate;
                     split[stages][end] = start;
diff --git a/scripts/assets/make_synthetic_gpt2_assets.py b/scripts/assets/make_synthetic_gpt2_assets.py
new file mode 100644
index 000000000..eebc349e5
--- /dev/null
+++ b/scripts/assets/make_synthetic_gpt2_assets.py
@@ -0,0 +1,126 @@
+#!/usr/bin/env python3
+"""Generate a small synthetic GPT-2 LLMC checkpoint and token file.
+
+The pipeline layout regression compares one reference run against several parallel
+layouts, so it only needs weights and tokens that are byte-identical across runs --
+not pretrained ones. Use this when the real llm.c starter pack cannot be downloaded
+(air-gapped machine, blocked mirror) or when a smaller/faster model is enough.
+
+    python3 scripts/assets/make_synthetic_gpt2_assets.py --out-dir data/gpt2-synthetic
+
+Both files use the formats the GPT-2 example already reads: the LLMC fp32 (version 3)
+checkpoint layout of example/gpt2/checkpoint_loader.cc and the uint16 token layout of
+example/common/tiny_shakespeare_dataset.cc.
+"""
+
+import argparse
+from pathlib import Path
+
+import numpy as np
+
+CHECKPOINT_MAGIC = 20240326
+CHECKPOINT_FP32_VERSION = 3
+TOKENS_MAGIC_UINT16 = 20240520
+TOKENS_VERSION = 1
+HEADER_INTS = 256
+
+
+def write_checkpoint(path: Path, args: argparse.Namespace) -> None:
+    rng = np.random.default_rng(args.seed)
+    layers, embd, vocab = args.n_layer, args.n_embd, args.padded_vocab_size
+
+    header = np.zeros(HEADER_INTS, dtype=np.int32)
+    header[0] = CHECKPOINT_MAGIC
+    header[1] = CHECKPOINT_FP32_VERSION
+    header[2] = args.block_size
+    header[3] = args.vocab_size
+    header[4] = layers
+    header[5] = args.n_head
+    header[6] = embd
+    header[7] = args.padded_vocab_size
+
+    def normal(*shape: int, std: float = 0.02) -> np.ndarray:
+        return rng.normal(0.0, std, size=shape).astype(np.float32)
+
+    def zeros(*shape: int) -> np.ndarray:
+        return np.zeros(shape, dtype=np.float32)
+
+    def ones(*shape: int) -> np.ndarray:
+        return np.ones(shape, dtype=np.float32)
+
+    # Same tensor order as llm.c: token/position embeddings, then one full pass over
+    # the layers per parameter, then the final norm.
+    tensors = [
+        normal(vocab, embd),
+        normal(args.block_size, embd),
+        ones(layers, embd),
+        zeros(layers, embd),
+        normal(layers, 3 * embd, embd),
+        zeros(layers, 3 * embd),
+        normal(layers, embd, embd),
+        zeros(layers, embd),
+        ones(layers, embd),
+        zeros(layers, embd),
+        normal(layers, 4 * embd, embd),
+        zeros(layers, 4 * embd),
+        normal(layers, embd, 4 * embd),
+        zeros(layers, embd),
+        ones(embd),
+        zeros(embd),
+    ]
+
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with path.open("wb") as out:
+        out.write(header.tobytes())
+        for tensor in tensors:
+            out.write(np.ascontiguousarray(tensor).tobytes())
+    print(f"checkpoint: {path} ({path.stat().st_size / (1 << 20):.1f} MiB)")
+
+
+def write_tokens(path: Path, args: argparse.Namespace) -> None:
+    rng = np.random.default_rng(args.seed + 1)
+
+    header = np.zeros(HEADER_INTS, dtype=np.int32)
+    header[0] = TOKENS_MAGIC_UINT16
+    header[1] = TOKENS_VERSION
+    header[2] = args.num_tokens
+
+    tokens = rng.integers(0, args.vocab_size, size=args.num_tokens, dtype=np.uint16)
+
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with path.open("wb") as out:
+        out.write(header.tobytes())
+        out.write(tokens.tobytes())
+    print(f"tokens: {path} ({args.num_tokens} tokens)")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--out-dir", default="data/gpt2-synthetic")
+    parser.add_argument("--n-layer", type=int, default=8)
+    # The GPT-2 loader keeps n_kv_head at the GPT2Config() default, and
+    # SanitizeGPT2Config() requires n_kv_head == n_head, so n_head must stay 12.
+    parser.add_argument("--n-head", type=int, default=12)
+    parser.add_argument("--n-embd", type=int, default=384)
+    parser.add_argument("--block-size", type=int, default=128)
+    parser.add_argument("--vocab-size", type=int, default=1024)
+    parser.add_argument("--padded-vocab-size", type=int, default=0,
+                        help="defaults to --vocab-size so that TP and non-TP runs read identical weights")
+    parser.add_argument("--num-tokens", type=int, default=262144)
+    parser.add_argument("--seed", type=int, default=20260919)
+    args = parser.parse_args()
+
+    if args.padded_vocab_size == 0:
+        args.padded_vocab_size = args.vocab_size
+    if args.padded_vocab_size < args.vocab_size:
+        raise SystemExit("--padded-vocab-size must be >= --vocab-size")
+    if args.n_embd % args.n_head:
+        raise SystemExit("--n-embd must be divisible by --n-head")
+
+    out_dir = Path(args.out_dir)
+    write_checkpoint(out_dir / "gpt2_synthetic.bin", args)
+    write_tokens(out_dir / "tokens_train.bin", args)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/scripts/suggest_pipeline_layout.py b/scripts/suggest_pipeline_layout.py
index fe6e61fb6..f6312ca22 100755
--- a/scripts/suggest_pipeline_layout.py
+++ b/scripts/suggest_pipeline_layout.py
@@ -47,10 +47,14 @@ def parse_profiler_records(paths: list[str], warmup_samples: int = 0) -> list[fl
     ]
 
 
-def balanced_partition(costs: list[float], stages: int) -> tuple[list[int], list[float]]:
+def balanced_partition(
+    costs: list[float], stages: int, embedding_cost: float = 0.0, lm_head_cost: float = 0.0
+) -> tuple[list[int], list[float]]:
     layers = len(costs)
     if stages <= 0 or stages > layers:
         raise ValueError("stages must satisfy 0 < stages <= number of layers")
+    if embedding_cost < 0 or lm_head_cost < 0:
+        raise ValueError("embedding and lm head costs must not be negative")
     prefix = [0.0]
     for cost in costs:
         prefix.append(prefix[-1] + cost)
@@ -58,9 +62,14 @@ def balanced_partition(costs: list[float], stages: int) -> tuple[list[int], list
     split = [[-1] * (layers + 1) for _ in range(stages + 1)]
     best[0][0] = 0.0
     for stage_count in range(1, stages + 1):
+        # The embedding always sits on stage 0 and the lm head on the last stage, so their
+        # cost joins the segment that closes those stages off.
+        leading = embedding_cost if stage_count == 1 else 0.0
         for end in range(stage_count, layers + 1):
+            trailing = lm_head_cost if stage_count == stages and end == layers else 0.0
             for start in range(stage_count - 1, end):
-                candidate = max(best[stage_count - 1][start], prefix[end] - prefix[start])
+                segment = prefix[end] - prefix[start] + leading + trailing
+                candidate = max(best[stage_count - 1][start], segment)
                 if candidate < best[stage_count][end]:
                     best[stage_count][end] = candidate
                     split[stage_count][end] = start
@@ -70,12 +79,20 @@ def balanced_partition(costs: list[float], stages: int) -> tuple[list[int], list
         start = split[stage + 1][end]
         counts[stage] = end - start
         end = start
+    return counts, stage_costs_of(costs, counts, embedding_cost, lm_head_cost)
+
+
+def stage_costs_of(
+    costs: list[float], counts: list[int], embedding_cost: float, lm_head_cost: float
+) -> list[float]:
     stage_costs = []
     start = 0
     for count in counts:
         stage_costs.append(sum(costs[start : start + count]))
         start += count
-    return counts, stage_costs
+    stage_costs[0] += embedding_cost
+    stage_costs[-1] += lm_head_cost
+    return stage_costs
 
 
 def main() -> None:
@@ -86,6 +103,10 @@ def main() -> None:
     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("--embedding-cost", type=float, default=0.0,
+                        help="extra cost the first stage carries for the embedding, in the same unit as the layers")
+    parser.add_argument("--lm-head-cost", type=float, default=0.0,
+                        help="extra cost the last stage carries for the final norm and lm head")
     parser.add_argument("--microbatches", type=int, default=1)
     parser.add_argument("--json-output", type=Path)
     args = parser.parse_args()
@@ -95,20 +116,20 @@ def main() -> None:
         if args.profiler_records
         else parse_numbers(args.costs or args.parameter_counts)
     )
-    counts, stage_costs = balanced_partition(costs, args.pipeline_parallel)
+    counts, stage_costs = balanced_partition(
+        costs, args.pipeline_parallel, args.embedding_cost, args.lm_head_cost
+    )
     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
+    uniform_costs = stage_costs_of(costs, uniform_counts, args.embedding_cost, args.lm_head_cost)
     bubble = (args.pipeline_parallel - 1) / (args.microbatches + args.pipeline_parallel - 1)
     result = {
         "partition": counts,
         "stage_costs": stage_costs,
         "maximum_stage_cost": max(stage_costs),
+        "embedding_cost": args.embedding_cost,
+        "lm_head_cost": args.lm_head_cost,
         "uniform_partition": uniform_counts,
         "uniform_stage_costs": uniform_costs,
         "uniform_maximum_stage_cost": max(uniform_costs),
diff --git a/tests/distributed/test_pipeline_layout.cc b/tests/distributed/test_pipeline_layout.cc
index dbec4b3a1..81cbe8907 100644
--- a/tests/distributed/test_pipeline_layout.cc
+++ b/tests/distributed/test_pipeline_layout.cc
@@ -56,6 +56,34 @@ TEST(PipelineLayoutTest, BalancesUserProvidedLayerCosts) {
     EXPECT_EQ(layout.stage_for_layer(5), 1);
 }
 
+TEST(PipelineLayoutTest, BalancesEmbeddingAndLmHeadCost) {
+    // Eight equal layers alone split 4/4, but a heavy lm head has to pull layers off the
+    // last stage and a heavy embedding off the first one.
+    const std::string layers = "1,1,1,1,1,1,1,1";
+    EXPECT_EQ(PipelineLayout::FromLayerCosts(8, 2, layers).layer_ranges(0),
+              (std::vector<std::pair<int, int>>{{0, 4}}));
+
+    const auto heavy_lm_head = PipelineLayout::FromLayerCosts(8, 2, layers + ",L:2");
+    EXPECT_EQ(heavy_lm_head.layer_ranges(0), (std::vector<std::pair<int, int>>{{0, 5}}));
+    EXPECT_EQ(heavy_lm_head.layer_ranges(1), (std::vector<std::pair<int, int>>{{5, 8}}));
+
+    const auto heavy_embedding = PipelineLayout::FromLayerCosts(8, 2, "E:2," + layers);
+    EXPECT_EQ(heavy_embedding.layer_ranges(0), (std::vector<std::pair<int, int>>{{0, 3}}));
+    EXPECT_EQ(heavy_embedding.layer_ranges(1), (std::vector<std::pair<int, int>>{{3, 8}}));
+
+    // Equal special-module costs cancel out, and a zero cost means "negligible".
+    EXPECT_EQ(PipelineLayout::FromLayerCosts(8, 2, "E:2," + layers + ",L:2").layer_ranges(0),
+              (std::vector<std::pair<int, int>>{{0, 4}}));
+    EXPECT_EQ(PipelineLayout::FromLayerCosts(8, 2, "e:0," + layers + ",l:0").layer_ranges(0),
+              (std::vector<std::pair<int, int>>{{0, 4}}));
+
+    // Every stage still owns at least one layer, so a dominant lm head can only shrink the
+    // last stage down to a single layer.
+    const auto four_stages = PipelineLayout::FromLayerCosts(8, 4, layers + ",L:3");
+    EXPECT_EQ(four_stages.stage_for_layer(7), 3);
+    EXPECT_EQ(four_stages.stage_for_layer(6), 2);
+}
+
 TEST(PipelineLayoutTest, SupportsArbitraryVirtualChunkOwnership) {
     const auto layout = PipelineLayout::FromChunkLayout(8, 2, "0:2,1:2,1:2,0:2");
 
@@ -98,6 +126,13 @@ TEST(PipelineLayoutTest, RejectsInvalidAutomaticLayoutInputs) {
     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(PipelineLayout::FromLayerCosts(3, 2, "N:1,1,1,1"), std::invalid_argument);
+    EXPECT_THROW(PipelineLayout::FromLayerCosts(3, 2, "E:1,E:1,1,1,1"), std::invalid_argument);
+    EXPECT_THROW(PipelineLayout::FromLayerCosts(3, 2, "L:-1,1,1,1"), std::invalid_argument);
+    EXPECT_THROW(PipelineLayout::FromLayerCosts(3, 2, "E:x,1,1,1"), std::invalid_argument);
+    EXPECT_THROW(PipelineLayout::FromLayerCosts(2, 2, "1.7e308,1,L:1.7e308"), std::invalid_argument);
+    // The tagged entries are extra, not a substitute for a per-layer cost.
+    EXPECT_THROW(PipelineLayout::FromLayerCosts(3, 2, "E:1,1,1"), 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);
diff --git a/tests/distributed/test_pipeline_layout_parallel_matrix.sh b/tests/distributed/test_pipeline_layout_parallel_matrix.sh
new file mode 100755
index 000000000..1dafb6e97
--- /dev/null
+++ b/tests/distributed/test_pipeline_layout_parallel_matrix.sh
@@ -0,0 +1,215 @@
+#!/usr/bin/env bash
+# Multi-GPU regression for custom pipeline layouts combined with DDP and TP.
+#
+# Runs a single-GPU reference and then every layout x parallelism combination the
+# layout feature claims to support, asserting that each one reproduces the reference
+# loss (and, where the parameters are not TP-sharded, every per-parameter gradient).
+#
+# Usage:
+#   tests/distributed/test_pipeline_layout_parallel_matrix.sh BUILD_DIR INPUT_BIN CHECKPOINT [NUM_GPUS]
+#
+# NUM_GPUS defaults to 8. Combinations that need more GPUs than are available are
+# skipped and reported as SKIP, so the script is still useful on a 2- or 4-GPU host.
+set -euo pipefail
+
+if [[ $# -lt 3 || $# -gt 4 ]]; then
+    echo "Usage: $0 BUILD_DIR INPUT_BIN CHECKPOINT [NUM_GPUS]" >&2
+    exit 2
+fi
+
+build_dir="$(realpath "$1")"
+input_bin="$(realpath "$2")"
+checkpoint="$(realpath "$3")"
+num_gpus="${4:-8}"
+source_dir="$(realpath "$(dirname "$0")/../..")"
+gpt2="$build_dir/gpt2"
+infini_run="$build_dir/infini_run"
+compare="$source_dir/scripts/precision_check/precision_compare.py"
+
+for path in "$gpt2" "$infini_run" "$input_bin" "$checkpoint" "$compare"; do
+    if [[ ! -e "$path" ]]; then
+        echo "Required test input does not exist: $path" >&2
+        exit 2
+    fi
+done
+
+# The layer count decides every partition below, so read it out of the LLMC header
+# instead of hard-coding one model size.
+layers="$(python3 - "$checkpoint" <<'PY'
+import struct, sys
+with open(sys.argv[1], "rb") as f:
+    header = struct.unpack("<256i", f.read(1024))
+magic, version, _, _, n_layer = header[0], header[1], header[2], header[3], header[4]
+assert magic == 20240326 and version == 3, f"not an fp32 LLMC checkpoint: {magic}/{version}"
+print(n_layer)
+PY
+)"
+if (( layers < 4 )); then
+    echo "This regression needs a checkpoint with at least 4 Transformer layers; got $layers" >&2
+    exit 2
+fi
+
+test_dir="$(mktemp -d /tmp/infinitrain-pipeline-matrix.XXXXXX)"
+trap 'rm -rf -- "$test_dir"' EXIT
+
+# Non-uniform partitions derived from the real layer count.
+pp2_head=$(( layers / 4 )); (( pp2_head > 0 )) || pp2_head=1
+pp2_partition="${pp2_head},$(( layers - pp2_head ))"
+pp4_partition="1,$(( layers - 3 )),1,1"
+quarter=$(( layers / 4 ))
+quarter_remainder=$(( layers - 3 * quarter ))
+chunk_layout="0:${quarter},1:${quarter},1:${quarter},0:${quarter_remainder}"
+# Equal layer costs plus an lm head worth two layers; the balanced split is the smallest
+# stage 0 size that minimises max(k, layers - k + 2).
+lm_head_costs="$(python3 -c "print(','.join(['1'] * $layers) + ',L:2')")"
+lm_head_head="$(python3 -c "
+layers = $layers
+print(min(range(1, layers), key=lambda k: (max(k, layers - k + 2), k)))")"
+megatron_layout="Et*${quarter}|t*${quarter}|t*${quarter}|t*${quarter_remainder}NL"
+
+common_args=(
+    --device=cuda
+    --input_bin="$input_bin"
+    --llmc_filepath="$checkpoint"
+    --batch_size=4
+    --sequence_length=64
+    # Must stay divisible by batch_size * sequence_length * DP size for every DP size below.
+    --total_batch_size=2048
+    --num_iteration=1
+    --freq_generate_txt=1000
+    --dtype=float32
+)
+
+gpu_list() { seq -s, 0 $(( $1 - 1 )); }
+
+extract_loss() { sed -n 's/.*train loss \([^ |]*\).*/\1/p' "$1" | tail -n 1; }
+
+assert_close() {
+    # assert_close NAME REFERENCE ACTUAL TOLERANCE
+    awk -v name="$1" -v reference="$2" -v actual="$3" -v tolerance="$4" 'BEGIN {
+        if (reference == "" || actual == "") {
+            printf "%s: missing loss (reference=%s actual=%s)\n", name, reference, actual > "/dev/stderr";
+            exit 1;
+        }
+        difference = reference - actual;
+        if (difference < 0) difference = -difference;
+        if (difference > tolerance) {
+            printf "%s: loss mismatch reference=%s actual=%s difference=%g\n",
+                   name, reference, actual, difference > "/dev/stderr";
+            exit 1;
+        }
+    }'
+}
+
+failures=0
+results=()
+
+echo "=== reference: single GPU, ${layers} layers ==="
+reference_grad="$test_dir/reference-grad"
+reference_log="$test_dir/reference.log"
+env GLOG_logtostderr=1 CUDA_VISIBLE_DEVICES=0 \
+    "$gpt2" "${common_args[@]}" --dump_gradients="$reference_grad" >"$reference_log" 2>&1
+reference_loss="$(extract_loss "$reference_log")"
+if [[ -z "$reference_loss" ]]; then
+    echo "Failed to extract the reference loss from $reference_log" >&2
+    exit 1
+fi
+find "$reference_grad" -type f -name '*.npy' -printf '%f\n' | sort >"$test_dir/reference-files"
+echo "reference loss: $reference_loss ($(wc -l <"$test_dir/reference-files") gradients)"
+
+run_case() {
+    # run_case NAME GPUS TOLERANCE COMPARE_GRADIENTS EXPECTED_LAYOUT_LINE -- <gpt2 flags...>
+    local name="$1" gpus="$2" tolerance="$3" compare_gradients="$4" expected_layout="$5"
+    shift 6 # drop the parsed fields and the "--" separator
+
+    if (( gpus > num_gpus )); then
+        echo "--- $name: SKIP (needs $gpus GPUs, have $num_gpus)"
+        results+=("SKIP  $name (needs $gpus GPUs)")
+        return 0
+    fi
+
+    local log="$test_dir/$name.log"
+    local grad_dir="$test_dir/$name-grad"
+    local extra=()
+    if [[ "$compare_gradients" == "grad" ]]; then extra=(--dump_gradients="$grad_dir"); fi
+
+    echo "--- $name: $gpus GPU(s)"
+    if ! env GLOG_logtostderr=1 CUDA_VISIBLE_DEVICES="$(gpu_list "$gpus")" \
+        "$infini_run" --nproc_per_node="$gpus" "$gpt2" "${common_args[@]}" "${extra[@]}" "$@" \
+        >"$log" 2>&1; then
+        echo "$name: training process failed, see $log" >&2
+        tail -n 20 "$log" >&2
+        results+=("FAIL  $name (training failed)")
+        failures=$(( failures + 1 ))
+        return 0
+    fi
+
+    local ok=1
+    if [[ -n "$expected_layout" ]] && ! grep -Fq "$expected_layout" "$log"; then
+        echo "$name: expected layout line not found: $expected_layout" >&2
+        ok=0
+    fi
+    if ! assert_close "$name" "$reference_loss" "$(extract_loss "$log")" "$tolerance"; then
+        ok=0
+    fi
+    if [[ "$compare_gradients" == "grad" ]]; then
+        find "$grad_dir" -type f -name '*.npy' -printf '%f\n' | sort >"$test_dir/$name-files"
+        if ! diff -u "$test_dir/reference-files" "$test_dir/$name-files" >"$test_dir/$name-files.diff"; then
+            echo "$name: gradient file set differs from the reference" >&2
+            head -n 20 "$test_dir/$name-files.diff" >&2
+            ok=0
+        elif ! python3 "$compare" --dir1 "$reference_grad" --dir2 "$grad_dir" --atol 1e-5 --rtol 0 \
+            >"$test_dir/$name-grad.log" 2>&1; then
+            echo "$name: per-parameter gradients differ from the reference" >&2
+            tail -n 10 "$test_dir/$name-grad.log" >&2
+            ok=0
+        fi
+    fi
+
+    if (( ok )); then
+        results+=("PASS  $name (loss $(extract_loss "$log"), $gpus GPU)")
+    else
+        results+=("FAIL  $name")
+        failures=$(( failures + 1 ))
+    fi
+}
+
+# Layout x parallelism matrix. TP shards parameters, so TP cases compare the loss only.
+run_case pp2_partition        2 1e-5 grad "stage 0: embedding layers[0,${pp2_head})" -- \
+    --pipeline_parallel=2 --pipeline_layer_partition="$pp2_partition"
+run_case pp4_partition        4 1e-5 grad "stage 1: layers[1,$(( layers - 2 )))" -- \
+    --pipeline_parallel=4 --pipeline_layer_partition="$pp4_partition"
+run_case pp2_ddp2             4 1e-5 grad "stage 0: embedding layers[0,${pp2_head})" -- \
+    --pipeline_parallel=2 --pipeline_layer_partition="$pp2_partition"
+run_case pp2_ddp4             8 1e-5 grad "stage 0: embedding layers[0,${pp2_head})" -- \
+    --pipeline_parallel=2 --pipeline_layer_partition="$pp2_partition"
+run_case pp2_tp2              4 1e-5 loss "stage 0: embedding layers[0,${pp2_head})" -- \
+    --pipeline_parallel=2 --tensor_parallel=2 --pipeline_layer_partition="$pp2_partition"
+run_case pp2_tp2_ddp2         8 1e-5 loss "stage 0: embedding layers[0,${pp2_head})" -- \
+    --pipeline_parallel=2 --tensor_parallel=2 --pipeline_layer_partition="$pp2_partition"
+run_case pp4_tp2              8 1e-5 loss "stage 1: layers[1,$(( layers - 2 )))" -- \
+    --pipeline_parallel=4 --tensor_parallel=2 --pipeline_layer_partition="$pp4_partition"
+run_case pp2_costs            2 1e-5 grad "stage 0: embedding layers[0,1)" -- \
+    --pipeline_parallel=2 --pipeline_layer_costs="$(python3 -c "print(','.join(['10'] + ['1'] * ($layers - 1)))")"
+run_case pp2_costs_lm_head    2 1e-5 grad "stage 0: embedding layers[0,${lm_head_head})" -- \
+    --pipeline_parallel=2 --pipeline_layer_costs="$lm_head_costs"
+run_case pp2_costs_lm_head_ddp2 4 1e-5 grad "stage 0: embedding layers[0,${lm_head_head})" -- \
+    --pipeline_parallel=2 --pipeline_layer_costs="$lm_head_costs"
+run_case pp2_vpp2_chunk       2 1e-5 grad "" -- \
+    --pipeline_parallel=2 --virtual_pipeline_parallel=2 --pipeline_chunk_layout="$chunk_layout"
+run_case pp2_vpp2_chunk_ddp2  4 1e-5 grad "" -- \
+    --pipeline_parallel=2 --virtual_pipeline_parallel=2 --pipeline_chunk_layout="$chunk_layout"
+run_case pp4_megatron         4 1e-5 grad "" -- \
+    --pipeline_parallel=4 --pipeline_model_parallel_layout="$megatron_layout"
+run_case pp4_megatron_ddp2    8 1e-5 grad "" -- \
+    --pipeline_parallel=4 --pipeline_model_parallel_layout="$megatron_layout"
+
+echo
+echo "=== summary (reference loss $reference_loss) ==="
+printf '%s\n' "${results[@]}"
+
+if (( failures )); then
+    echo "FAILED: $failures case(s) did not match the single-GPU reference" >&2
+    exit 1
+fi
+echo "PASS: every layout x parallelism combination matched the single-GPU reference"
diff --git a/tests/distributed/test_pipeline_layout_suggestion.py b/tests/distributed/test_pipeline_layout_suggestion.py
index c9102caf3..4dda6df51 100644
--- a/tests/distributed/test_pipeline_layout_suggestion.py
+++ b/tests/distributed/test_pipeline_layout_suggestion.py
@@ -18,6 +18,25 @@ def test_balances_contiguous_costs(self):
         self.assertEqual(counts, [1, 5])
         self.assertEqual(stage_costs, [10, 5])
 
+    def test_balances_embedding_and_lm_head_cost(self):
+        layers = [1] * 8
+        self.assertEqual(MODULE.balanced_partition(layers, 2)[0], [4, 4])
+
+        counts, stage_costs = MODULE.balanced_partition(layers, 2, lm_head_cost=2)
+        self.assertEqual(counts, [5, 3])
+        self.assertEqual(stage_costs, [5, 5])
+
+        counts, stage_costs = MODULE.balanced_partition(layers, 2, embedding_cost=2)
+        self.assertEqual(counts, [3, 5])
+        self.assertEqual(stage_costs, [5, 5])
+
+        # Equal costs on both ends cancel out; a zero cost means "negligible".
+        self.assertEqual(MODULE.balanced_partition(layers, 2, 2, 2)[0], [4, 4])
+        self.assertEqual(MODULE.balanced_partition(layers, 2, 0, 0)[0], [4, 4])
+
+        with self.assertRaises(ValueError):
+            MODULE.balanced_partition(layers, 2, lm_head_cost=-1)
+
     def test_reads_layer_profiler_records(self):
         records = """
 0 2026-08-21 TransformerLayer.0 cuda:0 12 100 1