深度学习的分布式训练与集合通信(一)
导言
随着近年来大模型的兴起,AI模型不断增大,巨大的计算量和数据量使得在单机单卡的环境下进行模型训练变得不再实际。分布式训练应运而生,人们使用多机多卡的计算集群来训练AI模型,这样一则可以增加计算的并行度,提升训练速度;二则可以消解单卡的存储压力,让数据分布式地存储在多张卡上。
分布式训练在给单卡计算和存储减负的同时,也增大了对卡间通信的需求。训练过程中,多卡需要进行频繁的数据交换与行为同步,集合通信或成为分布式训练性能优化上的瓶颈。分布式训练可能涉及多种集合通信操作,如AllGather,AllReduce等,了解这些操作是对其进行优化的基础。
本文以此为出发点,介绍常见的深度学习分布式训练的并行策略和背后使用到的集合通信操作。希望通过这篇文章,帮助读者理解分布式训练的原理,以及集合通信之于分布式训练的重要性和必要性。
在介绍集合通信操作的同时,会列举出其操作在HCCL(Huawei Collective Communication Library)中对应的接口表示,以方便读者进一步加深对HCCL的功能和使用场景理解。鉴于篇幅限制,将本文内容拆分成三个部分来讲述,本期将阐述第一部分的内容,欢迎持续关注:
- 第一部分:介绍两个方面的背景知识,一是模型训练的大体流程,二是集合通信操作的基本类型,然后介绍分布式训练中基础的数据并行(DP)的概念以及其通信模式;
- 第二部分:介绍分布式训练中的常见的三种模型并行(MP),即流水并行(PP)张量并行(TP)和专家并行(EP)的概念和通信模式;
- 第三部分:介绍一些进阶版的并行模式,如序列并行(SP),完全分片数据并行(FSDP),以及Zero系列并行等,介绍它们的概念以及通信模式。
模型训练流程
作为背景知识补充之一,先回顾一下深度学习模型训练流程。
在模型训练之前,有一些准备性工作,如数据预处理,模型选择,损失函数和优化方法的确定,以及超参数的预设等,这里不做详细展开。之后就可以开始模型的正式训练了,其过程大致可以分为以下几步:
- 初始化:对神经网络的权重和偏差进行随机的初始化赋值,将这些初始值作为学习过程的起点;
- 前向传播:输入数据被注入到神经网络中,并通过一系列乘加运算和激活函数,计算每层神经元的激活值,最终产生神经网络的预测输出;
- 损失计算:使用损失函数计算预测输出与实际目标输出之间的差异,量化预测值与真实值的偏差;
- 反向传播:计算损失函数相对于模型参数的梯度,其数值表明损失如何随着这些参数的微小变化而变化;这一步涉及到的数据除了网络中的所有参数外,还需要用到前向传播中计算出来的所有层神经元的激活值;
- 梯度下降:反向传播过程中计算出的梯度表示损失值上升最陡峭的方向,为了最小化损失,网络沿梯度的反方向来更新其参数,更新幅度由学习率控制;
- 迭代:对每一批训练数据(batch)重复步骤2到5多次(epoch),直到神经网络在训练数据上的性能达到令人满意的水平或收敛到一个解决方案。
图片来源:https://medium.com/data-science-365/overview-of-a-neural-networks-learning-process-61690a502fa
基础通信操作
对于通信操作,关注点主要有两个,一是数据的收发方是谁,二是收发的数据是怎样的。基于收发方的不同,通信模式被分成四类:一对一、一对多、多对一、多对多。其中,一对一的通信模式又被称为点到点通信(Point-to-Point,P2P),涉及到多个计算节点参与的一对多,多对一,和多对多的通信模式则属于集合通信(Collective Communication,CC)操作的范畴。
其次,基于收发数据的不同,例如是整块发还是分块发,发送的数据在接收端需不需要进行整合处理等,通信模式又被分成了更细致的门类,下文也会依次展开介绍。
另外,在本章的最后会列举出每一种通信操作在HCCL中对应的接口名称,方便后续大家查找和使用。
(1) 一对一
Send:将一个节点的数据发送到另外一个节点;
Receive:接收从另一个节点上发送来的数据;
(2) 一对多
Broadcast:将一个节点的数据广播到多个节点;
Scatter:将一个节点的数据切片分发到多个节点;
(3) 多对一
Gather:将多个节点上的数据收集到一个节点上;
Reduce:将多个节点上的数据收集到一个节点上并规约;
(4) 多对多
ReduceScatter:将多个节点上的数据按维度发送到相应节点上并进行数据规约操作;
AllGather:将每个节点上各自的数据发送到所有节点上;
AllReduce:将每个节点上的数据发送到所有节点上并进行数据规约操作;
AlltoAll:所有节点相互scatter和gather等量的数据;
AlltoAllV:所有节点相互scatter和gather不等量的数据。
对于上述各种通信操作,HCCL中都有其对应的接口(API)实现,以方便用户使用。API列举如下,想要获得更加详尽的HCCL接口信息,读者可以参考链接。
在了解了上述有关模型训练和通信操作的背景知识后,我们来看看分布式训练是如何利用多卡并行来共同完成大模型训练的,以及不同分布式训练策略背后的通信操作。
分布式训练的并行策略
什么是分布式训练?通俗易懂地说,就是将大模型训练这个涉及到庞大数据量和计算量的任务切成小份,分发给多个计算单元来共同完成(一个计算单元通常被称为一个计算节点或设备,物理上可以是一张NPU,GPU卡等)。
那什么是分布式训练的并行策略呢?简言之就是这个“切的刀法”。模型训练任务中包含着多个维度,我们可以依照不同维度将整个训练任务切分开,并行起来执行。不同的切分方法对应着不同的计算存储和通信模式,直接影响系统性能的好坏。如何进行这个并行维度的选择,是分布式训练并行策略所关注的问题。
目前常见的分布式训练并行策略主要有数据并行和模型并行,而模型并行又进一步分为流水并行,张量并行和专家并行等。除此之外,还有一些进阶版的并行策略,如序列并行,完全分片数据并行,以及Zero系列并行等,下面我们依次对其进行介绍。
1) 数据并行(DataParallelism,DP)
数据并行是指将一个批次(batch)的训练数据分成若干个小批次,分发给多个计算节点来进行训练的并行方式。
图片来源:https://medium.com/@aruna.kolluru/model-parallelism-390d32145a5a
数据并行需要首先将初始化的模型和优化器参数复制到所有节点上,然后将不同的小批次训练数据从存储设备装载到相应计算节点,随后执行上述模型训练流程中的步骤2)到4):前向传播,损失计算,反向传播。在各个计算节点基于自己分得的小批次数据计算得到网络参数的梯度后,梯度数据会在所有节点间被同步,随后各种进行参数优化步骤5),以此循环进行多批次的迭代。梯度同步的操作是为了确保所有节点在进行新批次数据迭代前有一样的网络参数。
由此可见,对于数据并行,集合通信操作发生在两处:其一是训练之初模型参数初始化后的模型广播,即将一个节点中初始化好的模型信息拷贝到所有节点,涉及到的通信操作是一对多的Broadcast;其二则是每一次各节点跑完小批次训练数据后的梯度同步,涉及到的通信操作是多对多的AllReduce,即将所有节点中的梯度信息整合平均后再同步给所有节点。
数据并行的通信操作总结成表格如下:
通信操作 |
通信数据 |
发生时间 |
通信次数 |
每次通信数据量 |
HCCL API |
Broadcast |
模型参数 |
参数初始化后 |
1次 |
模型参数量 |
|
AllReduce |
梯度 |
每批次数据计算出梯度后 |
每个Epoch通信1次 |
数据并行节点数x模型参数量 |
与之后要介绍的其他并行策略相比,数据并行单次通信数据较少且通信不频繁,因为其仅以Epoch迭代的频率进行通信,而不会在一个Epoch内的多次前向传播和反向传播的过程中进行通信。
数据并行的优势主要体现在:
- 它在处理大型数据集时特别有用,通过将一个批次的数据拆分成多个小批次进行并行学习,加快训练速度。
- 数据并行也可以通过调节并行的节点数量来灵活配置训练批次的大小(batch size),提升训练效果。
然而,数据并行也有使用的限制:
- 鉴于模型收敛速度的考虑,训练批次不能过大,因此数据并行的并行规模有自身的上限,且当并行的小批次数据被划分得太小时,通信的成本会大幅增加。
- 数据并行并不能直接帮助适应更大的模型,要想支持大模型训练,模型并行必不可少。
下期,将为大家带来模型并行的内容,包括分布式训练中的常见的三种模型并行方式,即流水并行(PP)张量并行(TP)和专家并行(EP),以及通信模式的介绍,感谢关注!
参考材料
- 点赞
- 收藏
- 关注作者
评论(0)