X

新加坡国立大学尤洋:高性能 AI 的突破丨GAIR 2023

过去数年,AI 模型的参数发生了极大变化。尤洋指出,从 2016 年至 2021 年 1 月,AI 模型的参数量是每 18 个月增长 40 倍;从 2018 年 1 月到 2021 年 1 月,AI 大语言模型的参数量每 18 个月增长 340 倍。而相形之下,2016 年 1 月至 2021 年 1 月间,GPU 的内存增长每 18 个月仅有 1.7 倍。

由此可见,训练成本高、周期长,是当前大模型发展最需要克服的难题。

针对这一问题,尤洋提出了 Colossal-AI 系统,从高效内存系统、N 维并行系统和大规模优化三个层次出发,以实现同样的设备条件下将数据移动的最小化,将 GPU 的吞吐量扩大至最高点。

尤洋还指出,现阶段的模型参数量以 10 万倍扩大、但层数增加不多,这或意味着:如今的 AI 发展可能不再是深度学习、而是进入了宽度学习时代。在模型变得更宽的情况下,面对大规模、长时间的 GPU 训练任务,大模型训练系统的核心将是如何实现 GPU 并行计算,以实现大模型训练越快越省钱的目标。

以下为尤洋的现场演讲内容,雷峰网作了不改变原意的编辑及整理:

AI 大模型的机遇和挑战

首先展示一张图片。图上的横坐标是时间,纵坐标是 AI 模型的参数量。

从 2016 年至 2021 年 1 月,AI 大模型的参数量大概每 18 个月增长 40 倍;从 2018 年 1 月到 2021 年 1 月,AI 模型的参数量每 18 个月增长 340 倍。

2016 年,当时世界上最好的模型是 ResNet-50,而今天最好的模型是 GPT-4。从架构上来看,虽然 OpenAI 没有对外公布 GPT-4 的架构,但对比 ResNet-50 的 50 层神经网络和 GPT-3 未达 100 层的架构,可以说 AI 模型近年来的层数并没有产生的太大的变化。

从 ResNet-50 到 GPT-4,虽然参数量大了 10 万倍左右,其实是每一层都变得更加宽了。包括 LLaMA-65B 版本,也是几十层的网络。

所以我们可能不是深度学习,而是进入了一个宽度学习时代。

可以看到,自 2019 年开始, Transformer 架构基本统一了 AI 大模型赛道,当前最高效的 AI 大模型均是 Transformer 架构。上图中的两条虚线,既展示了大模型参数的变化趋势,实际上也展现了 GPU 的变化趋势。

虽然现在英伟达的股价涨了很多倍,但包括英伟达在内的厂商,其 GPU 内存的增长速度远跟不上大模型的发展速度。

相较于过去六年模型参数量的增长速度变化,2016 年 1 月至 2021 年 1 月间,英伟达 GPU 的计算增长速度每 18 个月仅增长了 1.7 倍。

以 A100 80G 为例计算 GPT-3 训练所需的内存数,GPT-3 有大概 1750 亿参数,为方便计算取整数 2000 亿,等于 200 乘以 10 的 9 次方,每个单精度占用 4 个字节,仅参数就要占 800G 内存,梯度也占了 800G 内存。按照当前的优化方法储存一阶矩(first moment)、二阶矩(second moment)等信息均为 800G。也就是说,如果训练一个什么事情都不干大模型,至少需要几 T 的内存,单个 A100 GPU 仅 80G 内存远远不够,加上中间结果的 batch size 越大,内存开销也越大。

这也是为什么从内存角度上看,训练大模型首先需要有成千上万个 GPU 的原因。

举一个可量化的案例。PaLM 是一个 5400 亿的大模型,根据当前云计算市场价,训练 PaLM 需要承包至少 1000 个 GPU,成本约 900 多万美金。而 Meta 此前曾明确提到,LLaMA 需要使用到 2000 个 A100 GPU,并且用三周时间才可完成一次训练,由此计算可得出 LLaMA 单次训练的成本在 500 万美金。

但由于炼大模型并不仅限于一次训练,可能一个好的大模型产品迭代至少需要五六次,前期都是在试错。因此,据公开渠道分析, GPT-4 单次训练成本在 6000 万美金左右,且训练一次需要至少几个月时间。这也是为什么目前 ChatGPT 即便将其升级至最新版本,其底层还是 2021 年 9 月版本的模型。也就是说,从 2021 年 9 月至今,OpenAI 实质上并没有升级其产品,根本原因就在于,模型的每次训练不仅成本很高,训练周期也很长,因此大模型训练的代价高就非常严重。

设想一个问题,今天有一个 1 万亿参数以及 1000 亿参数的大模型,我们能否用什么方法,去检测万亿参数大模型比千亿参数大模型二者哪个效果更好?也即是目前常说的,模型参数量增大、到底要增大到什么程度?

到今天而言,我觉得这个问题暂时无法用科学回答。原因有几个。

首先,训练神经网络存在非凸优化的问题,目前训练所收敛的点多为局部最优解、而非全局最优解。因此,我们要验证神经网络训练到什么程度,在现有计算资源情况下是无法验证的。

第二个难度在于,大模型训练往往只训练一两个 epoch,而此前的 CNN 模型中,ResNet 训练有 90 个 epoch,甚至自监督学习训练有 1000 个 epoch,因此大模型只训练一两个 epoch 的方式,相当于只将数据集过了一两遍,收敛就更加不充分了。因此在训练成本如此之高的情况下,我们很难验证,一个 1 万亿参数的模型和 2 万亿参数的模型二者谁更好,因为它们潜力都没有能通过实验得到充分发挥。因此我认为,今天 AI 大模型是一个实验性学科,如何能高效提升这个实验的效率,降低成本,对整个行业的普及具有根本性的作用。

回到现实之中,为什么今天人人都在追求大模型?从数学逻辑上看,模型参数越大、效果越好,这是绝对的。

与此同时,成本也再不断攀高。目前训练大模型需要成百上千、甚至上万个 GPU,如何将上万个 GPU 的成本进一步降低,挑战非常大的。

在 20 年前,由于当时依靠的是主频的,所有的程序都是串行的,假设将硬件速度提高 10 倍,在一行代码都不用的更改的情况下,其运行速度也可以提升 10 倍。但到了如今,如果想将代码速度提升 10 倍,假定硬件已经增速 10 倍,但如果不优化代码,很可能速度反而会变慢。原因就在于,机器规模更大的情况下,比如 GPU 内部,GPU 内存和 CPU 之间的数据移动,或是 GPU 间的数据移动,再加上服务器实现的数据移动,会占据整个系统的大部分时间,把大部分时间都花在了数据移动上,模型的扩展性也会变得不好。

我认为,未来一个好的分布式软件和一个差的分布式软件,在上千个 GPU 上,甚至 500 个 GPU 上,其速度可能相差 10 倍。

Colossal-AI 如何运行?

基于上述的挑战,我们提出了大模型训练系统 Colossal-AI,提供优化方法,降低数据移动的代价,将模型扩展性效率提到最高。

一个具体的数据是,使用最简单的 PyTorch 训练 GPT-3,成本为 1000 万美金,英伟达经过优化后,用 Megatron 可将其成本减少至 300 万美金,而使用 Colossal-AI 后,成本可以降低到 130 万美金。可以看到,相同的设备条件下,数据移动的最小化将数据移动占比降低最低,能够把 GPU 吞吐量扩大至最高点。

针对上述问题,Colossal-AI 提出了三个层次。其他类似的软件也包括了这三个层次。

第一层是优化内存,先确保单个 GPU、单个服务器的内存效率最高,这是基础。

第二层是 N 维的并行。当前我们使用上千、上万个 GPU 时,其核心技术就是 ParallelComputing(并行计算)。从 1 个 GPU 扩到 10 个 GPU,因为其规模比较小,我们可以轻易获得 7 倍加速;从 10 个 到 100 个 GPU 时,往往可能只获得 4 倍加速,因为并行规模变大,它的通信代价变高了。而从 100 个 GPU 到 1000 个 GPU,因为通信代价进一步加高,很可能只获得 2 倍的加速。而从 1000 个 GPU 到 1 万个 GPU,如果软件运行情况不佳时,不仅可能无法加速,甚至还会更慢,因为设备将所有时间花耗在了更高密度的通信上。

其次是优化问题,未来 AI 大模型的发展方向我认为有两层,第一层是模型变得更加智能,设计出更好的结构,比如说从 BERT 到 GPT,或者从 ResNet 到 BERT等,都是在不断地尝试改变模型结构。

此外还有优化方法的进步,从 SGD 过渡到 MOMENTUM、ADAGRAD,到现在有 ADAM,未来又会有哪些更好的优化方法能够将效率提升 10 倍,这一点也非常重要。

具体到实际操作训练大模型的并行问题。

首先是数据并行,这是最简单、也是最高效的并行方法。数据并行指的是,假设现有 1 万张图片,每次循环处理 1000 张图片,如果有 10 个机器,每个机器分配 100 张,10 个循环即可完成所有图片的处理。

在数据并行的过程中需要进行汇总,每个机器用不同的数据获得不同梯度,机器在不同数据上学习不同的更改,并更新参数梯度,最后算出全局梯度,目前采用的是加和求平均的方式,效果已经非常好了。此前 Colossal-AI 在数据并行中的 LARS 方法,就为谷歌、 Meta、腾讯、索尼等公司,将 ImageNet 的训练时间从一小时缩短至一分钟。

数据并行是最基本的,同时也是最稳定的。将数据划分之后,假设未来有 1 万个 GPU,很容易发生的情况是,隔几个小时就有四五个 GPU 崩溃了,运维 1 万个 GPU 的集群很难,但数据并行的稳定之处在于,即便有 1 万个GPU 崩溃了十几个,但大体结果是不会变的,因为它是梯度加和求平均。

基于这个考虑,我认为数据并行是一个根本性的基础设施。

当然,仅用数据并行并不够,原因在于:数据并行有一个假设,必须将模型拷贝到每个 GPU 或服务器内,由服务器或 GPU 去交换梯度。但如果 GPU 仅 80G 内存时,万亿参数的模型则需要几十 T 的内存,这在 GPU 中是无法存放的,需要将模型切割至不同的 GPU 上再汇总结果。这种方法叫做模型并行。模型并行包括两种,第一种是张量并行( tensor paralism),即层内的模型并行。例如 GPT-3 的层数大概为八九十层,每层切割一次模型,将其层内计算分隔成多份,算完一层再算下一层,依次类推,这就是张量并行。

另一种方式则是 Pipeline Parallelism(流水线并行),介于数据并行和张量并行外的一种模型并行方式。通过构建几个数据 pipe(管道),每个数据 pipe 的数据点不同,相当于将一个大尺寸分割为多个小尺寸,通过这种方式进行 pipe 计算。假如有 10 个 pipe,10 个 pipe 代表十组不同的数据,第一个 pipe 计算第一层的数据,第二个 pipe 计算第二层……以此方式并行,类似我们盖楼一样,10 个工程队盖 1000 层楼,当第一个工程队在盖第一栋楼的第一层,第二个工程队盖第二栋楼的第二层,依此类推。

当楼数越多,楼和工程队之间的比值越高,效率也越高,相当于 10 个工程队在同时运转。其中每个工程队就相当于一个 GPU,每个楼就相当于一个 pipe,楼的层数相当于这个神经网络的层数,这就流水线并行的核心逻辑。

目前工业界已经做了相关的工作,除了 Colossal-AI 之外,还有英伟达的 TensorRT 和微软的 DeepSpeed,他们也是技术壁垒最高的两家公司。

但 Colossal-AI 与其不同之处是,Colossal-AI 专注于未来大模型的发展方向。可以看到,当前的模型还在变得更宽,而不是变得更深,张量并行将会更加重要,但它最大的弊端就在于,因为它是切割的是整个层,通信开销太大。这也是为什么英伟达 CEO 在 GTC 峰会上首次介绍 3D 并行时特别说明其通信开销太大的问题,只能放到一个服务器内去做。因此,Colossal-AI 主打 2D 张量并行和 2.5D 张量并行,将计算成本降低了一个数量级。

这就意味着用一维张量并行,1 万个机器里,每个机器都需要跟 9999 个机器打交道,而 2D 张量并行则是将其分成了各个子单元,每个机器只需要跟 96 个机器打交道。它的核心逻辑是,用一些 local synchronization(局部同步)去取代global synchronization(全局同步),以更多的局部通信去取代全局通信,这个过程中,设计调度是最困难的。

3D 张量并行也是同样,每升高一个维度,它的设计复杂度会高一个数量级,最终通信复杂度下降了。

在内存优化方面,目前 AI 大模型训练的内存开销很大,即便什么事情都不做,也需要几 T 的内存,如果不进行人工干预,一旦使用起来,可能需要几十 T 甚至是几百 T 的内存。

为了让模型预测效果更好,我们往往需要长序列数据,当前大模型的本质是通过一个单词的输出来预测下一个单词的概率,长序列成为刚需。对此, Colossal-AI 也推出了 Sequence Parallelism(序列并行)。

具体而言,在将序列进行切割后,会面临一个严重的问题是:在进行 attention score 时,每个 token 都需要跟全局序列中的其他 token 去评估,而切分后的服务器上只有部分 token,其他服务器上也会分布不同的 token,以至于每个服务器运行是需要同其他服务器打交道。

也就是说,假设今天屋子里 200 个人每人分别拿了一包零食,我希望每个人能品尝下其他所有人的零食,至少需要 200 个平方次交换,才能让每个人都尝到其他人的所有零食。那么最简单的方式是:所有人围一个圈,每个人将自己吃过的零食递给右手边的人,从自己的左手边获得零食,仅需 n-1 次,即 199 次的传递即可完成。从而降低了整个通信成本。

总结一下,目前 AI 大模型训练的核心技术栈,其实就是并行计算,因为我们要处理成百上千上万个 GPU 核心,把所有 GPU 并行利用起来。数据并行、张量并行、流水线并行以及数据序列并行是并行中较为核心的模块。

目前在内存优化方面,我们处于一个没有太多选择的环境,英伟达 GPU 是最好的,我们好像也没有其他更好的方案能够去取代它。但美中不足的是,英伟达 GPU 的内存有限,在这种情况下,我们能否思考如何利用 CPU 内存、NVMe内存,核心思想就是,GPU 内存放不下就移到 CPU 上,CPU 放不下就放到 NVMe 上,相当于在盖楼时,所需要的原材料自家楼下工地放不下,那我们就将其放到隔壁工厂。其技术的核心也在于最小化数据移动,即最小化 CPU、 GPU 之间的数据移动,最强化 CPU 和 NVMe 之间的数据移动,从而将吞吐力速度提升到最高。

走向开源

Colossal-AI 是一个开源软件,同时我们也做了一个商业化的平台,对没有 GPU 的用户,可以直接在平台上去训练部署自己的大模型。我们也提供了 LLaMA、PaLM、 GPT 等各种模型,最快仅需两三天就可完成一个模型微调。相比之前用户可能需要几周、乃至几个月来处理硬件、软件等基础设施,效率得到了极大提升。同时,Colossal-AI 也保护用户的隐私,平台不会保留、访问用户数据,这是 Colossal-AI 与 OpenAI ChatGPT 的本质区别。我们将数据上传至 Google Cloud 时,很多时候 Google 并没有碰我们的数据,但是 OpenAI GPT 会进行分析,AI 模型的不可解释性、训练不彻底等风险普遍存在。因此,未来也会有很多企业训练自己的大模型,Colossal-AI 做的,是最大化保护用户的隐私,同时提供大模型训练的工具。

在性能上,Colossal-AI 在同样的硬件上可以训练24 倍大的模型,相对于 DeepSpeed 的 3 倍加速,即便是一个低端的服务器,也可以借助 Colossal-AI 完成相应的模型训练。例如 LLaMA-65B 的训练,在 Colossal-AI 上使用同样的代码放,可以直接获得约 50% 的加速效率。

一个简单的比喻,比如说现在大模型是挖金子,英伟达是卖铲子的,那我们就是卖手套、卖衣服的,把挖金子的效率提到最高。

(雷峰网雷峰网)