本文约2200字,建议阅读6分钟
本文就介绍数据并行中的两个核心策略:DP和DDP。
在人工智能领域,大模型参数规模日益增长。这种增长使得传统单机训练面临两大挑战:显存容量限制(通常不超过80GB)和计算速度瓶颈(单卡算力约312 TFLOPS)。为此,数据并行(Data Parallelism)作为分布式训练的核心策略,成为支撑大模型训练的关键技术。本文就介绍数据并行中的两个核心策略:DP(Data Parallelism)和DDP(Distributed Data Parallelism)。
DP(Data Parallelism)数据并行
DP的架构设计中依赖主GPU(逻辑0卡)协调梯度聚合。图1描述了一个分布式机器学习训练的流程,其中涉及到一个主GPU(我们称之为“Master”)和多个工作GPU(我们称之为“Worker”)。DP策略可以通过torch.nn.DataParallel()语句实现,其步骤如下:
1.首先,Master GPU加载了整个数据集和模型。
2.Master GPU将数据集和模型分发给所有的Worker GPUs。
3.每个Worker GPU都会得到一部分数据和整个模型的副本。每个Worker GPU使用自己得到的数据和模型进行计算,这个过程叫做前向计算。计算完成后,每个Worker GPU都会产生一些输出。Master GPU收集所有Worker GPU的输出,并计算整个模型的损失(loss)。损失是衡量模型预测结果与实际结果差异的指标。
4.Master GPU计算损失后,会将这个损失分发给所有Worker GPU。每个Worker GPU会根据自己计算的输出和得到的损失,计算出梯度(Gradient)。梯度表示模型参数应该如何调整以减少损失。
5.所有Worker GPU计算出的梯度被发送回Master GPU,Master GPU将这些梯度汇总加和起来。
6.Master GPU使用汇总的梯度来更新模型的参数。最后,Master GPU将更新后的模型参数再次分发给所有的Worker GPU,以便进行下一轮的训练。
DP策略存在效率和通信问题,数据集首先需要被复制到主GPU中,然后再分割成多个部分,分发到不同的GPU上。此外,模型的权重参数只在主GPU上更新,这意味着在每次迭代前,都需要将数据同步到所有的GPU上。每次迭代后,所有GPU的输出结果还需要被汇总到主GPU上。这样的数据传输很快就成为了一个瓶颈,限制了整体的训练速度。同时,这也导致了主GPU和其他GPU之间的负载不均衡,比如主GPU可能使用了大量显存,而其他GPU的显存使用却很少,这可能导致主GPU的显存不足(OOM,即Out of Memory)。
图1 DP策略的示意图
DDP (Distributed Data Parallelism)分布式数据并行
受通信负载不均的影响,DP一般用于单机多卡场景。因此,DDP作为一种更通用的解决方案出现了,既能多机,也能单机,有效解决DP策略中存在的负载不均匀问题。DDP(Distributed Data Parallelism)利用了去中心化架构。在DDP中,不再有主GPU,每个GPU执行相同的任务,推理、损失函数计算、梯度计算等都可以在各个GPU上并行独立地完成,提高了训练的效率和速度。DDP实现并行训练的核心在于模型间的梯度同步,这是通过all-reduce(全归约)通信操作实现的,保证每个GPU会得到完全相同的梯度。接下来我们介绍目前最通用的all-Reduce方法:Ring-All Reduce(环状全归约)以便理解DDP中的数据通信操作。
Ring-All Reduce则分两大步骤实现该目标:Reduce-Scatter(归约分散)和All-Gather(全聚集)。
Reduce-Scatter(归约分散)
DP中使用的架构简单并易于使用,但随着GPU卡数的增加,服务器节点通信量随之线性增加。而环状通信算法被提出,每个节点仅与相邻的节点通信。环状通信算法的优点在于数据传输的通信量恒定,不随GPU数量增加而增加。图2展示了环状归约分散的基本原理。
图2展示了环状归约分散的基本原理,其基本原理如下:
图2 环状归约分散的基本原理。
1.初始状态:
o图中有三个节点(节点1、节点2、节点3),每个节点都有自己的数据(a1, b1, c1; a2, b2, c2; a3, b3, c3)。
2.第1步:
o每个节点首先与相邻的节点交换数据并进行归约操作。
o节点1将c1与节点2的c2相加,得到c1+c2。
o节点2将a2与节点3的a3相加,得到a2+a3。
o节点3将b1与节点1的b1相加,得到b1+b3。
o这一步完成后,每个节点都有部分归约的结果。
3.第2步:
o再次进行数据交换和归约操作。
o节点1将a1与节点2和节点3的a2+a3相加,得到a1+a2+a3。
o节点2将b1与节点1和节点3的b1+b2+b3相加,得到b1+b2+b3。
o节点3将c1与节点1和节点2的c1+c2相加,得到c1+c2+c3。
All-Gather(全聚集)
图3中展示了三个节点(节点1、节点2、节点3)在All-Gather操作中的初始状态和两个步骤的变化。
初始状态
节点1:拥有数据 a1+a2+a3,b1,c1+c2
节点2:拥有数据 a2+a3,b1+b2+b3,c2
节点3:拥有数据 a3,b1+b3,c1+c2+c3
在第一步中,每个节点开始发送并接收数据:
以节点1为例:从节点2接收数据,得到b1+b2+b3,同时c1+c2+c3发送给节点3
在第二步中,节点继续交换数据,直到每个节点都拥有所有其他节点的数据。
图3 全聚集的基本原理。
现在来分析一下DDP策略的通信时间。假设共有p个节点,假设每个节点的数据是V,则单独看归约分散和全聚集,每个单独过程里每个节点向外发送和接收供经过p-1个步骤,其发送和接收的数据量大约是
。假设每个节点发送和接收的带宽都是b,则整个过程(Reduce Scatter和All Gather)的总时间为
。当p足够大时,总时间为
,这意味着通信时间与节点数p无关,这是因为DDP把负载均衡到每一个节点中。
结论
数据并行策略和分布式数据并行策略能有效摆脱显存容量限制,突破计算速度瓶颈,在大模型训练中的应用日渐广泛。
编辑:王菁
王可汗,清华大学博士,人工智能算法研发工程师。对数据科学产生浓厚兴趣,对机器学习AI充满好奇。期待着在科研道路上,人工智能与工业界应用碰撞出别样的火花。希望结交朋友分享更多数据科学的故事,用数据科学的思维看待世界。
数据派研究部成立于2017年初,以兴趣为核心划分多个组别,各组既遵循研究部整体的知识分享和实践项目规划,又各具特色:
算法模型组:积极组队参加kaggle等比赛,原创手把手教系列文章;
调研分析组:通过专访等方式调研大数据的应用,探索数据产品之美;
系统平台组:追踪大数据&人工智能系统平台技术前沿,对话专家;
自然语言处理组:重于实践,积极参加比赛及策划各类文本分析项目;
制造业大数据组:秉工业强国之梦,产学研政结合,挖掘数据价值;
数据可视化组:将信息与艺术融合,探索数据之美,学用可视化讲故事;
网络爬虫组:爬取网络信息,配合其他各组开发创意项目。
点击文末“阅读原文”,报名数据派研究部志愿者,总有一组适合你~
转载须知
如需转载,请在开篇显著位置注明作者和出处(转自:数据派THUID:DatapiTHU),并在文章结尾放置数据派醒目二维码。有原创标识文章,请发送【文章名称-待授权公众号名称及ID】至联系邮箱,申请白名单授权并按要求编辑。
未经许可的转载以及改编者,我们将依法追究其法律责任。
数据派THU作为数据科学类公众号,背靠清华大学大数据研究中心,分享前沿数据科学与大数据技术创新研究动态、持续传播数据科学知识,努力建设数据人才聚集平台、打造中国大数据最强集团军。
新浪微博:@数据派THU
微信视频号:数据派THU
今日头条:数据派THU