<div class="tip inlineBlock share"> 译文链接: [How to Train Really Large Models on Many GPUs?](https://lilianweng.github.io/posts/2021-09-25-train-large/) **[Updated on 2022-03-13: add** [expert choice routing](https://lilianweng.github.io/posts/2021-09-25-train-large/#ec).] [Updated on 2022-06-10]: [Greg](https://gregbrockman.com/) **and I wrote a shorted and upgraded version of this post, published on OpenAI Blog:** [“Techniques for Training Large Neural Networks”](https://openai.com/blog/techniques-for-training-large-neural-networks/) </div> # 前言 近年来，我们在许多自然语言处理基准任务上取得了更好的结果，这得益于更大的预训练[语言模型](https://lilianweng.github.io/posts/2019-01-31-lm/)。如何训练大型和深层的神经网络是具有挑战性的，因为它需要大量的 GPU 内存和长时间的训练。 然而，单个 GPU 的内存有限，许多大模型的大小已经超过了单个 GPU 的容量。有几种并行性范例可以实现跨多个 GPU 进行模型训练，还有各种模型架构和内存节省设计，可以帮助实现训练非常大的神经网络。 --- # Training Parallelism 训练非常大的神经网络模型的主要瓶颈在于对GPU设备的巨大内存需求，已经远远超过了单 GPU 设备的承载能力。除了模型权重（例如数百亿个浮点数），通常还需要存储中间计算输出，如梯度和优化器状态（例如 Adam 中的动量和变化）。此外，训练大型模型通常需要大量的训练语料库，因此单个进程可能需要很长时间。 因此并行训练是必要的。并行性可以在不同的维度上进行，包括数据、模型架构和张量操作。 --- ## Data Parallelism 数据并行的最简单方式是将相同的模型权重复制到多个工作器(eg: GPU)中，并将一部分数据(eg: 均匀切分后)分配给每个工作器同时处理。 如果模型的大小大于单个 GPU 设备节点的内存，那么简单的动态规划方法无法很好地工作。像*GeePS* ([Cui et al. 2016](https://www.pdl.cmu.edu/PDL-FTP/CloudComputing/GeePS-cui-eurosys16.pdf))这样的方法会将暂时未使用的参数传输回 CPU，以便在模型太大无法适应一台机器时，与有限的 GPU 内存一起工作。数据交换传输在后台进行，不干扰训练计算。 At the end of each minibatch, workers need to synchronize gradients or weights to avoid staleness. There are two main synchronization approaches and both have clear pros & cons. 在每个minibatch结束时，各个工作器(GPU)需要同步梯度或权重以避免陈旧。有两种主要的同步方法，两者都有明显的优缺点。 1. 批同步并行（Bulk synchronous parallels , BSP）：工作节点在每个小批量结束时同步数据。这可以防止模型权重陈旧，并且具有良好的学习效率，但每台机器都必须停止并等待其他机器发送梯度。 2. 异步并行（Asynchronous parallel, ASP）：每个 GPU 工作器都异步处理数据，无需等待或停顿。然而，这很容易导致使用过时的权重，从而降低统计学习效率。即使它增加了计算时间，也可能无法加快收敛的训练时间。 在中间位置是每 $x$次全局同步梯度($x>1$)。这个功能在 Pytorch v1.5([Li et al. 2021](https://arxiv.org/abs/2006.15704))中被称为“gradient accumulation”在Distribution Data Parallel ([DDP](https://pytorch.org/tutorials/intermediate/ddp_tutorial.html))中。Bucketing gradients避免了即时的 `AllReduce` 操作，而是将多个梯度放入一个 `AllReduce` 中以提高吞吐量。可以基于计算图进行计算和通信调度优化。  图 1. Pytorch DDP 伪代码。（图片来源：Li 等人，2021 年） --- ## Model Parallelism 模型并行 (Model parallelism, MP)旨在解决模型权重无法适应单个节点的情况。计算和模型参数被分割到多台机器上。与数据并行不同，每个工作节点仅分配模型参数的一部分，从而减少了内存使用和计算量。 由于深度神经网络通常包含一系列垂直层，因此将大型模型按层进行分割，将一小组连续的层分组到一个工作节点上，这种方法看起来很直观。然而，通过顺序依赖性将每个数据批次通过多个工作节点运行的朴素实现会导致大量等待时间和计算资源的严重低效利用。  图2. 一个简答的模型并行设置，其中模型被垂直分割成4个分区。由于顺序依赖性，数据一次只能由一个工作人员处理，导致大量的空闲时间“冒泡”。（图片来源：Huang等人，2019年） --- ## Pipeline Parallelism 管道并行（Pipeline parallelism, PP）将模型并行与数据并行相结合，以减少低效的时间“泡沫”。其主要思想是将一个小批量数据分成多个微批量，并使每个阶段的工作器同时处理一个微批量。需要注意的是，每个微批量需要进行两次传递，一次前向传递和一次反向传递。工作器之间的通信只传输激活值（前向传递）和梯度（反向传递）。不同的方法在如何安排这些传递以及如何聚合梯度方面存在差异。分区（工作器）的数量也被称为管道深度。 在*GPipe* ([Huang et al. 2019](https://arxiv.org/abs/1811.06965))中，来自多个微-批次的梯度被聚合并在最后同步。同步梯度下降保证了学习的一致性和效率，无论工作器的数量如何。如图3所示，气泡仍然存在，但比图2中的要小得多。假设$m$均匀分割的微批次和 $d$个分区，假设每个微批次的前向和后向过程都需要一单位的时间，则气泡的比例为： $$ 1-\frac{2md}{(2m+2(d-1))*d} = \frac{d-1}{m+d-1} $$ GPipe论文观察到，如果微-批次的数量超过分区数量的4倍（应用激活重新计算时），气泡开销几乎可以忽略不计。  图3. GPipe中使用4个微-批次和4个分区的管道并行性示意图。GPipe在每个批次结束时同步地聚合和更新设备间的梯度。（图片来源：Huang等人，2019年） GPipe 在设备数量上实现了几乎线性的吞吐量加速，尽管如果模型参数在工作节点之间分布不均匀，这并不总是保证。 PipeDream ([Narayanan et al. 2019](https://cs.stanford.edu/~matei/papers/2019/sosp_pipedream.pdf))将每个工作节点安排为交替处理前向和后向传递（ `1F1B` ）。PipeDream 将每个模型分区命名为“stage”，每个stage的工作节点可以有多个副本来运行数据并行。在这个过程中，PipeDream 使用确定性的轮询负载均衡策略来分配多个阶段副本之间的工作，以确保同一小批次的前向和后向传递发生在同一个副本上。  图 4. PipeDream 中`1F1B`微-批量调度的示意图。（图片来源：Harlap 等人，2018 年） 由于 PipeDream 在所有工作节点之间没有批量结束的全局梯度同步，因此 `1F1B` 的本地实现很容易导致一个微批量的前向和后向传播使用不同版本的模型权重，从而降低学习效率。PipeDream 提出了几种设计来解决这个问题： * Weight stashing, 权重存储：每个工作节点跟踪多个模型版本，并确保在给定一个数据批次的前向和后向传播中使用相同版本的权重。 * Vertical sync, 垂直同步（可选）：模型权重的版本与激活和梯度一起在stage工作器之间流动。然后计算采用从前一个工作器传播的相应存储版本。这个过程保持了工作器之间的版本一致性。请注意，这是异步的，与GPipe不同。 在训练开始时，PipeDream首先对模型中每个层的计算内存成本和时间进行分析，然后优化将层划分为阶段的解决方案，这是一个动态规划问题。  图5. ILSVRC12上VGG16的结果。（顶部）准确率与时间的关系。整数表示阶段工作者的数量。ASP = Asynchronous parallel & BSP = Bulk synchronous parallels. (Bottom)（底部）不同并行配置的训练时间加速。直接流水线指的是没有数据并行的流水线并行。（图片来源：Harlap等人，2018年） 后来提出了两种 PipeDream 的变体，通过存储的模型版本来减少内存占用([Narayanan et al. 2021](https://arxiv.org/abs/2006.09503)).。 PipeDream-flush: 定期添加全局同步的管道刷新，就像 GPipe 一样。通过牺牲一点吞吐量，它极大地减少了内存占用（即只维护一个模型权重版本）。  图6. PipeDream-flush中的pipeline调度示意图。（图片来源：Narayanan等人，2021年） PipeDream-2BW仅保留两个模型权重版本，其中“2BW”是“双缓冲权重”的缩写。它每 $k$个微批次生成一个新的模型版本，$k$应该大于流水线深度 $d$， $k>d$ 。由于一些残留的反向传播仍依赖于旧版本，新更新的模型版本不能立即完全替代旧版本。总共只需要保存两个版本，因此内存成本大大降低。  图7. PipeDream-2BW中的pipeline scheduling示意图。（图片来源：Narayanan等人，2021年） --- ## Tensor Parallelism 模型并行和流水线并行都是将模型垂直分割。另一方面，我们可以将一个张量操作的计算水平分割到多个设备上，称为张量并行计算（Tensor parallelism, TP）。 让我们以 Transformer 为例，考虑到其受欢迎程度。Transformer 模型主要由多层 MLP 和自注意力块组成。Megatron-LM([Shoeybi et al. 2020](https://arxiv.org/abs/1909.08053)) 采用了一种简单的方式来并行化 MLP 和自注意力的层内计算。 A MLP layer in a transformer contains a GEMM (General matrix multiply) followed by an non-linear GeLU transfer. Let’s split weight matrix � by column: Transformer 中的 MLP 层包含一个 GEMM（通用矩阵乘法），后跟一个非线性的 GeLU 转换。让我们按列拆分权重矩阵 $A$： $$ \begin{equation} \begin{aligned} split(A) &= [A1, A2] \\ Y &= GeLU(X A) \\ [Y_1, Y_2] &= [GeLU(XA_1),GeLU(XA_2)] \end{aligned} \end{equation} $$ 根据上述分区，注意力块并行运行 GEMM，使用查询（ $Q$）、键（ $K$ ）和值权重（ $V$ ），然后将它们与另一个 GEMM 结合以生成注意力头结果。 $$ Attention(X,Q,K,V) = softmax(\frac{(XQ)(XK)^T}{\sqrt{d_k}})XV $$  图8. Megatron-LM 中提出的关键 Transformer 组件的张量并行性示意图。 (图片来源：Shoeybi 等人，2020 年) [Narayanan et al. (2021)](https://arxiv.org/abs/2104.04473)将pipeline, tensor, data parallelism以及 a new pipeline scheduling strategy结合，并将其方法命名为 PTD-P。每个工作节点不仅可以将连续的一组层（"模型块"）定位在一个设备上，还可以分配多个较小连续层子集的模型块给每个工作节点（例如，设备 1 有层 1、2、9、10；设备 2 有层 3、4、11、12；每个设备有两个模型块）。一个批次中的微批次数量应该能够被工作节点的数量整除（$m$）。如果每个工作节点有$v$个模型块，与 GPipe 调度相比，流水线气泡时间可以减少$v$倍。  Fig. 9. (Top) Default `1F1B` pipeline schedule as in PipeDream-flush. (Bottom) Interleaved 1F1B pipeline schedule. First model chunks are in dark colors and second chunks are in light colors. (Image source: Narayanan et al. 202)) 图 9.（顶部）默认的`1F1B`流水线调度，如 PipeDream-flush 所示。（底部）交错的`1F1B`流水线调度。第一个模型块使用深色表示，第二个模型块使用浅色表示。（图片来源：Narayanan 等人，202） --- # Mixture-of-Experts (MoE)  最近，专家混合（MoE）方法引起了很多关注，研究人员（主要来自 Google）试图推动模型规模的极限。这个想法的核心是集成学习：多个弱学习器的组合可以得到一个强学习器！ Within one deep neural network, ensembling can be implemented with a gating mechanism connecting multiple experts ([Shazeer et al., 2017](https://arxiv.org/abs/1701.06538)). The gating mechanism controls which subset of the network (e.g. which experts) should be activated to produce outputs. The paper named it “sparsely gated mixture-of-experts” (MoE) layer. 在一个深度神经网络中，可以通过连接多个专家的门控机制来实现集成（Shazeer等，2017）。门控机制控制着网络的哪个子集（例如哪些专家）应该被激活以产生输出。该论文将其命名为““sparsely gated mixture-of-experts”（MoE）层。 精确地说，一个MoE层包含 * $n$个前馈网络作为专家 $\{E_i\}_{i=1}^{n}$ * 可训练的门控网络$G$学习一个概率分布，以便将流量路由到少数选定的专家。 根据门控输出，不需要评估每个专家。当专家数量过多时，可以考虑使用两级分层 MoE。  图10. 混合专家（MoE）层的示意图。只有n个专家中的2个被选中并由门控网络激活。（图片来源：Shazeer等人，2017年） 一种简单的选择是将输入与可训练的权重矩阵$G_g$相乘，然后进行 softmax 操作: $G_\sigma = softmax(xW_g)$。然而，这会产生一个密集的控制向量用于门控，并且不会帮助节省计算资源，因为我们只有在$G^{(i)}(x)=0$时才需要评估一个专家。因此，MoE 层仅保留前$k$个值。它还在 中添加可调节的高斯噪声以改善负载平衡。这种机制被称为有噪声的 top-k 门控。  其中上标 $v^{(i)}$ 表示向量 $v$的第 i 维。函数 $topk(.,k)$ 通过将其他维度设置为$\infty$，选择具有最高值的前 $k$个维度。 To avoid the self-reinforcing effect that the gating network may favor a few strong experts all the time, [Shazeer et al. (2017)](https://arxiv.org/abs/1701.06538) proposed a soft constraint via an additional importance loss to encourage all the experts to have the same weights. It is equivalent to the square of the [coefficient of variation](https://en.wikipedia.org/wiki/Coefficient_of_variation) of batchwise average value per expert. 为了避免门控网络始终偏向少数强大的专家，[Shazeer et al. (2017)](https://arxiv.org/abs/1701.06538)提出了一种软约束，通过额外的重要性损失来鼓励所有专家具有相同的权重。它等同于批次平均值的变异系数的平方。 $$ L_{aux} = V_{aux} \cdot CV{(\sum\limits_{x \in X} G(x))}^2 $$ 其中$CV$是变异系数，损失权重 $w_{aux}$是一个需要调整的超参数。 由于每个专家网络只能获得训练样本的一部分（“缩小批次问题/The shrinking batch problem”），我们应该尽量使用尽可能大的批次大小来进行 MoE。然而，这受限于 GPU 内存。可以应用数据并行和模型并行来提高吞吐量。  图 11. 在 1-Billion-Word 语言建模基准测试中的测试困惑度。(左) 模型容量从左到右递增，包含 4、32、256、256、1024 和 4096 个专家。(右) 4 亿参数的 MoE 模型在不同计算预算下的性能。(图片来源：Shazeer 等人，2017) **GShard** ([Lepikhin et al., 2020](https://arxiv.org/abs/2006.16668)) 通过分片将 MoE Transformer 模型扩展到 6000 亿个参数。MoE Transformer 将每个其他前馈层替换为 MoE 层。分片的 MoE Transformer 只在多台机器上分片 MoE 层，而其他层则简单地复制。 GShard 中的门控函数$G$有几种改进设计： * *Expert capacity*, 专家容量：通过一个专家的令牌数量不应超过一个阈值，称为“专家容量”。如果一个令牌被路由到已达到其容量的专家，该令牌将被标记为“溢出”，并且门控输出将更改为零向量。 * *Local group dispatching*, 本地组调度：令牌被均匀分配到多个本地组，并在组级别上强制执行专家容量。 * *Auxiliary loss*, 辅助损失：动机与原始的MoE辅助损失类似。他们添加了一个辅助损失，以最小化数据分配给每个专家的比例的均方差。 * *Random routing*, 随机路由：以与其权重成比例的概率选择第二优的专家；否则，GShard将按照随机路由的方式进行，以增加一些随机性。  Fig. 12. Pseudo code of the group-level top-2 gating mechanism with auxiliary loss in GShard. (Image source: Lepikhin et al., 2020) 图 12. GShard 中带有辅助损失的group-level 的top-2门控机制的伪代码。（图片来源：Lepikhin 等人，2020 年） **Switch Transformer** ([Fedus et al. 2021](https://arxiv.org/abs/2101.03961)) 通过将密集前馈层替换为稀疏的开关 FFN 层，将模型规模扩展到数万亿个参数（！！）。在该层中，每个输入仅路由到一个专家网络。负载平衡的辅助损失是$loss_{aux}=w_{aux}\sum_{i=1}^n {f_{i}p_{i}}$给定 $n$个专家，其中 $f_i$ 是路由到第 $i$个专家的令牌的比例， $p_i$ 是由门控网络预测的第 $i$个专家的路由概率。  图 13. Switch transformer. sparse switch FFN 层位于蓝色框中。（图片来源：Fedus 等人，2021 年） 为了提高训练稳定性，switch transformer采用了以下设计： * *Selective precision*, 选择性精度。他们表明，仅将模型的局部部分选择性地转换为 FP32 精度可以提高稳定性，同时避免了 FP32 张量的昂贵通信成本。FP32 精度仅在路由器函数的主体内使用，并将结果重新转换为 FP16 精度。 * *Smaller initialization*,较小的初始化。权重矩阵的初始化采样自截断正态分布，均值为 $\mu=0$ ，标准差为 $\sigma=\sqrt{s/n}$ 。他们还建议将 Transformer 初始化比例参数 $s=1$减小到$s=0.1$ 。 * *Use higher expert dropout*, 使用更高的专家丢弃率。微调通常适用于小数据集。为了避免过拟合，每个专家内的丢弃率增加了相当大的量。有趣的是，他们发现增加所有层的丢弃率会导致性能下降。在论文中，他们在非专家层使用了 0.1 的丢弃率，但在专家 FF 层内使用了 0.4 的丢弃率。 Switch Transformer 论文总结了用于训练大模型的不同数据和模型并行策略，并用一个很好的插图进行了说明：  图 14. 对如何在多个 GPU 核心上分割（顶部）模型权重和（底部）数据的各种并行策略的示意图。在顶部一行中，每种颜色表示一个唯一的权重矩阵。在底部一行中，不同的颜色表示不同的标记集合。（图片来源：Fedus 等人，2021 年） GShard top-2 和 Switch Transformer top-1 都依赖于标记选择，其中每个标记选择最佳的一个或两个专家进行路由。它们都采用辅助损失来鼓励更平衡的负载分配，但这并不能保证最佳性能。此外，专家容量限制可能导致浪费的标记，因为如果一个专家达到其容量限制，这些标记将被丢弃。 **Export Choice (EC)** ([Zhou et al. 2022](https://arxiv.org/abs/2202.09368)) 路由使每个专家能够选择前 $k$ 个标记。这样，每个专家自然地保证了固定的容量，每个标记可以路由到多个专家。EC 可以实现完美的负载平衡，并且已经证明可以将训练收敛时间提高 2 倍。 给定 $e$ 个专家和输入矩阵 $X \in \mathbb{R} ^{n\times d}$ ，通过以下方式计算令牌到专家的关联得分：  一个令牌到专家的分配由三个矩阵 $I$, $G \in \mathbb{R} ^ {e \times k}$和 $P \in \mathbb{R} ^ {e \times k \times n}$表示。其中 $I[i,j]$注释了哪个令牌是第 $j$ 个专家的第 $i$ 次选择。门控矩阵 $G$ 存储了所选令牌的路由权重。 $P$ 是$I$ 的one-hot编码结果，用于生成门控 FFN 层的输入矩阵（ $P \cdot X \in \mathbb{R}^{e \times k \times d}$）。 $$ G,I = top-k(S^T, k); \ \ \ P=one-hot(I) $$ 一种export choice routing的正则化方法是限制每个令牌的最大专家数量。  每个条目中的 $A[i,j]$ ( 其中$A \in \mathbb{R}^{e\times n}$) 标记了第 $i$ 专家是否选择了第 $j$ 令牌。解决这个问题并不容易。本文使用了 Dykstra 的算法，该算法运行一系列多次迭代的计算步骤。在实验中，选择受限的专家选择导致微调性能略有下降。 参数$k$ 由$k=nc/e$ 确定，其中$n$ 是一个批次中的令牌总数，$c$ 是一个表示每个令牌使用的平均专家数量的容量因子。本文在大多数实验中使用$c=2$ ，但 $EC$ 与$c=1$ 仍然优于 top-1 令牌选择门控。有趣的是， $c=0.5$对训练性能只有轻微的影响。 $EC$ 的一个重大缺点是当批量大小太小时它无法工作，对于自回归文本生成也是如此，因为它需要知道未来的标记来进行顶部选择。 --- # Other Memory Saving Designs ## CPU Offloading 当GPU内存已满时，一个选择是将暂时未使用的数据转移到CPU，并在以后需要时读取它们([Rhu et al. 2016](https://arxiv.org/abs/1602.08124))。CPU卸载的想法很简单，但近年来由于它会导致训练时间变慢而变得不太流行。 ## Activation Recomputation **Activation recomputation** (also known as “activation checkpointing” or “gradient checkpointing”; [Chen et al. 2016](https://arvix.org/abs/1604.06174)) is a smart yet simple idea to reduce memory footprint at the cost of computation time. It reduces the memory cost of training a ℓ layer deep neural net to �(ℓ), which only additionally consumes an extra forward pass computation per batch. 激活重计算（也称为“activation checkpointing” 或 “gradient checkpointing”；[Chen et al. 2016](https://arvix.org/abs/1604.06174))是一种聪明而简单的想法，以减少内存占用为代价来节省计算时间。它将训练一个$l$层深度神经网络的内存成本降低到 $O(l)$，每个批次只额外消耗一个前向传递计算。 假设我们将一个$l$层网络均匀地划分为$d$个分区。只有在分区边界处保存和传递激活值。在分区内部层的中间激活值仍然需要用于计算梯度，因此在反向传递过程中需要重新计算。通过激活重计算，训练$M(l)$的内存成本为：  最低成本为$O(l)$，在$d=\sqrt{l}$  时。 激活重新计算技巧可以使内存成本相对于模型大小呈亚线性增长。  图 15. 不同内存节省算法的内存成本。Sharing：当不再需要中间结果时，回收使用的内存。Inplace：将输出直接保存在输入值的内存中。（图片来源：Chen 等人，2016 年） --- ## Mixed Precision Training [Narang & Micikevicius et al. (2018)](https://arxiv.org/abs/1710.03740)引入了一种使用半精度浮点（FP16）数值训练模型而不会丢失模型准确性的方法。  图16 一个层的混合精度训练过程。（图片来源：Narang＆Micikevicius 等人，2018 年） 避免在半精度下丢失关键信息的三种技术： * *Full-precision master copy of weights*, 权重的全精度主副本。维护一个完整精度（FP32）的模型权重副本，用于累积梯度。数字在前向和后向传递中舍入为半精度。动机是每个梯度更新（即梯度乘以学习率）可能太小，无法完全包含在 FP16 范围内（即 $2^{-24}$在 FP16 中变为零）。 * *Loss scaling*, 损失缩放。将损失放大以更好地处理梯度的小幅度（见图16）。放大梯度有助于将它们移动到可表示范围的右侧部分（包含较大值），保留原本会丢失的值。 * *Arithmetic precision*, 算术精度。对于常见的网络算术（例如向量点积、向量元素求和的归约），我们可以在FP32中累积部分结果，然后在保存到内存之前将最终输出保存为FP16。逐点操作可以在FP16或FP32中执行。  图 17. 完整精度下的梯度直方图。一旦模型切换到 FP16，左侧部分将被归零。（图片来源：Narang＆Micikevicius 等人，2018 年） 在他们的实验中，对于一些网络（例如图像分类，Faster R-CNN），不需要损失缩放，但对于其他网络（例如 Multibox SSD，大型 LSTM 语言模型）则是必需的。 --- ## Compression 中间结果通常会消耗大量内存，尽管它们只在前向传递和后向传递中需要一次。这两个使用之间存在明显的时间间隔。因此[Jain et al. (2018)](https://www.microsoft.com/en-us/research/uploads/prod/2018/04/fiddle-gist-isca18.pdf)提出了一种数据编码策略，用于在第一次传递后压缩中间结果，然后在后向传播时解码回来。 Their system *Gist* incorporates two encoding schemes: *Layer-specific lossless encoding*; focus on ReLU-Pool (“Binarize”) and ReLU-Conv (“Sparse storage and dense computation”) patterns. *Aggressive lossy encoding*; use delayed precision reduction (DPR). They observed that the first immediate use of feature maps should be kept at high precision but the second use can tolerate lower precision. 他们的系统 Gist 包含两种编码方案：*Layer-specific lossless encoding*:重点关注 ReLU-Pool（“二值化”）和 ReLU-Conv（“稀疏存储和密集计算”）模式。*Aggressive lossy encoding*；使用延迟精度降低（DPR）。他们观察到，特征图的第一次立即使用应保持高精度，但第二次使用可以容忍较低精度。 实验证明，Gist 可以在 5 个 SOTA 图像分类 DNN 中将内存成本降低 2 倍，平均降低 1.8 倍，仅有 4%的性能开销。 --- ## Memory Efficient Optimizer 优化器需要占用大量内存。以流行的 Adam 优化器为例，它在内部需要维护动量和方差，这两者与梯度和模型参数的规模相同。突然之间，我们需要保存 4 倍于模型权重的内存。 已经提出了几种优化器来减少内存占用。例如，*Adafactor* ([Shazeer et al. 2018](https://arxiv.org/abs/1804.04235))不像 Adam 那样存储完整的动量和变化，而是仅跟踪移动平均值的每行和每列总和，然后根据这些总和估计二阶矩。SM3* ([Anil et al. 2019](https://arxiv.org/abs/1901.11150)) 描述了一种不同的自适应优化方法，也大大减少了内存占用。 *ZeRO* (*Zero Redundancy Optimizer*; [Rajbhandari et al. 2019](https://arxiv.org/abs/1910.02054))根据对大型模型训练的两个主要内存消耗的观察，优化了用于训练大型模型的内存使用。 1. 大部分被模型状态占用，包括优化器状态（如 Adam 的动量和方差）、梯度和参数。混合精度训练需要大量内存，因为优化器需要保留 FP32 参数和其他优化器状态的副本，除了 FP16 版本。 2. 其它的内存被激活函数、临时缓冲区和无法使用的碎片化内存（在论文中称为残余状态）占用。 ZeRO 结合了两种方法，ZeRO-DP 和 ZeRO-R。ZeRO-DP 是一种增强的数据并行性，避免了模型状态的简单冗余。它通过动态通信调度将优化器状态、梯度和参数分割到多个数据并行进程中，以最小化通信量。ZeRO-R 通过分割激活函数的重新计算、恒定缓冲区大小和即时内存碎片整理来优化残余状态的内存消耗。 --- # Citation Cited as: > Weng, Lilian. (Sep 2021). How to train really large models on many GPUs? Lil’Log. https://lilianweng.github.io/posts/2021-09-25-train-large/. Or ``` @article{weng2021large, title = "How to Train Really Large Models on Many GPUs?", author = "Weng, Lilian", journal = "lilianweng.github.io", year = "2021", month = "Sep", url = "https://lilianweng.github.io/posts/2021-09-25-train-large/" } ``` # References [1] Li et al. [“PyTorch Distributed: Experiences on Accelerating Data Parallel Training”](https://arxiv.org/abs/2006.15704) VLDB 2020. [2] Cui et al. [“GeePS: Scalable deep learning on distributed GPUs with a GPU-specialized parameter server”](https://www.pdl.cmu.edu/PDL-FTP/CloudComputing/GeePS-cui-eurosys16.pdf) EuroSys 2016 [3] Shoeybi et al. [“Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism.”](https://arxiv.org/abs/1909.08053) arXiv preprint arXiv:1909.08053 (2019). [4] Narayanan et al. [“Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM.”](https://arxiv.org/abs/2104.04473) arXiv preprint arXiv:2104.04473 (2021). [5] Huang et al. [“GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism.”](https://arxiv.org/abs/1811.06965) arXiv preprint arXiv:1811.06965 (2018). [6] Narayanan et al. [“PipeDream: Generalized Pipeline Parallelism for DNN Training.”](https://cs.stanford.edu/~matei/papers/2019/sosp_pipedream.pdf) SOSP 2019. [7] Narayanan et al. [“Memory-Efficient Pipeline-Parallel DNN Training.”](https://arxiv.org/abs/2006.09503) ICML 2021. [8] Shazeer et al. [“The Sparsely-Gated Mixture-of-Experts Layer Noam.”](https://arxiv.org/abs/1701.06538) arXiv preprint arXiv:1701.06538 (2017). [9] Lepikhin et al. [“GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding.”](https://arxiv.org/abs/2006.16668) arXiv preprint arXiv:2006.16668 (2020). [10] Fedus et al. [“Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity.”](https://arxiv.org/abs/2101.03961) arXiv preprint arXiv:2101.03961 (2021). [11] Narang & Micikevicius, et al. [“Mixed precision training.”](https://arxiv.org/abs/1710.03740) ICLR 2018. [12] Chen et al. 2016 [“Training Deep Nets with Sublinear Memory Cost.”](https://arxiv.org/abs/1604.06174) arXiv preprint arXiv:1604.06174 (2016). [13] Jain et al. [“Gist: Efficient data encoding for deep neural network training.”](https://www.microsoft.com/en-us/research/uploads/prod/2018/04/fiddle-gist-isca18.pdf) ISCA 2018. [14] Shazeer & Stern. [“Adafactor: Adaptive learning rates with sublinear memory cost.”](https://arxiv.org/abs/1804.04235) arXiv preprint arXiv:1804.04235 (2018). [15] Anil et al. [“Memory-Efficient Adaptive Optimization.”](https://arxiv.org/abs/1901.11150) arXiv preprint arXiv:1901.11150 (2019). [16] Rajbhandari et al. [“ZeRO: Memory Optimization Towards Training A Trillion Parameter Models Samyam.”](https://arxiv.org/abs/1910.02054) arXiv preprint arXiv:1910.02054 (2019). [17] Zhou et al. [“Mixture-of-Experts with Expert Choice Routing”](https://arxiv.org/abs/2202.09368) arXiv preprint arXiv:2202.09368 (2022).