明敏 發自 凹非寺
量子位 | 公眾號 QbitAI
火爆全球的
Stable Diffusion
，預訓練和微調成本，正式被打下來了！
現在，預訓練的速度被提升到
6.5倍
，
成本直降85%
；
微調
只需RTX 2070/3050
即可搞定，硬體成本直降近
九成。
這就是國產專案
Colossal-AI最新開源
的加速方案。
開源地址：https：//github。com/hpcaitech/ColossalAI
今年，Stable Diffusion可謂火爆圈內圈外，而且最近憑藉免費開源還效果好的特點，進一步出圈，讓人看到了AIGC的巨大潛力和前景。
藝術圈和AI圈大神都爭相拿它創作，GitHub star數已達33。5k，背後公司估值甚至飆升至
10億美金
。
Stable Diffusion由此也成為擴散模型中最火出圈的專案。
BUT，它唯一的缺點是，
訓練成本真的很高
！
僅Stable Diffusion v1 版本的模型單次訓練便需要
150000個
A100 GPU Hour；
微調至少需要RTX 3090/4090這類高階消費級顯示卡。
其背後公司Stability AI為了維護運營和算力需求，擁有一個由4000塊A100組成的GPU叢集，光成本就超過
5000萬美元
（摺合人民幣約3。6億）。
實際上，由於擴散模型本身原理機制特殊，導致它本身對硬體視訊記憶體、計算速度的要求都很高，而且使用的資料庫也是千億量級的。
所以想要提升訓練速度、降低成本，還真不是件容易事。
那麼，Colossal-AI到底是如何實現的？
從視訊記憶體入手
首先來解釋，為什麼擴散模型對硬體的要求高。
擴散模型主要包括兩個過程：前向擴散和反向生成。
前向擴散
是對一張影象逐漸新增高斯噪聲，直至變成隨機噪聲。
反向生成
是去噪聲過程，將一個隨機噪聲使用多個U-Net進行逐漸去噪，直至生成一張影象，這也是擴散模型的訓練過程。
相比於傳統端到端深度學習模型，擴散模型的訓練過程顯然更為複雜。
其中Stable Diffusion除了擴散模型本身，還包含一個Frozen CLIP Textcoder 來輸入文字提示詞（text prompts）、一個自動編碼器（Autoencoder）將高解析度影象壓縮到潛在空間（Latent Space），並且還要在每個time step計算loss。
這樣一來，訓練方案對視訊記憶體開銷、計算速度都提出了極大挑戰。
針對如上問題，Colossal-AI分別從預訓練、微調兩個步驟上進行了最佳化。
預訓練階段，主要做的就是
降低視訊記憶體開銷
。
一般來說batch size越大，訓練速度越快，擴散模型亦是如此。
透過ZeRO、Gemini、Chunk-based記憶體管理等策略以及Flash Attention模組最佳化Cross-attention計算，Colossal-AI可以讓使用者在
8GB視訊記憶體
的普通消費級顯示卡（如RTX 2070/3050）上訓練擴散模型。
如果在A100上，則能將單卡Batch Size直接擴大到
256
。
對比Stable-diffusion-v1 FP32的Distributed Data Parallel （DDP） ，訓練可以
提速6.5倍
。
由此可以大幅縮短訓練時長，從而降低高達數百萬美元的訓練成本。
在微調上
，對於大多數AIGC玩家而言，都傾向於選擇使用開源的預訓練模型權重來進行微調個性化下游任務。
一方面是由於擴散模型本身複雜，另一方面是Stable Diffusion預訓練採用的是LAION-5B資料集，包含5850億個圖片文字對，需要240TB儲存空間。
但現有的很多開源微調方案中，使用的訓練並行方式主要為DDP，這導致訓練過程中佔用的視訊記憶體很多。
即使微調也
至少需要RTX 3090/4090
這類最高階的消費級顯示卡。
加之，現階段很多開源的訓練框架都沒有給出完整的訓練配置與指令碼，還需要自己來做一些補全和除錯工作。
基於現狀，Colossal-AI提出了首個
同時開源完整訓練配置引數和訓練指令碼的方案
。
這種方案一方面可以讓使用者不再侷限於現有的權重推理，更加靈活地訓練出針對新下游任務的最新版細分模型，擴大應用範圍；
另一方面，由於引入了視訊記憶體最佳化等技術，方案支援在普通個人電腦的
單張顯示卡
（如
GeForce RTX 2070/3050 8GB
）上快速完成微調任務。
要知道，RTX 2070現在的價格可比3090/4090便宜了將近9成。
此外，對於速度不敏感的任務，還可以進一步使用
Colossal-AI NVMe
，即利用低成本的硬碟空間降低視訊記憶體消耗。
最佳化秘籍是啥？
所以視訊記憶體開銷是透過哪些技術降下去的？
主要基於三個方面：
ZeRO+Gemini
Flash Attention
FP16、acitvation checkpoint等常見最佳化技術
先來看
ZeRO+Gemini
。
ZeRO是零冗餘最佳化器。
這種方法主要透過切分最佳化器狀態、梯度、模型引數，使GPU僅儲存當前計算所需的部分，從而來消除資料並行、模型並行中存在的記憶體冗餘。
和經典的資料並行策略相比，它能極大提高記憶體使用效率，同時不犧牲計算粒度和通訊效率。
與此同時，Colossal-AI還引入了
Chunk機制
進一步提升ZeRO效能。
將運算順序上連續的一組引數存入一個Chunk中（Chunk即一段連續的記憶體空間），每個Chunk的大小相同。
Chunk方式組織記憶體可以保證PCI-e和GPU-GPU之間網路頻寬的高效利用，減小了通訊次數，同時避免潛在的記憶體碎片。
△Chunk機制
Gemini
是Colossal-AI的異構記憶體空間管理器。
它支援將最佳化器從GPU解除安裝到CPU，以節省GPU記憶體佔用。
打破GPU記憶體牆的限制，可以實現同時利用GPU記憶體和CPU記憶體（由 CPU DRAM 或 NVMe SSD記憶體組成），從而進一步擴充套件可訓練模型規模。
△透過ZeRO+Gemini提升硬體的模型容量
第二方面是
Flash Attention機制
。
LDM（Latent Diffusion Models）是改善擴散模型效率的一種方法，它透過在模型框架中引入交叉注意力機制（Cross-Attention）來實現多模態訓練，使得擴散模型可以更靈活地實現對class-condition， text-to-image， layout-to-image的支援。
但是Cross-Attention層對比原始擴散模型的CNN層增加了額外的計算開銷，極大增加了訓練成本。
相比之下，
Flash Attention是針對長序列Attention的加速版本
。
它使用Flatten來減少GPU高頻寬記憶體（HBM）之間的記憶體讀/寫次數，Flash Attention同時針對塊狀稀疏的Attention，設計了一個近似的注意力演算法，比現有任何attention方法都要快。
它能將Attention的速度
提升104%
，將端到端訓練的峰值視訊記憶體
減少23%
。
△ Flash Attention
除了以上兩方面，Colossal-AI還集成了FP16、activation checkpoint等常見最佳化技術。
例如，activate checkpoint透過用計算換取記憶體來工作。
它避免儲存整個計算圖的所有中間啟用用於反向計算，在檢查點部分不儲存中間啟用，而是在反向傳遞中重新計算它們，進一步降低了視訊記憶體。
而FP16在基本不影響精度前提下，將原本的32位浮點數運算轉為16位，降低視訊記憶體使用，提升計算效率。
上手也超簡單
到了實操環節，Colossal-AI
只需少量程式碼
，即可快捷啟動擴散模型訓練。
不同於常見PyTorch開源專案，Stable Diffusion是基於PyTorch Lighting搭建的。
PyTorch Lightning為流行的深度學習框架PyTorch提供了簡潔易用、靈活高效的高階介面，為廣大AI研究人員提供了簡潔易用的高層次抽象，從而使深度學習實驗更易於閱讀和再現，
受PyTorch Lightning的邀請，Colossal-AI已整合作為PyTorch Lightning的官方大模型解決方案。
以Stable Diffusion為例，啟動程式碼如下（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也對OPT、HuggingFace等熱門模型和社群提供了良好支援及最佳化。
低成本微調
在微調方面，Colossal-AI提供了基於HuggingFace上開源的Stable Diffusion模型權重進行微調的功能。
使用者只需簡單修改Dataloader載入自己的微調資料集並讀取預訓練權重，再修改引數配置yaml檔案並執行訓練指令碼，便可在個人電腦微調屬於自己的個性化模型。
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同時支援原生Stable Diffusion推理管道。
在完成訓練或精調後，只需直接呼叫diffuser庫並載入自己儲存的模型引數，即可直接進行推理，
無需進行其他改動
。
這樣可以方便新使用者熟悉推理流程，並讓習慣使用原版框架的使用者快速上手。
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
實際上，如上提到的Colossal-AI不只能用來加速訓練擴散模型，它面向各種大模型。
透過多維自動並行、異構記憶體管理、大規模最佳化庫、自適應任務排程等方法，Colossal-AI可以高效快速部署AI大模型訓練和推理，降低應用成本。
自開源以來，Colossal-AI多次在GitHub及Paper With Code熱榜上位列世界第一。
經國際專家的嚴格評審，
Colossal-AI已成功入選為SC、AAAI、PPoPP等國際AI與HPC頂級會議的官方教程。
目前，其相關解決方案已經成功在自動駕駛、雲計算、零售、醫藥、晶片等行業知名廠商落地應用。
例如針對生物醫藥行業的蛋白質結構預測模型AlphaFold，基於Colossal-AI的最佳化方案，FastFold成功將
單張GPU可推理氨基酸序列長度擴增到一萬，覆蓋99.9999%的蛋白質。
只用筆記本上的消費級顯示卡就能解析90%蛋白質。
而且還能對訓練、推理進行全流程並行加速，目前已助力多家新型藥物研發企業縮短開發流程，降低開發成本。
傳送門
開源地址：
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/
— 完 —
量子位 QbitAI · 頭條號簽約
關注我們，第一時間獲知前沿科技動態
