AI训练的最大障碍不是算力,而是“内存墙”

二者增长速度相差巨大

子豪 发自 凹非寺
量子位 报道 | 公众号 QbitAI

AI训练的计算量每年都在大幅增长,最近有研究指出,AI训练未来的瓶颈不是算力,而是GPU内存

AI加速器通常会简化或删除其他部分,以提高硬件的峰值计算能力,但是却难以解决在内存和通信上的难题。

无论是芯片内部、芯片间,还是AI加速器之间的通信,都已成为AI训练的瓶颈。

Transformer模型中的参数数量(红色)呈现出2年240倍的超指数增长,而单个GPU内存(绿色)仅以每2年2倍的速度扩大。

AI训练的最大障碍不是算力,而是“内存墙”

△多年来SOTA模型的参数数量(红点)以及AI加速器存储容量(绿点)的演变

训练AI模型的内存需求,通常是参数数量的几倍。因为训练需要存储中间激活,通常会比参数(不含嵌入)数量增加3-4倍的内存。

于是,AI训练不可避免地撞上了“内存墙”(Memory Wall),内存墙不仅是内存容量,也包括内存传输带宽。

在很多情况下,数据传输的容量和速度,都没有触摸到内存墙。

AI训练的最大障碍不是算力,而是“内存墙”

△训练不同神经网络模型所需的内存量

从图中可以看出,每当GPU内存容量增加时,开发人员就会设计出新模型;

2019年GPT-2所需的内存容量,已经是2012年的AlexNet的7倍以上;

自谷歌团队在2017年提出Transformer,模型所需的内存容量开始大幅增长。

为什么不能靠多GPU堆显存

那么,为了摆脱单个硬件的有限内存容量和带宽,是否可以将训练扩展到多个AI加速器,使用分布式内存呢?

事实上,这样做也会面临内存墙的问题,神经网络加速器之间移动数据的通信瓶颈,甚至比芯片上的数据移动还慢且低效。

与单系统内存的情况类似,扩展带宽的技术难题还尚未被攻克。仅在很少的通信和数据传输的情况下,横向扩展才适用于计算密集型问题。

AI训练的最大障碍不是算力,而是“内存墙”

从图中可以看出,20年间,运算设备的算力提高了90,000倍;

虽然存储器从DDR发展到GDDR6x,能够用于显卡、游戏终端和高性能运算,接口标准也从PCIe1.0a升级到NVLink3.0;

但是和算力的提高幅度相比,通讯带宽的增长只有30倍,可以说非常缓慢。

由于算力和内存之间的差距越来越大,训练更大的模型也会更加困难。

怎样突破“内存墙”

怎样解决内存限制问题?作者从三个方面进行了分析。

训练算法的改进

训练神经网络模型的一大挑战,就是要进行蛮力超参数调整。虽然可以通过二阶随机优化方法来实现,不过目前的方法却增加了3-4倍的内存占用,这一点仍需解决。

微软的Zero方法(一种万亿级模型参数训练方法),实现了在相同内存下,通过去除多余的优化状态变量,来训练8倍大的模型。

也可以在传递过程中只存储或检查激活的子集,而不保存所有激活,以此将内存占用减少5倍,不过需要增加20%的计算量。

此外,从单精度算法到半精度(FP16)算法的进展,使硬件计算能力提高了10倍以上,可以进一步研究适合INT8精读训练的优化算法。

高效部署

最新的SOTA模型(例如:GPT-3)需要分布式内存部署,这是一个很大的挑战。可以通过降低精度或删除其冗余参数,来压缩这些模型,以进行推理。

在训练或推理过程中,可以降低至INT4精度,模型占用空间和延迟能够减少8倍。不过,想要将训练精度降低到远低于FP16,仍然很困难。

而删除冗余参数,则可能导致准确率下降。当前的方法能够修剪30%的具有结构化稀疏性的神经元,以及80%的具有非结构化稀疏性的神经元,以保证对准确性的影响最小。

AI加速器的设计

虽然很难同时提高存储带宽和峰值计算能力,但是可以牺牲峰值计算,以获得更好的带宽。

在带宽受限问题上,CPU的性能要比GPU好得多,但是与相比GPU相比,CPU的峰值计算能力要小一个数量级左右。

因此,可以研究一种在二者之间的另一种架构,实现更高效的缓存。

研究数据可戳下方链接查看~

原文链接:

https://medium.com/riselab/ai-and-memory-wall-2cb4265cb0b8
https://github.com/amirgholami/ai_and_memory_wall

版权所有,未经授权不得以任何形式转载及使用,违者必究。