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

語言: CN / TW / HK
時間 2023-01-04 16:03:23
最近，AI大模型連續火爆出圈，人工智慧生成模型（AIGC）的熱度尚未褪去，聊天機器人ChatGPT便引發全網熱議，兩週吸引百萬使用者。還有卷趴一半程式設計師的AlphaCode，生成全新蛋白質的ESM2等，不斷探索AI大模型落地的新領域。面對大模型帶來的技術革命，連谷歌都拉響“紅色警報”，擔心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 模型庫。
 

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，進一步方便使用者使用。

訓練

為了加快訓練速度，降低訓練成本，使用更大的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的相關研究與應用落地。
 
     
     
 

背後視訊記憶體優化

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記憶體牆的限制，進一步擴充套件了可訓練模型規模。

一行程式碼快速上手

作為PyTorch Lightning的官方合作伙伴， 僅需一行程式碼即可呼叫Colossal-AI的上述視訊記憶體優化。
  
from lightning.pytorch import trainer, LightningModule
from lightning.pytorch.strategies import ColossalAIStrategy

Mystrategy = ColossalAIStrategy(use_chunk=True, enable_distributed_storage=True, placement_policy=auto)

trainer = Trainer(accelerator="gpu", devices=4, precision=16, strategy=Mystrategy)
trainer.fit(model)

DreamBooth微調

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

推理

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

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記憶體佔用峰值，達到理論上的較優水平。相較於將模型按層切分的“類流水線”分佈方式，模型並行可以提高非密集請求下的視訊記憶體使用效率。
 

一行程式碼自動並行

大模型的分散式混合部署是一個非常複雜的問題，目前常見的分散式大模型訓練方案，都依賴使用者人工反覆嘗試以及系統專家的經驗來進行配置部署。然而這對絕大多數AI開發者來說十分不友好，因為不希望將時間精力花費在研究分散式系統和試錯上。
 
Colossal-AI提供了業界急需的一套高效易用 自動並行系統。相比現有其他手動配置複雜並行策略和修改模型的解決方案，Colossal-AI 僅需增加一行程式碼，提供cluster資訊以及單機訓練模型即可獲得分散式訓練能力，並且 原生支援包括Hugging Face，Timm等熱門AI模型庫。
  
# wrap the model using auto_engine
model, optimizer = auto_engine(model, optimizer, cluster_info)
# normal training loop
...
因此，Colossal-AI可以 極大地降低AI開發者的使用分散式技術訓練和微調大模型門檻。同時，自動並行系統可以從更細粒度搜索並行策略， 找到更加高效的並行方案。
 
     

Graph Tracing

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

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

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

分散式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的擴充套件性更強，這在大規模叢集上進行訓練時，可以展現出很大的優勢。

與activation checkpoint結合

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

     

關於Colossal-AI

     


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

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