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

子豪 发自 凹非寺 量子位 报道

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

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

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

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

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

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

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

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

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

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

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

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

为什么不能靠多GPU堆显存

那么，为了摆脱单个硬件的有限内存容量和带宽，是否可以将训练扩展到多个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-2cb4265cb0b8https://github.com/amirgholami/ai_and_memory_wall