本文經(jīng)AI新媒體量子位（公眾號ID:QbitAI）授權(quán)轉(zhuǎn)載，轉(zhuǎn)載請聯(lián)系出處。

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

話不多說，直接看結(jié)果：

• Stable Diffusion 2.0訓(xùn)練/微調(diào)/推理，顯存消耗最多可節(jié)省5.6倍，使硬件成本直降至1/46，一行代碼即可啟用；
• 1750億參數(shù)大模型BLOOM單機推理，顯存消耗節(jié)省4倍，硬件成本直降至十幾分之一。
• 一行代碼實現(xiàn)自動搜索最佳并行策略，顯著降低分布式訓(xùn)練上手門檻，原生支持Hugging Face、Timm等熱門AI模型庫。

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

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

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

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

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

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

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

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

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

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

Stable Diffusion2.0低成本訓(xùn)練/微調(diào)/推理

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

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

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

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

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

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

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

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

具體內(nèi)容有以下幾點：

• 顯存消耗可節(jié)省5.6倍，硬件成本最多降至1/46
• 支持DreamBooth單GPU快速個性化微調(diào)
• 推理顯存消耗節(jié)省2.5倍

而且該方案也將于近期合并進(jìn)入Hugging Face，進(jìn)一步方便用戶使用。

訓(xùn)練

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

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

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

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

背后顯存優(yōu)化

Flash Attention

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

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

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

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

ZeRO + Gemini

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

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

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

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

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

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

一行代碼快速上手

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

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, precisinotallow=16, strategy=Mystrategy)
trainer.fit(model)

DreamBooth微調(diào)

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

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

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

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

推理

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

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

model = replace_module(model)

用RTX3090推理1750億BLOOM模型

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

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

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

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

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

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

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

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

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

一行代碼自動并行

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

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

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

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

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

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

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

Graph Tracing

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

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

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

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

細(xì)粒度分布式訓(xùn)練策略搜索

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

現(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對集群進(jìn)行了抽象管理。

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

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

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

• Colossal-AI的device mesh可以profiling到集群性能指標(biāo)，對不同的通信算子進(jìn)行耗時估算。
• Colossal-AI的shape consistency會貪心地搜索sharding spec間的轉(zhuǎn)換方式，而不是樸素地逐dimension進(jìn)行轉(zhuǎn)換，這樣能找到更高效的轉(zhuǎn)換路徑，進(jìn)而使得sharding spec間的轉(zhuǎn)換通信開銷更小。
• 加入了all_to_all操作，使得Colossal-AI的擴(kuò)展性更強，這在大規(guī)模集群上進(jìn)行訓(xùn)練時，可以展現(xiàn)出很大的優(yōu)勢

與activation checkpoint結(jié)合

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

相比于大部分將最大顯存壓縮作為目標(biāo)的技術(shù)方案，Colossal-AI的搜索目標(biāo)是在顯存預(yù)算以內(nèi)，找到最快的activation checkpoint方案。

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

關(guān)于Colossal-AI

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

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

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

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

傳送門

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