当前飞桨集合通信模式已经可以支持文心 ERNIE 千亿语言模型的训练能力，其 Sharding-DP 策略更是在近期助力文心 ERNIE 的多项任务分数刷新 GLUE 榜单。而这个 Sharding-DP 策略正是飞桨集合通信模式为了训练 ERNIE 这样的大规模复杂模型所支持的多种并行策略中的一种。那么飞桨是使用哪些策略成功支持文心 ERNIE 千亿语言模型训练的呢？这些策略是如何工作的呢？接下来将为大家详细介绍。
ERNIE 千亿级模型采用 100 多层 Transformer 网络结构，计算复杂，训练需要占用 T 级显存资源，如果想用更少的机器高效训练，必须采取一系列性能优化和显存优化措施。
首先看如何性能优化。我们通过一个公式来看哪些因素可以影响训练速度，在固定的硬件环境下：
总训练速度 ∝ 单卡速度 * 卡数 * 多卡加速比
其中单卡速度由数据读取和计算速度决定；多卡加速比由计算 / 通信效率决定。显而易见，这三个是关键因素。除了单卡可以使用的算子融合、混合精度之类的基础性能优化策略之外，分布式训练还引入一系列并行策略。并行策略的核心思想是将数据和计算有关的图 / 算子切分到不同设备上，同时尽可能降低设备间通信所需的代价，合理使用多台设备资源，实现高效的并发调度训练，最大化提升训练速度。常见并行策略有数据并行 DP（Data Parallel）、Layer 间并行（流水线并行 PP，Pipeline Parallel）、Layer 内并行（模型并行 MP，Model Parallel）。如下表所示。我们从设备资源和计算 / 通信效率来分析三种策略的优缺点：
数据并行训练加速比最高，但要求每个设备上都备份一份模型，显存占用比较高。为此我们的改进方案是分组参数切片数据并行策略（具体原理后文介绍），兼容了 MP+DP 的优势，但缺点是通信量大。
模型并行，通信占比高，适合在机器内做模型并行且支持的模型类型有限。
流水线并行，训练设备容易出现空闲状态，加速效率没有 DP 高；但能减少通信边界支持更多的层数，适合在机器间使用。
其次看显存问题，通过下表分析的显存占用来源可以看出，上述的并行策略同样可以很好地应对不同来源的显存占用，更多的层数可以通过流水线并行和分组参数切分策略来解决；某层参数很大可以通过模型并行来解决；其次飞桨还提供一些其它灵活的优化方式，例如每层输出占用的显存，可以通过重计算和 Offload 来解决。
综上所述，针对性能优化和显存优化，几种并行策略都有用武之地，但是同时也有各自的局限性，所以如果想高效训练千亿模型，需要这几种策略相互组合，取长补短，发挥各自的优势。
那么如何组合呢？飞桨研发人员首先在单机内使用模型并行和分组参数切片组合的 2D 策略，这么选择的原因是这两个策略通信量较大，适合使用机器内的卡间通信；然后为了承载千亿规模模型，再叠加流水线并行策略，使用多台机器共同分担；最后为了做到高效，在外层又叠加了数据并行来增加并发数量，提升整体训练速度。这样业内首个 4D 混合并行策略就诞生了。
图 8 4D 混合并行策略示意图
更多详情，请点击阅读原文
内容中包含的图片若涉及版权问题，请及时与我们联系删除
评论
沙发等你来抢