AI畫畫模型成本被打下來了！預(yù)訓(xùn)練成本直降85%，微調(diào)只需單張RTX 2070，這個(gè)國產(chǎn)開源項(xiàng)目又上新了

明敏 發(fā)自 凹非寺

量子位 | 公眾號(hào) QbitAI

火爆全球的Stable Diffusion，預(yù)訓(xùn)練和微調(diào)成本，正式被打下來了！

現(xiàn)在，預(yù)訓(xùn)練的速度被提升到6.5倍，成本直降85%；

微調(diào)只需RTX 2070/3050即可搞定，硬件成本直降近九成。

這就是國產(chǎn)項(xiàng)目Colossal-AI最新開源的加速方案。

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

今年，Stable Diffusion可謂火爆圈內(nèi)圈外，而且最近憑借免費(fèi)開源還效果好的特點(diǎn)，進(jìn)一步出圈，讓人看到了AIGC的巨大潛力和前景。

藝術(shù)圈和AI圈大神都爭(zhēng)相拿它創(chuàng)作，GitHub star數(shù)已達(dá)33.5k，背后公司估值甚至飆升至10億美金。

Stable Diffusion由此也成為擴(kuò)散模型中最火出圈的項(xiàng)目。

BUT，它唯一的缺點(diǎn)是，訓(xùn)練成本真的很高！

僅Stable Diffusion v1 版本的模型單次訓(xùn)練便需要150000個(gè)A100 GPU Hour；

微調(diào)至少需要RTX 3090/4090這類高端消費(fèi)級(jí)顯卡。

其背后公司Stability AI為了維護(hù)運(yùn)營和算力需求，擁有一個(gè)由4000塊A100組成的GPU集群，光成本就超過5000萬美元（折合人民幣約3.6億）。

實(shí)際上，由于擴(kuò)散模型本身原理機(jī)制特殊，導(dǎo)致它本身對(duì)硬件顯存、計(jì)算速度的要求都很高，而且使用的數(shù)據(jù)庫也是千億量級(jí)的。

所以想要提升訓(xùn)練速度、降低成本，還真不是件容易事。

那么，Colossal-AI到底是如何實(shí)現(xiàn)的？

從顯存入手

首先來解釋，為什么擴(kuò)散模型對(duì)硬件的要求高。

擴(kuò)散模型主要包括兩個(gè)過程：前向擴(kuò)散和反向生成。

前向擴(kuò)散是對(duì)一張圖像逐漸添加高斯噪聲，直至變成隨機(jī)噪聲。

反向生成是去噪聲過程，將一個(gè)隨機(jī)噪聲使用多個(gè)U-Net進(jìn)行逐漸去噪，直至生成一張圖像，這也是擴(kuò)散模型的訓(xùn)練過程。

相比于傳統(tǒng)端到端深度學(xué)習(xí)模型，擴(kuò)散模型的訓(xùn)練過程顯然更為復(fù)雜。

其中Stable Diffusion除了擴(kuò)散模型本身，還包含一個(gè)Frozen CLIP Textcoder 來輸入文本提示詞（text prompts）、一個(gè)自動(dòng)編碼器（Autoencoder）將高分辨率圖像壓縮到潛在空間（Latent Space），并且還要在每個(gè)time step計(jì)算loss。

這樣一來，訓(xùn)練方案對(duì)顯存開銷、計(jì)算速度都提出了極大挑戰(zhàn)。

針對(duì)如上問題，Colossal-AI分別從預(yù)訓(xùn)練、微調(diào)兩個(gè)步驟上進(jìn)行了優(yōu)化。

預(yù)訓(xùn)練階段，主要做的就是降低顯存開銷。

一般來說batch size越大，訓(xùn)練速度越快，擴(kuò)散模型亦是如此。

通過ZeRO、Gemini、Chunk-based內(nèi)存管理等策略以及Flash Attention模塊優(yōu)化Cross-attention計(jì)算，Colossal-AI可以讓用戶在8GB顯存的普通消費(fèi)級(jí)顯卡（如RTX 2070/3050）上訓(xùn)練擴(kuò)散模型。

如果在A100上，則能將單卡Batch Size直接擴(kuò)大到256。

對(duì)比Stable-diffusion-v1 FP32的Distributed Data Parallel (DDP) ，訓(xùn)練可以提速6.5倍。

由此可以大幅縮短訓(xùn)練時(shí)長，從而降低高達(dá)數(shù)百萬美元的訓(xùn)練成本。

在微調(diào)上，對(duì)于大多數(shù)AIGC玩家而言，都傾向于選擇使用開源的預(yù)訓(xùn)練模型權(quán)重來進(jìn)行微調(diào)個(gè)性化下游任務(wù)。

一方面是由于擴(kuò)散模型本身復(fù)雜，另一方面是Stable Diffusion預(yù)訓(xùn)練采用的是LAION-5B數(shù)據(jù)集，包含5850億個(gè)圖片文本對(duì)，需要240TB儲(chǔ)存空間。

但現(xiàn)有的很多開源微調(diào)方案中，使用的訓(xùn)練并行方式主要為DDP，這導(dǎo)致訓(xùn)練過程中占用的顯存很多。

即使微調(diào)也至少需要RTX 3090/4090這類最高端的消費(fèi)級(jí)顯卡。

加之，現(xiàn)階段很多開源的訓(xùn)練框架都沒有給出完整的訓(xùn)練配置與腳本，還需要自己來做一些補(bǔ)全和調(diào)試工作。

基于現(xiàn)狀，Colossal-AI提出了首個(gè)同時(shí)開源完整訓(xùn)練配置參數(shù)和訓(xùn)練腳本的方案。

這種方案一方面可以讓用戶不再局限于現(xiàn)有的權(quán)重推理，更加靈活地訓(xùn)練出針對(duì)新下游任務(wù)的最新版細(xì)分模型，擴(kuò)大應(yīng)用范圍；

另一方面，由于引入了顯存優(yōu)化等技術(shù)，方案支持在普通個(gè)人電腦的單張顯卡（如GeForce RTX 2070/3050 8GB）上快速完成微調(diào)任務(wù)。

要知道，RTX 2070現(xiàn)在的價(jià)格可比3090/4090便宜了將近9成。

此外，對(duì)于速度不敏感的任務(wù)，還可以進(jìn)一步使用Colossal-AI NVMe，即利用低成本的硬盤空間降低顯存消耗。

優(yōu)化秘籍是啥？

所以顯存開銷是通過哪些技術(shù)降下去的？

主要基于三個(gè)方面：

• ZeRO+Gemini
• Flash Attention
• FP16、acitvation checkpoint等常見優(yōu)化技術(shù)

先來看ZeRO+Gemini。

ZeRO是零冗余優(yōu)化器。

這種方法主要通過切分優(yōu)化器狀態(tài)、梯度、模型參數(shù)，使GPU僅保存當(dāng)前計(jì)算所需的部分，從而來消除數(shù)據(jù)并行、模型并行中存在的內(nèi)存冗余。

和經(jīng)典的數(shù)據(jù)并行策略相比，它能極大提高內(nèi)存使用效率，同時(shí)不犧牲計(jì)算粒度和通信效率。

與此同時(shí)，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)存碎片。

Gemini是Colossal-AI的異構(gòu)內(nèi)存空間管理器。

它支持將優(yōu)化器從GPU卸載到CPU，以節(jié)省GPU內(nèi)存占用。

打破GPU內(nèi)存墻的限制，可以實(shí)現(xiàn)同時(shí)利用GPU內(nèi)存和CPU內(nèi)存（由 CPU DRAM 或 NVMe SSD內(nèi)存組成），從而進(jìn)一步擴(kuò)展可訓(xùn)練模型規(guī)模。

第二方面是Flash Attention機(jī)制。

LDM（Latent Diffusion Models）是改善擴(kuò)散模型效率的一種方法，它通過在模型框架中引入交叉注意力機(jī)制（Cross-Attention）來實(shí)現(xiàn)多模態(tài)訓(xùn)練，使得擴(kuò)散模型可以更靈活地實(shí)現(xiàn)對(duì)class-condition, text-to-image, layout-to-image的支持。

但是Cross-Attention層對(duì)比原始擴(kuò)散模型的CNN層增加了額外的計(jì)算開銷，極大增加了訓(xùn)練成本。

相比之下，Flash Attention是針對(duì)長序列Attention的加速版本。

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

它能將Attention的速度提升104%，將端到端訓(xùn)練的峰值顯存減少23%。

除了以上兩方面，Colossal-AI還集成了FP16、activation checkpoint等常見優(yōu)化技術(shù)。

例如，activate checkpoint通過用計(jì)算換取內(nèi)存來工作。

它避免存儲(chǔ)整個(gè)計(jì)算圖的所有中間激活用于反向計(jì)算，在檢查點(diǎn)部分不保存中間激活，而是在反向傳遞中重新計(jì)算它們，進(jìn)一步降低了顯存。

而FP16在基本不影響精度前提下，將原本的32位浮點(diǎn)數(shù)運(yùn)算轉(zhuǎn)為16位，降低顯存使用，提升計(jì)算效率。

上手也超簡(jiǎn)單

到了實(shí)操環(huán)節(jié)，Colossal-AI只需少量代碼，即可快捷啟動(dòng)擴(kuò)散模型訓(xùn)練。

不同于常見PyTorch開源項(xiàng)目，Stable Diffusion是基于PyTorch Lighting搭建的。

PyTorch Lightning為流行的深度學(xué)習(xí)框架PyTorch提供了簡(jiǎn)潔易用、靈活高效的高級(jí)接口，為廣大AI研究人員提供了簡(jiǎn)潔易用的高層次抽象，從而使深度學(xué)習(xí)實(shí)驗(yàn)更易于閱讀和再現(xiàn)，

受PyTorch Lightning的邀請(qǐng)，Colossal-AI已集成作為PyTorch Lightning的官方大模型解決方案。

以Stable Diffusion為例，啟動(dòng)代碼如下(Python)：

from colossalai.nn.optimizer import HybridAdam
from lightning.pytorch import trainer
class MyDiffuser(LightningModule):
…
def configure_sharded_model(self) -> None:
# create your model here
self.model = construct_diffuser_model(…)
…
def configure_optimizers(self):
# use the specified optimizer
optimizer = HybridAdam(self.model.parameters(), self.lr)
…
model = MyDiffuser()
trainer = Trainer(accelerator=”gpu”, devices=1, precision=16, strategy=”colossalai”)
trainer.fit(model)

除此之外，Colossal-AI和PyTorch Lightning也對(duì)OPT、HuggingFace等熱門模型和社區(qū)提供了良好支持及優(yōu)化。

低成本微調(diào)

在微調(diào)方面，Colossal-AI提供了基于HuggingFace上開源的Stable Diffusion模型權(quán)重進(jìn)行微調(diào)的功能。

用戶只需簡(jiǎn)單修改Dataloader載入自己的微調(diào)數(shù)據(jù)集并讀取預(yù)訓(xùn)練權(quán)重，再修改參數(shù)配置yaml文件并運(yùn)行訓(xùn)練腳本，便可在個(gè)人電腦微調(diào)屬于自己的個(gè)性化模型。

model:
target: ldm.models.diffusion.ddpm.LatentDiffusion
params:
your_sub_module_config:
target: your.model.import.path
params:
from_pretrained: ‘your_file_path/unet/diffusion_pytorch_model.bin’
…
lightning:
trainer:
strategy:
target: pytorch_lightning.strategies.ColossalAIStrategy
params:
…
python main.py –logdir /your_log_dir -t -b config/train_colossalai.yaml

快速推理

Colossal-AI同時(shí)支持原生Stable Diffusion推理管道。

在完成訓(xùn)練或精調(diào)后，只需直接調(diào)用diffuser庫并加載自己保存的模型參數(shù)，即可直接進(jìn)行推理，無需進(jìn)行其他改動(dòng)。

這樣可以方便新用戶熟悉推理流程，并讓習(xí)慣使用原版框架的用戶快速上手。

from diffusers import StableDiffusionPipeline
pipe = StableDiffusionPipeline.from_pretrained(
“your_ColoDiffusion_checkpoint_path”
).to(“cuda”)
image = pipe(‘your prompt’, num_inference_steps=50)[“sample”][0]
image.save(‘file path’)

One More Thing

實(shí)際上，如上提到的Colossal-AI不只能用來加速訓(xùn)練擴(kuò)散模型，它面向各種大模型。

通過多維自動(dòng)并行、異構(gòu)內(nèi)存管理、大規(guī)模優(yōu)化庫、自適應(yīng)任務(wù)調(diào)度等方法，Colossal-AI可以高效快速部署AI大模型訓(xùn)練和推理，降低應(yīng)用成本。

自開源以來，Colossal-AI多次在GitHub及Paper With Code熱榜上位列世界第一。

經(jīng)國際專家的嚴(yán)格評(píng)審，Colossal-AI已成功入選為SC、AAAI、PPoPP等國際AI與HPC頂級(jí)會(huì)議的官方教程。

目前，其相關(guān)解決方案已經(jīng)成功在自動(dòng)駕駛、云計(jì)算、零售、醫(yī)藥、芯片等行業(yè)知名廠商落地應(yīng)用。

例如針對(duì)生物醫(yī)藥行業(yè)的蛋白質(zhì)結(jié)構(gòu)預(yù)測(cè)模型AlphaFold，基于Colossal-AI的優(yōu)化方案，F(xiàn)astFold成功將單張GPU可推理氨基酸序列長度擴(kuò)增到一萬，覆蓋99.9999%的蛋白質(zhì)。

只用筆記本上的消費(fèi)級(jí)顯卡就能解析90%蛋白質(zhì)。

而且還能對(duì)訓(xùn)練、推理進(jìn)行全流程并行加速，目前已助力多家新型藥物研發(fā)企業(yè)縮短開發(fā)流程，降低開發(fā)成本。

傳送門

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

參考鏈接：
[1]https://medium.com/@yangyou_berkeley/diffusion-pretraining-and-hardware-fine-tuning-can-be-almost-7x-cheaper-85e970fe207b
[2]https://github.com/CompVis/stable-diffusion
[3]https://arxiv.org/abs/2205.14135
[4]https://arxiv.org/abs/2112.10752
[5]https://openai.com/blog/triton/