明敏 整理自 MEET2023
量子位 | 公眾號 QbitAI
AI模型急速增大與硬體算力緩慢增長的矛盾，已愈加尖銳。
尤其在AI落地進入深水區的背景下，
如何降本增效
成為了行業熱點議題。
但具體到實際動作，該怎麼做？
在MEET 2023智慧未來大會現場，潞晨科技創始人、新加坡國立大學校長青年教授
尤洋
指明瞭一個方向：
未來，我們迫切需要一個可擴充套件性的高效計算基礎設施。
而且，潞晨科技已經用實際行動給出論證，它們推出的Colossal-AI系統，在過去一年裡迅速成長為開源加速方案中的明星專案，GitHub上攬星超過7k。
為了完整體現尤洋的分享及思考，在不改變原意的基礎上，量子位對他的演講內容進行了編輯整理。
關於MEET 智慧未來大會：MEET大會是由量子位主辦的智慧科技領域頂級商業峰會，致力於探討前沿科技技術的落地與行業應用。今年共有數十家主流媒體及直播平臺報道直播了MEET2023大會，吸引了超過300萬行業使用者線上參會，全網總曝光量累積超過2000萬。
演講要點
AI落地正面臨一個大問題：大模型/最先進AI技術訓練成本太高了。
我們迫切需要一個可擴充套件性的高效計算基礎設施。
Colossal-AI由三部分組成：高效記憶體管理系統、自動N維並行技術、大規模最佳化技術
Colossal-AI積極融入全球生態系統，為PyTorch、Hugging Face、PyTorch Lighting提供支援
（以下為尤洋演講全文）
AI大模型落地成本太高了
我今天分享的主題為
《Colossal-AI：一種全新深度學習系統，面向未來各種大模型應用場景的低成本落地》
。
在此先簡單介紹一下Colossal-AI的團隊成員。
我自己是新加坡國立大學校長青年教授，在加州大學伯克利分校獲得博士學位，並很榮幸獲得了IEEE Early Career Excellence Award。
另一位主要成員是
James Demmel教授
，他是UC伯克利前計算機系主任和EECS院長，也是UC伯克利傑出教授，美國科學院、工程院院士。
還有幾位是負責Colossal-AI開發推廣的朋友
（如圖所示）
。
下面進入正題——
Colossal-AI到底想解決什麼樣的問題？
可以先來看看這張圖。
橫座標為時間，縱座標為AI模型的引數量。
圖中有3條線，第一條從2016-2021年，模型大小從200多萬增長到了1。6萬億，翻了成千上萬倍，量化起來是
每18個月翻40倍
。
如果從2018年算起就更誇張了，每18個月翻了
340倍
。
與之形成鮮明對比的是，
GPU視訊記憶體增長相對有限
，每18個月只翻了
1.7倍
。
在未來，訓練AI大模型是一種必然。
但目前如OpenAI等頂尖AI公司，他們訓練AI的時候用了成百上千個GPU才能完成這件事。
核心在於：
第一模型計算量太大
第二引數量很大，導致視訊記憶體需求也很大
即使用模型微調或推理，也需要幾千GB的記憶體，這就是今天要面對的問題。
未來，我們迫切需要一個可擴充套件性的高效計算基礎設施。
所以我們就專注這一點，打造了Colossal-AI這個系統。
另外，中小型企業的需求也十分重要。比如GPT-2、Transformer，都是谷歌、微軟等科技巨頭的產物，中小型企業到底是什麼情況？
右邊這張圖可以說明情況。
橫座標是時間，縱座標是AI模型的引數量。
從這裡可以看出，
中小型企業使用的模型量也呈指數級上升
。
這有一個很有意思的現象，不論企業大小，大家普遍都在把模型做得更大，核心的原因在於模型更大效能可以更好。
如上左圖表明，橫座標體現了模型引數不斷增大，縱座標則是效能表現。
總結起來，AI在落地方面正面臨著一個很大的問題，大模型或者是最先進的AI技術訓練成本太高了。
Stability AI每年光計算就需要花費約2000萬美元，這顯然是小公司、科研單位難以承受的。
即使只拿大模型去做微調和推理，本地應用的挑戰也很大，需要好幾千GB的記憶體。
最終，想把大模型部署到生產線上，如果一個企業或者單位從零開始自己做，需要的人力為70人左右，而在歐美地區光是養活70個人，成本就需要2000萬美金。
所以說大模型的訓練成本提高、和自身複雜性使得整個生態變得更加困難。
何為Colossal-AI？
由此，我們打造了Colossal-AI整個系統。
它主要包括3個部分：
高效記憶體管理系統
（因為大模型本質上還是太吃記憶體，無法直接容納）
自動N維並行技術，提升計算效率
大規模最佳化技術提升擴充套件性
從三方面把AI模型訓練部署的效能提到最高。
目前來看，目標是希望使用者只需要在自己的單機筆記本上寫好程式碼，透過Colossal-AI能夠無縫地部署到雲端或者是超級計算機上。
我來簡單談一下背景。
目前訓練大模型大概有3種並行方式：資料並行、張量並行以及流水線並行。
資料並行是指，比如有1萬個資料表把它分到10臺機器上，每臺機器獲得1000個數據。
張量並行，是在層內劃分資料。
流水線並行，是在層與層之間去劃分資料。
對於我們而言，之前在資料並行上已經做了很多工作。
比如說，我們設計的
LARS、LAMB方法
，先後幫助騰訊、索尼、谷歌等公司把模型訓練時間從1小時縮短到1分鐘，我們的LAMB方法，也幫助谷歌把BERT訓練時間從3天縮短到76分鐘。
Colossal-AI的解決方案，首先是支援了上述主流並行方案，然後我們創新地打造了2D張量並行、2。5D張量並行以及3D張量並行，以及提出了資料序列並行，還提供了降低視訊記憶體消耗的異構記憶體管理和大規模並行最佳化，把它們整合起來提供一套
自動並行的解決方案
。
其實AI工程師、研究員不需要理解背後的技術細節，只需要提供模型資訊和計算資源，就可以自動地把計算資源能力發揮到最大化，同時完成虛擬模型訓練和自動部署，輕鬆低成本應用AI大模型。
在此簡單講一下張量並行。
一維張量並行是最簡單的層內劃分，二維張量並行也是層內劃分，但邏輯上把這些資料移動按二維的方式進行分佈。
這樣的好處是，透過我們精心的設計算法，可以使得傳輸資料通訊的時候，
每一個機器只用和同維度的機器打交道就行了
。
原本一維情況下，1臺機器需要和1萬臺機器打交道，現在只需要跟數百個機器打交道。
三維的效果更明顯，在100萬臺機器的情況下，每臺機器都需要跟其他每一臺機器打交道，即1000000-1臺。但現在只需要100萬的根號3次方，即100個機器打交道就行了，進一步降低了它的通訊複雜度。
這張圖可以展示二維張量並行和三維張量並行以及跟其他並行方式的區別。
與此同時，我們也提出了
資料序列並行
。
原因在於，未來真實AI執行中，往往是長序列資料。這些長資料非常棘手，為什麼？
因為本身模型應用已經很大了，如果在蛋白質摺疊或者5G、元宇宙場景，長資料再來增加記憶體開銷的話，會直接導致系統崩潰。
為了減緩系統的崩潰，我們引入了資料序列並行，主要有兩個點：
對MLP層比較友好。因為它本身就是相互獨立的。更麻煩的是要在Transformer中算Attention score的時候，需要所有token。
這裡有一個問題是，如果用比較簡單的方法，每個機器需要跟所有機器打交道，顯然得不償失。
比如有p個GPU，通訊代價就變成了p的平方，100個GPU通訊需要每步迴圈1萬個訊息，不划算。
由此，我們設計了
環狀self-attention分析演算法
，核心思路是每次從左手鄰居獲得一個k，透過p-1次就可以把這件事解決，將通訊次數從1萬降低到99。
同時也提供
自動並行技術
。
自動並行的作用，使得使用者只需把模型寫好，不需要知道平行計算分散式系統下層的各種軟體基礎設施。
因為這些東西的學習負擔太高了，使用者只需告訴我們有多少個機器、多少資源、訓練多少模型，就能自動把模型劃分到它的計算資源上，從而幫助完成訓練、部署。
另一個方面是
記憶體最佳化技術
。
眾所周知GPU記憶體是非常有限的，假如只用GPU記憶體，比如說用PyTorch baseline訓練GPT-3，至少需要幾百個GPU才能完成。
畢竟每個GPU都只有16/40/80G記憶體。
所以我們需要充分利用各個層級的記憶體。
GPU記憶體放不下，就把部分資料遷移到CPU上，CPU放不下再放到NVMe硬碟上。
但要知道，資料移動開銷是遠大於計算的，也就是GPU和CPU之間的資料傳輸，遠慢於CPU和NVMe之間的資料移動。所以要如何儘可能多地把資料移出CPU、GPU，同時能夠減少移動量。
這就需要
動態的智慧系統
，所以我們打造了Colossal-AI系統有4個模組去統計tensor資訊、確定tensor狀態，可以實時移動合適的tensor。
總之，透過這種方式可以大幅提升效率。
這有一個比較，右上角圖中紅色代表普通PyTorch。可以看出在6G記憶體的情況下，用PyTorch只能訓練1億引數，但是用Colossal-AI可以訓練15億引數，
擴大了15倍
。
相同硬體下，在GPT-2上的效果也更好，使用PyTorch本來只能訓練1億引數，透過Colossal-AI可以擴大到120億。
加速效果還可以比Megatron-LM更快。
像Stable Diffusion這樣的新模型上，Colossal-AI也能取得
6.5倍加速
，並可以把成本降到原本的1/7。
Colossal-AI也支援推理，比如在一臺伺服器上可以很好把1750億引數的OPT部署起來。
目前，Colossal-AI在開源後的增長情況還是非常好的，和之前比較成功的開源軟體Grafana、Spark等相比，至少在前13個月Colossal-AI的增長速度是遠超它們的，目前達到7k顆GitHub Star。
Colossal-AI的使用者也是遍佈於全球的，包括中國、美國、歐洲、印度、東南亞等。
與此同時，Colossal-AI還在積極融入全球生態系統。
比如目前世界上第一大AI生態系統是Facebook的PyTorch，它最好的模型是OPT，在OPT的官網上直接指向了Colossal-AI。
第二大生態系統Hugging Face也找我們做了最佳化。
第三大AI生態系統PyTorch Lighting依賴於Colossal-AI，他們的使用者能直接應用Colossal-AI提升效率。
如下目前開源社群上Colossal-AI的使用者，都是一些國內外大型企業。
Colossal-AI還有很多具體行業的落地應用，比如蛋白質摺疊，可以大幅提升效率。
Colossal-AI現已對外開源，程式碼公開免費，可以自由使用。
如果有問題，可以訪問Colossal-AI的GitHub主頁進行提問，我們免費提供技術諮詢，也有微信群可以交流討論。
我今天的演講就到這裡，有問題也可以隨時聯絡我交流，感謝大家關注。
（最後，如果想回看大會全程，請點選閱讀原文）
— 完 —
量子位 QbitAI · 頭條號簽約