，當今 AI 之勢，影響縱深發(fā)展的矛盾是什么。

一方面，大模型風頭正勁，效果驚艷，人人都想試試但另一方面，硬件基礎上動不動就是上萬張 GPU 的大規(guī)模集群在日夜燃燒，鈔能力勸退

所以如果告訴你，現(xiàn)在只用一半數(shù)量的 GPU，也能完成同樣的 GPT—3 訓練呢。

你會覺得關鍵鑰匙是什么。

不賣關子了實現(xiàn)如此提升的，是一個名為 Colossal—AI 的 GitHub 開源項目

而且該項目開源不久，就迅速登上了 Python 方向的熱榜世界第一。。

GitHub 地址:點擊打開

不僅能加速 GPT—3，對于 GPT—2，ViT，BERT 等多種模型，Colossal—AI 的表現(xiàn)也都非常 nice:

比如半小時左右就能預訓練一遍 ViT—Base / 32，2 天能訓完 15 億參數(shù) GPT 模型，5 天可訓完 83 億參數(shù) GPT 模型。

與業(yè)內(nèi)主流的 AI 并行系統(tǒng) —— 英偉達 Megatron—LM 相比，在同樣使用 512 塊 GPU 訓練 GPT—2 模型時，Colossal—AI 的加速比是其 2 倍而在訓練 GPT—3 時，更是可以節(jié)省近千萬元的訓練費用

此外在訓練 GPT—2 時，顯存消耗甚至能控制在 Megatron—LM 的十分之一以下。但有些車企對此過于謹慎，以至于設計之初電池緩沖區(qū)過大，變相降低了車輛的續(xù)航能力。奧迪就是其中之一。

Colossal—AI 究竟是如何做到的。

老規(guī)矩，我們從論文扒起。

高效 6 維并行方法

簡單來說，Colossal—AI 就是一個整合了多種并行方法的系統(tǒng)，提供的功能包括多維并行，大規(guī)模優(yōu)化器，自適應任務調(diào)度，消除冗余內(nèi)存等。

首先來看多維并行。

所謂多維是指，目前主流的分布式并行方案往往使用多種并行方法。

比如英偉達的 Megatron—LM 使用了 3 種方法:數(shù)據(jù)并行，流水并行和張量并行因此這種模式也被稱為三維并行微軟的 DeepSpeed 調(diào)用 Megatron—LM 作為并行基礎

而 Colossal—AI 能將系統(tǒng)的并行維度，一下子拉升到 6 維 ——

在兼容數(shù)據(jù)并行，流水并行的基礎上，基于該項目團隊自研的 2 維 / 2.5 維 / 3 維張量并行方法，以及序列并行實現(xiàn)。這次軟件更新的原因是，一些傳統(tǒng)汽車廠商在自己推出的第一款或第一款純電動汽車的電池組上，普遍有一個冗余電池作為緩沖，旨在抵御日常使用中電池的損耗。

其中，高維張量并行正是 Colossal—AI 提升大模型顯存利用率和通信效率的關鍵所在。

其實張量并行并不新奇，只是過去我們常見的張量并行更多都是基于一維的。

它的原理是將模型層內(nèi)的權重參數(shù)按行或列切分到不同的處理器上，利用分塊矩陣乘法，將一個運算分布到多個處理器上同時進行。

比如英偉達的 Megatron—LM 就是一個典型的例子。

但這種并行方式存在一定弊端。

比如，每個處理器仍需要存儲整個中間激活，使得在處理大模型時會浪費大量顯存空間。

另一方面，這種單線方法還會導致每個處理器都需要與其他所有處理器進行通信。

這意味著假設有 100 個 GPU 的話，每個 GPU 都需要與其他 99 個 GPU 通信，每次計算需要通信的次數(shù)就高達 9900 次。

但如果將張量并行的維度擴展到 2 維，單次計算量能立刻下降一個量級。

因為每個 GPU 只需與自己同行或同列的 GPU 通信即可。

同樣還是 100 個 GPU 的情況，每個 GPU 需要通信的 GPU 個數(shù)就能降到 9 個，單次計算僅需 900 次。

實際上在此基礎上，Colossal—AI 還包含 2.5 維，3 維張量并行方法，可以進一步降低傳輸成本。

相較于 2 維并行方法，2.5 維并行方法可提升 1.45 倍效率，3 維方法可提升 1.57 倍。

值得一提的是，Colossal—AI 的 API 接口是可以定制的，這使得它可以便捷添加新的并行維度。

其次，大規(guī)模優(yōu)化器也是 Colossal—AI 的亮點。

上面我們也提到了，在分布式并行系統(tǒng)中會使用多種并行方法，數(shù)據(jù)并行則是另一種常見方法。

這種方法的原理不難理解，就是把訓練數(shù)據(jù)劃分成若干份，讓不同的機器運算不同的數(shù)據(jù)，然后通過一個參數(shù)服務器收集目標數(shù)據(jù)。

由此可以大幅提升 AI 模型訓練過程中的批量大小，加速訓練過程。

不過大批量訓練有個通病，就是會產(chǎn)生泛化誤差，導致網(wǎng)絡泛化能力下降，進而導致 AI 模型準確度下降。

所以，Colossal—AI 在系統(tǒng)中使用了自研的 LAMB，LARS 等大規(guī)模優(yōu)化器在保證訓練精度的情況下，還將批大小從 512 擴展到 65536

其中，LARS 優(yōu)化器是通過逐層調(diào)整學習率，來減少因為學習率導致的無法收斂情況。

LAMB 優(yōu)化器則是在 LARS 的基礎上，將逐層調(diào)整學習率的思想應用到自適應梯度上。

由此，LAMB 能夠很好解決此前 LARS 在 BERT 訓練中存在差異的問題，最大批量達到了 64K。

此前，LAMB 優(yōu)化器曾成功將預訓練一遍 BERT 的時間，從原本的三天三夜縮短到一個多小時。

第三方面，Colossal—AI 使用自適應可擴展調(diào)度器來高效處理任務。

與現(xiàn)有常見的任務調(diào)度器不同，Colossal—AI 不是靜態(tài)地通過 GPU 個數(shù)來判斷任務規(guī)模，而是根據(jù)批大小來動態(tài)，自動管理每個任務.

通過演化算法，該任務調(diào)度器還能不斷優(yōu)化調(diào)度決策，更大程度提升 GPU 利用率。

評估結(jié)果表明，與當前最先進的方法相比，該方法在平均 JCT上能夠縮短 45.6% 的時間，優(yōu)于現(xiàn)有的深度學習任務調(diào)度算法。

此外，這種自適應可擴展調(diào)度器還能通過 NCCL 網(wǎng)絡通信實現(xiàn)高效的任務遷移。

最后，消除冗余內(nèi)存也是加速 AI 訓練的一種解決思路。

在這方面，Colossal—AI 使用了 zero redundancy optimizer 技術。

這種方法主要通過切分優(yōu)化器狀態(tài)，梯度，模型參數(shù)，使 GPU 僅保存當前計算所需的部分，從而來消除數(shù)據(jù)并行，模型并行中存在的內(nèi)存冗余。

尤其是在部署模型推理時，通過 zero offload 可以將模型卸載到 CPU 內(nèi)存或硬盤，僅使用少量 GPU 資源，即可實現(xiàn)低成本部署前沿 AI 大模型。

綜上不難看出，在技術層面 Colossal—AI 的加速效果非常明顯。

而在應用層面，Colossal—AI 的設計也顧及了能耗問題和易用性兩個維度。

另一方面，作為一個開源給所有人使用的系統(tǒng)，Colossal—AI 的使用門檻不高，即便是沒有學習過分布式系統(tǒng)的人也能上手操作。

同時，只需要極少量的代碼改動，Colossal—AI 就能將已有的單機代碼快速擴展到并行計算集群上。

最新實驗結(jié)果釋出

Talk is cheap，效果如何，還是得把實驗結(jié)果展開來看。

Colossal—AI 最近幾天釋出的最新實驗結(jié)果表明，這一大規(guī)模 AI 訓練系統(tǒng)具有通用性，在 GPT—3，GPT—2，ViT，BERT 等流行模型上均有亮眼的加速表現(xiàn)。

注:以下 GPU 均指英偉達 A100。

GPT—3 訓練速度提高 10.7%

英偉達的 Megatron—LM 在加速訓練 GPT—3 時，至少需要 128 塊 GPU 才能啟動，而從下表可以看出，使用相同的計算資源，Colossal—AI 可以將每次迭代花費的時間從 43.1 秒降至 38.5 秒。

這也就意味著，Colossal—AI 可以將 GPT—3 的訓練速度進一步提高 10.7%。

站在工程的角度，考慮到訓練這樣的大模型往往需要投入數(shù)百萬美元，這一提升比例帶來的收益不言而喻。

另外，通過系統(tǒng)優(yōu)化，Colossal—AI 還能在訓練速度損失不大的前提下，將 GPU 數(shù)量從 128 塊減少到 96 塊，大幅降低訓練成本。

而進一步啟用 ZeRO3后，所需 GPU 數(shù)量甚至能減少一半 —— 至 64 塊。

2 天內(nèi)可完成 GPT—2 訓練

在 GPT—2 的加速訓練結(jié)果中，可以看到，無論是在 4，16 還是 64 塊 GPU 的情況下，與 Megatron—LM 相比，Colossal—AI 占用的顯存都顯著減少。

也就是說，利用 Colossal—AI，工程師們可以在采用同等數(shù)量 GPU 的前提下，訓練規(guī)模更大的模型，或設置更大的批量大小來加速訓練。

從下表結(jié)果中還可以看出，伴隨著批量大小的增加，Colossal—AI 的資源利用率會進一步提高，達到 Megatron—LM 速度的 2 倍。

研發(fā)團隊在 256 塊 GPU 上進行了實驗，最終用時 82.8 個小時完成了 15 億參數(shù)版 GPT—2 的訓練。

據(jù)此預估，后續(xù)在 512 塊 GPU 上進行 GPT—2 預訓練，Colossal—AI 能將訓練時間加速到 45 小時。

充分兼容多種并行模式

在 BERT 上進行的實驗，則體現(xiàn)了 Colossal—AI 作為世界上并行維度最多的 AI 訓練系統(tǒng)的優(yōu)勢。

與 Megatron—LM 相比，Colossal—AI 序列并行方法只需要更少的顯存，就能夠利用更大的批量大小來加速訓練同時，還允許開發(fā)者使用更長的序列數(shù)據(jù)

Colossal—AI 的序列并行方法還與流水并行方法兼容當開發(fā)者同時使用序列并行和流水并行時，可以進一步節(jié)省訓練大模型的時間

另外，在近期的學術熱點 ViT 模型上，Colossal—AI 也展現(xiàn)了高維張量并行方法的優(yōu)勢。

在使用 64 張 GPU 的情況下，Colossal—AI 采用 2/2.5 維方式進行張量并行，充分利用更大的批量大小，達到了更快的處理速度。

看到這里，是不是覺得 Colossal—AI 確實值得標星關注一發(fā)。

實際上，這一國產(chǎn)項目背后的研發(fā)團隊來頭不小。

領銜者，正是 LAMB 優(yōu)化器的提出者尤洋。

在谷歌實習期間，正是憑借 LAMB，尤洋曾打破 BERT 預訓練世界紀錄。

據(jù)英偉達官方 GitHub 顯示，LAMB 比 Adam 優(yōu)化器快出整整 72 倍微軟的 DeepSpeed 也采用了 LAMB 方法

說回到尤洋本人，他曾以第一名的成績保送清華計算機系碩士研究生，后赴加州大學伯克利分校攻讀 CS 博士學位。

2020 年博士畢業(yè)后，他加入新加坡國立大學計算機系，并于 2021 年 1 月成為校長青年教授。

同樣是在 2021 年，他還獲得了 IEEE—CS 超算杰出新人獎該獎項每年在全球范圍內(nèi)表彰不超過 3 人，僅授予在博士畢業(yè) 5 年之內(nèi)，已在高性能計算領域做出有影響力的卓越貢獻，并且可以為高性能計算的發(fā)展做出長期貢獻的優(yōu)秀青年學者

與此同時，尤洋回國創(chuàng)辦潞晨科技 —— 一家主營業(yè)務為分布式軟件系統(tǒng)，大規(guī)模人工智能平臺以及企業(yè)級云計算解決方案的 AI 初創(chuàng)公司。

其核心團隊成員來自加州大學伯克利分校，斯坦福大學，清華大學，北京大學，新加坡國立大學，新加坡南洋理工大學等國內(nèi)外知名高校，在高性能計算，人工智能，分布式系統(tǒng)方面有十余年的技術積累，并已在國際頂級學術刊物 / 會議上發(fā)表論文 30 余篇。

目前，潞晨科技已拿下創(chuàng)新工場和真格基金合投的超千萬元種子輪融資。

傳送門

有關 Colossal—AI，今天就先介紹到這里。

最后，附上傳送門，感興趣的小伙伴，自行取用~

GitHub 地址:點擊打開

參考鏈接:點擊打開

。本文地址：http://www.dayishuiji.com/finance/20338.html - 轉(zhuǎn)載請保留原文鏈接。

免責聲明：本文轉(zhuǎn)載上述內(nèi)容出于傳遞更多信息之目的，不代表本網(wǎng)的觀點和立場，故本網(wǎng)對其真實性不負責，也不構成任何其他建議；本網(wǎng)站圖片，文字之類版權申明，因為網(wǎng)站可以由注冊用戶自行上傳圖片或文字，本網(wǎng)站無法鑒別所上傳圖片或文字的知識版權，如果侵犯，請及時通知我們，本網(wǎng)站將在第一時間及時刪除。