feat: add configurable pipeline layouts

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。测试使用临时目录，
+退出时自动清理梯度和日志。

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%。