简单学点大模型-分布式训练

最新推荐文章于 2024-08-26 11:36:51 发布

我的薯条呢

最新推荐文章于 2024-08-26 11:36:51 发布

阅读量134

点赞数

文章标签：分布式语言模型人工智能自然语言处理

本文链接：https://blog.csdn.net/2201_75734742/article/details/132990400

版权

近年来，深度学习被广泛应用到各个领域，包括计算机视觉、语言理解、语音识别、广告推荐等。在这些不同的领域中，⼀个共同的特点就是模型规模越来越大，比如 GPT-3 模型的参数量达到1750 亿。即使用1024张 80 GB 的 A100，那么完整训练 GPT-3 的时长都需要1个月。

模型规模的扩大，对硬件（算力、内存）的发展提出要求。

训练 AI 模型时候所需要的内存一般比模型参数量还要多几倍，与之相比，AI 硬件上的内存大小仅仅是以每两年翻2倍的速率在增长。目前AI硬件更专注于提高硬件的峰值算力，但是通常以简化或者删除其他部分（例如内存的分层架构）为代价，这已不足以应付最新AI模型的训练。

目前 NLP 中的 SOTA Transformer 类模型的算力需求，以每两年750倍的速率增长，模型参数数量则以每两年240倍的速率增长。相比之下，硬件算力峰值的增长速率为每两年3.1倍，DRAM 还有硬件互连带宽增长速率则都为每两年1.4倍，已经逐渐被需求甩在身后。

硬件内存、算力的增长速度跟不上模型训练规模扩大的速度，这时就出现了“内存墙”问题。内存墙问题不仅与内存容量大小相关，也包括内存的传输带宽。这涉及到多个级别的内存数据传输。

单机多卡训练，即并行训练。神经网络的训练不仅需要多个设备进行计算，还涉及到设备之间的数据传输，只有协调好集群中的计算与通信，才能做高效的分布式训练。

假设神经网络中某⼀层是做矩阵乘法，其中输入的数据 $x$ 的形状为 $4\times 5$ ，模型参数 $\omega$ 的形状为 $5\times 8$ ，那么，矩阵乘法输出形状为 $4\times 8$ 。

单机单卡的训练中，以上矩阵乘法，先计算得到 $out$ ，并将 $out$ 传递给下⼀层，并最终计算得到的损失 $loss$ ，然后在反向传播过程中，得到 $\frac{\partial loss}{\partial \omega }$ ，用于更新 $\omega$ 。（计算模型的损失更新模型）

分布式训练中，依据是切分 $x$ 还是 $\omega$ 的不同，分为“数据并行”和“模型并行”策略。

数据并行，就是将数据 $x$ 进行切分，而每个设备上的模型 $\omega$ 是相同的。如下图所示， $x$ 被按照第0维度平均切分到2个设备上，两个设备上都有完整的模型 $\omega$ 。

这样，在两台设备上，分别得到的输出，都只是逻辑上输出的⼀半（形状为 $2\times 8$ ），将两个设备上的输出拼接到⼀起，才能得到逻辑上完整的输出。

注意，因为数据被分发到了2个设备上，因此反向传播过程，各自设备上得到的 $\frac{\partial loss}{\partial \omega }$ 会不⼀样。因此在数据并行策略下，在反向传播过程中，需要对各个设备上的梯度进行AllReduce，以确保各个设备上的模型始终保持⼀致。

当数据集较大，模型较小时，由于反向过程中为同步梯度产生的通信代价较小，此时选择数据并⾏⼀般比较有优势，常见的视觉分类模型，如 ResNet50，比较适合采用数据并行。

AllReduce其实是一类算法，目标是高效得将不同机器中的数据整合（reduce）之后再把结果分发给各个机器。在深度学习应用中，数据往往是一个向量或者矩阵，通常用的整合则有Sum、Max、Min等。下图展示了AllReduce在有四台机器，每台机器有一个长度为四的向量时的输入和输出。

所谓的模型并行，就是每个设备上的数据 $x$ 是完整的、⼀致的，而模型 $\omega$ 被切分到了各个设备上，每个设备只拥有模型的⼀部分，所有计算设备上的模型拼在⼀起，才是完整的模型。

如下图所示， $\omega$ 被按照第1维度平均切分到2个设备上，两个设备上都有完整的数据 $x$ 。两个设备上的输出也需要通过拼接才能得到逻辑上的输出。

模型并行的好处是，省去了多个设备之间的梯度 AllReduce；但是，由于每个设备都需要完整的数据输入，因此，数据会在多个设备之间进行广播，产生通信代价。比如，上图中的最终得到的，如果它作为下⼀层网络的输入，那么它就需要被广播发送到两个设备上。语言模型，如 BERT，常采用模型并行。

当神经网络过于巨大，无法在一个设备上存放时，除了上述的模型并行的策略外，还可以选择流水并行。流水并行指将网络切为多个阶段，并分发到不同的计算设备上，各个计算设备之间以“接力”的方式完成训练。

如下图，展示了一个逻辑上的4层网络（T1 至 T4)是如何做流水并行的。

4层网络被切分到2个计算设备上，其中 GPU0 上进行 T1 与 T2 的运算，GPU1 上进行 T3 与 T4 的计算。

GPU0 上完成前两层的计算后，它的输出被当作 GPU1 的输入，继续进行后两层的计算。

网络的训练中，也可以将多种并行策略混用。

以 GPT-3 为例，以下是它训练时的设备并行方案：它首先被分为 64 个阶段，进行流水并行。每个阶段都运行在 6 台 DGX-A100 主机上。在6台主机之间，进行的是数据并行训练；每台主机有 8 张 GPU 显卡，同一台机器上的8张 GPU 显卡之间是进行模型并行训练。

关注