僅用 480 塊 GPU 跑出萬億參數！全球首個「低碳版」巨模型 M6 來了

繼今年 3 月阿里達摩院釋出國內首個千億參數多模態大模型 M6（MultiModality-to-MultiModality MultitaskMega-transformer，以下簡稱 M6） 之後，6 月 25 日，達摩院宣佈對 M6 進行全新升級，帶來「低碳版」巨模型 M6，在全球範圍內首次大幅降低了萬億參數超大模型訓練能耗，更加符合業界對低碳、高效訓練 AI 大模型的迫切需求。

通過一系列突破性的技術創新，達摩院團隊僅使用 480 卡 GPU，即訓練出了規模達人類神經元 10 倍的萬億參數多模態大模型 M6，與英偉達、谷歌等海外公司實現萬億參數規模相比，能耗降低超八成、效率提升近 11 倍。

在本文中，我們將從萬億參數多模態大模型 M6 所帶來的創新突破為起點，分享其背後所採用的 MoE 架構原理和實現，以及達摩院對 MoE 架構的探索與發現。

何為大模型？

大模型將成下一代人工智慧基礎設施，在 AI 圈內已成共識。與生物體神經元越多往往越聰明類似，參數規模越大的 AI 模型，往往擁有更高的智慧上限，訓練大模型或將讓人類在探索通用人工智慧上更進一步。然而，大模型算力成本也相當高昂，很大程度阻礙了學界、工業界對大模型潛力的深入研究。

對此，達摩院與阿里雲機器學習 PAI 平臺、EFLOPS 計算叢集等團隊改進了 MoE（Mixture-of-Experts）框架，大大擴增了單個模型的承載容量。同時，通過加速線性代數、混合精度訓練、半精度通訊等優化技術，大幅提升了萬億模型訓練速度，且在效果接近無損的前提下有效降低了所需計算資源。

萬億參數的 M6-MoE 模型

大模型研究的一大技術挑戰是，模型擴展到千億及以上參數的超大規模時，將很難放在一臺機器上。如果使用模型+流水並行的分散式策略，一方面在程式碼實現上比較複雜，另一方面由於前向和反向傳播 FLOPs 過高，模型的訓練效率將非常低，在有限的時間內難以訓練足夠的樣本。

為了解決這一難題，達摩院智慧計算實驗室團隊採用了 Mixture-of-Experts（MoE）技術方案，該技術能夠在擴展模型容量並提升模型效果的基礎上，不顯著增加運算 FLOPs，從而實現高效訓練大規模模型的目的。

普通 Transformer 與 MoE 的對比如下圖所示。在經典的資料並行 Transformer 分散式訓練中，各 GPU 上同一 FFN 層使用同一份參數。當使用圖中最右側所示的 MoE 策略時，則不再將這部分參數在 GPU 之間共享，一份 FFN 參數被稱為 1 個 expert，每個 GPU 上將存放若干份參數不同的 experts。

在前向過程中，對於輸入樣本的每個 token，達摩院團隊使用 gate 機制為其選擇分數最高的 k 個 experts，並將其 hiddenstates 通過 all-to-all 通訊傳送到這些 experts 對應的 GPU 上進行 FFN 層計算，待計算完畢後傳送回原 GPU，k 個 experts 的輸出結果根據 gate 分數加權求和，再進行後續運算。為了避免部分 experts 在訓練中接收過多 tokens 從而影響效率，MoE 往往設定一個 capacity 超參指定每個 expert 處理 token 的最大數量，超出 capacity 的 token 將在 FFN 層被丟棄。不同的 GPU 輸入不同的訓練資料分片。通過這種 expert 並行的策略，模型的總參數和容量大大擴增。由於單個樣本經過 gate 稀疏啟用後只使用有限個 experts，每個樣本所需要的計算量並沒有顯著增加，這帶來了突破千億乃至萬億規模的可能性。

在 MoE 模型的具體實現上，谷歌的工作依賴 mesh tensorflow 和 TPU。達摩院則使用阿里雲自研框架 Whale 開發萬億 M6-MoE 模型。將 FFN 層改造成 expert 並行，達摩院主要使用了 Whale 的運算元拆分功能。在實現基本 MoE 策略的基礎上，達摩院團隊進一步整合 Gradient checkpointing、XLA 優化、混合精度訓練、半精度通訊等訓練效率優化技術，並採用了 Adafactor 優化器，成功在 480 張 NVIDIA V100-32GB 上完成萬億模型的訓練。

在訓練中，他們採用絕對值更小的初始化，並且適當減小學習率，保證了訓練的穩定性，實現正常的模型收斂，而訓練速度也達到了約480 samples/s。通過對比1000億、2500億和10000億參數規模的模型收斂曲線（如下圖所示），達摩院團隊發現參數規模越大確實能帶來效果上的進一步提升。

然而，值得注意的是，目前擴參數的方式還是橫向擴展（即增加expert數和intermediate size），而非縱向擴展（即擴層數），未來該團隊也將進一步探索縱向擴展，尋求模型深度與寬度的最優平衡。

探索 MoE，進一步提升模型效果！

除了規模擴展外，達摩院對 MoE 架構開展了更進一步的探索研究，觀察哪些因素對 MoE 模型的效果和效率影響較大。利用 MoE 架構擴大模型規模的一大關鍵是 expert 並行。而在 expert 並行中，幾大因素決定著模型的計算和通訊，包括負載均衡策略，topk 策略及對應的 capacity 設計等。

在 M6 團隊對負載均衡在 MoE 實驗的觀察過程中，他們考慮到負載均衡的問題，通過採用啟發式的方法解決該問題，如上述的 expert capacity 和對應的 residual connection 的方法。Google 的 Gshard 和 SwitchTransformer 沿用了 MoE 原文經典的做法加入了 auxiliaryload balancing loss。

目前還沒有相關工作觀察負載均衡的情況究竟有多嚴重，以及它是不是真的會影響模型的效果。達摩院團隊在小規模的 M6 模型上進行了對 auxiliary loss 的消融實驗，觀察到該 loss 對最終模型效果影響甚微，甚至沒有帶來正向效果，然而它確實對 loadbalance這個問題非常有效。如下圖所示：

上圖彩色曲線線表示各個層的 expert 接收有效 token 的變異係數隨著訓練進行的變化，灰色曲線表明訓練階段的 log PPL。圖中變異係數 CV 表明每一層 expert 負載均衡情況，各曲線表明其隨著訓練步數的變化。

不難發現，訓練初期所有模型均有較嚴重的負載不均衡問題，剛開始少數的 expert 接收了絕大部分的 token，導致很多 token 直接被丟棄，但它們均能實現快速下降，尤其具備 auxiliary loss 的模型 CV 能降低到 0.3 左右，也可觀察到在該水平下均衡程度很高，每個 expert 都能接收大量有效 token。然而與之相反，不加 auxiliary loss 的模型表現非常不同，有的層甚至在訓練後期出現 CV 的飆升。但不管對比訓練階段的 log PPL，還是對比下游語言模型任務的 PPL，不帶 auxiliary loss 的模型都表現更優。這一定程度上反映其實負載均衡對最終效果的影響並不大。

達摩院 M6 團隊進一步探索了關鍵的 top-k gating 策略 k 值和 capacity（C） 的選擇。首先，他們簡單地將 k 值擴大，發現k值越大其實效果越好。但考慮到選用不同的 k 值，C 則對應根據公式進行調整。通過對 C 調整到 k=1 的水平，觀察不同 k 值的 MoE 模型的表現，達摩院團隊觀察到 k 值更大模型依然表現越好，儘管 k 值增加帶來的優勢逐漸不太明顯。

但 k 值的增加根據 Gshard top-2 gating 的實現，除了存在實現層面上一定的冗餘和困難外，迴圈 argmax 的操作也會導致速度變慢。此外，第二個 expert 的行為會受到第一個 expert 的影響，讓訓練和測試存在差異。

達摩院團隊用 expert prototyping 的簡單方式替代，相較 baseline 實現了效果提升，且未顯著增加計算成本。expert prototyping，即將 expert 分成 k 組，在每組中再進行 top-k 的操作（通常採用top-1，便於理解），然後將k組的結果進行組合，也稱之為k top-1。這種方式實現上更直接簡便，並且允許組和組之間並行做top-k操作，更加高效。

達摩院團隊觀察到，在不同規模的模型上，expert prototyping都能取得比baseline更好的效果，同時速度和計算上也相比top-k更有優勢。且其在更大規模的模型上優勢變得更大，在百億模型下游imagecaptioning任務上甚至能觀察到優於top-k的表現：

因此達摩院團隊將該方法推廣到萬億參數M6超大模型，並對應和上述的萬億baseline做了對比。目前，萬億參數模型訓練了大約3萬步，已經顯著優於同等規模的基線模型，呈現約5倍的收斂加速。

沿著這個方向，值得做的工作還有很多：考慮到分組的特性，應當讓組和組之間產生足夠的差異，讓每個組選出來的experts儘可能實現組合的效果等等。

M6：首個實現商業化落地的多模態大模型

隨著萬億參數 M6 模型的落地，阿里達摩院在超大規模預訓練模型領域邁上新的臺階，且 M6 巨模型也成為國內首個實現商業化落地的多模態大模型。

在商業應用層面，M6 擁有超越傳統 AI 的認知和創造能力，它擅長繪畫、寫作、問答，在電商、製造業、文學藝術等諸多領域擁有廣泛應用前景。其中以 AI 領域為例，M6 將作為 AI 助理設計師正式上崗阿里新制造平臺犀牛智造，通過結合潮流趨勢進行快速設計、試穿效果模擬，有望大幅縮短快時尚新款服飾設計週期。隨著實踐經驗的增長，M6 設計的能力還將不斷進化。

結合阿里的電商背景，M6 團隊希望通過 M6 大模型優異的文字到圖片生成能力，和電商領域產業鏈深度融合，挖掘潛在的應用價值。此前 OpenAI DALL·E 生成圖片清晰度達 256×256，M6 則將圖片生成清晰度提升至 1024×1024。

（以上為 M6 生成的高清服裝設計圖的示例）

解放設計師雙手，以下為 M6 參與新款服裝設計的流程圖：

除文生圖外，M6 還可以在工業界直接落地圖生文能力，能夠快速為商品等圖片提供描述文案。該能力目前已在淘寶、支付寶部分業務上試應用。同時，多模態大模型為精準的跨模態搜尋帶來可能。目前M6已建立從文字到圖片的匹配能力，未來，或將建立從文字到視訊內容的認知能力，為搜尋形態帶來變革。

達摩院資深演算法專家楊紅霞表示，「接下來，M6 團隊將繼續把低碳 AI 做到極致，推進應用進一步落地，並探索對通用大模型的理論研究。」

參考文獻

[1] Lepikhin,Dmitry, et al. "Gshard: Scaling giant models with conditional computationand automatic sharding." ICLR, 2021.

[2] Fedus, William, Barret Zoph, and Noam Shazeer. "Switch Transformers:Scaling to Trillion Parameter Models with Simple and Efficient Sparsity."arXiv preprint arXiv:2101.03961 (2021).

[3] Shazeer, Noam, and Mitchell Stern. "Adafactor: Adaptive learning rateswith sublinear memory cost." International Conference on Machine Learning.PMLR, 2018.

[4] Wang, Ang, et al. "Whale: A Unified Distributed TrainingFramework." arXiv preprint arXiv:2011.09208 (2020).

Windows 11 正式官宣：全新 UI、支援安卓 App、應用商店 0 抽成！
提前洩露小米10至尊版，博主賠償100萬；傳位元組跳動實習生刪除GB以下所有機器學習模型；快手宣佈7月將取消大小周|極客頭條
防毒軟體 McAfee 創始人獄中身亡，75 年傳奇人生畫下句號