微信开源「派大星」：4000元游戏电脑能带动7亿参数GPT！

金磊 发自 凹非寺

量子位 报道 | 公众号 QbitAI

一听到训练大模型，是不是第一感觉就是贵、烧钱、玩不起？

但我说，一台4000多块钱的游戏电脑，谁都能训练上亿参数的大模型呢？

别不信，这是真的。

而这就归功于微信AI团队，最近推出的一款利器——派大星。

但不是你印象中的那个派大星昂~

微信AI的派大星（PatricStar），其实是一个超大预训练模型训练系统。

要知道，在这个领域中，以往都是微软DeepSeed独占鳌头。

此次微信一出手，可以说是直接秒杀了微软：

在 8xV100 和 240GB CPU 内存节点上，训练了一个120 亿参数的 GPT 模型，是当前最佳方案DeepSpeed模型规模上限的1.5 倍。

但毕竟针对的是大模型，“烧钱”是出了名的难题。

而微信AI的派大星就显得相当的亲民了。

即使在700美元的个人游戏电脑上，它也可以训练一个7亿参数的 GPT 模型！

现在，人人都可以在家训练大模型了！

划重点：已开源！

为什么要搞派大星？

大规模预训练模型，已然成为技术发展中的新潮流。

以BERT、GPT为代表的预训练模型的出现，可以说是自然语言处理（NLP）领域的里程碑事件。

NLP，正在进入了预训练时代。

那么像派大星这样的训练系统，真的有必要吗？

答案是肯定的。

从技术角度来看，预训练模型（PTM）通常使用一个堆叠了多个Transformer结构神经网络，在大量文本上预训练通用语言特征表示。

然后，通过微调将学到的知识转移到不同的下游任务。

预训练模型使用大量来自互联网的文本数据，可以捕获自然语言的细微特征，并在下游任务上获得非常惊艳的表现效果。

于是，AI社区的共识是采用预训练模型，作为特定NLP任务的主干，而不是在与任务相关的数据集上从头开始训练模型。

预训练模型的力量源泉，是它拥有的数以亿计的参数规模，这对运行它的计算和内存资源都提出了巨大的要求。

因此，预训练模型训练仍是一小部分人的游戏。

所有发表百亿级模型训练成果的团队，所采用的的设备都是如DGX型号的AI超级计算机。

它的一个节点就配置了8张GPU，1.5TB内存，3.84TB SSDs，还使用NVLink作为高速通信网络。

目前最大的预训练模型Megatron-Turing，包含5300亿参数，其预训练过程就是在560个DGX A100节点的集群上完成的。

这种配置在大多数工业界数据中心都是遥不可及的。

而通过像派大星这样的训练系统，便可以让这种“遥不可及”变得“唾手可得”，让大模型可以普惠到更多的开发人员，实现PTM的“共同富裕”。

再从绿色AI角度来看，预训练模型的预训练的过程是极其烧钱和有害环境的。

比如，从头训练型一次万亿级别的预训练模型要烧掉154万人民币，耗电所产生的碳排放相当于数十辆小汽车从出厂到报废的碳排放总和。

出于社会利益最大化考虑，预训练技术未来的产业形态，应该是中心化的：

少部分财力雄厚的机构，用超大规模集群承担预训练阶段的计算和环境开销；大多数从业人员在小规模、相对简陋的硬件上针对自身业务进行微调。

前者只需要相对少量的计算和碳排放，而后者的诉求却被当前的预训练软件所严重忽略。

现如今，派大星的到来，让大规模预训练模型的训练变得“多快好省”了起来。

而且不仅是对于机构，更是有益于个人开发者。

……

那么派大星的效果，具体又怎样呢？

不是魔改，是从头搭建，性能达SOTA

值得一提的是，派大星并不是基于DeepSpeed的魔改，代码是团队从头开始搭建起来的。

派大星框架非常直观的一个特点，便是简单易用，而且还是可以兼容其他并行方案的那种。

例如，开发者可以使用几行代码端到端的加速PyTorch的训练过程。

from patrickstar.runtime import initialize_engine

config = {
"optimizer": {
"type": "Adam",
"params": {
"lr": 0.001,
"betas": (0.9, 0.999),
"eps": 1e-6,
"weight_decay": 0,
"use_hybrid_adam": True,
},
},
"fp16": { # loss scaler params
"enabled": True,
"loss_scale": 0,
"initial_scale_power": 2 ** 3,
"loss_scale_window": 1000,
"hysteresis": 2,
"min_loss_scale": 1,
},
"default_chunk_size": 64 * 1024 * 1024,
"release_after_init": True,
"use_cpu_embedding": False,
}

def model_func():
# MyModel is a derived class for torch.nn.Module
return MyModel(...)

model, optimizer = initialize_engine(model_func=model_func, local_rank=0, config=config)

...

for data in dataloader:
optimizer.zero_grad()

loss = model(data)
model.backward(loss)
optimizer.step()

接下来，我们一起看一下派大星的性能效果。

上图便展示了DeepSpeed stage3，PyTorch系统在 1、2、4、8 个 GPU 上的性能（y轴通过对数方式重新缩放）。

这些点代表在一个 GPU 上使用 4、8、16、32 和 64 批大小测试的最佳结果。

（注：圆点周围的值表示派大星在吞吐量及其对DeepSpeed的加速；deeps是DeepSpeed仅使用数据并行的效果，我们接下来称之为DeepSpeed-DP，deeps-mpX 是 DeepSpeed使用X路的模型并行结果；模型的计量单位是B表示十亿Billon。）*

PyTorch 仅适用于 1B 模型大小的情况，派大星在8个GPU上比PyTorch快1.37倍，在 1、2、4 个 GPU 情况下与 PyTorch 相似。

使用相同的零冗余优化器来实现数据并行，派大星在大多数情况下（14 个中有 12 个）优于 DeepSpeed-DP，并且数据并行方式训练8B和12B之间模型大小的唯一解决方案。

不难看出，尤其是针对小模型，改进是非常明显了（0.90x-1.49x）。

而在增加模型大小时，派大星不会显着降低计算效率。

此外，派大星在增加 GPU 数量时显示出超线性可扩展性。

若是将派大星与模型并行解决方案进行了比较，又会是怎样的结果？

例如在上图中，还比较了DeepSpeed在8个GPU卡上使用Zero-DP方案叠加2路模型并行和4路模型并行的性能。

派大星在所有测试用例上实现了最大的模型规模120亿参数，以及最佳的性能效率。

在模型并行的帮助下，DeepSpeed将模型规模扩展到了80亿参数。

但是，MP引入了更多的通信开销；性能明显低于派大星和 DeepSpeed-DP。

……

效果是有够惊艳的了，但接下来的一个问题便是：

关键技术是什么？

破局者：异构训练

或许你会说了，让数据并行不就完事了吗？

事实却并非如此。

对于预训练模型来说，最常用的数据并行技术不适用，这是因为模型数据无法再容纳在单个 GPU 的内存中。

GPU硬件的存储规模上限，像一堵墙一样限制住了PTM的可训练规模，因此从业人员通常称之为”GPU内存墙”现象。

近两年来，通过利用并行训练在多个 GPU 内存之间分配模型数据，例ZeRO-DP、模型并行、流水线并行尝试使 PTM 大小突破内存墙。

但是，使用这些技术又需要不断扩大GPU规模，也意味着更高设备的投入，那么此局怎么破？

异构训练技术，了解一下。

它不仅可以显著提升单GPU训练模型的规模，而且可以和并行训练技术正交使用。

异构训练通过在CPU和GPU中，容纳模型数据并仅在必要时将数据移动到当前设备来利用 GPU 内存、CPU 内存（由 DRAM 或 NVMe 内存组成）。

其他方案如数据并行、模型并行、流水线并行，都在异构训练基础上进一步扩展GPU规模。

预训练模型在训练期间，存在必须管理的两种类型训练数据：

• 模型数据由参数、梯度和优化器状态组成，其规模与模型结构定义相关；
• 非模型数据主要由算子生成的中间张量组成，根据训练任务的配置动态变化，例如批量大小。

模型数据和非模型数据相互竞争GPU内存。

然而，目前最佳的异构训练方案DeepSpeed的Zero-Offload/Infinity，仍存在很大优化空间。

在不考虑非模型数据的情况下，DeepSpeed在CPU和GPU内存之间静态划分模型数据，并且它们的内存布局对于不同的训练配置是恒定的。

这种静态分区策略会导致几个问题。

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

其次，当数据以张量为粒度的不同内存空间之间传输时通信效率低下，并且当你可以预先将模型数据放置在目标计算设备上时，一些CPU-GPU通信量是不必要的。

因此DeepSpeed在微信的数据中心单GPU只能运行60亿参数的模型，而且效率十分低下，远不如在DGX上的报告结果130亿参数。

派大星则通过以细粒度的方式管理模型数据，以更有效地使用异构内存来克服这些缺点。

它将模型数据张量组织成块，即相同大小的连续内存块。

块在异构内存空间中的分布在训练期间根据它们的张量状态动态编排。

通过重用不共存的块，派大星还比DeepSpeed的方案进一步降低了模型数据的内存占用。

派大星使用预热迭代来收集运行时模型数据可用 GPU 内存的统计数据。

基于收集到的统计数据的有效块驱逐策略和设备感知算子放置策略，为的就是减少 CPU-GPU 数据移动量。

最后，使用零冗余优化器（ZeroReduencyOptimizer）的Zero-DP数据并行方法，通过块的集合GPU 通信来使用扩展到多个GPU。

团队介绍

这项研究主要由腾讯微信AI团队和新加坡国立大学团队共同完成。

论文一作是来自微信AI的高级工程师Jiarui Fang，清华大学博士毕业。

其主要工作是通过创新并行计算技术提升在线和离线NLP任务的运算效率。

他曾经还曾开源过一款Tranformer模型推理加速工具TurboTransformer。

……

那么最后，你是否也想训练一个专属的大模型呢？戳下方链接试试吧~

派大星开源地址：
https://github.com/Tencent/PatrickStar

论文地址：
https://arxiv.org/abs/2108.05818