AIGC落地門檻被打下來(lái)了：硬件預(yù)算一口氣降至1/46，低成本上手Stable Diffusion2.0，一行代碼自動(dòng)并行

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

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

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

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

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

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

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

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

言下之意無(wú)非四個(gè)字：燒不起了。

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

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

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

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

開(kāi)源地址：
https://github.com/hpcaitech/ColossalAI

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

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

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

（畢竟1.0都還沒(méi)玩明白呢）

但話說(shuō)回來(lái)，還是老問(wèn)題，AIGC模型落地的成本高啊。

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

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

Stable Diffusion 2.0基于簡(jiǎn)單易用的PyTorch Lightning框架搭建。

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

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

• 顯存消耗可節(jié)省5.6倍，硬件成本最多降至1/46
• 支持DreamBooth單GPU快速個(gè)性化微調(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)練硬件門檻。

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

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

由此可以讓更多用戶在消費(fèi)級(jí)GPU上，低成本地開(kāi)展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是針對(duì)長(zhǎng)序列 attention 的加速版本，使用 Flatten 來(lái)減少 GPU 高帶寬內(nèi)存（HBM）之間的內(nèi)存讀 / 寫次數(shù)，F(xiàn)lash attention 同時(shí)針對(duì)塊狀稀疏的 attention，設(shè)計(jì)了一個(gè)近似的注意力算法，比任何現(xiàn)有的近似 attention 方法都要快。

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

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

ZeRO + Gemini

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

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

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

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

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

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

一行代碼快速上手

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

fromlightning.pytorchimporttrainer,LightningModule
fromlightning.pytorch.strategiesimportColossalAIStrategy
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微調(diào)

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

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

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

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

推理

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

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

model=replace_module(model)

用RTX3090推理1750億BLOOM模型

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

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

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

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

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

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

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

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

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

一行代碼自動(dòng)并行

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

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

然而，這對(duì)于絕大多數(shù)AI開(kāi)發(fā)者來(lái)說(shuō)很不友好，因?yàn)榇蠹叶疾幌Ｍ堰^(guò)多時(shí)間精力花費(fèi)在研究分布式系統(tǒng)和試錯(cuò)上。

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

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

#wrapthemodelusingauto_engine
model,optimizer=auto_engine(model,optimizer,cluster_info)
#normaltrainingloop
…

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

Graph Tracing

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

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

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

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

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

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

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

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

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

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

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

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

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

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

與activation checkpoint結(jié)合

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

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

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

關(guān)于Colossal-AI

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

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

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

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

傳送門

開(kāi)源地址：
https://github.com/hpcaitech/ColossalAI

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