硬體預算最高直降46倍!低成本上手AIGC和千億大模型,一行程式碼自動並行,Colossal-AI再升級

最近,AI 大模型連續火爆出圈,人工智慧生成模型(AIGC)的熱度尚未褪去,聊天機器人 ChatGPT 便引發全網熱議,兩週吸引百萬使用者。還有卷趴一半程式設計師的 AlphaCode,生成全新蛋白質的 ESM2 等,不斷探索 AI 大模型落地的新領域。面對大模型帶來的技術革命,連Google都拉響 「紅色警報」,擔心 ChatGPT 砸掉自己搜尋引擎的飯碗。

作為當下最火熱的開源 AI 大模型解決方案,Colossal-AI 已收穫 Github Star 七千多顆,此前在 Stable Diffusion、GPT-3、AlphaFold 等大模型上展現卓越性能優勢。針對 AI 大模型落地成本高昂這一痛點,Colossal-AI 本次更新聚焦於降低大模型應用成本以及增強易用性,包括:

  • Stable Diffusion 2.0 低成本訓練 / 微調 / 推理,最多可降低視訊記憶體消耗 5.6 倍,使硬體成本直降 46 倍,一行程式碼即可使用;
  • 1750 億參數 BLOOM 模型單機推理,視訊記憶體消耗降低 4 倍,使硬體成本降低 10 餘倍;
  • 一行程式碼實現自動搜尋最佳並行策略,顯著降低分散式訓練上手門檻,原生支持 Hugging Face,Timm 等熱門 AI 模型庫。

開源地址:

https://github.com/hpcaitech/ColossalAI

1. Stable Diffusion 2.0 低成本訓練 / 微調 / 推理

AIGC(人工智慧生成內容)目前是 AI 領域最熱門的話題之一,還被 Science 評為 2022 年度十大科學突破。短短數月,爆紅 AI 圈的 Stable Diffusion 已升級到 2.0 版本,正如部分使用者所言:「我從未見過任何技術發展如此之快。V1 還沒整透徹,V2 就來了」。

然而,高昂的成本也一定程度上限制了 AIGC 產業的大規模普及。以 Stable Diffusion 為例,其背後的 Stability AI維護超過 4000 個英偉達 A100 的 GPU 集群,並已為此支出超過 5000 萬美元的運營成本。面對快速迭代的模型、演算法和下游任務,如何降低應用成本成為 AIGC 真正走向落地的核心問題。

Stable Diffusion 2.0 基於簡單易用的 PyTorch Lightning 框架搭建。作為 PyTorch Lightning 的官方大模型解決方案,Colossal-AI 在第一時間跟進,開源了更為高效,硬體門檻更低的全套訓練 / 微調 / 推理方案:

  • 可減少訓練視訊記憶體消耗 5.6 倍,硬體成本最高降低 46 倍;
  • 支持 DreamBooth 單 GPU 快速個性化微調;
  • 推理視訊記憶體消耗降低 2.5 倍。

該方案也將於近期合併進入當下最火熱的 AI 模型社區 Hugging Face,進一步方便使用者使用。

1.1 訓練

為了加快訓練速度,降低訓練成本,使用更大的 batch size 已成為被廣泛使用的有效手段。但 GPU 有限的視訊記憶體容量,嚴重限制了 batch size 大小,推高了訓練硬體門檻。

Colossal-AI 通過一系列視訊記憶體最佳化技術和支持 Stable Diffusion 2.0,使 Stable Diffusion 平均在每個 GPU 上使用大 batch size 16 訓練的視訊記憶體需求從 64.5GB 降低 5.6 倍至 11.6GB,還可擴展至單 GPU 或多 GPU 並行。相比使用最先進的 A100 80GB,目前僅需 3060 等消費級顯示卡即可滿足需求,硬體成本最高直降 46 倍。更多使用者可以在消費級 GPU 上,低成本地開展 Stable Diffusion 的相關研究與應用落地。

1.1.1 背後視訊記憶體最佳化
1.1.1 背後視訊記憶體最佳化

1.1.1 背後視訊記憶體最佳化

Flash Attention:早在 Stable Diffusion 1.0 版本,Colossal-AI 就率先引入了 Flash Attention 技術,成功將 attention 的速度提升 104%,將端到端訓練的峰值視訊記憶體減少 23%。Flash Attention 是針對長序列 attention 的加速版本,使用 Flatten 來減少 GPU 高頻寬記憶體(HBM)之間的記憶體讀 / 寫次數,Flash attention 同時針對塊狀稀疏的 attention,設計了一個近似的注意力演算法,比任何現有的近似 attention 方法都要快。在 Stable Diffusion 1.0 版本,整個 Diffusion Model 只有少量 attention 層,Flash attention 還沒有體現出其性能優勢。在 Stable Diffusion 2.0 中,由於將大量卷積層替換為 attention 層,進一步發揮了 Flash Attention 的視訊記憶體最佳化潛力。

ZeRO + Gemini:Colossal-AI 支持使用零冗餘最佳化器 (ZeRO) 的方法來消除記憶體冗餘,與經典的資料並行性策略相比,可極大提高記憶體使用效率,同時不犧牲計算粒度和通訊效率。此外,Colossal-AI 還引入了 Chunk 機制進一步提升 ZeRO 的性能。運算順序上連續的一組參數存入一個 Chunk 中(Chunk 即一段連續的記憶體空間),每個 Chunk 的大小相同。Chunk 方式組織記憶體可以保證 PCI-e 和 GPU-GPU 之間網路頻寬的高效利用,減小了通訊次數,同時避免潛在的記憶體碎片。

而 Colossal-AI 的異構記憶體空間管理器 Gemini 支持將最佳化器狀態從 GPU 卸載到 CPU,以節省 GPU 記憶體佔用。可以同時利用 GPU 記憶體、CPU 記憶體(由 CPU DRAM 或 NVMe SSD 記憶體組成)來突破單 GPU 記憶體牆的限制,進一步擴展了可訓練模型規模。

1.1.2 一行程式碼快速上手

作為 PyTorch Lightning 的官方合作伙伴,僅需一行程式碼即可調用 ColossalAI 的上述視訊記憶體最佳化。

from lightning.pytorch import trainer, LightningModulefrom lightning.pytorch.strategies import ColossalAIStrategyMystrategy = ColossalAIStrategy(use_chunk=True,enable_distributed_storage=True, placement_policy=auto)trainer = Trainer(accelerator="gpu", devices=4, precision=16,strategy=Mystrategy)trainer.fit(model)

1.2 DreamBooth 微調

DreamBooth 是一種用於個性化文字到圖像模型(例如 Stable Diffusion)的方法,只需要幾張(3~5 張)期望主題的圖像,即可生成指定主題的一系列圖像。使用者只需直接運行檔案 train_dreambooth_colossalai.py,即可在該微調任務上充分發揮 Colossal-AI 的視訊記憶體最佳化,個性化快速微調自己的圖文模型,極大地降低了使用門檻。

1.3 推理

1.3 推理

由於模型推理對數值精度不敏感,這為實現低精度的低成本推理提供了可能。對於 Stable Diffusion 2.0 模型,可以通過添加一行程式碼,支持模型的 Int8 量化推理,視訊記憶體消耗降低 2.5 倍,僅需 3.1GB,且不造成顯著性能損失。

model = replace_module(model)

2. 1750 億 BLOOM 模型低成本推理

隨著模型規模的不斷增大,模型推理的視訊記憶體及記憶體佔用也變成了不可忽視的關鍵因素。以 Hugging Face 發佈的 1750 億參數開源模型 BLOOM 為例,如果直接使用常見的 FP32/FP16 進行推理,在單節點 8 張 GPU 使用模型並行,每張 GPU 需要消耗至少 87.5GB/43.8GB 的視訊記憶體。如此大的視訊記憶體佔用,即使是最先進的 8 卡 A100(80GB/40GB)伺服器,也無法直接部署推理服務,而多節點推理又會帶來沉重的額外成本和通訊開銷。

Colossal-AI 實現了高效的 Int8 量化和模型並行推理,可以將 1750 億參數的 BLOOM 等大模型的推理服務,部署到 3090/4090 等消費級顯示卡的 8 卡伺服器,同時不產生顯著的 CPU 記憶體佔用提升及性能損耗。相比原有的 A100 方案,可降低硬體部署成本十餘倍

在 Colossal-AI 中,通過對模型進行了 Int8 量化,可將模型總體視訊記憶體佔用從 352.3GB(FP16)降低到 185.6GB, 同時使用 Colossal-AI 的模型並行技術,將每張顯示卡的佔用減少到了 23.2GB。在模型並行中,為了不增加 CPU 記憶體佔用,Colossal-AI 在主進程中進行模型的量化和切分工作,其餘的每個進程中分別使用 lazy_init 獲得幾乎不佔視訊記憶體和記憶體的 meta model,再通過 gloo backend 在進程之間傳遞模型參數。通過上述方案,在沒有分段載入模型參數的情況下,便可以使得 CPU 記憶體佔用峰值,達到理論上的較優水平。相較於將模型按層切分的 「類流水線」 分佈方式,模型並行可以提高非密集請求下的視訊記憶體使用效率。

3. 一行程式碼自動並行

大模型的分散式混合部署是一個非常複雜的問題,目前常見的分散式大模型訓練方案,都依賴使用者人工反覆嘗試以及系統專家的經驗來進行配置部署。然而這對絕大多數 AI 開發者來說十分不友好,因為不希望將時間精力花費在研究分散式系統和試錯上。

Colossal-AI 提供了業界急需的一套高效易用自動並行系統。相比現有其他手動配置複雜並行策略和修改模型的解決方案,Colossal-AI僅需增加一行程式碼,提供 cluster 資訊以及單機訓練模型即可獲得分散式訓練能力,並且原生支持包括 Hugging Face,Timm 等熱門 AI 模型庫。

# wrap the model using auto_enginemodel, optimizer = auto_engine(model, optimizer, cluster_info)# normal training loop...
因此,Colossal-AI 可以

因此,Colossal-AI 可以極大地降低 AI 開發者的使用分散式技術訓練和微調大模型門檻。同時,自動並行系統可以從更細粒度搜尋並行策略,找到更加高效的並行方案

3.1 Graph Tracing

Colossal-AI 是首個基於 PyTorch 框架使用靜態圖分析的自動並行系統。PyTorch 作為一個動態圖框架,獲取其靜態的執行計劃是機器學習系統領域被長期研究的問題。Colossal-AI 使用基於 torch.FX Tracer 的 ColoTracer,在 tracing 過程中推導並記錄了每個 tensor 的元資訊,例如 tensor shape,dims,dtype 等,可以為後續的自動並行策略搜尋提供幫助。因此 Colossal-AI 具有更好的模型泛化能力,而不是依靠模型名或手動修改來適配並行策略。

3.2 細粒度分散式訓練策略搜尋

Colossal-AI 會在滿足記憶體預算的限制下,以最快運行時間為目標,為每個 op 進行策略搜尋,最終得到真實訓練時的策略,包括每個 tensor 的切分策略,不同計算節點間需要插入的通訊運算元類型,是否要進行運算元替換等。現有系統中的張量並行,資料並行,NVIDIA 在 Megatron-LM 等並行系統中使用的 column 切分和 row 切分並行等混合並行,都是自動並行可以搜尋到的策略的子集。除了這些可以手動指定的並行方式外,Colossal-AI 的自動並行系統有能力為每個 op 指定獨特的並行方式,因此有可能找到比依賴專家經驗和試錯配置的手動切分更好的並行策略。

3.3 分散式 tensor 與 shape consistency 系統

與 PyTorch 最新發布的 DTensor 類似,Colossal-AI 也使用了 device mesh 對集群進行了抽象管理。具體來說,Colossal-AI 使用 sharding spec 對 tensor 的分散式儲存狀態進行標註,使用 shape consistency manager 自動地對同一 tensor 在不同 sharding spec 間進行轉換。這讓 Colossal-AI 的通用性和易用性極大地提升,藉助 shape consistency manager 可以沒有負擔地切分 tensor,而不用擔心上游 op 的 output 與下游的 input 在集群中的儲存方式不同。

相較於 PyTorch DTensor,Colossal-AI 有以下 3 個優勢:

  1. Colossal-AI 的 device mesh 可以 profiling 到集群性能指標,對不同的通訊運算元進行耗時估算。
  2. Colossal-AI 的 shape consistency 會貪心地搜尋 sharding spec 間的轉換方式,而不是樸素地逐 dimension 進行轉換,這樣能找到更高效的轉換路徑,進而使得 sharding spec 間的轉換通訊開銷更小。
  3. 加入了 all_to_all 操作,使得 Colossal-AI 的擴展性更強,這在大規模集群上進行訓練時,可以展現出很大的優勢。

3.4 與 activation checkpoint 結合

作為大模型訓練中必不可少的視訊記憶體壓縮技術,Colossal-AI 也提供了對於 activation checkpoint 的自動搜尋功能。相比於大部分將最大視訊記憶體壓縮作為目標的技術方案,Colossal-AI 的搜尋目標是在視訊記憶體預算以內,找到最快的 activation checkpoint 方案。同時,為了避免將 activation checkpoint 的搜尋一起建模到 SPMD solver 中導致搜尋時間爆炸,Colossal-AI 做了 2-stage search 的設計,因此可以在合理的時間內搜尋到有效可行的分散式訓練方案。

4. 關於 Colossal-AI

4. 關於 Colossal-AI

4. 關於 Colossal-AI

面向大模型時代的通用深度學習系統 Colossal-AI,可實現高效快速部署 AI 大模型訓練和推理,降低 AI 大模型應用成本。自開源以來,Colossal-AI 已經多次在 GitHub 熱榜位列世界第一,獲得 GitHub Star 超七千顆,併成功入選 SC、AAAI、PPoPP 等國際 AI 與 HPC 頂級會議的官方教程。

Colossal-AI 相關解決方案已成功在自動駕駛、雲端運算、零售、醫藥、晶片等行業知名廠商落地應用,廣受好評。例如,近期火爆的 ChatGPT 尚未開源,且不具備聯網功能。Colossal-AI 已成功幫助某世界 500 強企業,開發具備線上搜尋引擎能力增強的聊天機器人模型。

開源地址:

https://github.com/hpcaitech/ColossalAI

參考連結:

https://www.hpc-ai.tech/blog/colossal-ai-0-2-0

相關文章