PP自动并行算法
背景与挑战
流水线并行是将模型网络层切分成多个stage，再把stage映射到不同的设备上，使得不同设备并行计算神经网络的不同部分。流水线并行大大缓解了单卡内存瓶颈问题，并通过多卡之间的流水训练提高了硬件的利用率。流水线并行成为了当前大模型训练最常用的并行方式之一。然而当前流水线并行在内存消耗和性能方面并非最优，主要存在两大问题：
1）内存不均衡：当前流水线常用调度模式（1F1B）下，靠近模型前面层的流水线stage的内存占用远多于后面的stage内存占用，并且内存占用差距有2~3倍，总体上可训练的模型规模受限于PP-Stage 0的显存消耗。
2）流水线气泡：流水线1F1B调度策略在每个设备上交替进行小批次数据的前向后向计算，由于各流水设备之间计算负载不均衡或者网络通信的波动，导致设备与设备之间存在等待（流水线气泡），影响训练性能。
解决方案
本系统基于在线profiling+PP建模搜索，通过使能内存优化模块、性能优化模块分别最大化流水线并行训练的内存和性能。内存优化模块旨在通过自动寻找流水线并行中stage的最优层分布和细粒度重计算模块，均匀分配每个卡上的显存，优化存在显存瓶颈的PP-Stages，降低峰值内存；性能优化模块采用mbs序列和前向反向调度序列自动寻优和多流异步通信机制，压缩流水线气泡，提升训练性能。
内存优化模块
基于在线profiling+PP建模搜索，自动构建出最优的内存排布方案均衡化各个stage之间的内存开销，降低峰值内存的同时最小化端到端训练时间，具备较好的易用性和泛化性。具体而言，在层分布和细粒度重计算的联合搜索空间自动寻优内存排布方案： ① PP层分布切分：采用不均匀层切分策略，自动搜索最优层切分方式，均衡化每个卡消耗的显存，从而优化存在显存瓶颈的PP-Stages，降低峰值内存。 ② 细粒度重计算：利用流水线气泡时间来做重计算，保证性能不劣化，通过自动寻优细粒度的重计算策略，进一步降低峰值内存。
性能优化模块
在满足训练峰值内存开销不超过设备最大内存容量的条件下，通过自动寻找流水线并行中最优的mbs序列及前向反向调度序列，最小化端到端训练时间。 ① 动态mbs：在给定的gbs下，自动搜索最优mbs序列。通过小mbs加速流水线的启动与冷却，压缩气泡时间，稳态阶段自动寻找最高效的mbs进行计算，缩短稳态阶段计算时间，提升端到端训练性能。 ② 前反向调度：通过调整流水线并行过程中前反向计算的顺序，结合多流异步通信机制，压缩流水线稳态气泡，提升训练性能。
PP自动并行系统如下图所示：
使用场景
该系统主要用于开启流水线并行的训练场景，使用PP自动并行系统可有效优化内存不足或流水线气泡占比过大的问题。 使用条件:
--pipeline-model-parallel-size >= 2；- 内存、性能优化模块不能同时使用。
使用方法
（1）当内存不足时，可启用PP自动并行内存优化模块，请首先在训练脚本中添加 --automated-pipeline 标志启用功能。
（2）当流水线气泡过大导致训练性能不优时，可启用PP自动并行性能优化模块，请首先在训练脚本中添加 --automated-pipeline-perf 标志启用功能。
使用效果
PP自动并行内存优化模块收益：LLaMA2-7B，LLaMA-13B，LLaMA2-70B等使用流水线并行PP配置训练的模型，叠加本算法后平均峰值内存减少11.5%，平均性能劣化小于1%。性能优化模块收益：LLaMA2-7B，LLaMA-13B，LLaMA3-8B等使用流水线并行PP配置训练的模型，叠加本算法后平均性能提升7.6%。