从分布式训练到大模型训练

要了解大模型训练难,我们得先看看从传统的分布式训练,到大模型的出现,需要大规模分布式训练的原因。接着第二点去了解下大规模训练的挑战。


从分布式训练到大规模训练

常见的训练方式是单机单卡,也就是一台服务器配置1块AI芯片,这是最简单的训练方式。随着数据量的增加,希望加快模型的训练速度,于是出现了单机多卡,多块AI芯片并行,以一台机器上配置8块AI芯片为例,把数据切分成8份,分别在8块AI芯片上都跑一次BP算法,计算出梯度,然后所有AI芯片上计算出的梯度进行平均,更新模型参数。这样的话,以前一次BP只能训练1个batch的数据,现在就是8个batch。

image

模型和数据规模的增大,意味着训练时间的增长。为了提升模型训练的速度,可以增加计算资源来缩短训练时间,于是出现了分布式训练。所谓分布式训练,物理上就是使用多台机器,每台机器上都有多块AI芯片,网络模型运行在不同机器的不同AI芯片上,以加快整体训练速度。简单来理解,实质就是将单卡的负载拆到了多卡上。

image

来看看引入分布式训练的两个原因:

1)第一个是数据规模大,导致训练时间很长。在单机8卡情况下,对于COCO在115k规模的数据集,假设训练resnet152模型需要40个小时。训练Open Image dataset v4图片数据集在1,740k规模下,则相对来说可能需要40天时间。对于生命宝贵、毕业延期的炼丹同学来说,需要不断的调整网络结构、尝试新的参数,修改下游任务等工作,单机单卡和单机多卡的训练速度是无法接受的,于是很有必要使用分布式进行训练。

2)第二点就是分布式训练可能带来精度上的提升。先回忆一下,为什么要用SGD来优化模型,随机梯度下降的“随机”是指每次从数据集里面随机抽取一个小batch数据来进行计算真实值和模型预测值之间的误差,然后反向传播。之所以只选一个小batch,一是因为随机数据产生的小batch梯度方向,理论和实践证明,一定程度上可以模拟整个数据的梯度下降方向,二是因为AI芯片的内存有限。但是在实际情况下,小batch梯度并不足够代替整个数据集的梯度方向,每次小batch在BP算法求解出来的梯度方向,与整体数据集并不完全一致。这样就会导致优化迭代(训练)过程不断震荡。使用分布式训练,可以使用更大的batch size,避免优化过程中的震荡。

数据并行:在分布式训练场景,因为数据规模增大,可以通过数据并行,可以对输入数据进行切分,每块AI芯片只需要处理一部分数据。优点是计算通信可以重叠,通信需求低,并行效率高;缺点是无法运行大模型。

image

模型并行:同时,模型每一层的参数量增大、模型的层数越深,可以通过模型并行修改网络层内的计算方式,将单层的计算负载和内存负载切分到多块AI芯片上。缺点是通信和计算是串行的,对通信的要求高。

image

流水线并行:另外还可以利用流水并行,将不同的网络层放在不同的AI芯片上运行,进而进一步将计算负载和内存负载切分至多块AI芯片。优点是计算的通信可以重叠,通信需求低;但是存在流水线间空闲bubble,需要和重计算配合使用。

image

通过对数据、网络模型、以及运行的流水线进行切分,分布式训练可以有效地减少单卡的资源负载,一方面提升了训练任务的计算和内存吞吐量,另一方面使得原来单机单卡无法训练的任务,变成可能。

现在一切看上去都是那么的美好,有分布式训练就够了。理论上,AI芯片数量越多,模型训练越快。但是,随着训练数据集规模的进一步增长,数据并行就会出现局限性:当训练资源扩大到一定规模时,由于通信瓶颈的存在,增加计算资源的边际效应,就会越来越明显,甚至增加资源也没办法进行加速,这个比例称为加速比。假设单设备的吞吐量为T,有n个设备系统的吞吐量应为nT,设这个系统实际达到的吞吐量为T_n。最终的加速比(scale factor)为:

scale factor=T_n/nT

image

理论上,我们希望无论有设备数n为多少,加速比越接近1越好。

另外,随着网络模型的不断增大,会出现模型需要的内存急剧膨胀,算子的增加也会使得单卡就算进行模型切分和流水线切分,也难以在合理的时间内完成一个step的训练,最终导致分布式训练不再适用于这类型的任务。这也是大模型的出现,对系统的挑战。

这个时候,面对大模型,我们需要引入大规模训练技术,在解决内存的同时,也不会被计算资源给限制住,让算法开发人员可以方便进行高效的分布式调试调优和编写代码。愉快地训练大模型。


大规模训练的挑战

相比普通的分布式训练,大规模训练在技术上,需要考虑的问题更加复杂。

首先,面对单卡无法装载的大模型,如何利用多卡来突破内存限制的瓶颈是个问题;其次,大规模训练会用到大量的计算资源,大量计算资源间如何通信、协作是另一个难题;最后,如何平衡各类层出不穷的大规模训练技术,使得众多技术形成一个完整高效的训练方案,更是一大学问。

下面,将大规模训练技术面临的挑战分为四个部分:内存、通讯、计算和调优。

image

内存墙

模型训练无可避免的问题就是内存墙。

下面举两个例子,在进行ResNet152对120G的ImageNet数据进行训练时,在使用Batch Size=128,Input Image Size=512x512的前提下,训练过程占用了16G的内存,在GPU一般具有16G内存的情况下,没办法把模型放在一个一块AI芯片上。想进一步加大每次执行数据的Batch size,模型训练的内存占用也会随之增长,最后高于AI芯片的内存容量,触碰到了内存墙,模型无法训练。

以2000亿参数的鹏程盘古大模型为例,2000亿的参数内存占用就消耗了754GB,训练过程因为会有权重、激活、优化器状态、再加上自动微分所产生临时变量,需要3500GB的内存,一个大模型训练就需要100多块具有32G内存的AI芯片。

要理解内存墙出现的本质,就需要了解内存增长的本质。在ResNet50的1轮迭代,打开看看内存占用变化:

1)网络模型开始计算的时候,内存占用不断增加,直到达到峰值1.2GB。

2)峰值过后内存开始逐渐释放,内存占用慢慢降到320M。

3)1个step计算结束后,仍有一部分内存驻留,内存保持在320M。

image

可以发现,模型训练对内存的占用可以分为两部分:

1)一部分是模型自身的权重参数、优化器状态信息,由于是比较固定的所以称为静态参数。

2)另一部分是模型在前向计算和反向传播的时候,会产生诸如前向输出Forward Output、梯度输出Output Gradient、算子计算时候的临时变量等,这部分内存会在反向传播时逐渐释放掉,因此被称为动态参数。

从这里面可以发现,在计算过程的内存峰值,是否遇到内存墙,主要是由动态内存决定。即使绞尽脑汁,降低静态内存意义并不是很大,关键在于降低内存占用的峰值。

静态内存

以ResNet50为例,静态内存占用大约1.2G:

ResNet50 = 250M(权重参数) + 900M(优化器参数和动量) ~ 1.2G

在计算过程中,神经网络模型每一层的卷积或者全连接计算,都会把权重W_m长期保存下来,用作网络的权重参数更新。另外针对诸如ADAM的优化器,会存储优化器的动量等信息,用于优化器计算。

一块AI芯片有16G,最大能塞满20+亿参数的模型,但是这时候已经没有额外空间,留给动态内存进行分配啦。所以说降低动态内存的峰值才是人间正道。

image

动态内存

静态内存比较好理解,下面主要看看动态内存。动态内存一般指的在自动微分Autogard过程,产生的中间变量。

神经网络层数非常多,在计算过程中,神经网络模型每一层的前向输出Forward Output都需要保存下来给反向传播时候使用;而在神经网络的反向传播计算过程中,需要记录下权重梯度Weight Gradient(W_g),给优化器使用,在后续的step中更新梯度的信息;另外还有梯度输出Output Gradient同样在反向传播的过程中,继续在反向传播的过程中,作为上一层网络模型的输入。

当然,这里面还会有一些正向和反向的算子中间变量需要消耗内存的,由此,上述的O_f、O_g、W_g、OP累积下来,会造成巨大的内存开销。从ResNet50的例子中,我们可以看出,动态内存是静态内存的X倍(,静态内存占用了1.2G内存,而动态内存则占用了25G内存,比例非常不平衡)。

image

静态内存和动态内存都可能造成内存墙的问题。相互独立但又相互制约,任意一侧的增大都会导致留给另一侧的显存空间变小,所以单单对一侧做优化是不够的,必须同时优化静态和动态内存。

为了能够将大模型运行起来,需要使用模型并行、流水并行等技术,但是这些技术又会降低AI芯片的计算吞吐量。

通讯墙

大模型通过模型并行、流水线并行切分到AI集群后,通讯便成了主要的性能瓶颈。

从模型切分的角度来考虑,大模型再怎么切分其仍然是一个模型,因而模型切分到不同的机器后,仍然需要通信来对分布在不同机器的参数进行总体聚合。

在通讯方式的角度来看,参数聚合(All Reduce)就会对通讯提出很多需求,例如使用同步的更新策略,还是使用异步的更新策略,如何对模型局部变量进行更新等。

另外,在网络通讯的角度,由于专用的AI加速芯片中内存与计算单元ALU非常进,使得片内带宽很大,计算速度很快,但是在集群中的网络传输速度远远不能匹配专用的AI加速芯片运算速率。

直接用更大的带宽不能解决这些问题吗?

从图中,我们可以看到,随着网络带宽从1Gbps到100Gbps,利用率从接近100%下降到40%左右。随着网络带宽的增加,带宽的利用率将会越来越低,高带宽所带来的收益会遇到瓶颈。

image

随着机器规模的扩大,基于同步的All Reduce通讯聚合方式,会因为大量的AI芯片和服务器之间频繁进行同步,出现水桶效应,也就是最慢的一路通讯,将会决定整个AI集群的通讯的高度。

如果采用目前比较流行的Ring-AllReduce的通信聚合方式,当通讯的环越大,通讯的延长将会不断地被扩大。另外网络协议的多次握手的方式,诸如此类的开销会导致训练无法有效利用带宽。

总体而言,通讯过程中,需要综合考虑数据参数量、计算量、计算类型、数据样本量、集群带宽拓扑和通讯策略等不同的因素,才能设计出性能较优的切分策略,最大化利用通讯效率,提高通讯比。

性能墙

性能墙呢主要是指计算资源利用率的问题。随着大模型的提出,对算力需求更加迫切,理论上在4K的集群上每块卡快1分钟,总体就快了68个小时。大模型会增加对算力的需求,但是随着大模型引入各项分布式并行技术的同时,会降低计算资源的利用率。

计算效率在MMlab的贴总结得非常好,我们同样从底至上分为算子层(Operator Level)、图层(Graph Level)和任务层(Task Level)三个方面。

算子层(Operator Level)

在Operator Level的问题大模型和普通模型遇到的问题都是类似,主要是指对算子级别的优化。

1)其中比较突出的问题是模型小算子过多,可以通过算子融合进行优化,例如对Conv+Bn+Act三个算子进行融合成一个大算子。但是还有很多算子融合策略无法穷举,如何通过计算编译技术,对小算子进行融合。

2)第二点就是算子实现不够高效,类似于卷积CONV算子针对2x2和3x3 Kernel大小的使用Winograd算法代替。可是这种算法是人为提出了,更多的高效Kernel算子,如何利用编译技术,去寻找最好的多线程并行计算。

3)最后一点就是内存局部性差,对算子和内存的开销进行分析,可以对计算时算子输出有相同的shape进行复用。

图层(Graph Level)

Graph Level指如何对计算图进行优化,进而加速大规模训练。

一般而言,如果没有遇到控制流和大模型,一个模型只有一个图。但是当遇到大模型的时候,由于一块AI芯片的内存没办法满足大模型的内存需求,因此需要对算子和模型进行切分。这时候涉及两方面的问题。

1)如何搜索和切分处计算效率更高的子图,分配到不同的机器上进行计算。

2)数据在通讯和内存之间如何增加overlay重叠部分,提高单卡整体的计算率。

任务层Task Level

大集群大模型的分布式混合了各种并行训练策略,在任务层级,性能的瓶颈从计算转移到了通信,如何降低通信量,缩减通信域规模,尽可能把通信时延隐藏在计算中,是大规模训练的核心关注点。

大规模训练技术中,不仅要求AI芯片的计算性能足够强悍,同时也依赖于AI框架的大规模分布式训练的运行和调度效率,以及在分布式并行等各种优化手段的权衡。

调优墙

在数千节点的集群上进行模型开发,听到就头皮发麻。我平时把EffectNet网络模型魔改,还不一定能够收敛,调一次参数再训练一次,单机多卡一个星期就过去了。那要是大模型训练了1个月遇到Loss跑飞了怎么办?

所以在数千节点的集群上,需要考虑到提升算法工程师分布式调试调优的效率,另外还要考虑降低工程师对大模型进行并行切分的难度。除了对人的考虑,还要对硬件集群的管理,需要保证计算的正确性、性能、可用性。要是有一台机器坏了,如何快速恢复训练中的参数。


总结一句话,大模型训练考验的是算法、数据、框架、资源调度等全栈和全流程的综合能力。希望MindSpore能够在大模型引领业界。

最后编辑于
©著作权归作者所有,转载或内容合作请联系作者
  • 序言:七十年代末,一起剥皮案震惊了整个滨河市,随后出现的几起案子,更是在滨河造成了极大的恐慌,老刑警刘岩,带你破解...
    沈念sama阅读 203,271评论 5 476
  • 序言:滨河连续发生了三起死亡事件,死亡现场离奇诡异,居然都是意外死亡,警方通过查阅死者的电脑和手机,发现死者居然都...
    沈念sama阅读 85,275评论 2 380
  • 文/潘晓璐 我一进店门,熙熙楼的掌柜王于贵愁眉苦脸地迎上来,“玉大人,你说我怎么就摊上这事。” “怎么了?”我有些...
    开封第一讲书人阅读 150,151评论 0 336
  • 文/不坏的土叔 我叫张陵,是天一观的道长。 经常有香客问我,道长,这世上最难降的妖魔是什么? 我笑而不...
    开封第一讲书人阅读 54,550评论 1 273
  • 正文 为了忘掉前任,我火速办了婚礼,结果婚礼上,老公的妹妹穿的比我还像新娘。我一直安慰自己,他们只是感情好,可当我...
    茶点故事阅读 63,553评论 5 365
  • 文/花漫 我一把揭开白布。 她就那样静静地躺着,像睡着了一般。 火红的嫁衣衬着肌肤如雪。 梳的纹丝不乱的头发上,一...
    开封第一讲书人阅读 48,559评论 1 281
  • 那天,我揣着相机与录音,去河边找鬼。 笑死,一个胖子当着我的面吹牛,可吹牛的内容都是我干的。 我是一名探鬼主播,决...
    沈念sama阅读 37,924评论 3 395
  • 文/苍兰香墨 我猛地睁开眼,长吁一口气:“原来是场噩梦啊……” “哼!你这毒妇竟也来了?” 一声冷哼从身侧响起,我...
    开封第一讲书人阅读 36,580评论 0 257
  • 序言:老挝万荣一对情侣失踪,失踪者是张志新(化名)和其女友刘颖,没想到半个月后,有当地人在树林里发现了一具尸体,经...
    沈念sama阅读 40,826评论 1 297
  • 正文 独居荒郊野岭守林人离奇死亡,尸身上长有42处带血的脓包…… 初始之章·张勋 以下内容为张勋视角 年9月15日...
    茶点故事阅读 35,578评论 2 320
  • 正文 我和宋清朗相恋三年,在试婚纱的时候发现自己被绿了。 大学时的朋友给我发了我未婚夫和他白月光在一起吃饭的照片。...
    茶点故事阅读 37,661评论 1 329
  • 序言:一个原本活蹦乱跳的男人离奇死亡,死状恐怖,灵堂内的尸体忽然破棺而出,到底是诈尸还是另有隐情,我是刑警宁泽,带...
    沈念sama阅读 33,363评论 4 318
  • 正文 年R本政府宣布,位于F岛的核电站,受9级特大地震影响,放射性物质发生泄漏。R本人自食恶果不足惜,却给世界环境...
    茶点故事阅读 38,940评论 3 307
  • 文/蒙蒙 一、第九天 我趴在偏房一处隐蔽的房顶上张望。 院中可真热闹,春花似锦、人声如沸。这庄子的主人今日做“春日...
    开封第一讲书人阅读 29,926评论 0 19
  • 文/苍兰香墨 我抬头看了看天上的太阳。三九已至,却和暖如春,着一层夹袄步出监牢的瞬间,已是汗流浃背。 一阵脚步声响...
    开封第一讲书人阅读 31,156评论 1 259
  • 我被黑心中介骗来泰国打工, 没想到刚下飞机就差点儿被人妖公主榨干…… 1. 我叫王不留,地道东北人。 一个月前我还...
    沈念sama阅读 42,872评论 2 349
  • 正文 我出身青楼,却偏偏与公主长得像,于是被迫代替她去往敌国和亲。 传闻我的和亲对象是个残疾皇子,可洞房花烛夜当晚...
    茶点故事阅读 42,391评论 2 342