So-Large-LM-Task07: 大模型的分布式训练
输入“/”快速插入
So-Large-LM-Task07: 大模型的分布式训练
飞书用户6850
1月26日创建
Parallelism / 大模型的分布式训练
本章的 Lecture 是手写体,辨识相当有难度,因此对教程进行了替换(替换内容更为详细)。该 Lecture 文本可见我本人的肉眼识别版:
CS324 Parallelism Lecture
本节源自OneFlow:
https://docs.oneflow.org/master/parallelism/01_introduction.html
为什么分布式训练越来越流行
模型规模越来越大,训练时间越来越长。
模型规模的扩大,对硬件(算力、内存)的发展提出要求。然而,因为
内存墙
的存在,单一设备的算力及容量,受限于物理定律,持续提高芯片的集成越来越困难,难以跟上模型扩大的需求。
为了解决算力增速不足的问题,人们考虑用多节点集群进行分布式训练,以提升算力,分布式训练势在必行。
常见的并行策略
神经网络的训练并不是单纯地“把原来一个设备做的事情分给多个设备各自做”,它不仅需要计算,还涉及到
数据传输
,只有协调好集群的计算与
通信
,才能做高效的分布式训练。
以最简单的矩阵乘法为例解释 Data Parallelism 和 Model Parallelism 的区别。
设神经网络某一层中做矩阵乘法,其
输入
的 shape 为 (4, 5),权重
的 shape 为 (5, 8) ,那么输出
的 shape 为 (4, 8)
在单机单卡的训练中,该矩阵乘法运算得到输出
,随后被继续前向传递,最后得到 loss
。随后将
进行反向传播以更新 NN 的参数,在参数
处,其
。
在分布式训练中,根据切分
还是切分
,分为 Data Parallelism 和 Model Parallelism。
(把 x 想像成一个 batch)
Data Parallelism / 数据并行
顾名思义,Data Parallelism 将 data
进行切分,而每个设备上的模型参数
是完整的、一致的。
如下图所示,
在其第 0 维度上被平均切分到 2 个设备上,两个设备上都有完整的
。正因如此,两台设备上运算的结果的 shape 也只是原本结果 shape 的平均切分 (2, 8),将两个设备上的输出再进行第 0 维度上的拼接,才能得到原本的结果。
将 data 进行切分带来了新的问题。data 被分发到了 2 个设备上,所以在反向传播过程中,各自设备上计算的梯度
不一样,如果在各个设备上分别使用各自计算的梯度进行更新,会导致 2 个设备上的模型参数
不一致,这样的话分布式训练就失去了意义,其不能得到统一的、最优的参数。
因此,在 Data Parallelism 策略下,反向传播的过程中不能让各个设备分别使用各自计算的梯度进行更新,而要让其进行
AllReduce
(用所有设备上的梯度信息计算一个全局的梯度信息),以确保各个设备上的模型参数
始终保持一致。
当数据集较大,网络较小时,由于反向传播过程中同步梯度信息的通信代价较小,此时选择 Data Parallelism 一般比较有优势,常见的 CV 分类模型,如 ResNet50,就比较适合采用 Data Parallelism 策略。
Model Parallelism / 模型并行
(这里描述的 Model Parallelism 也被称为 Tensor Parallelism)
当网络较大时,Data Parallelism 同步梯度信息的通信代价很大,甚至网络本身可能巨大到无法存放到单一的计算设备中。这时,可以采用 Model Parallelism 策略解决问题。
顾名思义,Model Parallelism 将模型参数
进行切分,而每个设备上的数据
是完整的、一致的。
如下图所示,
在其第 1 维度上被平均切分到 2 个设备上,两个设备上都有完整的
。正因如此,将两个设备上的输出再进行第 1 维度上的拼接,才能得到原本的结果。
Model Parallelism 的好处是,省去了多个设备之间梯度信息的 AllReduce(每个参数都是用同样的输入数据产生的 loss 优化的)。但是,由于每个设备的 NN 都需要完整的输入数据,因此,输入数据需要在各个设备间进行广播,产生通信代价。例如,上图中最终得到的输出
,如果它是下一层网络的输入,那么它就需要被广播发送到这 2 个设备上。
NLP 中的 LLM 模型,如 BERT,常采用 Model Parallelism。
Pipeline Parallelism / 流水并行
(这里描述的是最基本的 pipeline parallelism,严格来说是串行的,称为bubble,但是后续的改进让他成为了并行,就像操作系统的轮转一样)
当网络较大时,除了使用 Model Parallelism 策略,还可以选择 Pipeline Parallelism。
Pipeline Parallelism 指的是将网络切分为多个阶段,分发到不同设备上,各个设备之间以“接力”的方式完成训练。
如下图,展示了一个逻辑上的 4 层网络(T1 到 T4,即 Layer1 -> Layer4)进行 Pipeline parallel 的过程。该 4 层网络在运行阶段上被平均切分到 2 个设备上,其中 GPU0 上进行 T1 与 T2 的运算,GPU1 上进行 T3 与 T4 的运算。GPU0 完成 T1 和 T2 的计算后,其输出被当作 GPU1 点输入,继续进行 T3 和 T4 的计算。