AIGC落地門檻被打下來了：硬件預算一口氣降至1/46

明敏 發(fā)自 凹非寺量子位 | 公眾號 QbitAI

從AI畫畫到NLP大模型，AIGC的落地成本被一次性打下來了！

話不多說，直接看結果：

• Stable Diffusion 2.0訓練/微調/推理，顯存消耗最多可節(jié)省5.6倍，使硬件成本直降至1/46，一行代碼即可啟用；

• 1750億參數(shù)大模型BLOOM單機推理，顯存消耗節(jié)省4倍，硬件成本直降至十幾分之一。

• 一行代碼實現(xiàn)自動搜索最佳并行策略，顯著降低分布式訓練上手門檻，原生支持Hugging Face、Timm等熱門AI模型庫。

要知道，在AIGC爆火的另一面，居高不下的成本困擾著整個行業(yè)。

上周，首批AI畫畫公司之一StockAI被迫宣布關閉平臺。原因無他，創(chuàng)始人表示：

公司驅動成本太高了，目前的收入難以為繼。

即便如ChatGPT身后有OpenAI和微軟撐腰，也在平臺開放幾周后發(fā)出公告，開始限制用戶每日使用次數(shù)。

言下之意無非四個字：燒不起了。

總而言之，降低AI大模型落地成本，是目前行業(yè)內亟需解決的問題。

與此同時，開源AI大模型解決方案Colossal-AI在過去一年迅速躥紅，GitHub上已經(jīng)收獲7k+Star。

如上降本方案，便出自它之手。

具體是怎么實現(xiàn)的？往下看，

開源地址：https://github.com/hpcaitech/ColossalAI

相較于1.0版本，Stable Diffusion 2.0不僅提升了模型生成圖像像素，還引入了Depth2img模型、text-guided修復模型等，功能更加完善。

這波上新其實讓用戶們既驚喜又措手不及。

（畢竟1.0都還沒玩明白呢）

但話說回來，還是老問題，AIGC模型落地的成本高啊。

以Stable Diffusion為例，其背后的Stability AI維護超過 4000 個英偉達 A100 的 GPU 集群，并已為此支出超過 5000 萬美元的運營成本。

面對快速迭代的模型、算法和下游任務，如何降低應用成本成為AIGC真正走向落地的核心問題。

Stable Diffusion 2.0基于簡單易用的PyTorch Lightning框架搭建。

作為PyTorch Lightning的官方大模型解決方案，Colossal-AI第一時間進行跟進。

具體內容有以下幾點：

• 顯存消耗可節(jié)省5.6倍，硬件成本最多降至1/46

• 支持DreamBooth單GPU快速個性化微調

• 推理顯存消耗節(jié)省2.5倍

訓練

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

通過一系列顯存優(yōu)化技術，Colossal-AI使得Stable Diffusion平均在每個GPU上使用大batch size 16訓練的顯存需求，從64.5GB降低到了11.6GB、節(jié)省5.6倍，還可擴展至單GPU或多GPU并行。

相比使用最先進的A100 80GB，目前僅需3060等消費級顯卡即可滿足需求，硬件成本最高直降至1/46。

由此可以讓更多用戶在消費級GPU上，低成本地開展Stable Diffusion的相關研究與應用落地。

背后顯存優(yōu)化

Flash Attention

早在Stable Diffusion 1.0版本，Colossal-AI就率先引入了Flash Attention技術，成功將attention的速度提升 104%，將端到端訓練的峰值顯存減少 23%。

Flash Attention是針對長序列 attention 的加速版本，使用 Flatten 來減少 GPU 高帶寬內存（HBM）之間的內存讀 / 寫次數(shù)，F(xiàn)lash attention 同時針對塊狀稀疏的 attention，設計了一個近似的注意力算法，比任何現(xiàn)有的近似 attention 方法都要快。

在Stable Diffusion 1.0版本，整個Diffusion Model只有少量attention層，F(xiàn)lash attention還沒有體現(xiàn)出其性能優(yōu)勢。

在Stable Diffusion 2.0中，由于將大量卷積層替換為attention層，進一步發(fā)揮了Flash Attention的顯存優(yōu)化潛力。

ZeRO + Gemini

Colossal-AI支持使用零冗余優(yōu)化器（ZeRO）的方法來消除內存冗余，與經(jīng)典的數(shù)據(jù)并行性策略相比，可極大提高內存使用效率，同時不犧牲計算粒度和通信效率。

此外，Colossal-AI 還引入了Chunk機制進一步提升ZeRO的性能。

運算順序上連續(xù)的一組參數(shù)存入一個Chunk中（Chunk即一段連續(xù)的內存空間），每個Chunk的大小相同。

Chunk方式組織內存可以保證PCI-e和GPU-GPU之間網(wǎng)絡帶寬的高效利用，減小了通信次數(shù)，同時避免潛在的內存碎片。

而Colossal-AI的異構內存空間管理器Gemini支持將優(yōu)化器狀態(tài)從GPU卸載到CPU，以節(jié)省GPU內存占用。

可以同時利用GPU內存、CPU內存（由CPU DRAM或NVMe SSD內存組成）來突破單GPU內存墻的限制，進一步擴展了可訓練模型規(guī)模。

一行代碼快速上手

作為PyTorch Lightning的官方合作伙伴，僅需一行代碼即可調用Colossal-AI的上述顯存優(yōu)化。

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)

DreamBooth微調

在推出Stable Diffusion 2.0加速方案的同時，Colossal-AI還“順手”發(fā)布了DreamBooth模型的微調方案。

這是谷歌在今年8月發(fā)布的模型。它只需3-5張圖片，再加上文字表述，就能讓指定物體遷移到其他場景或風格中去。

和Dall-E 2、Imagen等最大的不同是，DreamBooth能對選定對象忠實還原。

方案中，用戶只需直接運行文件train_dreambooth_colossalai.py，即可在該微調任務上充分發(fā)揮Colossal-AI的顯存優(yōu)化，個性化快速微調自己的圖文模型，極大降低使用門檻。

推理

由于模型推理對數(shù)值精度不敏感，這為實現(xiàn)低精度的低成本推理提供了可能。

對于Stable Diffusion 2.0模型，可以通過添加一行代碼，支持模型的Int8量化推理，顯存消耗節(jié)省2.5倍，僅需3.1GB，且不造成顯著性能損失。

model = replace_module(model)

AI畫畫爆火的另一邊，NLP大模型的趨勢也還在延續(xù)。

今年7月，Hugging Face發(fā)布了1750億參數(shù)開源模型BLOOM，它動用384塊A100煉成。

如果直接使用常見的FP32/FP16進行推理，在單節(jié)點8張GPU使用模型并行，每張GPU需要消耗至少87.5GB/43.8GB的顯存。

如此大的顯存占用，即使是最先進的8卡A100（80GB/40GB）服務器，也無法直接部署推理服務，而多節(jié)點推理又會帶來沉重的額外成本和通信開銷。

基于這一現(xiàn)狀，Colossal-AI實現(xiàn)了高效的Int8量化和模型并行推理，可以將1750億參數(shù)的BLOOM等大模型的推理服務，部署到3090/4090等消費級顯卡的8卡服務器，同時不產生顯著的CPU內存占用提升及性能損耗。

相比原有的A100方案，可將硬件部署成本降低到原有的十幾分之一。

通過對模型進行Int8量化，Colossal-AI可將模型總體顯存占用從352.3GB（FP16）降低到185.6GB， 同時使用Colossal-AI的模型并行技術，將每張顯卡的占用減少到了23.2GB。

在模型并行中，為了不增加CPU內存占用，Colossal-AI在主進程中進行模型的量化和切分工作，其余的每個進程中分別使用lazy_init獲得幾乎不占顯存和內存的meta model，再通過gloo backend在進程之間傳遞模型參數(shù)。

通過上述方案，在沒有分段加載模型參數(shù)的情況下，便可以使得CPU內存占用峰值，達到理論上的較優(yōu)水平。相較于將模型按層切分的“類流水線”分布方式，模型并行可以提高非密集請求下的顯存使用效率。

大模型的分布式混合部署是一個非常復雜的問題。

目前常見的分布式大模型訓練方案，都依賴于用戶人工反復嘗試以及系統(tǒng)專家基于經(jīng)驗進行配置部署。

然而，這對于絕大多數(shù)AI開發(fā)者來說很不友好，因為大家都不希望把過多時間精力花費在研究分布式系統(tǒng)和試錯上。

由此，Colossal-AI的高效易用自動并行系統(tǒng)，可以說是解大家燃眉之急了。

僅需增加一行代碼，它就能提供cluster信息以及單機訓練模型即可獲得分布式訓練能力，并且原生支持包括Hugging Face，Timm等熱門AI模型庫。

# wrap the model using auto_enginemodel, optimizer = auto_engine(model, optimizer, cluster_info)# normal training loop...

因此，Colossal-AI可以極大地降低AI開發(fā)者的使用分布式技術訓練和微調大模型門檻。同時，自動并行系統(tǒng)可以從更細粒度搜索并行策略，找到更加高效的并行方案。

Graph Tracing

Colossal-AI是首個基于PyTorch框架使用靜態(tài)圖分析的自動并行系統(tǒng)。

PyTorch作為一個動態(tài)圖框架，獲取其靜態(tài)的執(zhí)行計劃是機器學習系統(tǒng)領域被長期研究的問題。

Colossal-AI使用基于torch.FX Tracer的ColoTracer，在tracing過程中推導并記錄了每個tensor的元信息，例如tensor shape，dims，dtype等，可以為后續(xù)的自動并行策略搜索提供幫助。

因此Colossal-AI具有更好的模型泛化能力，而不是依靠模型名或手動修改來適配并行策略。

細粒度分布式訓練策略搜索

Colossal-AI會在滿足內存預算的限制下，以最快運行時間為目標，為每個op進行策略搜索，最終得到真實訓練時的策略，包括每個tensor的切分策略，不同計算節(jié)點間需要插入的通信算子類型，是否要進行算子替換等。

現(xiàn)有系統(tǒng)中的張量并行、數(shù)據(jù)并行，NVIDIA在Megatron-LM等并行系統(tǒng)中使用的column切分和row切分并行等混合并行，都是自動并行可以搜索到的策略的子集。

除了這些可以手動指定的并行方式外，Colossal-AI的自動并行系統(tǒng)有能力為每個op指定獨特的并行方式，因此有可能找到比依賴專家經(jīng)驗和試錯配置的手動切分更好的并行策略。

分布式tensor與shape consistency系統(tǒng)

與PyTorch最新發(fā)布的DTensor類似，Colossal-AI也使用了device mesh對集群進行了抽象管理。

具體來說，Colossal-AI使用sharding spec對tensor的分布式存儲狀態(tài)進行標注，使用shape consistency manager自動地對同一tensor在不同sharding spec間進行轉換。

這讓Colossal-AI的通用性和易用性極大地提升，借助shape consistency manager可以沒有負擔地切分tensor，而不用擔心上游op的output與下游的input在集群中的存儲方式不同。

相較于PyTorch DTensor，Colossal-AI有以下3個優(yōu)勢：

• Colossal-AI的device mesh可以profiling到集群性能指標，對不同的通信算子進行耗時估算。

• Colossal-AI的shape consistency會貪心地搜索sharding spec間的轉換方式，而不是樸素地逐dimension進行轉換，這樣能找到更高效的轉換路徑，進而使得sharding spec間的轉換通信開銷更小。

• 加入了all_to_all操作，使得Colossal-AI的擴展性更強，這在大規(guī)模集群上進行訓練時，可以展現(xiàn)出很大的優(yōu)勢

與activation checkpoint結合

作為大模型訓練中必不可少的顯存壓縮技術，Colossal-AI也提供了對于activation checkpoint的自動搜索功能。

相比于大部分將最大顯存壓縮作為目標的技術方案，Colossal-AI的搜索目標是在顯存預算以內，找到最快的activation checkpoint方案。

同時，為了避免將activation checkpoint的搜索一起建模到SPMD solver中導致搜索時間爆炸，Colossal-AI做了2-stage search的設計，因此可以在合理的時間內搜索到有效可行的分布式訓練方案。

通用深度學習系統(tǒng)Colossal-AI面向大模型時代，它可實現(xiàn)高效快速部署AI大模型訓練和推理，降低AI大模型應用成本。

自開源以來，Colossal-AI已經(jīng)多次在GitHub熱榜位列世界第一，獲得GitHub Star超七千顆，并成功入選SC、AAAI、PPoPP等國際AI與HPC頂級會議的官方教程。

Colossal-AI相關解決方案已成功在自動駕駛、云計算、零售、醫(yī)藥、芯片等行業(yè)知名廠商落地應用，廣受好評。

例如近期爆火的ChatGPT，尚未開源且不具備聯(lián)網(wǎng)功能。Colossal-AI已成功幫助某世界500強企業(yè)，開發(fā)具備在線搜索引擎能力增強的聊天機器人模型。

開源地址：https://github.com/hpcaitech/ColossalAI

參考鏈接：https://www.hpc-ai.tech/blog/colossal-ai-0-2-0