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

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

別不信，這是真的。

而這就歸功于微信AI團隊，最近推出的一款利器——派大星。

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

微信AI的派大星（PatricStar），其實是一個超大預訓練模型訓練系統(tǒng)。

要知道，在這個領域中，以往都是微軟DeepSeed獨占鰲頭。

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

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

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

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

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

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

劃重點：已開源！

為什么要搞派大星？

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

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

NLP，正在進入了預訓練時代。

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

答案是肯定的。

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

然后，通過微調將學到的知識轉移到不同的下游任務。

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

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

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

因此，預訓練模型訓練仍是一小部分人的游戲。

所有發(fā)表百億級模型訓練成果的團隊，所采用的的設備都是如DGX型號的AI超級計算機。

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

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

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

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

再從綠色AI角度來看，預訓練模型的預訓練的過程是極其燒錢和有害環(huán)境的。

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

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

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

前者只需要相對少量的計算和碳排放，而后者的訴求卻被當前的預訓練軟件所嚴重忽略。

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

而且不僅是對于機構，更是有益于個人開發(fā)者。

……

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

不是魔改，是從頭搭建，性能達SOTA

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

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

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

frompatrickstar.runtimeimportinitialize_engineconfig={"optimizer":{"type":"Adam","params":{"lr":0.001,"betas":(0.9,0.999),"eps":1e-6,"weight_decay":0,"use_hybrid_adam":True,},},"fp16":{#lossscalerparams"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,}defmodel_func():#MyModelisaderivedclassfortorch.nn.ModulereturnMyModel(...)model,optimizer=initialize_engine(model_func=model_func,local_rank=0,config=config)...fordataindataloader:optimizer.zero_grad()loss=model(data)model.backward(loss)optimizer.step()

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

上圖便展示了DeepSpeed stage3，PyTorch系統(tǒng)在 1、2、4、8 個 GPU 上的性能（y軸通過對數(shù)方式重新縮放）。

這些點代表在一個 GPU 上使用 4、8、16、32 和 64 批大小測試的最佳結果。

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

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

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

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

而在增加模型大小時，派大星不會顯著降低計算效率。

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

若是將派大星與模型并行解決方案進行了比較，又會是怎樣的結果？

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

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

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

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

……

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

關鍵技術是什么？

破局者：異構訓練

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

事實卻并非如此。

對于預訓練模型來說，最常用的數(shù)據(jù)并行技術不適用，這是因為模型數(shù)據(jù)無法再容納在單個 GPU 的內存中。

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

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

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

異構訓練技術，了解一下。

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

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

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

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

模型數(shù)據(jù)由參數(shù)、梯度和優(yōu)化器狀態(tài)組成，其規(guī)模與模型結構定義相關；

非模型數(shù)據(jù)主要由算子生成的中間張量組成，根據(jù)訓練任務的配置動態(tài)變化，例如批量大小。

模型數(shù)據(jù)和非模型數(shù)據(jù)相互競爭GPU內存。

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

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

這種靜態(tài)分區(qū)策略會導致幾個問題。

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

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

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

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

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

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

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

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

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

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

團隊介紹

這項研究主要由騰訊微信AI團隊和新加坡國立大學團隊共同完成。

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

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

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

……

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

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

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

本文來自微信公眾號“量子位”（ID:QbitAI），作者：金磊，36氪經(jīng)授權發(fā)布。

關鍵詞： 能帶 開源 參數(shù)