什么？单块 GPU也能训练大模型了？

还是 20 系就能拿下的那种？？？

没开玩笑，事实已经摆在眼前：

RTX 2060 6GB 普通游戏本能训练15 亿参数模型；

RTX 3090 24GB 主机直接单挑180 亿参数大模型；

Tesla V100 32GB 连240 亿参数都能拿下。

相比于 PyTorch 和业界主流的 DeepSpeed 方法，提升参数容量能达到 10 多倍。

而且这种方法完全开源，只需要几行代码就能搞定，修改量也非常少。

这波操作真是直接腰斩大模型训练门槛啊，老黄岂不是要血亏。

那么，搞出如此大名堂的是何方大佬呢？

它就是国产开源项目 Colossal-AI。

自开源以来，曾多次霸榜 GitHub 热门第一。

△开源地址：https://github.com/hpcaitech/ColossalAI

主要做的事情就是加速各种大模型训练，GPT-2、GPT-3、ViT、BERT 等模型都能搞定。

比如能半小时左右预训练一遍 ViT-Base/32，2 天训完 15 亿参数 GPT 模型、5 天训完 83 亿参数 GPT 模型。

同时还能省 GPU。

比如训练 GPT-3 时使用的 GPU 资源，可以只是英伟达 Megatron-LM 的一半。

那么这一回，它又是如何让单块 GPU 训练百亿参数大模型的呢？

我们深扒了一下原理 ~

高效利用 GPU+CPU 异构内存

为什么单张消费级显卡很难训练 AI 大模型？

显存有限，是最大的困难。

当今大模型风头正盛、效果又好，谁不想上手感受一把？

但动不动就 "CUDA out of memory"，着实让人遭不住。

目前，业界主流方法是微软 DeepSpeed 提出的ZeRO（Zero Reduency Optimizer）。

它的主要原理是将模型切分，把模型内存平均分配到单个 GPU 上。

数据并行度越高，GPU 上的内存消耗越低。

这种方法在 CPU 和 GPU 内存之间仅使用静态划分模型数据，而且内存布局针对不同的训练配置也是恒定的。

由此会导致两方面问题。

第一，当 GPU 或 CPU 内存不足以满足相应模型数据要求时，即使还有其他设备上有内存可用，系统还是会崩溃。

第二，细粒度的张量在不同内存空间传输时，通信效率会很低；当可以将模型数据提前放置到目标计算设备上时，CPU-GPU 的通信量又是不必要的。

目前已经出现了不少 DeepSpeed 的魔改版本，提出使用电脑硬盘来动态存储模型，但是硬盘的读写速度明显低于内存和显存，训练速度依旧会被拖慢。

针对这些问题，Colossal-AI 采用的解决思路是高效利用GPU+CPU 的异构内存。

具体来看，是利用深度学习网络训练过程中不断迭代的特性，按照迭代次数将整个训练过程分为预热和正式两个阶段。

预热阶段，监测采集到非模型数据内存信息；

正式阶段，根据采集到的信息，预留出下一个算子在计算设备上所需的峰值内存，移动出一些 GPU 模型张量到 CPU 内存。

大概逻辑如下所示：

这里稍微展开说明下，模型数据由参数、梯度和优化器状态组成，它们的足迹和模型结构定义有关。

非模型数据由 operator 生成的中间张量组成，会根据训练任务的配置（如批次大小）动态变化。

它俩常干的事呢，就是抢 GPU 显存。

所以，就需要在 GPU 显存不够时 CPU 能来帮忙，与此同时还要避免其他情况下内存浪费。

Colossal-AI 高效利用 GPU+CPU 的异构内存，就是这样的逻辑。

而以上过程中，获取非模型数据的内存使用量其实非常难。