主要观点总结
本文主要介绍了针对大规模语言模型流水线并行训练中的负载均衡问题，提出一种基于符号抽象的全局线性规划算法，旨在同时均衡计算和内存负载，提高模型训练效率。文章详细描述了流水线并行的典型编排方式、技术背景和挑战，以及自动负载均衡算法的实现原理和应用流程。
关键观点总结
关键观点1: 流水线并行训练中的挑战和痛点
随着模型规模的增大，流水线并行训练面临内存和计算负载不均衡的问题，导致训练效率低下。手工调优需要大量时间和经验，且难以找到最优策略。
关键观点2: 基于符号代价模型和线性规划的自动负载均衡算法
通过对模型的性能和内存进行分解，构建线性规划问题，在分钟级时间内求解出最优的layer分配和重计算策略，提高模型训练效率。
关键观点3: Dryrun工具的应用
Dryrun工具能够模拟实际执行过程，提供准确的内存预估，结合自动负载均衡工具，实现单卡仿真大集群并行策略的调优，降低调优资源消耗，提升调优效率。
关键观点4: 昇思MindSpore提供的解决方案
昇思MindSpore团队提出了一种全局求解方法，通过自动生成最优的stage-layer分配和重计算策略，解决流水线并行策略配置对模型端到端性能的影响。该方案通过结合自动流水负载均衡工具和Dryrun集群内存仿真工具，实现低成本高效模型调优。
正文
图1
DeepSeek-V3 671B 大集群调优
大模型训练中为了保障训练的效果（loss的收敛和最终的评测效果），通常会控制global batch size不能太大。在限定global batch size的情况下，流水并行的micro size数也受到了限制，尤其是在大集群训练时，数据并行dp的维度数非常大。根据bubble计算公式（p-1）/m，当m过小时，bubble相对较大，设备空等时间较多，影响训练性能。针对interleave场景的调优难题，自动负载工具做到了对内存排布和计算时间的建模，用户可以任意给定chunk数，工具将给出此条件下的最优解。同时工具自带模拟器，方便用户对比各种不同的并行策略。
下图通过流水并行模拟器展示了超大集群的千亿参数Llama模型训练的调优实例。
图2 超大集群千亿参数规模Llama调优
图2中共有三种流水线负载均衡配置，分别是对layer进行平均切分，且不开启重计算；平均切分开启全部重计算以及经过自动负载均衡调优后的策略。
最终通过超大集群实测，自动负载均衡调优相比于专家调优提升性能14%。
# 02
图1展示了最简单的流水并行切分，将模型的层进行了切分，分布在不同的设备上，虽然节省了内存占用，但是可以很明显的看到非常大的bubble，在实际训练时性能会非常差，因此在实际训练中基本上不会直接使用native pipeline。谷歌为了提高流水线的效率，提出了Gpipe，如图2所示，设备上的训练数据被分成了更细的micro batch，不同的micro batch可以并行计算，这样就压缩了训练过程中的bubble，提升了训练性能。
图3 Native Pipeline Parallelism
图4 GPipe
在这种范式下，每个device上的内存压力是一样的，都需要在反向之前存储一个step中所有layer的激活值（activation/Acti.）。因此，GPipe虽然通过流水线的方式把模型分割到了不同的设备上，但其所采取的流水线编排方式在调大micro batch时，内存占用过高，影响了其扩展性。因为通常来说，micro batch越大，可以流水并行的计算部分越多，bubble越小。
图5 1F1B & interleaved pipeline
PipeDream在GPipe batch-forward-batch-backward的编排基础上提出了一种新的编排方式，1-forward-1-backward(1F1B)(图5上半部分），流水线在一个前向处理完成后便可以开始计算反向，这样激活值在反向执行后就可以及时释放，相比于Gpipe提高了内存使用效率，增加micro batch的数量不再会对设备上的峰值内存有影响。在这种1F1B的编排方式之下，流水线的最小bubble为（p-1）/m，p和m分别是流水线stage数和micro batch数。性能持平GPipe的同时进一步节省了内存压力。以下表为例，当micro batch数量为8的时候，stage = 4的流水线并行最少也能为设备节省50%的内存。
因为在峰值内存方面的优势，业界的流水线并行普遍采用了1F1B的编排方式。经过深入分析可以发现，1F1B编排模式下，设备的内存压力相比于GPipe而言减少了很多，然而流水线中每个stage的内存压力是不平衡的（见上表最右列）。每个设备上的峰值内存由进入稳定状态之前1F1B累积的激活值决定，第i个stage累计的激活值带来的内存开销可以通过公式计算，计算结果显示在流水线中靠前的设备（即i数值较小的）峰值内存更高。
Megatron团队在2021年提出了流水线交织（interleaved pipeline）的排布方式（图5下半部分），这是一种基于1F1B的优化，通过把一个设备上原本连续的layer继续分割成多个chunk，进一步细化了流水线的颗粒度，进一步压缩bubble到（p-1）/vm，v是每个stage切分出的chunk数。当stage = 4时，如果micro batch数为16，理论上1F1B下bubble最小占比为18%，而分出两个chunk之后，bubble的占比极限下可以进一步被压缩到9%左右。不过，流水线交织并不能解决内存不平衡，只是引入了新变量chunk，导致调优又新增了一个维度，复杂度几何级上升。
