微信開源「派大星」：4000元游戲電腦能帶動(dòng)7億參數(shù)GPT！

金磊 發(fā)自 凹非寺

量子位 報(bào)道 | 公眾號(hào) QbitAI

一聽到訓(xùn)練大模型，是不是第一感覺就是貴、燒錢、玩不起？

但我說，一臺(tái)4000多塊錢的游戲電腦，誰都能訓(xùn)練上億參數(shù)的大模型呢？

別不信，這是真的。

而這就歸功于微信AI團(tuán)隊(duì)，最近推出的一款利器——派大星。

但不是你印象中的那個(gè)派大星昂~

微信AI的派大星（PatricStar），其實(shí)是一個(gè)超大預(yù)訓(xùn)練模型訓(xùn)練系統(tǒng)。

要知道，在這個(gè)領(lǐng)域中，以往都是微軟DeepSeed獨(dú)占鰲頭。

此次微信一出手，可以說是直接秒殺了微軟：

在 8xV100 和 240GB CPU 內(nèi)存節(jié)點(diǎn)上，訓(xùn)練了一個(gè)120 億參數(shù)的 GPT 模型，是當(dāng)前最佳方案DeepSpeed模型規(guī)模上限的1.5 倍。

但畢竟針對(duì)的是大模型，“燒錢”是出了名的難題。

而微信AI的派大星就顯得相當(dāng)?shù)?strong>親民了。

即使在700美元的個(gè)人游戲電腦上，它也可以訓(xùn)練一個(gè)7億參數(shù)的 GPT 模型！

現(xiàn)在，人人都可以在家訓(xùn)練大模型了！

劃重點(diǎn)：已開源！

為什么要搞派大星？

大規(guī)模預(yù)訓(xùn)練模型，已然成為技術(shù)發(fā)展中的新潮流。

以BERT、GPT為代表的預(yù)訓(xùn)練模型的出現(xiàn)，可以說是自然語言處理（NLP）領(lǐng)域的里程碑事件。

NLP，正在進(jìn)入了預(yù)訓(xùn)練時(shí)代。

那么像派大星這樣的訓(xùn)練系統(tǒng)，真的有必要嗎？

答案是肯定的。

從技術(shù)角度來看，預(yù)訓(xùn)練模型（PTM）通常使用一個(gè)堆疊了多個(gè)Transformer結(jié)構(gòu)神經(jīng)網(wǎng)絡(luò)，在大量文本上預(yù)訓(xùn)練通用語言特征表示。

然后，通過微調(diào)將學(xué)到的知識(shí)轉(zhuǎn)移到不同的下游任務(wù)。

預(yù)訓(xùn)練模型使用大量來自互聯(lián)網(wǎng)的文本數(shù)據(jù)，可以捕獲自然語言的細(xì)微特征，并在下游任務(wù)上獲得非常驚艷的表現(xiàn)效果。

于是，AI社區(qū)的共識(shí)是采用預(yù)訓(xùn)練模型，作為特定NLP任務(wù)的主干，而不是在與任務(wù)相關(guān)的數(shù)據(jù)集上從頭開始訓(xùn)練模型。

預(yù)訓(xùn)練模型的力量源泉，是它擁有的數(shù)以億計(jì)的參數(shù)規(guī)模，這對(duì)運(yùn)行它的計(jì)算和內(nèi)存資源都提出了巨大的要求。

因此，預(yù)訓(xùn)練模型訓(xùn)練仍是一小部分人的游戲。

所有發(fā)表百億級(jí)模型訓(xùn)練成果的團(tuán)隊(duì)，所采用的的設(shè)備都是如DGX型號(hào)的AI超級(jí)計(jì)算機(jī)。

它的一個(gè)節(jié)點(diǎn)就配置了8張GPU，1.5TB內(nèi)存，3.84TB SSDs，還使用NVLink作為高速通信網(wǎng)絡(luò)。

目前最大的預(yù)訓(xùn)練模型Megatron-Turing，包含5300億參數(shù)，其預(yù)訓(xùn)練過程就是在560個(gè)DGX A100節(jié)點(diǎn)的集群上完成的。

這種配置在大多數(shù)工業(yè)界數(shù)據(jù)中心都是遙不可及的。

而通過像派大星這樣的訓(xùn)練系統(tǒng)，便可以讓這種“遙不可及”變得“唾手可得”，讓大模型可以普惠到更多的開發(fā)人員，實(shí)現(xiàn)PTM的“共同富?！?。

再?gòu)?strong>綠色AI角度來看，預(yù)訓(xùn)練模型的預(yù)訓(xùn)練的過程是極其燒錢和有害環(huán)境的。

比如，從頭訓(xùn)練型一次萬億級(jí)別的預(yù)訓(xùn)練模型要燒掉154萬人民幣，耗電所產(chǎn)生的碳排放相當(dāng)于數(shù)十輛小汽車從出廠到報(bào)廢的碳排放總和。

出于社會(huì)利益最大化考慮，預(yù)訓(xùn)練技術(shù)未來的產(chǎn)業(yè)形態(tài)，應(yīng)該是中心化的：

少部分財(cái)力雄厚的機(jī)構(gòu)，用超大規(guī)模集群承擔(dān)預(yù)訓(xùn)練階段的計(jì)算和環(huán)境開銷；大多數(shù)從業(yè)人員在小規(guī)模、相對(duì)簡(jiǎn)陋的硬件上針對(duì)自身業(yè)務(wù)進(jìn)行微調(diào)。

前者只需要相對(duì)少量的計(jì)算和碳排放，而后者的訴求卻被當(dāng)前的預(yù)訓(xùn)練軟件所嚴(yán)重忽略。

現(xiàn)如今，派大星的到來，讓大規(guī)模預(yù)訓(xùn)練模型的訓(xùn)練變得“多快好省”了起來。

而且不僅是對(duì)于機(jī)構(gòu)，更是有益于個(gè)人開發(fā)者。

……

那么派大星的效果，具體又怎樣呢？

不是魔改，是從頭搭建，性能達(dá)SOTA

值得一提的是，派大星并不是基于DeepSpeed的魔改，代碼是團(tuán)隊(duì)從頭開始搭建起來的。

派大星框架非常直觀的一個(gè)特點(diǎn)，便是簡(jiǎn)單易用，而且還是可以兼容其他并行方案的那種。

例如，開發(fā)者可以使用幾行代碼端到端的加速PyTorch的訓(xùn)練過程。

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系統(tǒng)在 1、2、4、8 個(gè) GPU 上的性能（y軸通過對(duì)數(shù)方式重新縮放）。

這些點(diǎn)代表在一個(gè) GPU 上使用 4、8、16、32 和 64 批大小測(cè)試的最佳結(jié)果。

（注：圓點(diǎn)周圍的值表示派大星在吞吐量及其對(duì)DeepSpeed的加速；deeps是DeepSpeed僅使用數(shù)據(jù)并行的效果，我們接下來稱之為DeepSpeed-DP，deeps-mpX 是 DeepSpeed使用X路的模型并行結(jié)果；模型的計(jì)量單位是B表示十億Billon。）*

PyTorch 僅適用于 1B 模型大小的情況，派大星在8個(gè)GPU上比PyTorch快1.37倍，在 1、2、4 個(gè) GPU 情況下與 PyTorch 相似。

使用相同的零冗余優(yōu)化器來實(shí)現(xiàn)數(shù)據(jù)并行，派大星在大多數(shù)情況下（14 個(gè)中有 12 個(gè)）優(yōu)于 DeepSpeed-DP，并且數(shù)據(jù)并行方式訓(xùn)練8B和12B之間模型大小的唯一解決方案。

不難看出，尤其是針對(duì)小模型，改進(jìn)是非常明顯了（0.90x-1.49x）。

而在增加模型大小時(shí)，派大星不會(huì)顯著降低計(jì)算效率。

此外，派大星在增加 GPU 數(shù)量時(shí)顯示出超線性可擴(kuò)展性。

若是將派大星與模型并行解決方案進(jìn)行了比較，又會(huì)是怎樣的結(jié)果？

例如在上圖中，還比較了DeepSpeed在8個(gè)GPU卡上使用Zero-DP方案疊加2路模型并行和4路模型并行的性能。

派大星在所有測(cè)試用例上實(shí)現(xiàn)了最大的模型規(guī)模120億參數(shù)，以及最佳的性能效率。

在模型并行的幫助下，DeepSpeed將模型規(guī)模擴(kuò)展到了80億參數(shù)。

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

……

效果是有夠驚艷的了，但接下來的一個(gè)問題便是：

關(guān)鍵技術(shù)是什么？

破局者：異構(gòu)訓(xùn)練

或許你會(huì)說了，讓數(shù)據(jù)并行不就完事了嗎？

事實(shí)卻并非如此。

對(duì)于預(yù)訓(xùn)練模型來說，最常用的數(shù)據(jù)并行技術(shù)不適用，這是因?yàn)槟Ｐ蛿?shù)據(jù)無法再容納在單個(gè) GPU 的內(nèi)存中。

GPU硬件的存儲(chǔ)規(guī)模上限，像一堵墻一樣限制住了PTM的可訓(xùn)練規(guī)模，因此從業(yè)人員通常稱之為”GPU內(nèi)存墻”現(xiàn)象。

近兩年來，通過利用并行訓(xùn)練在多個(gè) GPU 內(nèi)存之間分配模型數(shù)據(jù)，例ZeRO-DP、模型并行、流水線并行嘗試使 PTM 大小突破內(nèi)存墻。

但是，使用這些技術(shù)又需要不斷擴(kuò)大GPU規(guī)模，也意味著更高設(shè)備的投入，那么此局怎么破？

異構(gòu)訓(xùn)練技術(shù)，了解一下。

它不僅可以顯著提升單GPU訓(xùn)練模型的規(guī)模，而且可以和并行訓(xùn)練技術(shù)正交使用。

異構(gòu)訓(xùn)練通過在CPU和GPU中，容納模型數(shù)據(jù)并僅在必要時(shí)將數(shù)據(jù)移動(dòng)到當(dāng)前設(shè)備來利用 GPU 內(nèi)存、CPU 內(nèi)存（由 DRAM 或 NVMe 內(nèi)存組成）。

其他方案如數(shù)據(jù)并行、模型并行、流水線并行，都在異構(gòu)訓(xùn)練基礎(chǔ)上進(jìn)一步擴(kuò)展GPU規(guī)模。

預(yù)訓(xùn)練模型在訓(xùn)練期間，存在必須管理的兩種類型訓(xùn)練數(shù)據(jù)：

• 模型數(shù)據(jù)由參數(shù)、梯度和優(yōu)化器狀態(tài)組成，其規(guī)模與模型結(jié)構(gòu)定義相關(guān)；
• 非模型數(shù)據(jù)主要由算子生成的中間張量組成，根據(jù)訓(xùn)練任務(wù)的配置動(dòng)態(tài)變化，例如批量大小。

模型數(shù)據(jù)和非模型數(shù)據(jù)相互競(jìng)爭(zhēng)GPU內(nèi)存。

然而，目前最佳的異構(gòu)訓(xùn)練方案DeepSpeed的Zero-Offload/Infinity，仍存在很大優(yōu)化空間。

在不考慮非模型數(shù)據(jù)的情況下，DeepSpeed在CPU和GPU內(nèi)存之間靜態(tài)劃分模型數(shù)據(jù)，并且它們的內(nèi)存布局對(duì)于不同的訓(xùn)練配置是恒定的。

這種靜態(tài)分區(qū)策略會(huì)導(dǎo)致幾個(gè)問題。

首先，當(dāng)GPU內(nèi)存或CPU內(nèi)存不足以滿足其相應(yīng)的模型數(shù)據(jù)要求時(shí)，即使當(dāng)時(shí)其他設(shè)備上仍有可用內(nèi)存，系統(tǒng)也會(huì)崩潰。

其次，當(dāng)數(shù)據(jù)以張量為粒度的不同內(nèi)存空間之間傳輸時(shí)通信效率低下，并且當(dāng)你可以預(yù)先將模型數(shù)據(jù)放置在目標(biāo)計(jì)算設(shè)備上時(shí)，一些CPU-GPU通信量是不必要的。

因此DeepSpeed在微信的數(shù)據(jù)中心單GPU只能運(yùn)行60億參數(shù)的模型，而且效率十分低下，遠(yuǎn)不如在DGX上的報(bào)告結(jié)果130億參數(shù)。

派大星則通過以細(xì)粒度的方式管理模型數(shù)據(jù)，以更有效地使用異構(gòu)內(nèi)存來克服這些缺點(diǎn)。

它將模型數(shù)據(jù)張量組織成塊，即相同大小的連續(xù)內(nèi)存塊。

塊在異構(gòu)內(nèi)存空間中的分布在訓(xùn)練期間根據(jù)它們的張量狀態(tài)動(dòng)態(tài)編排。

通過重用不共存的塊，派大星還比DeepSpeed的方案進(jìn)一步降低了模型數(shù)據(jù)的內(nèi)存占用。

派大星使用預(yù)熱迭代來收集運(yùn)行時(shí)模型數(shù)據(jù)可用 GPU 內(nèi)存的統(tǒng)計(jì)數(shù)據(jù)。

基于收集到的統(tǒng)計(jì)數(shù)據(jù)的有效塊驅(qū)逐策略和設(shè)備感知算子放置策略，為的就是減少 CPU-GPU 數(shù)據(jù)移動(dòng)量。

最后，使用零冗余優(yōu)化器（ZeroReduencyOptimizer）的Zero-DP數(shù)據(jù)并行方法，通過塊的集合GPU 通信來使用擴(kuò)展到多個(gè)GPU。

團(tuán)隊(duì)介紹

這項(xiàng)研究主要由騰訊微信AI團(tuán)隊(duì)和新加坡國(guó)立大學(xué)團(tuán)隊(duì)共同完成。

論文一作是來自微信AI的高級(jí)工程師Jiarui Fang，清華大學(xué)博士畢業(yè)。

其主要工作是通過創(chuàng)新并行計(jì)算技術(shù)提升在線和離線NLP任務(wù)的運(yùn)算效率。

他曾經(jīng)還曾開源過一款Tranformer模型推理加速工具TurboTransformer。

……

那么最后，你是否也想訓(xùn)練一個(gè)專屬的大模型呢？戳下方鏈接試試吧~

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

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