過去數(shù)年，AI 模型的參數(shù)發(fā)生了極大變化。尤洋指出，從 2016 年至 2021 年 1 月，AI 模型的參數(shù)量是每 18 個(gè)月增長(zhǎng) 40 倍；從 2018 年 1 月到 2021 年 1 月，AI 大語言模型的參數(shù)量每 18 個(gè)月增長(zhǎng) 340 倍。而相形之下，2016 年 1 月至 2021 年 1 月間，GPU 的內(nèi)存增長(zhǎng)每 18 個(gè)月僅有 1.7 倍。

由此可見，訓(xùn)練成本高、周期長(zhǎng)，是當(dāng)前大模型發(fā)展最需要克服的難題。

針對(duì)這一問題，尤洋提出了 Colossal-AI 系統(tǒng)，從高效內(nèi)存系統(tǒng)、N 維并行系統(tǒng)和大規(guī)模優(yōu)化三個(gè)層次出發(fā)，以實(shí)現(xiàn)同樣的設(shè)備條件下將數(shù)據(jù)移動(dòng)的最小化，將 GPU 的吞吐量擴(kuò)大至最高點(diǎn)。

尤洋還指出，現(xiàn)階段的模型參數(shù)量以 10 萬倍擴(kuò)大、但層數(shù)增加不多，這或意味著：如今的 AI 發(fā)展可能不再是深度學(xué)習(xí)、而是進(jìn)入了寬度學(xué)習(xí)時(shí)代。在模型變得更寬的情況下，面對(duì)大規(guī)模、長(zhǎng)時(shí)間的 GPU 訓(xùn)練任務(wù)，大模型訓(xùn)練系統(tǒng)的核心將是如何實(shí)現(xiàn) GPU 并行計(jì)算，以實(shí)現(xiàn)大模型訓(xùn)練越快越省錢的目標(biāo)。

以下為尤洋的現(xiàn)場(chǎng)演講內(nèi)容，雷峰網(wǎng)作了不改變?cè)獾木庉嫾罢恚?/p>

AI 大模型的機(jī)遇和挑戰(zhàn)

首先展示一張圖片。圖上的橫坐標(biāo)是時(shí)間，縱坐標(biāo)是 AI 模型的參數(shù)量。

從 2016 年至 2021 年 1 月，AI 大模型的參數(shù)量大概每 18 個(gè)月增長(zhǎng) 40 倍；從 2018 年 1 月到 2021 年 1 月，AI 模型的參數(shù)量每 18 個(gè)月增長(zhǎng) 340 倍。

2016 年，當(dāng)時(shí)世界上最好的模型是 ResNet-50，而今天最好的模型是 GPT-4。從架構(gòu)上來看，雖然 OpenAI 沒有對(duì)外公布 GPT-4 的架構(gòu)，但對(duì)比 ResNet-50 的 50 層神經(jīng)網(wǎng)絡(luò)和 GPT-3 未達(dá) 100 層的架構(gòu)，可以說 AI 模型近年來的層數(shù)并沒有產(chǎn)生的太大的變化。

從 ResNet-50 到 GPT-4，雖然參數(shù)量大了 10 萬倍左右，其實(shí)是每一層都變得更加寬了。包括 LLaMA-65B 版本，也是幾十層的網(wǎng)絡(luò)。

所以我們可能不是深度學(xué)習(xí)，而是進(jìn)入了一個(gè)寬度學(xué)習(xí)時(shí)代。

可以看到，自 2019 年開始， Transformer 架構(gòu)基本統(tǒng)一了 AI 大模型賽道，當(dāng)前最高效的 AI 大模型均是 Transformer 架構(gòu)。上圖中的兩條虛線，既展示了大模型參數(shù)的變化趨勢(shì)，實(shí)際上也展現(xiàn)了 GPU 的變化趨勢(shì)。

雖然現(xiàn)在英偉達(dá)的股價(jià)漲了很多倍，但包括英偉達(dá)在內(nèi)的廠商，其 GPU 內(nèi)存的增長(zhǎng)速度遠(yuǎn)跟不上大模型的發(fā)展速度。

相較于過去六年模型參數(shù)量的增長(zhǎng)速度變化，2016 年 1 月至 2021 年 1 月間，英偉達(dá) GPU 的計(jì)算增長(zhǎng)速度每 18 個(gè)月僅增長(zhǎng)了 1.7 倍。

以 A100 80G 為例計(jì)算 GPT-3 訓(xùn)練所需的內(nèi)存數(shù)，GPT-3 有大概 1750 億參數(shù)，為方便計(jì)算取整數(shù) 2000 億，等于 200 乘以 10 的 9 次方，每個(gè)單精度占用 4 個(gè)字節(jié)，僅參數(shù)就要占 800G 內(nèi)存，梯度也占了 800G 內(nèi)存。按照當(dāng)前的優(yōu)化方法儲(chǔ)存一階矩（first moment）、二階矩（second moment）等信息均為 800G。也就是說，如果訓(xùn)練一個(gè)什么事情都不干大模型，至少需要幾 T 的內(nèi)存，單個(gè) A100 GPU 僅 80G 內(nèi)存遠(yuǎn)遠(yuǎn)不夠，加上中間結(jié)果的 batch size 越大，內(nèi)存開銷也越大。

這也是為什么從內(nèi)存角度上看，訓(xùn)練大模型首先需要有成千上萬個(gè) GPU 的原因。

舉一個(gè)可量化的案例。PaLM 是一個(gè) 5400 億的大模型，根據(jù)當(dāng)前云計(jì)算市場(chǎng)價(jià)，訓(xùn)練 PaLM 需要承包至少 1000 個(gè) GPU，成本約 900 多萬美金。而 Meta 此前曾明確提到，LLaMA 需要使用到 2000 個(gè) A100 GPU，并且用三周時(shí)間才可完成一次訓(xùn)練，由此計(jì)算可得出 LLaMA 單次訓(xùn)練的成本在 500 萬美金。

但由于煉大模型并不僅限于一次訓(xùn)練，可能一個(gè)好的大模型產(chǎn)品迭代至少需要五六次，前期都是在試錯(cuò)。因此，據(jù)公開渠道分析， GPT-4 單次訓(xùn)練成本在 6000 萬美金左右，且訓(xùn)練一次需要至少幾個(gè)月時(shí)間。這也是為什么目前 ChatGPT 即便將其升級(jí)至最新版本，其底層還是 2021 年 9 月版本的模型。也就是說，從 2021 年 9 月至今，OpenAI 實(shí)質(zhì)上并沒有升級(jí)其產(chǎn)品，根本原因就在于，模型的每次訓(xùn)練不僅成本很高，訓(xùn)練周期也很長(zhǎng)，因此大模型訓(xùn)練的代價(jià)高就非常嚴(yán)重。

設(shè)想一個(gè)問題，今天有一個(gè) 1 萬億參數(shù)以及 1000 億參數(shù)的大模型，我們能否用什么方法，去檢測(cè)萬億參數(shù)大模型比千億參數(shù)大模型二者哪個(gè)效果更好？也即是目前常說的，模型參數(shù)量增大、到底要增大到什么程度？

到今天而言，我覺得這個(gè)問題暫時(shí)無法用科學(xué)回答。原因有幾個(gè)。

首先，訓(xùn)練神經(jīng)網(wǎng)絡(luò)存在非凸優(yōu)化的問題，目前訓(xùn)練所收斂的點(diǎn)多為局部最優(yōu)解、而非全局最優(yōu)解。因此，我們要驗(yàn)證神經(jīng)網(wǎng)絡(luò)訓(xùn)練到什么程度，在現(xiàn)有計(jì)算資源情況下是無法驗(yàn)證的。

第二個(gè)難度在于，大模型訓(xùn)練往往只訓(xùn)練一兩個(gè) epoch，而此前的 CNN 模型中，ResNet 訓(xùn)練有 90 個(gè) epoch，甚至自監(jiān)督學(xué)習(xí)訓(xùn)練有 1000 個(gè) epoch，因此大模型只訓(xùn)練一兩個(gè) epoch 的方式，相當(dāng)于只將數(shù)據(jù)集過了一兩遍，收斂就更加不充分了。因此在訓(xùn)練成本如此之高的情況下，我們很難驗(yàn)證，一個(gè) 1 萬億參數(shù)的模型和 2 萬億參數(shù)的模型二者誰更好，因?yàn)樗鼈儩摿Χ紱]有能通過實(shí)驗(yàn)得到充分發(fā)揮。因此我認(rèn)為，今天 AI 大模型是一個(gè)實(shí)驗(yàn)性學(xué)科，如何能高效提升這個(gè)實(shí)驗(yàn)的效率，降低成本，對(duì)整個(gè)行業(yè)的普及具有根本性的作用。

回到現(xiàn)實(shí)之中，為什么今天人人都在追求大模型？從數(shù)學(xué)邏輯上看，模型參數(shù)越大、效果越好，這是絕對(duì)的。

與此同時(shí)，成本也再不斷攀高。目前訓(xùn)練大模型需要成百上千、甚至上萬個(gè) GPU，如何將上萬個(gè) GPU 的成本進(jìn)一步降低，挑戰(zhàn)非常大的。

在 20 年前，由于當(dāng)時(shí)依靠的是主頻的，所有的程序都是串行的，假設(shè)將硬件速度提高 10 倍，在一行代碼都不用的更改的情況下，其運(yùn)行速度也可以提升 10 倍。但到了如今，如果想將代碼速度提升 10 倍，假定硬件已經(jīng)增速 10 倍，但如果不優(yōu)化代碼，很可能速度反而會(huì)變慢。原因就在于，機(jī)器規(guī)模更大的情況下，比如 GPU 內(nèi)部，GPU 內(nèi)存和 CPU 之間的數(shù)據(jù)移動(dòng)，或是 GPU 間的數(shù)據(jù)移動(dòng)，再加上服務(wù)器實(shí)現(xiàn)的數(shù)據(jù)移動(dòng)，會(huì)占據(jù)整個(gè)系統(tǒng)的大部分時(shí)間，把大部分時(shí)間都花在了數(shù)據(jù)移動(dòng)上，模型的擴(kuò)展性也會(huì)變得不好。

我認(rèn)為，未來一個(gè)好的分布式軟件和一個(gè)差的分布式軟件，在上千個(gè) GPU 上，甚至 500 個(gè) GPU 上，其速度可能相差 10 倍。

Colossal-AI 如何運(yùn)行？

基于上述的挑戰(zhàn)，我們提出了大模型訓(xùn)練系統(tǒng) Colossal-AI，提供優(yōu)化方法，降低數(shù)據(jù)移動(dòng)的代價(jià)，將模型擴(kuò)展性效率提到最高。

一個(gè)具體的數(shù)據(jù)是，使用最簡(jiǎn)單的 PyTorch 訓(xùn)練 GPT-3，成本為 1000 萬美金，英偉達(dá)經(jīng)過優(yōu)化后，用 Megatron 可將其成本減少至 300 萬美金，而使用 Colossal-AI 后，成本可以降低到 130 萬美金?？梢钥吹?，相同的設(shè)備條件下，數(shù)據(jù)移動(dòng)的最小化將數(shù)據(jù)移動(dòng)占比降低最低，能夠把 GPU 吞吐量擴(kuò)大至最高點(diǎn)。

針對(duì)上述問題，Colossal-AI 提出了三個(gè)層次。其他類似的軟件也包括了這三個(gè)層次。

第一層是優(yōu)化內(nèi)存，先確保單個(gè) GPU、單個(gè)服務(wù)器的內(nèi)存效率最高，這是基礎(chǔ)。

第二層是 N 維的并行。當(dāng)前我們使用上千、上萬個(gè) GPU 時(shí)，其核心技術(shù)就是 ParallelComputing（并行計(jì)算）。從 1 個(gè) GPU 擴(kuò)到 10 個(gè) GPU，因?yàn)槠湟?guī)模比較小，我們可以輕易獲得 7 倍加速；從 10 個(gè) 到 100 個(gè) GPU 時(shí)，往往可能只獲得 4 倍加速，因?yàn)椴⑿幸?guī)模變大，它的通信代價(jià)變高了。而從 100 個(gè) GPU 到 1000 個(gè) GPU，因?yàn)橥ㄐ糯鷥r(jià)進(jìn)一步加高，很可能只獲得 2 倍的加速。而從 1000 個(gè) GPU 到 1 萬個(gè) GPU，如果軟件運(yùn)行情況不佳時(shí)，不僅可能無法加速，甚至還會(huì)更慢，因?yàn)樵O(shè)備將所有時(shí)間花耗在了更高密度的通信上。

其次是優(yōu)化問題，未來 AI 大模型的發(fā)展方向我認(rèn)為有兩層，第一層是模型變得更加智能，設(shè)計(jì)出更好的結(jié)構(gòu)，比如說從 BERT 到 GPT，或者從 ResNet 到 BERT等，都是在不斷地嘗試改變模型結(jié)構(gòu)。

此外還有優(yōu)化方法的進(jìn)步，從 SGD 過渡到 MOMENTUM、ADAGRAD，到現(xiàn)在有 ADAM，未來又會(huì)有哪些更好的優(yōu)化方法能夠?qū)⑿侍嵘?10 倍，這一點(diǎn)也非常重要。

具體到實(shí)際操作訓(xùn)練大模型的并行問題。

首先是數(shù)據(jù)并行，這是最簡(jiǎn)單、也是最高效的并行方法。數(shù)據(jù)并行指的是，假設(shè)現(xiàn)有 1 萬張圖片，每次循環(huán)處理 1000 張圖片，如果有 10 個(gè)機(jī)器，每個(gè)機(jī)器分配 100 張，10 個(gè)循環(huán)即可完成所有圖片的處理。

在數(shù)據(jù)并行的過程中需要進(jìn)行匯總，每個(gè)機(jī)器用不同的數(shù)據(jù)獲得不同梯度，機(jī)器在不同數(shù)據(jù)上學(xué)習(xí)不同的更改，并更新參數(shù)梯度，最后算出全局梯度，目前采用的是加和求平均的方式，效果已經(jīng)非常好了。此前 Colossal-AI 在數(shù)據(jù)并行中的 LARS 方法，就為谷歌、 Meta、騰訊、索尼等公司，將 ImageNet 的訓(xùn)練時(shí)間從一小時(shí)縮短至一分鐘。

數(shù)據(jù)并行是最基本的，同時(shí)也是最穩(wěn)定的。將數(shù)據(jù)劃分之后，假設(shè)未來有 1 萬個(gè) GPU，很容易發(fā)生的情況是，隔幾個(gè)小時(shí)就有四五個(gè) GPU 崩潰了，運(yùn)維 1 萬個(gè) GPU 的集群很難，但數(shù)據(jù)并行的穩(wěn)定之處在于，即便有 1 萬個(gè)GPU 崩潰了十幾個(gè)，但大體結(jié)果是不會(huì)變的，因?yàn)樗翘荻燃雍颓笃骄?/p>

基于這個(gè)考慮，我認(rèn)為數(shù)據(jù)并行是一個(gè)根本性的基礎(chǔ)設(shè)施。

當(dāng)然，僅用數(shù)據(jù)并行并不夠，原因在于：數(shù)據(jù)并行有一個(gè)假設(shè)，必須將模型拷貝到每個(gè) GPU 或服務(wù)器內(nèi)，由服務(wù)器或 GPU 去交換梯度。但如果 GPU 僅 80G 內(nèi)存時(shí)，萬億參數(shù)的模型則需要幾十 T 的內(nèi)存，這在 GPU 中是無法存放的，需要將模型切割至不同的 GPU 上再匯總結(jié)果。這種方法叫做模型并行。模型并行包括兩種，第一種是張量并行（ tensor paralism），即層內(nèi)的模型并行。例如 GPT-3 的層數(shù)大概為八九十層，每層切割一次模型，將其層內(nèi)計(jì)算分隔成多份，算完一層再算下一層，依次類推，這就是張量并行。

另一種方式則是 Pipeline Parallelism（流水線并行），介于數(shù)據(jù)并行和張量并行外的一種模型并行方式。通過構(gòu)建幾個(gè)數(shù)據(jù) pipe（管道），每個(gè)數(shù)據(jù) pipe 的數(shù)據(jù)點(diǎn)不同，相當(dāng)于將一個(gè)大尺寸分割為多個(gè)小尺寸，通過這種方式進(jìn)行 pipe 計(jì)算。假如有 10 個(gè) pipe，10 個(gè) pipe 代表十組不同的數(shù)據(jù)，第一個(gè) pipe 計(jì)算第一層的數(shù)據(jù)，第二個(gè) pipe 計(jì)算第二層......以此方式并行，類似我們蓋樓一樣，10 個(gè)工程隊(duì)蓋 1000 層樓，當(dāng)?shù)谝粋€(gè)工程隊(duì)在蓋第一棟樓的第一層，第二個(gè)工程隊(duì)蓋第二棟樓的第二層，依此類推。

當(dāng)樓數(shù)越多，樓和工程隊(duì)之間的比值越高，效率也越高，相當(dāng)于 10 個(gè)工程隊(duì)在同時(shí)運(yùn)轉(zhuǎn)。其中每個(gè)工程隊(duì)就相當(dāng)于一個(gè) GPU，每個(gè)樓就相當(dāng)于一個(gè) pipe，樓的層數(shù)相當(dāng)于這個(gè)神經(jīng)網(wǎng)絡(luò)的層數(shù)，這就流水線并行的核心邏輯。

目前工業(yè)界已經(jīng)做了相關(guān)的工作，除了 Colossal-AI 之外，還有英偉達(dá)的 TensorRT 和微軟的 DeepSpeed，他們也是技術(shù)壁壘最高的兩家公司。

但 Colossal-AI 與其不同之處是，Colossal-AI 專注于未來大模型的發(fā)展方向?？梢钥吹剑?dāng)前的模型還在變得更寬，而不是變得更深，張量并行將會(huì)更加重要，但它最大的弊端就在于，因?yàn)樗乔懈畹氖钦麄€(gè)層，通信開銷太大。這也是為什么英偉達(dá) CEO 在 GTC 峰會(huì)上首次介紹 3D 并行時(shí)特別說明其通信開銷太大的問題，只能放到一個(gè)服務(wù)器內(nèi)去做。因此，Colossal-AI 主打 2D 張量并行和 2.5D 張量并行，將計(jì)算成本降低了一個(gè)數(shù)量級(jí)。

這就意味著用一維張量并行，1 萬個(gè)機(jī)器里，每個(gè)機(jī)器都需要跟 9999 個(gè)機(jī)器打交道，而 2D 張量并行則是將其分成了各個(gè)子單元，每個(gè)機(jī)器只需要跟 96 個(gè)機(jī)器打交道。它的核心邏輯是，用一些 local synchronization（局部同步）去取代global synchronization（全局同步），以更多的局部通信去取代全局通信，這個(gè)過程中，設(shè)計(jì)調(diào)度是最困難的。

3D 張量并行也是同樣，每升高一個(gè)維度，它的設(shè)計(jì)復(fù)雜度會(huì)高一個(gè)數(shù)量級(jí)，最終通信復(fù)雜度下降了。

在內(nèi)存優(yōu)化方面，目前 AI 大模型訓(xùn)練的內(nèi)存開銷很大，即便什么事情都不做，也需要幾 T 的內(nèi)存，如果不進(jìn)行人工干預(yù)，一旦使用起來，可能需要幾十 T 甚至是幾百 T 的內(nèi)存。

為了讓模型預(yù)測(cè)效果更好，我們往往需要長(zhǎng)序列數(shù)據(jù)，當(dāng)前大模型的本質(zhì)是通過一個(gè)單詞的輸出來預(yù)測(cè)下一個(gè)單詞的概率，長(zhǎng)序列成為剛需。對(duì)此， Colossal-AI 也推出了 Sequence Parallelism（序列并行）。

具體而言，在將序列進(jìn)行切割后，會(huì)面臨一個(gè)嚴(yán)重的問題是：在進(jìn)行 attention score 時(shí)，每個(gè) token 都需要跟全局序列中的其他 token 去評(píng)估，而切分后的服務(wù)器上只有部分 token，其他服務(wù)器上也會(huì)分布不同的 token，以至于每個(gè)服務(wù)器運(yùn)行是需要同其他服務(wù)器打交道。

也就是說，假設(shè)今天屋子里 200 個(gè)人每人分別拿了一包零食，我希望每個(gè)人能品嘗下其他所有人的零食，至少需要 200 個(gè)平方次交換，才能讓每個(gè)人都嘗到其他人的所有零食。那么最簡(jiǎn)單的方式是：所有人圍一個(gè)圈，每個(gè)人將自己吃過的零食遞給右手邊的人，從自己的左手邊獲得零食，僅需 n-1 次，即 199 次的傳遞即可完成。從而降低了整個(gè)通信成本。

總結(jié)一下，目前 AI 大模型訓(xùn)練的核心技術(shù)棧，其實(shí)就是并行計(jì)算，因?yàn)槲覀円幚沓砂偕锨先f個(gè) GPU 核心，把所有 GPU 并行利用起來。數(shù)據(jù)并行、張量并行、流水線并行以及數(shù)據(jù)序列并行是并行中較為核心的模塊。

目前在內(nèi)存優(yōu)化方面，我們處于一個(gè)沒有太多選擇的環(huán)境，英偉達(dá) GPU 是最好的，我們好像也沒有其他更好的方案能夠去取代它。但美中不足的是，英偉達(dá) GPU 的內(nèi)存有限，在這種情況下，我們能否思考如何利用 CPU 內(nèi)存、NVMe內(nèi)存，核心思想就是，GPU 內(nèi)存放不下就移到 CPU 上，CPU 放不下就放到 NVMe 上，相當(dāng)于在蓋樓時(shí)，所需要的原材料自家樓下工地放不下，那我們就將其放到隔壁工廠。其技術(shù)的核心也在于最小化數(shù)據(jù)移動(dòng)，即最小化 CPU、 GPU 之間的數(shù)據(jù)移動(dòng)，最強(qiáng)化 CPU 和 NVMe 之間的數(shù)據(jù)移動(dòng)，從而將吞吐力速度提升到最高。

走向開源

Colossal-AI 是一個(gè)開源軟件，同時(shí)我們也做了一個(gè)商業(yè)化的平臺(tái)，對(duì)沒有 GPU 的用戶，可以直接在平臺(tái)上去訓(xùn)練部署自己的大模型。我們也提供了 LLaMA、PaLM、 GPT 等各種模型，最快僅需兩三天就可完成一個(gè)模型微調(diào)。相比之前用戶可能需要幾周、乃至幾個(gè)月來處理硬件、軟件等基礎(chǔ)設(shè)施，效率得到了極大提升。同時(shí)，Colossal-AI 也保護(hù)用戶的隱私，平臺(tái)不會(huì)保留、訪問用戶數(shù)據(jù)，這是 Colossal-AI 與 OpenAI ChatGPT 的本質(zhì)區(qū)別。我們將數(shù)據(jù)上傳至 Google Cloud 時(shí)，很多時(shí)候 Google 并沒有碰我們的數(shù)據(jù)，但是 OpenAI GPT 會(huì)進(jìn)行分析，AI 模型的不可解釋性、訓(xùn)練不徹底等風(fēng)險(xiǎn)普遍存在。因此，未來也會(huì)有很多企業(yè)訓(xùn)練自己的大模型，Colossal-AI 做的，是最大化保護(hù)用戶的隱私，同時(shí)提供大模型訓(xùn)練的工具。

在性能上，Colossal-AI 在同樣的硬件上可以訓(xùn)練24 倍大的模型，相對(duì)于 DeepSpeed 的 3 倍加速，即便是一個(gè)低端的服務(wù)器，也可以借助 Colossal-AI 完成相應(yīng)的模型訓(xùn)練。例如 LLaMA-65B 的訓(xùn)練，在 Colossal-AI 上使用同樣的代碼放，可以直接獲得約 50% 的加速效率。

一個(gè)簡(jiǎn)單的比喻，比如說現(xiàn)在大模型是挖金子，英偉達(dá)是賣鏟子的，那我們就是賣手套、賣衣服的，把挖金子的效率提到最高。

（雷峰網(wǎng)(公眾號(hào)：雷峰網(wǎng))雷峰網(wǎng)）

雷峰網(wǎng)原創(chuàng)文章，未經(jīng)授權(quán)禁止轉(zhuǎn)載。詳情見轉(zhuǎn)載須知。