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

[why1te]

总结

本 PR 为 Pipeline Parallel（流水线并行）增加统一的布局描述与解析能力，使 GPT-2 和 LLaMA 3 可以在同一套模型构建、权重加载和调度流程中使用以下三种布局：

• 默认均匀布局；
• 自定义每个逻辑 Chunk 的 Transformer 层数；
• 显式指定每个 Chunk 所属的物理 Stage。

未设置新参数时，程序继续使用原有的均匀分层与交错分配规则，默认行为保持不变。

设计方案

1. 统一布局总结

类型	作用
LayerRange	表示全局 Transformer 层的左闭右开区间 [begin, end)
PipelineChunkSpec	表示显式布局输入中的 stage_id 和层数
PipelineChunkLayout	保存 Chunk 的全局编号、本地编号、Stage 归属和层范围
PipelineStageLayout	保存一个 Stage 的本地 Chunk 列表和特殊模块标记
PipelineLayout	保存完整、只读的 Pipeline 布局，并提供统一查询接口

PipelineLayout 提供三种构建接口：

接口	行为
BuildUniformLayout()	均匀划分层数，并按全局 Chunk 编号交错分配 Stage
BuildCustomLayout()	使用用户给出的 Chunk 层数，保持默认交错归属
BuildChunkLayout()	同时使用用户给出的层数和显式 Stage 归属

主要查询接口包括 GetStage()、GetChunk()、GetStageForLayer()、GetLocalLayerIndex()、GetInputStage()、GetOutputStage() 和 GetMaxLocalChunks()。

2. 参数解析与布局创建

新增两个互斥参数：

参数	含义
--pipeline_layer_partition	按全局 Chunk 顺序指定每个 Chunk 的层数
--pipeline_chunk_layout	按全局 Chunk 顺序指定 stage_id:layer_count

ParsePipelineLayoutRequest() 负责格式与维度检查；取得真实模型层数后，ResolvePipelineLayout() 调用对应布局工厂。随机初始化和 checkpoint 加载两条入口最终都向模型、Loader、Pipeline 包装器和调度器传递同一个布局对象。

3. 模型构建与权重加载

• TransformerModel 通过 ConvertLayoutToStageInfo() 复用现有模型构建逻辑，不维护第二套 Transformer 实现；
• Embedding 跟随第一个逻辑 Chunk；Final Norm 和 LM Head 跟随最后一个逻辑 Chunk；
• GPT-2 和 LLaMA 3 Loader 根据全局层号查询 owner，并将全局层号转换为目标 Stage 内的本地连续层号；
• 支持一个 Stage 持有多个不连续的全局层区间。

4. Pipeline 包装、调度与通信

• 调度任务同时记录全局 Chunk、本地 Chunk 和物理 Stage；
• 相邻 Chunk 属于不同 Stage 时使用异步收发通信；
• 相邻 Chunk 属于同一 Stage 时直接传递本地激活值，不经过跨 Stage 通信；
• 反向传播一次穿过连续的本地 Chunk，避免重复执行；
• 对跨 Stage 后再次返回旧 Stage 的显式布局，使用能够保证通信继续推进的 GPipe 顺序，避免双方同时发送导致等待；
• 不涉及 Stage 回访的旧布局继续保持原任务顺序。

5. 兼容性设计

• 未提供新参数时，默认布局、模型层范围和调度顺序保持不变；
• checkpoint 文件格式不变；
• 保留旧 GetStageInfo() 和旧构造函数作为兼容入口；
• 精度检查器将本地层名还原为全局层名，使不同布局可以按相同名称比较输出和梯度。

使用示例

PP=2、每个 Stage 一个 Chunk，按 7,5 自定义分层：

--pipeline_parallel=2 \
--virtual_pipeline_parallel=1 \
--pipeline_layer_partition=7,5

PP=2、最大每个 Stage 两个 Chunk，同时自定义分层与归属：

--pipeline_parallel=2 \
--virtual_pipeline_parallel=2 \
--pipeline_chunk_layout=0:4,0:3,1:3,1:2 \
--freq_generate_txt=0 \
--val_loss_every=0 \
--sample_every=0

第二个示例生成四个逻辑 Chunk，层数为 4,3,3,2，owner 为 0,0,1,1。Embedding 位于 Stage 0，Final Norm 和 LM Head 位于 Stage 1。

测试结果

1. 构建与单元测试

• GPT-2 和 LLaMA 3 的 CPU、CUDA 目标均构建成功；
• Pipeline 相关 CTest：32/32 通过；
• 命令行参数用例：16/16 通过。

单元测试覆盖布局构建与非法输入、参数解析、特殊模块归属、GPT-2/LLaMA Loader、默认调度兼容、任意 owner、本地 Chunk 链、循环通信推进和输入维度检查。

2. 双卡端到端正确性

测试项	结果
不同布局持续训练	均完成 10 步
精度采集训练	正常退出
模块输出与输出梯度	通过，每组比较 254 个张量
参数梯度、loss、输入与标签	通过，每组比较 149 个参数
数值有效性	2,140 条张量记录均无 NaN/Inf
默认路径兼容性	逐步 loss 最大绝对差约 1e-6
运行稳定性	训练均正常退出，无超时或 CUDA/NCCL 错误

严格 FP32 比较使用 atol=1e-5, rtol=0。候选 D（4,3,3,2 / 0011）还与默认 vPP 布局进行了单独复核：两次训练最终 loss 均为 5.356190，254/254 个模块输出与输出梯度全部通过。

3. 性能与 Pipeline 负载

测试固定 GPT-2 124M、FP32、PP=2、vPP=2、训练批次大小 4、序列长度 128、4 个 microbatch 和每步 2048 个 token。每组独立运行 4 次、每次 100 步，丢弃前 20 步预热数据，全部样本均保留。

布局	层数 / owner	吞吐中位数（token/s）	相对 A	Stage 0/1 计算时间（ms/步）	Stage 差值	流水线空泡
A 默认布局	3,3,3,3 / 0101	33,066.1	基线	30.57 / 38.18	7.61 ms	46.23%
B 仅调整 owner	3,3,3,3 / 0011	41,159.4	+24.48%	30.71 / 38.17	7.46 ms	37.17%
C 仅调整分层	4,3,3,2 / 0101	34,311.1	+3.77%	35.27 / 33.47	1.80 ms	44.23%
D 同时调整	4,3,3,2 / 0011	43,518.2	+31.61%	35.25 / 33.48	1.77 ms	30.64%

流水线空泡表示 Stage 没有执行非 NCCL 模型计算的时间比例，不等同于 GPU 完全空闲。接收等待归因是根据任务语义和 NCCL 调用顺序得到的保守下界。性能结果仅适用于当前机器、模型、参数和四个候选，不表示全局最优。