大规模分布式训练简介

故事一说

古有曹冲称象，为人津津乐道。而“大规模分布式训练”堪称机器学习领域的的“曹冲称象”。
三国时曹操等大臣面临的问题是大象的重量超过了称的极限。如果我们把故事中的对象替换一下：

• 大象替换成类似GPT3这类的参数超万亿的大模型
• 称替换成类似GPU-A100,这类的机器学习处理器
• 重量替换成机器训练大模型所需的内存

那什么是“大规模分布式训练”呢？

翻译翻译就是：大模型预训练所需的内存超过了单张深度学习处理器的内存。于是机器学习的工程师们纷纷化身曹冲，提出肢解大模型而又不尽量影响模型运行的的方案！

比较尴尬的是， 现代工程师们找不到一只可以容纳整个模型的船（也许2077年，nvidia发布A1000或者H1000，就是我们要找的船），于是工程师只好另辟蹊径，拿多杆秤称一头象了。如何协调多杆秤工作，就是一个技术活了。
简而言之：大规模分布式训练的目标是协调多台机器简单高效的运行大规模的模型。

问题剖析

假设运行模型所需的内存为 Requried_Momory，单卡的内存 Device_Memory。

现状是 Required_Memory >> Device_Memory，这是大规模分布式训练面临的核心矛盾，供需严重失衡。

解决这个矛盾的就是让这个不等式变成 Required_Memory <= Device_Memory，使供需平衡。

解决方案无非就是两个思路：节流和开源。
节流，即尝试减少 Required_Memory。开源就是增大 Device_Memory。

模型内存分析

弄清问题的本质是供需不平衡，接下来我们需要厘清这这两个变量是怎么产生的。
Device_Memory : 由硬件决定的，当芯片型号确定时，可以认定这是个常量。
Required_Memory: 根据笔者了解，模型训练中内存消耗主要来自两块，分别是

• 神经元的中间值
• 参数值（包括权值，梯度值和优化器的状态值）

大规模分布式训练的解决方案

节流

假设要训练一个参数量为10亿的模型，即参数量大小固定了，无法进行节流，节流的主意只能打在了神经元的中间值上。

常见的节流策略如下：

1 激活重计算（Activation-Recompute）

在网络中标记少量的算子 ，前向计算只保留这些被标记的算子的输出结果（激活值）， 其余前向算子的输出结果直接被释放，这样就可以极大减少激活值消耗的内存。当反向更新梯度需要前向算子的输出时，利用被标记的算子重新计算获取。
激活重计算是一个以时间换空间的策略。

开源

开源的目标是扩大 Device_Memory，但 Device_Memory 受硬件限制，无法改变。但脑袋一动，单卡的 Device_Memory 是固定的，可以多张卡协作，共同运行模型，这样就相当于扩大了 Device_Memory。

常见的开源策略如下：

1 数据并行（Data Parallel）

传统的数据并行（DDP），在反向计算完成后，每张卡上的梯度和权值，优化器的状态值是完全一致的。即各张卡之间进行的是完全一致的计算。因此可以将这部分的计算分散到各张卡上，每张卡仅保存部分的梯度，梯度和状态值。
针对大模型的数据并行策略，微软在DeepSpeed 框架提出了零冗余优化器（ZeRO Redunrancy Optimizer, 以下简称ZeRO）实现了节省模型训练过程中所需的内存。

2 模型并行（Model Parallel）

将模型中特定子图中的权值均匀的分配到多张卡上，每张卡进行部分计算，然后通过节点间通信获取完整的输出。从而有效降低单卡上的内存消耗。
模型并行策略的详细介绍：大规模分布式训练–模型并行策略

流水并行（Pipeline Parallel）

将网络切成若干子网络，每个子网络单独的运行在一张卡上，上一层的输出作为下一层的输入。

3 负载均衡（ZeRO-Offload）（不推荐）
在深度学习中，GPU 显存（Device Memory）的特点是昂贵、高速且容量小，而 CPU 的主存（Host Memory）的特点是便宜、相对低速和大容量； 那么将前向计算中的一些
暂时用不到的 activation 临时换出到 CPU 主存上，等到反向计算需要时再换入到 GPU 显存里，通过这种方式也可以节省显存。

负载均衡是一种将CPU的内存加入Device_Memory的方式，但这种换入换出也是有时间成本的。因此也是一种以时间换空间的策略。大批量的频繁的进行换入换出操作，很容易让带宽成为瓶颈，不推荐使用。

开源节流，以期供需平衡

通过以上介绍，可以看出每个方案都不是无损的，完美的。节流的方案会带来性能上的损失。开源的方案很氪金；同时卡数越多，通信开销越大。
同时需要注意的是，对于一个大规模的模型，很难依靠单个策略解决内存不足问题，需要多个方案组合使用，开源节流一起支棱起来。因此需要研究人员结合模型特征和硬件条件，平衡各方案的利弊，找到最符合模型的解决方案。