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 主頁進行提問，我們免費提供技術咨詢，也有微信群可以交流讨論。
我今天的演講就到這裡，有問題也可以随時聯系我交流，感謝大家關注。
