Diffusion預訓練成本降低6.5倍，微調硬體成本降低7倍！ Colossal-AI完整開源方案低成本加速AIGC產業落地

如何更好、更快、更便宜地實現訓練、微調 AIGC 模型，已成為 AIGC 商業化和應用程式爆發的最大痛點。

Colossal-AI 基於在大模型民主化的專業技術積累，開源完整Stable Diffusion 預訓練和個性化微調方案，預訓練時間加速和經濟成本降低6.5 倍，個性化微調硬體成本降低7倍！在個人電腦的 RTX 2070/3050 上即可快速完成微調任務流程，讓 Stable Diffusion 等 AIGC 模式的觸手可及。

開源位址：

https://github.com/hpcaitech/ColossalAI

火爆的AIGC 賽道與高成本

AIGC（AI- Generated Content 人工智慧生成內容）是當前AI 領域最熱門的話題之一，尤其是伴隨著Stable Diffusion、Midjourney、NovelAI、DALL-E 等為代表的文本生成圖像的跨模態應用湧現，AIGC 更是火爆出圈，廣受關注。

Stable Diffusion 產生圖像

由於AIGC 激發了大量行業需求，它已被視為下一波AI 浪潮的重要方向之一，業界廣泛期望出現基於AIGC 在文字、音訊、圖像視訊、遊戲、元宇宙等技術場景的新技術革命和殺手級應用。 AIGC 在相關場景的成功商業化落地，潛在的數萬億美元市場，更是讓相關新創公司成為資本寵兒，如Stability AI、Jasper 等成立僅一兩年便已獲得上億美元融資，晉升獨角獸行列。

AI 模型規模與效能的同步成長

但高昂的硬體需求和訓練成本仍嚴重阻礙 AIGC 產業的快速發展。 AIGC 應用的出色表現通常建立在 GPT-3 或 Stable Diffusion 等大模型之上，並針對特定下游任務和應用進行微調。以大火的Stable Diffusion 為例，儘管其背後的Stability AI 成立不久，卻維護了超過4000 個英偉達A100 的GPU 集群，並已為此支出超過5000 萬美元的運營成本，僅Stable Diffusion v1 版本的模型單次訓練便需要150000 個A100 GPU Hour。

Diffusion model

Diffusion model(擴散模型）的想法最早在2015 年的論文Deep Unsupervised Learning using Nonequilibrium Thermodynamics 被提出，2020的論文Denoising Diffusion Probabilistic Models （DDPM）將其推到了一個新的高度，之後基於擴散模型的DALL-E 2， Imagen， Stable Diffusion 在生成任務上取得了遠超生成對抗網絡（GAN）、變微分自動編碼器（VAE）、自回歸模型（AR）等傳統生成模型的效果。

擴散模型包含兩個過程：前向擴散過程和反向生成過程，前向擴散過程是對一張圖像逐漸添加高斯噪聲直至變成隨機噪音，而反向生成過程是去噪聲過程，將一個隨機噪音使用多個U-Net 進行逐漸去噪音直至生成一張圖像，這也是擴散模型訓練的部分。

Latent Diffusion model

比較傳統端對端的深度學習模型，擴散模型的訓練過程無疑更為複雜，以Stable Diffusion 為例，除了擴散模型本身，還有一個Frozen CLIP Textcoder 來輸入text prompts，以及一個Autoencoder 實作將高解析度影像壓縮到潛在空間（Latent Space），並在每個time step 計算loss。這對訓練方案的顯存開銷，計算速度都提出了更大的挑戰。

更低成本－預訓練加速與少資源微調

預訓練最佳化

對於預訓練而言，一般batch size 越大，訓練速度也越快，Diffusion model 也是類似的。 Colossal- AI 透過ZeRO，Gemini,Chunk-based 記憶體管理等策略以及Flash Attention 模組優化Cross-attention 計算，大大降低了Diffusion model 的訓練的顯存開銷，使用戶在10G 顯存的消費級顯示卡（如RTX3080）上就可以訓練Diffusion model，在A100 這樣的專用顯示卡上最大可以直接支援單卡Batch Size 256 的訓練, 對比stable-diffusion-v1-1 的FP32 的DistributedDataParallel (DDP) 訓練可以提速6.5 倍。這意味著數百萬美元的訓練成本可降低 6.5 倍，大幅降低 AIGC 產業訓練成本和入場門檻！

Acceleration of Colossal-AI to Stable Diffusion

個人化微調最佳化

由於Stable Diffusion 的預訓練採用的LAION-5B 資料集共5850 億個圖片文字對，需要240TB 儲存空間，再結合模型的複雜性，顯然完整預訓練的成本極高：Stable Diffusion 的Stability 團隊花費超過5000 萬美元部署了4,000 塊A100 GPU。對於大多數 AIGC 玩家而言，更切實的選擇是使用開源的預訓練模型權重來進行微調個人化下游任務。

但其他現有的開源finetune 方案中使用的訓練並行方式主要為DDP，這導致訓練過程中顯存佔用極大，即使微調也需要至少使用RTX 3090 或4090 最高端的消費級顯示卡才能啟動。同時，現階段開源的許多訓練框架並沒有給予完整的訓練配置與腳本，需要使用者花費額外時間進行煩瑣的補全與調試。

不同於其他解決方案，Colossal-AI 是首個同時開源完整的訓練配置參數和訓練腳本的方案，讓使用者可以隨時訓練出針對新下游任務的最新版細分模型，使用更加靈活且應用範圍更廣。而由於Colossal-AI 引入顯存優化等技術，僅在普通個人電腦的單張消費級顯示卡上（如GeForce RTX 2070/3050 8GB），即可快速完成微調任務流程，相比RTX 3090 或4090 可降低約7 倍硬體成本，大大降低了使用Stable Diffusion 等AIGC 模型的門檻和成本，使用戶不再局限於現有的權重推理，方便快速完成個人化客製化服務。對於速度不敏感的任務，還可以進一步使用 Colossal-AI NVMe，即利用低成本的硬碟空間降低記憶體消耗。

Memory Reduction of Colossal-AI to Stable Diffusion

背後優化技術

ZeRO Gemini

Colossal-AI 支援使用零冗餘優化器(ZeRO) 的方法來消除記憶體冗餘，與經典的資料並行性策略相比，可極大提高記憶體使用效率，同時不犧牲計算粒度和通訊效率。

Colossal-AI 引進了 Chunk 機制，我們可以進一步提升 ZeRO 的效能。運算順序上連續的一組參數存入一個 Chunk 中（Chunk 即一段連續的記憶體空間），每個 Chunk 的大小都相同。 Chunk 方式組織記憶體可以確保 PCI-e 和 GPU-GPU 之間網路頻寬的高效利用，減少了通訊次數，同時避免潛在的記憶體碎片。

Chunk 機制

此外，Colossal-AI 的異質記憶體空間管理器Gemini 支援將最佳化器狀態從GPU 卸載到CPU ，以節省GPU記憶體佔用。可以同時利用 GPU 記憶體、CPU 記憶體（由 CPU DRAM 或 NVMe SSD 記憶體組成）來突破單 GPU 記憶體牆的限制，進一步擴展了可訓練模型規模。

透過ZeRO Gemini 提升硬體的模型容量

Flash Attention

LDM(Latent Diffusion Models) 透過在模型架構中引入cross-attention(交叉注意力層) 來實現多模態訓練，使得Diffusion model 可以更靈活地實現對class-condition, text-to-image, layout-to-image 的支援。然而 cross-attention 層對比原始 Diffusion model 的 CNN 層增加了額外的計算開銷，大大增加了訓練成本。

Colossal-AI 透過引入 Flash attention 機制，成功將 attention 的速度提升 104%，將端對端訓練的峰值顯存減少 23%。 Flash attention 是針對長序列attention 的加速版本，使用Flatten 來減少GPU 高頻寬記憶體（HBM）之間的記憶體讀/ 寫次數, Flash attention 同時針對塊狀稀疏的attention，設計了一個近似的注意力演算法，比任何現有的近似attention 方法都要快。

其他優化

Colossal-AI 也整合了 FP16、activation checkpoint 等常見優化技術。例如，activate checkpoint 透過用計算換取記憶體來運作。它避免儲存整個計算圖的所有中間激活用於反向計算，在檢查點部分不保存中間激活，而是在反向傳遞中重新計算它們，進一步降低了顯存。而 FP16 在基本上不影響精確度前提下，將原本的 32 位元浮點數運算轉為 16 位元，降低記憶體使用，提升運算效率。

快速上手使用

不同於常見的 PyTorch 開源項目，目前火熱的 stable diffusion 是基於 PyTorch Lightning 搭建的。 PyTorch Lightning 為流行的深度學習框架PyTorch 提供了簡潔易用、靈活高效的高級接口，為廣大AI 研究人員提供了簡潔易用的高層次抽象，從而使深度學習實驗更易於閱讀和再現，已在GitHub上收穫了20.5k 顆Star。

受 PyTorch Lightning 的邀請，Colossal-AI 已整合作為 PyTorch Lightning 的官方大模型解決方案。由於兩者的強強聯合，現在 AI 研究者可以更有效率地訓練和使用 diffusion 模型。以訓練 stable diffusion model 為例，只需少量程式碼即可快速啟動。

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

上述針對Diffusion 為代表的AIGC 訓練優化突破基於面向大模型時代的通用深度學習系統Colossal-AI，它透過高效多維自動並行、異質記憶體管理、大規模優化庫、自適應任務調度等實現高效快速部署AI 大模型訓練和推理，降低AI 大模型應用成本。自開源以來，Colossal-AI 已經多次在 GitHub 及 Papers With Code 熱榜位列世界第一，與眾多已有數萬 star 的明星開源項目一起受到海內外關注！經國際專家的嚴格評審，Colossal-AI 已成功入選為 SC、AAAI、PPoPP 等國際 AI 與 HPC 頂級會議的官方教程。

Colossal-AI 應用：更好的蛋白質結構預測解決方案

Colossal-AI 相關解決方案已成功在自動駕駛、雲端運算、零售、醫藥、晶片等業界知名廠商落地應用，廣受好評。例如，針對生物醫藥產業的蛋白質結構預測模型AlphaFold，基於Colossal-AI 的優化方案FastFold 成功將單張GPU 可推理的最大氨基酸序列長度突破至一萬，覆蓋了99.9999% 的蛋白質，僅用筆記型電腦上的消費級顯示卡即可解析90% 蛋白質。也能進一步對訓練、推理進行全流程並行加速，已協助多家新型藥物研發企業縮短開發流程，並降低研發成本。

開源位址：

https://github.com/hpcaitech/ColossalAI

以上是Diffusion預訓練成本降低6.5倍，微調硬體成本降低7倍！ Colossal-AI完整開源方案低成本加速AIGC產業落地的詳細內容。更多資訊請關注PHP中文網其他相關文章！