導讀 隨著 Chat-GPT 的普及，如今大模型已成為不可阻擋的趨勢。本文從 Colossal-AI 這個行業軟體入手，介紹如何提高大規模 AI 訓練、微調和部署的效率。
今天的介紹會圍繞下面五點展開：
1. 大模型訓練的挑戰
2. N 維並行系統
3. 高效記憶體管理系統
4. 傑出效能和應用案例
5. Q&A
分享嘉賓｜尤洋博士 新加坡國立大學 校長青年教授、潞晨科技 創始人兼董事長
編輯整理｜王吉東
內容校對｜李瑤
出品社群｜DataFun
01
大模型訓練的挑戰
1. 大模型發展現狀及問題
上圖中展示了大模型的發展歷程，其中縱座標是 AI 模型的引數量。在過去幾年間，AI 大模型的引數量飛速增長，大約每 18 個月時間就會增長 40 倍。例如 2016 年，世界最好的大模型 ResNet-50，引數量約為 2000 萬；而到 2020 年的 GPT-3 模型，引數量已達到 1750 億；到今天，根據 OpenAI 透露的訊息，MOE 混合專家系統這種架構的大模型引數量大約 1.7 萬億；2021 年初 Google 的 Switch Transformer 模型採用了相同的混合專家系統架構，引數量已達到 1.6 萬億。
近來，OpenAI 在訓練的 GPT-5 模型，預計將達到 100 萬億引數。此外，Inception 等創業公司甚至認為未來模型的引數量將超過 GPT-5 模型 100 倍。
隨著模型規模的指數級增長，未來可能需要使用成千甚至上萬個 GPU。因此，背後的軟體基礎設施至關重要，不管是 AI 框架還是分散式系統，都將起到關鍵作用。
未來的 AI 訓練，其核心在於擴充套件性和高效性。所謂擴充套件性，即從 1 個 GPU 擴充套件到 10 個、100 個、1000 個、1 萬個甚至更多的 GPU，而在保證擴充套件性的同時能否始終保持百分之八九十的擴充套件性效率？這就是所謂的擴充套件性和高效性。
2. Why Do We Believe in Large Models?
我們投入了大量精力去完善基礎設施，就是為未來 AI 大模型而準備的。從源頭上說，我們之所以看好大模型的發展，原因主要有兩點：其一，在模型更大的情況下，確實能夠獲得更優的效能，其 AI 智慧程度也更高；其二，大模型的“大”字其實很難定義（今天，100 億引數屬於大模型，明天，可能 1 萬億引數才是大模型），但是有一點是肯定的，那就是未來的小模型也會呈現指數級上升。比如我們今天覺得 Bert 已經成為了一個小模型，兩年之後可能 GPT-3 也變成了我們眼中的小模型；再過三年，可能今天的 GPT 子版本 ChatGPT 也變成小模型了。
如上圖所示，橫座標為時間，縱座標為模型引數量。例如 2015 年，中小型企業一般會使用 LeNet；到 2016 年，中小型企業已經能用 ResNet-50 了；到 2020 年的時候，中小型企業甚至可以使用 Bert、GPT-2 等模型；預計到明年，中小型企業會陸續使用 LLAMA 這類模型（引數量大約 700 億）。如上文所說，小模型實際上也在指數級上升。
因此所有 AI 模型的訓練和微調都需要分散式，因為至少需要幾十甚至上百張 GPU。近期，英偉達的 GH 200 透過將 CPU 和 GPU 連線在一起，使用更高速的網路將其串聯。未來，GPU 訓練會更加便捷，就像今天的伺服器一樣，廠家自行採購一套實體伺服器或雲端伺服器，即可實現大模型的訓練。
大模型的訓練往往需要很高的並行度，那麼如何提高模型訓練的效率呢？
首先可以關注雲上計算的價格。如上圖，橫座標是時間，縱座標是單位美元能買到的浮點數計算次數（T flops 次數）。從圖中可以看出，從 2015 年到 2020 年，該次數呈指數級上升了 45 倍。因此，絕對的算力實際上是在變便宜的，即單位美元買到的浮點數計算次數是呈指數上升的。因此，如果能將算力挖掘出來，預計到 2026 年，GPT-3 的訓練成本將從現在的 300 萬美金降低到 7 萬美金左右。
如果進一步最佳化模型的並行訓練效率，設計出更好的演算法，使其收斂速度變得更快，這個訓練成本可能會降低至兩三萬美金（甚至可能會降到 1 萬美金以下）。
3. 大模型的挑戰
然而現在，訓練成本仍然是很高的，例如谷歌的 PaLM 模型，按照雲上的價格，需要 900 萬美金。另一方面，如果一個公司想從零開始訓練一個大模型，至少需要投入五六十個人，帶來了較高的人力成本，增加了模型落地的困難。
對於大部分的業務型企業，企業本身希望快速實現對大模型的訓練、部署和微調，這些對大模型的落地是非常關鍵的因素，而上述現狀對大模型工作在這類企業的開展就會變得非常困難。
4. Colossal-AI
基於上述原因，我們打造了 Colossal-AI 這個工具，希望幫助更多企業實現大模型的快速訓練、微調和部署。
Colossal-AI 包括三個部分：
- 高效的記憶體管理系統
- N 維分散式系統
- 低延遲推理系統
利用以上三個部分，三位一體地實現大模型的高效訓練。不管使用者使用什麼樣的程式設計介面或者框架，透過 Colossal-AI 的核心都可以在 CPU、GPU、NPU 上獲得最好的效能。
上圖回顧了 Colossal-AI 開源社群發展的情況，橫座標是時間，縱座標是 Github 星數，可以看出 Colossal-AI 目前的增長速度已經遠遠超過了同時期的 Spark、ClickHouse、MongoDB、Kafka 等。
其背後的原因，除了 Colossal-AI 自身優秀的開源設計，還在於這一波 AI 的浪潮預計比上一波大資料的浪潮更大，未來 AI 的價值可能會貫穿到企業的整個生命週期，因此未來的機會是非常大的，降低模型的訓練難度和訓練成本，是非常重要的一件事情。
上圖是 Colossal-AI 的使用者分佈圖，可以看出，目前在中國、美國、歐洲、印度、東南亞都有很多的使用者。
同時，很榮幸，Colossal-AI 被收錄到很多國際頂尖 AI 會議的官方教程，如今年的 Neural NIPS（之前叫 NIPS，該會議是目前 AI 的第一技術峰會，每年 OpenAI、谷歌等頂尖技術人員，包括像圖靈獎得主等，都會到該會議上演講），Colossal-AI 很榮幸成為該會議的官方教程。此外，Colossal-AI 曾被收錄進了 AAAI、CVPR 等頂尖 AI 會議的官方教程。
目前，MLPerf 是衡量 AI 晶片的公認指標，是英偉達、英特爾、谷歌、阿里巴巴等公司共同發起的基準，現在稱為 mlcomments。現在 MLPerf 的官方 Github 也依賴 Colossal-AI，因此 Colossal-AI 也很榮幸地成為了 MLPerf 標準的一部分。
此外，在全球 AI 生態系統中，目前 Colossal-AI 也處於非常理想的位置。例如，目前世界第一大 AI 生態 PyTorch，其官方生態直接依賴於 Colossal-AI 的分散式；此外，世界第二大 AI 生態 Hugging Face 的很多應用也是基於 Colossal-AI 的。目前開源 AI 大模型由 Facebook 所主導，LLAMA 和PyTorch 都是 Facebook 開發的，目前 Facebook 官方大模型 OPT 有一個連結就指向了 Colossal-AI。
02
N 維並行系統
未來大模型的訓練，會涉及處理成百上千個 GPU 叢集內部的關係問題，即如何把並行效率提升到極致。未來使用 AI 超級計算機（如使用 256 個 GH 200 組裝而成），雖然其頻寬很高，通訊傳輸效率很高，但仍然需要足夠的並行度才能將算力發揮出來。
對於其他科技巨頭，如華為，使用類似英偉達的 GH 200 這種超級計算機，其在平行計算方面，主要使用的方式包括資料並行、流水線並行和張量並行，而第一個實現這項技術的還是英偉達。
GPT-3 模型是 2020 年 6 月份訓練出來的。在隨後的 GTC 峰會上，英偉達黃總介紹了他們是如何訓練 GPT-3 的，當時提出 3-D 並行這一概念。他們用了 3072 個 GPU（每個伺服器有 8 個 GPU），每個伺服器內使用 Tensor parallelism 進行張量並行，然後 64 個伺服器形成一個小組，小組內部使用流水線並行，而 6 個小組之間使用資料並行，這樣就是 8x64x6=3072。
這樣的配置，主要考慮的核心點就是張量並行的通訊代價太高，因此在伺服器內部使用流入線並行的方式，通訊代價相對較低。在伺服器之間，為了擴充套件到更大，又引入了資料並行的方式。而未來 GH200 組成的超級計算機，其網路頻寬可能會非常高，晶片之間可以直接交換，這樣就可以直接使用張量並行的方式（雖然通訊代價高，但如果網路傳輸非常好，則可以很好地解決這個問題）。
回到今天，對於張量並行，最開始實現的是英偉達的 microtrain，然而其並沒有把張量並行最佳化到極致。
張量的本質就是矩陣，矩陣的切片是非常靈活且容易實現的。然而理論上每個機器都需要跟其他所有機器打交道，這裡就會涉及到 global synchronization，而我們的想法是能不能用更多的 local synchronization 來取代 global synchronization，這樣會獲得指數級的速度增長。因此我們提到了 2 維張量並行、3 維張量並行，甚至 2.5 維張量並行，其實質就是上述思想。
第二點，回顧大模型的變化趨勢會發現，2016 年的 Resnet 50 模型，有 2000 萬引數，包含 50 層；到 2020 年，GPT-3 模型有將近 2000 億引數，只包含不到 100 層；而今天的 GPT-4 模型也只有 120 層。模型的引數規模越來越大，而層數卻沒有大的變化，說明模型的層變得更寬了。在這樣的趨勢下，未來對張量並行的要求將會更高，因此張量並行未來仍有很大的最佳化空間。對於流水線並行，其好處在於通訊代價很低，但是它總是有 bubble 的開銷，會限制並行的效率，其 pipeline parallelism 只能在一定規模內使用，無法無限擴充套件。
第三點，我們需要新的並行 sequence parallelism。以對話類大模型為例，比如 ChatGPT，其 sequence Lens（contact Lens）是 8000（Bert 是 128512）。GPT 的核心原理是使用上下文資訊去預測下一個詞，因此，contact lens（即 sequel less）越長，預測的下一個單詞就越準確。有些創業公司甚至想把 contact lens 做到 100 萬，然而這樣會帶來巨大的記憶體開銷壓力（後文會介紹這一部分的最佳化）。
國內的一些通用大模型公司在 AI 基礎設施這部分的關注點主要在於從千卡擴充套件到萬卡。然而隨著 GPU 數量的擴充套件，系統穩定性會下降，即使 OpenAI、微軟這樣的公司，其萬卡叢集的利用率也只有 30% 左右。如果能把穩定性的問題解決好，下一步就是擴充套件性的問題。
從 100 卡擴充套件到 1000 卡，已經帶來了很大的困難，其效率可能會從 70% 直接降到 30% 左右。而從 1 千卡到 1 萬卡，如果沒有足夠的平行計算效率，其訓練可能會變得更慢。
這就引出了另一個問題：要將 AI 訓練的基礎設施增長到萬卡訓練的規模，哪方面是欠缺的？首先 pipeline parallelism 只能在區域性使用，因此無法擴充套件到萬卡；如果再加上 tensor parallelism 以及 sequence parallelism，未來還需要進一步最佳化資料並行。
資料並行的好處是，base size 增大 10 倍，epoch 數儘量固定，則計算量也是恆定的。如果 epoch 數量能夠固定下來，batch size 增大 10 倍，則迴圈次數就減少 10 倍；如果同時實現資料並行，其通訊就是梯度加和求平均，這樣單次通訊量又減少 10 倍；因此資料並行在萬卡訓練中也會發揮出更大的作用。假如把 batch size 擴大 10 倍，在保證 accuracy 的基礎上從千卡擴充套件到萬卡，其通訊量會減少 10 倍，因此會更加高效。
2017 年曾有一篇論文《training ImageNet in one hour》，使用一小時訓練 ImageNet。這一篇論文可以認為是整個分散式AI訓練的開山之作。如今我們可能會覺得他的技術很簡單，但其是第一個使用百卡來訓練深度神經網路的一個高效案例。
當時只用了資料並行，而裝置也只是 Ethernet，甚至沒有使用 Infinity band，但其擴充套件性也能超過 90%，這說明 data parallelism 的通訊效率還是非常高的。因此，對於未來的萬卡訓練，這個 data parallelism 也會非常關鍵。此外，我們仍然需要少部分的 pipeline parallelism、比較高效的 tensor parallelism 以及 sequence parallelism，也需要更大的 data parallelism。
前文提到的 data parallelism，只要將 epoch 數恆定，將 batch size 增大 10 倍，迴圈數就會減少 10 倍，而通訊就是梯度加和求平均，因此還是非常 robust 的。假如使用 1000 個機器進行訓練，即使有 5-10 個機器崩潰，也不會影響 data pair reserve 結果的一致性，因為只要樣本分佈在統計學上接近正態分佈或其他較理想的分佈，那麼使用 1000 個 T5 加權平均和用 995 個 T5 加權平均的結果不會相差太大。
實際上 DIP LSM 業務存在一些自身問題，當 batch side 開得過大時，迴圈次數反而增長更大（例如 by side 增大了 1000 倍，迴圈次數增長了 2000 倍）。這就需要在 1/1000 的迴圈內尋找區域性最優解，所以變得更具挑戰性。因為這種情況下很容易陷入 sharp minimum，導致 accuracy 下降嚴重，從而得不償失。因此 data parallelism 的挑戰就是要設計出更好的最佳化方法。
Colossal-AI 中 loss 和 lamb 也做了很好的實現：首先在保持 icon 的情況下，以 VIP 為例，把 batch size 從 128 擴到 32K，同時計算時間從 733 小時降低到 40 分鐘，因此可以看出其擴充套件性也是線性的。
之前我們的一些方法曾幫助騰訊、索尼、谷歌、富士通等公司將音訊拆下來，訓練時間從一小時縮短到了一分鐘，並且曾幫助谷歌把 batch 訓練時間從 3 天縮至 76 分鐘。Data parallelism 可能在未來的萬卡訓練中發揮出更大作用。
下面介紹 pipeline parallelism 和 tensor parallelism。
Megatron 一維的方法非常簡潔，但是其通訊代價較高，因為其本質是 global synchronization，每一層都做的話計算量是非常恐怖的；這一點不同於 data parallelism，後者是算完一層後直接同步去通訊。
而 Tensor parallelism 則是算完這一層才進行下一層，因此沒有 overlap 的機會，導致每個層內都使用 together or reduce 的代價變得過高。
因此前文提到使用更多的 local synchronization 取代 global synchronization。
以 2D 方法為例，將每個機器視為一個網格，每次訓練都先和同列的機器打交道，這樣只需要跟根號 p 個機器打交道即可。同理，3D 方法就是將其分成三個維度，每次只和同個維度的機器打交道。
從上圖可以看出，2D 方法的效率還是挺高的，且規模越大，其效果越好。
上圖可以看出，引數記憶體的開銷是差不多的。在通訊層面，bandwidth latency 開銷，可以看出維度的增加確實提高了通訊效率。
未來的長序列也非常有挑戰性，現在訓練大模型的記憶體開銷已經很高了，假如訓練一個大概 2000 億引數的 AI 模型（如 GPT-3），假定使用 LP32，則需要 200 乘以 10 的 9 次方；假如每個引數是 LP32 化，就是 4 個 Byte，則引數首先就會耗費 1000G 的記憶體（梯度最佳化器同理）。因此，在模型開始訓練之前就需要上 T 的記憶體，而 GPU 一般只有 100G 的記憶體。該問題的嚴重性在於，如果使用更長的序列，則序列中每兩個 token 之間都要計算 attention score，會產生很多中間記憶體的需求。因此序列越長，記憶體開銷就越大，當達到一定的臨界值之後，將會呈指數級上升。
因此，直觀上，需要將序列劃分成若干個不同的部分，將每一部分分配給不同 GPU 上不同的伺服器，進而控制其記憶體開銷上限。然而，在不同的 GPU 存放不同的 token，每個 token 都需要和其他 token 計算 attention score，這樣會帶來 p2 的通訊操作，造成非常高的通訊代價。基於此，我們提出了 Ring Self-Attention 方法。
該方法基於百度的 ring or reduce，其核心思想如下：比如 50 個人，每個人抱了一大包零食，現在需要每個人都嘗一下其他所有人的零食，這樣將會需要 2500 次的重新傳輸，但是如果大家手拉手圍成一個圈，每次從左手邊獲得一包零食，再將其傳給右手邊的一個人，這樣只需要批檢一次自操作即可。
Ring Self-Attention，每次從左邊鄰居獲得一個 key，再將其傳遞給右邊鄰居，進而計算出一個 tension score，迴圈往復，這樣只需批檢一次重新操作，即可算出所有資訊。
上圖中展示了一些很直觀的比較，可以看出，throughput 相比之前的方法要快 50% 左右，能夠同時在同樣裝置下訓練的 sequence 長度比之前增長 5%。
03
高效記憶體管理系統
現在的AI架構主要還是基於 GPU 伺服器去做訓練，每個 GPU 的記憶體非常有限，因為記憶體只有 100G 左右，而訓練一個千億引數的模型至少需要上 T 的記憶體。CPU 上也有很多記憶體，GPU 放不下就放 CPU，再放不下就放 NVMe。
然而這樣帶來的問題就是會嚴重降低效率，因為 GPU 內的計算遠快於 GPU 和 CPU 之間，以及 CPU 和 NVMe 之間的傳輸。因此使用這種多級的記憶體就需要最小化 GPU 和 CPU 之間，以及與 CPU 和 NVMe 之間的資料移動。
我們的方案是將引數分成了塊狀，每次選取有用的塊，暫時用不上的塊先存放到 CPU 上，透過這種方式動態地進行儲存，將效率提到最高。使用這種方式，只要有一塊記憶體，都會放滿最常用的引數或者梯度。
04
傑出效能和應用案例
與 PyTorch 相比，Colossal-AI 可以訓練 20 倍大小的模型，在同樣裝置條件下可以獲得 3 倍的加速。從下圖的比較可以看出，在相同裝置條件下，可訓練的引數量呈指數級上升。
從 Throughput 的圖可以看出，Colossal-AI 與 DeepSpeed 相比，throughput 也是更高的。例如，在 Stable Diffusion 上可以獲得約 6 倍的加速。同樣在 inference 上打 Colossal-AI 的話也可以獲得 40% 左右的加速，而加速之後，其記憶體得到了大幅的下降，但是對圖的質量是沒有任何影響的。
從上圖可以看出，Stable Diffusion 的出圖效果還是比較理想的。首先跟 PyTorch DDP 相比，Colossal-AI 在推理上可以獲得 40% 左右的加速，在訓練上可以獲得 7 倍左右的加速。在同樣的裝置條件下，可以存放 10 倍大小的模型。
最近 Colossal-AI 也對 LLAMA 模型進行了最佳化，特別是在訓練環節，獲得了 2 倍左右的加速。
以上是 Colossal-AI 訓練平臺的介紹，可以實現一站式的 AI 模型訓練、微調和部署。如果感興趣，歡迎下載並檢視 Colossal-AI 的開原始碼，同時歡迎關注公眾號，並加入微信群。
05
Q&A
Q1：大模型場景中，不同模型的結構差別比較大。以 Transformer 結構為例，包括 encoder-decoder 聯合架構，也有純 encoder 或純 decoder 的架構。那麼在訓練最佳化和推理最佳化的過程中，每一種模型可能都需要進行 specific 的最佳化；如何看待 overhead 呢？
A1：對於 overhead，常會聽到抱怨模型太多了。我們這個最佳化每種最佳化只能對到某個模型上，然後做像 confusion 這種。這個情況其實是往好的方向發展，因為近期逐漸都收斂到了 LLAMA、GPT、Falcon 這幾種常見的模型，上述模型彼此間的相似度較高。例如未來 GPT-4 將相當於 GPT 加上 MOE；而加上 MOE 之後，對運算元的下方是沒有任何影響的。正因為未來大模型的試錯成本越來越高，反而會導致未來收斂成少數常見的大模型，這是有利的地方。回到這個問題，如何通用地解決這類問題，其實個人感覺並沒有太好的辦法，只能儘量將通用的需求做好，在此基礎上做最佳化。並且更嚴重的一個問題就是現在國產晶片有很多力量，就是想想未來佈局一些新的事情，這種條件下我們可能就需要更多的最佳化。我覺得短期之內假如說有一個新的晶片或者有一個新的模型，讓它直接跑起來，還是有挑戰性的，所以說可能還真的需要一些最佳化。
Q2：相較於應用層，在模型框架層最佳化過程中，遇到最大的困難點是什麼？
A2：近期全球算力緊張，受限於算力，只能透過和企業合作，幫助企業訓練模型來驗證框架的穩定性，這樣反而不利於這個框架的發展。因此，算力資源對模型框架的發展是非常關鍵的。
Q3：和 deep speed Metro LM 相比，Colossal-AI 最大的技術優勢是什麼？
A3：每個團隊的技術背景都不太一樣，例如 deep speed 主要 做sharding Zero 記憶體最佳化方面的工作，而 Megtron 是英偉達在晶片背景條件下進行的整體最佳化。潞晨科技以前是做高效能計算的，因此會認為 AI 模型的分散式訓練本質上就是資料移動的最佳化，因此核心思路在於通訊的最佳化，從而實現上百、上千個 GPU 具有良好的擴充套件性。由於不同團隊的基礎背景不一樣，因此對同一問題的看法也會不一樣。
以上就是本次分享的內容，謝謝大家。
