WAIC 2023 | 新加坡國立大學(xué)尤洋教授 ：AI大模型的挑戰(zhàn)與系統(tǒng)優(yōu)化
在 WAIC 2023 AI 開發(fā)者論壇上，新加坡國立大學(xué)校長青年教授、潞晨科技創(chuàng)始人兼董事長尤洋發(fā)表了主題演講 《Colossal-AI：AI 大模型的挑戰(zhàn)與系統(tǒng)優(yōu)化》。演講中，尤洋首先介紹了大模型「巨人癥」問題所在以及開源軟件 Colossal-AI 社區(qū)目前發(fā)展情況。接下來，尤洋介紹了 Colossal-AI 背后的一些技術(shù)細(xì)節(jié)，主要是訓(xùn)練大模型的并行系統(tǒng)以及內(nèi)存優(yōu)化方面的努力。最后展示了 Colossal-AI 在產(chǎn)業(yè)應(yīng)用上取得的成果。
以下為尤洋在 WAIC 2023 AI 開發(fā)者論壇上的演講內(nèi)容，機器之心進行了不改變原意的編輯、整理：今天的演講主要介紹我們公司做的開源軟件 Colossal-AI 的一些技術(shù)原理和應(yīng)用。 首先簡單介紹一下我本人。我在加州大學(xué)伯克利分校獲得博士學(xué)位，現(xiàn)在在新加坡國立大學(xué)任教，同時創(chuàng)立了潞晨科技。公司另外一位核心成員 James Demmel 教授是美國科學(xué)院工程院院士，也是加州大學(xué)伯克利分校前計算機系主任兼院長。今天的演講分四個部分。第一部分簡單介紹大模型的挑戰(zhàn)以及目前 Colossal-AI 社區(qū)的發(fā)展情況，接下來兩個部分介紹一些技術(shù)細(xì)節(jié)，最后一部分介紹具體應(yīng)用上的效果。首先給大家展示一張圖：橫坐標(biāo)是時間，縱坐標(biāo)是 AI 模型的參數(shù)量。過去六年中，最好的 AI 模型參數(shù)量已經(jīng)上升了 1 萬倍左右。比如，2016 年 ResNet-50 只有 2000 萬參數(shù)，2020 年 GPT-3 已經(jīng)達(dá)到 1750 億參數(shù)的規(guī)模。據(jù)說 GPT-4 也是用的混合專家系統(tǒng)，跟谷歌 Switch Transformer、智源的「悟道」都是同一種技術(shù)。Switch Transformer 參數(shù)規(guī)模大概 1.6 萬億，據(jù)說 GPT-4 有 16 個專家（模型），每個專家（模型）有千億左右（參數(shù)）。所以說，過去六、七年 —— 從 ResNet-50 到 GPT-4—— 最好模型的參數(shù)量剛好大了 10 萬倍左右。但是，以 GPT-3 為例，模型構(gòu)造沒有到 100 層，ResNet-50 也是 50 層左右，層數(shù)基本上沒變化，模型不是變得更深而是變得更寬，大了 1-10 萬倍左右，也給 GPU 內(nèi)存造成更大壓力，但 GPU 內(nèi)存每 18 月只增長 1.7 倍，這就需要對下一代人工智能基礎(chǔ)設(shè)施進行優(yōu)化或者重建。所以，我們打造了 Colossal-AI 系統(tǒng)。這是 Colossal-AI 系統(tǒng)結(jié)構(gòu)圖，包括三個層次。第一個層次是內(nèi)存管理系統(tǒng)，因為大模型太吃內(nèi)存。第二部分是 N-Dim 并行技術(shù)（N 維并行技術(shù)）。據(jù)說 OpenAI 已經(jīng)用 10 萬張 GPU 卡訓(xùn)練大模型。前兩天一家美國創(chuàng)業(yè)公司融資了 13 億美金，背后基礎(chǔ)設(shè)施據(jù)說已經(jīng)有 2 萬張 GPU 卡。未來，從 1 個 GPU 到 10、100、10000 個 GPU ，自動擴展效率會對訓(xùn)練系統(tǒng)產(chǎn)生根本性影響，所以我們打造了 N 維并行系統(tǒng)。 第三部分是低延遲的推理系統(tǒng)，也是目前 Colossal-AI 的主要模塊。模型訓(xùn)練好后要服務(wù)用戶，用戶每調(diào)用一次模型就是做一次推理，這跟成本有直接關(guān)系。所以，推理的延遲要很低，成本要降到最低。雖然 Colossal-AI 開源社區(qū)只推出了 20 個月左右，但發(fā)展速度非?？?。（下圖）橫坐標(biāo)是時間，縱坐標(biāo)是 GitHub 上的星數(shù)，可以看出 Colossal-AI 增長速度遠(yuǎn)超于傳統(tǒng)開源軟件。 Colossal-AI 增速也遠(yuǎn)超與 Colossal-AI 類似軟件，比如 DeepSpeed。目前 Colossal-AI 用戶遍布全球。中國、美國、歐洲、印度、東南亞都有很多用戶。在全球 AI 生態(tài)系統(tǒng)中也都發(fā)揮了更重要的作用。目前世界公認(rèn)的第一大 AI 生態(tài)是 PyTorch ，它的生態(tài)圖中有一個鏈接指向了 Colossal-AI。Lightning AI 是 PyTorch 第一大子生態(tài)，直接依賴于 Colossal-AI。 HuggingFace 世界公認(rèn)的第二大 AI 生態(tài)，很多應(yīng)用也是基于 Colossal-AI 做的。現(xiàn)在開源生態(tài)做的最好的是 PyTorch 也由 Facebook 主導(dǎo)，我們的開源大模型很多基于 Facebook 的 LLaMA。OPT 也是 Facebook 的一個很重要模型。OPT 官網(wǎng)截圖顯示，它也有一個鏈接指向了 Colossal-AI，OPT 用戶也可以通過 Colossal-AI 進行優(yōu)化。第二部分講一下技術(shù)細(xì)節(jié)。假設(shè)未來需要用 1 萬或者 10 萬張 GPU 卡訓(xùn)練大模型，怎么充分利用這些 GPU 卡呢？與二十年前相比，今天特別是下層系統(tǒng)軟件面臨的問題有很大不同。20 年前，英特爾芯片快了 3 倍，因為都是串行代碼，一行代碼都不用改就可以直接快 3 倍。但在今天，提升算力主要靠并行計算，不管是單個 GPU 內(nèi)的多線程還是多個 GPU，甚至于兩三周之前英偉達(dá) CEO 黃仁勛在演講中說到，要把成百上千個 GPU 用高速網(wǎng)絡(luò)連接在一起（現(xiàn)在已經(jīng)是這樣的狀況了）。如何最大化算法或者說上層應(yīng)用并行度，將對機器的性能發(fā)揮產(chǎn)生非常實質(zhì)性影響。在同樣條件下，好的分布式軟件導(dǎo)致速度上出現(xiàn) 5-10 倍的差距也很正常，所以，并行系統(tǒng)非常重要。 目前，訓(xùn)練 AI 大模型的并行系統(tǒng)主要有三個。最開始也是黃仁勛在 GTC 中介紹的，包括數(shù)據(jù)并行、模型并行和流水線并行。模型并行是指層內(nèi)并行。剛才也說了，現(xiàn)在處在寬度學(xué)習(xí)時代，而不是深度學(xué)習(xí)時代。從 ResNet-50 到 GPT-3，層的數(shù)量基本沒有變化，但參數(shù)量大了 1 萬倍，變寬了很多。在層的寬度越來越大、每層計算量越來越大的情況下，張量并行有可能發(fā)揮更加重要的作用。當(dāng)然，張量并行最大的弊端是它的通信開銷太大，這也是黃仁勛在 2021 年演講里提到把所有張量并行都放在服務(wù)器內(nèi)的原因，跨服務(wù)器的通信開銷太大，得不償失。所以，如何優(yōu)化張量并行的通信就非常重要。我們又打造了二維張量并行、三維張量并行、2.5 倍張量并行，去最小化數(shù)據(jù)移動。另一個是數(shù)據(jù)序列并行。我們前段時間看到一些工作，看能不能把 sequence 增長到 4K 的量級。比如對 ChatGPT 應(yīng)用而言，sequence 越長，一定概率上，預(yù)測效果也會更好，因為上下文得到了更多信息。但是，Transformer 的架構(gòu)決定了每個 token 都需要跟其他的 token 算 Attention score，內(nèi)存開銷非常大。數(shù)據(jù)序列并行主要優(yōu)化的也是 sequence 劃分之后，如何盡量減少它們的數(shù)據(jù)移動。當(dāng)然，最傳統(tǒng)的還是數(shù)據(jù)并行。我會一一介紹這三種并行系統(tǒng)。先介紹一下最簡單的數(shù)據(jù)并行。最開始我們訓(xùn)練 ResNet、AlexNet 甚至 BERT，很大程度上可以把模型下放到單個芯片，只要單個芯片能把模型放下去，芯片和芯片之間做數(shù)據(jù)并行就可以了。比如，將一堆數(shù)據(jù)分 10 份，每個 GPU 放十分之一。這樣情況下，工作就很理想了 —— 只需要增大 batch size 10 倍，假定能把循環(huán)次數(shù)減少 10 倍，這樣 batch size 線性增大、循環(huán)次數(shù)線性減少就變成了一個很理想的擴展性問題，這樣我們可以保持 Epoch 不變，但總體運行時間顯著減少。數(shù)據(jù)并行的通信，其實就是每次循環(huán)的時候大概需要交換一下梯度就可以了。當(dāng)然，梯度大小和參數(shù)大小是一樣的，對傳統(tǒng) CNN 而言，它們參數(shù)也不太多，所以，通信也非常友好，數(shù)據(jù)并行之前也是最常用的一種方式。但是，大模型時代面臨的問題是循環(huán)次數(shù)逐漸減少情況下，收斂變得越來越困難 —— 假如之前用 1 萬次循環(huán)做收斂，batch size 增加 10 倍就需要在 1000 次循環(huán)中完成收斂。所以，數(shù)據(jù)并行的難度在于它的 accuracy 有時候會掉很多。因此需要看能不能設(shè)計出更好的優(yōu)化方法，去解決這個問題。我們團隊打造了 LARS 和 LAMB 的方法，也發(fā)表在 ICLR 論文上，現(xiàn)在引用次數(shù)也快 700 次。最近，我們基于 LAMB 方法做了另外的工作，今年年初也獲得了 ACL 杰出論文，另一篇獲得 AAAI 2023 杰出論文。其實都是基于這個工作線，優(yōu)化訓(xùn)練的收斂性和效率。包括谷歌，F(xiàn)acebook SEER，Deepmind BYOL 等使用了 LARS 和 LAMB 方法。我們也幫助谷歌把 BERT 訓(xùn)練時間從三天縮減到 76 分鐘。說完數(shù)據(jù)并行，再說一下模型并行。模型并行就像蓋樓 —— 要蓋一個很寬的樓，每層 1 萬平，直接用 10 個工程隊把它分成 10 份，每個工程隊一份，同步蓋完一層后，再蓋下一層。所以，層內(nèi)并行的問題是通信開銷太大。第三個是流水線并行。并行的字面意思可以理解為，算完這一層算下一層，第 N 層要依賴于第 N-1 層的結(jié)果，第 N+1 層也要依賴于第 N 層的結(jié)果，前后依賴。比如，現(xiàn)在要蓋 1000 棟樓，有 20 個工程隊。第一個工程隊先蓋第一棟樓第一層，蓋好之后，第二個工程隊才能入場蓋第二層。這時，第一個工程隊移步第二棟樓，開始蓋第一層，直到蓋完 20 棟樓的第一層時，所有 20 個工程隊才能都進入施工現(xiàn)場。流水線并行的層數(shù)和 GPU 數(shù)量之間關(guān)系，就像工程隊數(shù)量（類似 GPU 數(shù)量）與樓層數(shù)（相當(dāng)于流水線并行的層數(shù)）關(guān)系一樣密切。數(shù)據(jù)流水線數(shù)和 GPU 之間比值越大，并行效率越高。20 個工程隊蓋 1000 棟樓，很多時候并行度可以達(dá)到 20，但用 20 個工程隊去蓋 30 棟樓，效率就不是很高。所以，流水線并行本質(zhì)上還是需要將 batch size 擴得很大，因為不同流水線進行了不同計算，本質(zhì)上它的并行也來源于數(shù)據(jù)并行。這是三種傳統(tǒng)并行方式。GPT-3 剛出來的時候，OpenAI 用了 1 萬個 GPU 訓(xùn)練 GPT-3，但是英偉達(dá)優(yōu)化之后發(fā)現(xiàn) 3072 個 GPU 就夠：把 64 個服務(wù)器分成一組，每個服務(wù)器內(nèi)有 8 個 GPU，8 個 GPU 做張量并行。64 個服務(wù)器之間就做流水線并行，每個組有 64 個服務(wù)器，組之間再做數(shù)據(jù)并行，因為流水線并行的傳輸代價非常小，只傳層和層之間的信息。我們之前的工作做了很多張量并行的優(yōu)化。張量并行未來發(fā)展空間也是最大的，因為模型的層變得更寬，通信開銷也非常高昂。而二維張量并行、三維張量并行的核心思想是用更多的局部序列化替換全局序列化 ，用更多局部通信換取全局通信，從而降低通信成本，提高效率。這幾種張量并行方式的訓(xùn)練效率差異非常大，特別是二維和三維張量并行。可以看一下最后兩列（最右側(cè)兩列）：第一個是通信帶寬上的代價，第二個是通信延遲的代價。二維、三維、2.5 維張量并行可以大幅度降低通信和內(nèi)存開銷。數(shù)據(jù)序列并行非常重要。前段時間有一個美國創(chuàng)業(yè)公司發(fā)了一個大新聞，可以把單次數(shù)據(jù)序列的 input sequence length (輸入序列長度) 做到 4K Token。這個事情很重要，這張圖可以說明一些問題：橫坐標(biāo)反映數(shù)據(jù)序列長度，縱坐標(biāo)預(yù)測下一個詞的精度 —— 數(shù)據(jù)序列越長我們預(yù)測的越準(zhǔn)確，因為它有更多的上下文信息。雖然不好的地方在于計算變得更加復(fù)雜，但長的 Sequences 是算法側(cè)需求，下層的基礎(chǔ)設(shè)施也應(yīng)該盡可能去滿足。這兩天我又看到一篇論文，有團隊聲稱把數(shù)據(jù)序列甚至能做到 100 萬（當(dāng)然需要再求證）。數(shù)據(jù)序列變長的需求還是會長期存在，越長確實效果就越好。這張圖展示的是為什么變長之后，給下層實現(xiàn)、基礎(chǔ)設(shè)施帶來很多麻煩。現(xiàn)在基于 Transformer 的架構(gòu)，每個 Sequence 都有很多 token，每兩個 token 之間都需要算 Attention score，導(dǎo)致內(nèi)存開銷非常大，特別是 Sequence 長度到一定級別后，內(nèi)存開銷會指數(shù)級上漲。Transformer 結(jié)構(gòu)內(nèi)存壓力本來很大。比如訓(xùn)練 GPT-3 ，如果按照 2000 億參數(shù)量計算，假設(shè)用單精度，每個參數(shù)要占 4 個字節(jié)，僅參數(shù)就要占 800G 內(nèi)存，梯度要占 800G 內(nèi)存，等等。何況 Transformer Attention score 等更是指數(shù)級上升，內(nèi)存壓力非常大。因此，如果數(shù)據(jù)序列想變得更長，就需要并行，把不同 Token 劃分到不同的 GPU 上。另一個問題是如何減少 GPU 之間的數(shù)據(jù)移動。100 個 GPU，所有都需要跟其他 GPU 交換數(shù)據(jù)，這是 Transformer 的 Attention 機制決定的。我們打造了環(huán)狀 self-attention 通信算法，盡量減少數(shù)據(jù)之間移動，每次循環(huán)只需要跟左右鄰居打交道，把通信復(fù)雜度從 P 平方降低到 P-1（P，指 GPU 數(shù)量或服務(wù)器數(shù)量）。時間關(guān)系，實驗結(jié)果就不介紹了，感興趣的話可以看一下我們的官網(wǎng)，Colossal-AI 做了很多優(yōu)化。比如接下來要介紹的內(nèi)存優(yōu)化。其實就是減少 GPU 和 CPU 之間、CPU 與 NVME 硬件之間的數(shù)據(jù)移動。這一點在未來非常重要。因為 GPU 內(nèi)存有限，我們需要用 CPU 甚至 NVME 硬件，但它們之間的數(shù)據(jù)移動可能比計算慢上千倍，如何最小化它們之間的數(shù)據(jù)移動，以更低成本容納大模型將非常重要。快速過一下一些應(yīng)用效果。在一些重要應(yīng)用比如 Stable diffusion 上，Colossal-AI 取得了很好的加速。我們可以把內(nèi)存開銷降低 5.6 倍，硬件成本降低 46 倍。Colossal-AI 經(jīng)過了很多業(yè)界檢驗。用 Colossal-AI 加速過的效果圖是沒有任何損失的。 我們通過 Colossal-AI 做了 ColossalChat，是世界上首個開源的最接近 ChatGPT 原始技術(shù)方案，具備完整 RLHF 流程的低成本 ChatGPT 復(fù)現(xiàn)方案。只有 70 億或 100 億的參數(shù)模型可以達(dá)到更好效果。相對于業(yè)界標(biāo)準(zhǔn)，Colossal-AI 可以將推理上獲得 30-50% 的提速，在訓(xùn)練上獲得大概 7 倍加速。包括可以把 ChatGPT 訓(xùn)練成本從 300 萬美金降低到 140 萬美金左右。接下來是一些 demo 展示和與斯坦福羊駝的對比，以及一些更進一步分布式 PPO 和 evaluation 的工作。當(dāng)然，由于我們做的是計算優(yōu)化，因此不會限制于某個特定行業(yè)或者模型，具備良好的通用性。例如對于生物醫(yī)藥行業(yè)的蛋白質(zhì)預(yù)測模型 AlphaFold2，我們也可以優(yōu)化提升約 10 倍的訓(xùn)練推理速度。如果感興趣可以看一下 Colossal-AI 官網(wǎng)，加我們的微信群，也非常歡迎大家在 GitHub 提 issue，或者提交 PR 共同建設(shè) Colossal-AI 項目。我的介紹就到這里，謝謝。
本站僅提供存儲服務(wù)，所有內(nèi)容均由用戶發(fā)布，如發(fā)現(xiàn)有害或侵權(quán)內(nèi)容，請
點擊舉報。
