如何更好、更快和更便宜地實(shí)現(xiàn)訓(xùn)練、微調(diào) AIGC 模型，已成為 AIGC 商業(yè)化和應(yīng)用爆發(fā)的最大痛點(diǎn)。

Colossal-AI 基于在大模型民主化的專業(yè)技術(shù)積累，開源完整 Stable Diffusion 預(yù)訓(xùn)練和個(gè)性化微調(diào)方案，預(yù)訓(xùn)練時(shí)間加速和經(jīng)濟(jì)成本降低 6.5 倍，個(gè)性化微調(diào)硬件成本降低 7 倍！在個(gè)人電腦的 RTX 2070/3050 上即可快速完成微調(diào)任務(wù)流程，讓 Stable Diffusion 等 AIGC 模型的觸手可及。

開源地址：

https://github.com/hpcaitech/ColossalAI

火爆的 AIGC 賽道與高昂成本

AIGC（AI-Generated Content 人工智能生成內(nèi)容）是當(dāng)前 AI 領(lǐng)域最熱門的話題之一，尤其是伴隨著 Stable Diffusion、Midjourney、NovelAI、DALL-E 等為代表的文本生成圖像的跨模態(tài)應(yīng)用涌現(xiàn)，AIGC 更是火爆出圈，廣受關(guān)注。

Stable Diffusion 生成圖像

由于 AIGC 激發(fā)了大量行業(yè)需求，它已被視為下一波 AI 浪潮的重要方向之一，業(yè)界廣泛期望出現(xiàn)基于 AIGC 在文本、音頻、圖像視頻、游戲、元宇宙等技術(shù)場(chǎng)景的新技術(shù)革命和殺手級(jí)應(yīng)用。AIGC 在相關(guān)場(chǎng)景的成功商業(yè)化落地，潛在的數(shù)萬(wàn)億美元市場(chǎng)，更是讓相關(guān)初創(chuàng)公司成為資本寵兒，如 Stability AI、Jasper 等成立僅一兩年便已獲得上億美元融資，晉升獨(dú)角獸行列。

AI 模型規(guī)模與性能的同步增長(zhǎng)

但高昂的硬件需求和訓(xùn)練成本仍嚴(yán)重阻礙著 AIGC 行業(yè)的快速發(fā)展。AIGC 應(yīng)用的出色表現(xiàn)通常建立在 GPT-3 或 Stable Diffusion 等大模型之上，并針對(duì)特定下游任務(wù)和應(yīng)用進(jìn)行微調(diào)。以大火的 Stable Diffusion 為例，盡管其背后的 Stability AI 成立不久，卻維護(hù)了超過(guò) 4000 個(gè)英偉達(dá) A100 的 GPU 集群，并已為此支出超過(guò) 5000 萬(wàn)美元的運(yùn)營(yíng)成本，僅 Stable Diffusion v1 版本的模型單次訓(xùn)練便需要 150000 個(gè) A100 GPU Hour。

Diffusion model

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

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

Latent Diffusion model

對(duì)比傳統(tǒng)端到端的深度學(xué)習(xí)模型，擴(kuò)散模型的訓(xùn)練過(guò)程無(wú)疑更為復(fù)雜，以 Stable Diffusion 為例，除了擴(kuò)散模型本身，還有一個(gè) Frozen CLIP Textcoder 來(lái)輸入 text prompts，以及一個(gè) Autoencoder 實(shí)現(xiàn)將高分辨率圖像壓縮到潛在空間（Latent Space），并在每個(gè) time step 計(jì)算 loss。這對(duì)訓(xùn)練方案的顯存開銷，計(jì)算速度都提出了更大的挑戰(zhàn)。

更低成本——預(yù)訓(xùn)練加速與少資源微調(diào)

預(yù)訓(xùn)練優(yōu)化

對(duì)于預(yù)訓(xùn)練而言，一般 batch size 越大，訓(xùn)練速度也越快，Diffusion model 也是類似的。Colossal- AI 通過(guò) ZeRO，Gemini, Chunk-based 內(nèi)存管理等策略以及 Flash Attention 模塊優(yōu)化 Cross-attention 計(jì)算，極大地降低了 Diffusion model 的訓(xùn)練的顯存開銷，使用戶在 10G 顯存的消費(fèi)級(jí)顯卡（如 RTX3080）上就可以訓(xùn)練 Diffusion model，在 A100 這樣的專用顯卡上最大可以直接支持單卡 Batch Size 256 的訓(xùn)練, 對(duì)比 stable-diffusion-v1-1 的 FP32 的 DistributedDataParallel (DDP) 訓(xùn)練可以提速 6.5 倍。這意味著數(shù)百萬(wàn)美元的訓(xùn)練成本可降低 6.5 倍，極大降低 AIGC 行業(yè)訓(xùn)練成本和入場(chǎng)門檻！

Acceleration of Colossal-AI to Stable Diffusion

個(gè)性化微調(diào)優(yōu)化

由于 Stable Diffusion 的預(yù)訓(xùn)練采用的 LAION-5B 數(shù)據(jù)集共 5850 億個(gè)圖片文本對(duì)，需要 240TB 儲(chǔ)存空間，再結(jié)合模型的復(fù)雜性，顯然完整預(yù)訓(xùn)練的成本極高：Stable Diffusion 的 Stability 團(tuán)隊(duì)花費(fèi)超過(guò) 5000 萬(wàn)美元部署了 4,000 塊 A100 GPU。對(duì)于大多數(shù) AIGC 玩家而言，更切實(shí)的選擇是使用開源的預(yù)訓(xùn)練模型權(quán)重來(lái)進(jìn)行微調(diào)個(gè)性化下游任務(wù)。

但其他現(xiàn)有的開源 finetune 方案中使用的訓(xùn)練并行方式主要為 DDP，這導(dǎo)致訓(xùn)練過(guò)程中顯存占用極大，即使微調(diào)也需要至少使用 RTX 3090 或 4090 最高端的消費(fèi)級(jí)顯卡才能啟動(dòng)。同時(shí)，現(xiàn)階段開源的很多訓(xùn)練框架并沒(méi)有給出完整的訓(xùn)練配置與腳本，需要用戶花費(fèi)額外時(shí)間進(jìn)行煩瑣的補(bǔ)全和調(diào)試。

不同于其他解決方案，Colossal-AI 是首個(gè)同時(shí)開源完整的訓(xùn)練配置參數(shù)和訓(xùn)練腳本的方案，讓用戶可以隨時(shí)訓(xùn)練出針對(duì)新下游任務(wù)的最新版細(xì)分模型，使用更加靈活且應(yīng)用范圍更廣。而且由于 Colossal-AI 引入顯存優(yōu)化等技術(shù)，僅在普通個(gè)人電腦的單張消費(fèi)級(jí)顯卡上（如 GeForce RTX 2070/3050 8GB），即可快速完成微調(diào)任務(wù)流程，相比 RTX 3090 或 4090 可降低約 7 倍硬件成本，大大降低了使用 Stable Diffusion 等 AIGC 模型的門檻和成本，使用戶不再局限于現(xiàn)有的權(quán)重推理，方便快捷完成個(gè)性化定制服務(wù)。對(duì)于速度不敏感的任務(wù)，還可以進(jìn)一步使用 Colossal-AI NVMe，即利用低成本的硬盤空間降低顯存消耗。

Memory Reduction of Colossal-AI to Stable Diffusion

背后優(yōu)化技術(shù)

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)存碎片。

Chunk 機(jī)制

此外，Colossal-AI 的異構(gòu)內(nèi)存空間管理器 Gemini 支持將優(yōu)化器狀態(tài)從 GPU 卸載到 CPU ，以節(jié)省 GPU 內(nèi)存占用?？梢酝瑫r(shí)利用 GPU 內(nèi)存、CPU 內(nèi)存（由 CPU DRAM 或 NVMe SSD 內(nèi)存組成）來(lái)突破單 GPU 內(nèi)存墻的限制，進(jìn)一步擴(kuò)展了可訓(xùn)練模型規(guī)模。

通過(guò) ZeRO + Gemini 提升硬件的模型容量

Flash Attention

LDM(Latent Diffusion Models) 通過(guò)在模型架構(gòu)中引入 cross-attention(交叉注意力層) 來(lái)實(shí)現(xiàn)多模態(tài)訓(xùn)練，使得 Diffusion model 可以更靈活地實(shí)現(xiàn)對(duì) class-condition, text-to-image, layout-to-image 的支持。然而 cross-attention 層對(duì)比原始 Diffusion model 的 CNN 層增加了額外的計(jì)算開銷，極大增加了訓(xùn)練成本。

Colossal-AI 通過(guò)引入 Flash attention 機(jī)制，成功將 attention 的速度提升 104%，將端到端訓(xùn)練的峰值顯存減少 23%。Flash attention 是針對(duì)長(zhǎng)序列 attention 的加速版本，使用 Flatten 來(lái)減少 GPU 高帶寬內(nèi)存（HBM）之間的內(nèi)存讀 / 寫次數(shù), Flash attention 同時(shí)針對(duì)塊狀稀疏的 attention，設(shè)計(jì)了一個(gè)近似的注意力算法，比任何現(xiàn)有的近似 attention 方法都要快。

其他優(yōu)化

Colossal-AI 還集成了 FP16、activation checkpoint 等常見(jiàn)優(yōu)化技術(shù)。例如，activate checkpoint 通過(guò)用計(jì)算換取內(nèi)存來(lái)工作。它避免存儲(chǔ)整個(gè)計(jì)算圖的所有中間激活用于反向計(jì)算，在檢查點(diǎn)部分不保存中間激活，而是在反向傳遞中重新計(jì)算它們，進(jìn)一步降低了顯存。而 FP16 在基本不影響精度前提下，將原本的 32 位浮點(diǎn)數(shù)運(yùn)算轉(zhuǎn)為 16 位，降低顯存使用，提升計(jì)算效率。

快速上手使用

不同于常見(jiàn)的 PyTorch 開源項(xiàng)目，當(dāng)前火熱的 stable diffusion 是基于 PyTorch Lightning 搭建的。PyTorch Lightning 為流行的深度學(xué)習(xí)框架 PyTorch 提供了簡(jiǎn)潔易用、靈活高效的高級(jí)接口，為廣大 AI 研究人員提供了簡(jiǎn)潔易用的高層次抽象，從而使深度學(xué)習(xí)實(shí)驗(yàn)更易于閱讀和再現(xiàn)，已在 GitHub 上收獲了 20.5k 顆 Star。

受 PyTorch Lightning 的邀請(qǐng)，Colossal-AI 已集成作為 PyTorch Lightning 的官方大模型解決方案。得益于兩者的強(qiáng)強(qiáng)聯(lián)合，現(xiàn)在 AI 研究者們可以更加高效地訓(xùn)練和使用 diffusion 模型。以訓(xùn)練 stable diffusion model 為例，僅需少量代碼即可快捷啟動(dòng)。

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)

Colossal-AI 為了滿足用戶通過(guò)較少資源短時(shí)間訓(xùn)練出可以生成有自己風(fēng)格的模型的需求，提供了基于 HuggingFace 上開源的 Stable Diffusion 模型權(quán)重進(jìn)行微調(diào)的功能。用戶只需簡(jiǎn)單修改 Dataloader 載入自己的微調(diào)數(shù)據(jù)集并讀取預(yù)訓(xùn)練權(quán)重，簡(jiǎ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 庫(kù)并加載自己保存的模型參數(shù)即可直接進(jìn)行推理，無(wú)需進(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

上述針對(duì) Diffusion 為代表的 AIGC 訓(xùn)練優(yōu)化突破基于面向大模型時(shí)代的通用深度學(xué)習(xí)系統(tǒng) Colossal-AI，它通過(guò)高效多維自動(dòng)并行、異構(gòu)內(nèi)存管理、大規(guī)模優(yōu)化庫(kù)、自適應(yīng)任務(wù)調(diào)度等實(shí)現(xiàn)高效快速部署 AI 大模型訓(xùn)練和推理，降低 AI 大模型應(yīng)用成本。自開源以來(lái)，Colossal-AI 已經(jīng)多次在 GitHub 及 Papers With Code 熱榜位列世界第一，與眾多已有數(shù)萬(wàn) star 的明星開源項(xiàng)目一起受到海內(nèi)外關(guān)注！經(jīng)國(guó)際專家的嚴(yán)格評(píng)審，Colossal-AI 已成功入選為 SC、AAAI、PPoPP 等國(guó)際 AI 與 HPC 頂級(jí)會(huì)議的官方教程。

Colossal-AI 應(yīng)用：更好的蛋白質(zhì)結(jié)構(gòu)預(yù)測(cè)解決方案

Colossal-AI 相關(guān)解決方案已成功在自動(dòng)駕駛、云計(jì)算、零售、醫(yī)藥、芯片等行業(yè)知名廠商落地應(yīng)用，廣受好評(píng)。例如，針對(duì)生物醫(yī)藥行業(yè)的蛋白質(zhì)結(jié)構(gòu)預(yù)測(cè)模型 AlphaFold，基于 Colossal-AI 的優(yōu)化方案 FastFold 成功將單張 GPU 可推理的最大氨基酸序列長(zhǎng)度突破至一萬(wàn)，覆蓋了 99.9999% 的蛋白質(zhì)，僅用筆記本電腦上的消費(fèi)級(jí)顯卡即可解析 90% 蛋白質(zhì)。還能進(jìn)一步對(duì)訓(xùn)練、推理進(jìn)行全流程并行加速，已助力多家新型藥物研發(fā)企業(yè)縮短開發(fā)流程，降低研發(fā)成本。

開源地址：

https://github.com/hpcaitech/ColossalAI