PP自动并行

问题分析

靠近模型前面的流水线stage的内存占用多于模型后面的stage内存占用，并且内存占用差距有2~3倍，总体上模型规模受限于PP-Stage 0的显存。当前通过离线建模搜索手动配置PP层的分布可以缓解显存不平衡的问题，但存在多模型的泛化性问题，端到端训练效果无法得到保证。

解决方案

本算法通过自动寻找流水线并行中stage的最优网络层分布和细粒度重计算模块，均匀分配每个卡上的显存，优化存在显存瓶颈的PP-stages，降低峰值内存。

解决思路

基于在线建模搜索自动构建出最优的内存排布方案使得各个stage之间的内存相对均衡，降低峰值内存的同时最小化端到端训练时间，具备更好的易用性和泛化性。通过数学理论建模+在线profiling的方式。从PP维度，在层分布和细粒度重计算的联合搜索空间自动寻优： ① PP层分布：搜索最优层切分方式，均匀分配每个卡上的显存，优化存在显存瓶颈的PP-stages，降低峰值内存。 ② 细粒度重计算模块：通过自动寻优重计算策略，进一步降低峰值内存，同时保证性能不劣化。

PP自动并行1.0流程如下图所示：

使用场景

该特性主要用于训练过程中显存不足的场景，使用PP自动并行可有效降低显存的占用。 **注意：**使用条件：--pipeline-model-parallel-size >= 2

使用方法

启用PP自动并行，请首先在训练脚本中添加 --automated-pipeline 标志开启PP自动并行策略。

使用效果

LLaMA2-7B，LLaMA-13B，LLaMA2-70B等使用流水线并行PP配置训练的模型，叠加本算法后平均峰值内存减少5%，平均性能劣化小于1%。