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

Colossal-AI 基于在大模型民主化的專業(yè)技術(shù)積累，開源完整 Stable Diffusion 預(yù)訓(xùn)練和個性化微調(diào)方案，預(yù)訓(xùn)練時(shí)間加速和經(jīng)濟(jì)成本降低 6.5 倍，個性化微調(diào)硬件成本降低 7 倍！在個人電腦的 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ù)場景的新技術(shù)革命和殺手級應(yīng)用。AIGC 在相關(guān)場景的成功商業(yè)化落地，潛在的數(shù)萬億美元市場，更是讓相關(guān)初創(chuàng)公司成為資本寵兒，如 Stability AI、Jasper 等成立僅一兩年便已獲得上億美元融資，晉升獨(dú)角獸行列。

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

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

Diffusion model

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

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

Latent Diffusion model

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

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

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

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

Acceleration of Colossal-AI to Stable Diffusion

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

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

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

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

Memory Reduction of Colossal-AI to Stable Diffusion

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

ZeRO + Gemini

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

Colossal-AI 引入了 Chunk 機(jī)制，我們可以進(jìn)一步提升 ZeRO 的性能。運(yùn)算順序上連續(xù)的一組參數(shù)存入一個 Chunk 中（Chunk 即一段連續(xù)的內(nèi)存空間），每個 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)存組成）來突破單 GPU 內(nèi)存墻的限制，進(jìn)一步擴(kuò)展了可訓(xùn)練模型規(guī)模。

通過 ZeRO + Gemini 提升硬件的模型容量

Flash Attention

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

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

其他優(yōu)化

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

快速上手使用

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

受 PyTorch Lightning 的邀請，Colossal-AI 已集成作為 PyTorch Lightning 的官方大模型解決方案。得益于兩者的強(qiáng)強(qiáng)聯(lián)合，現(xiàn)在 AI 研究者們可以更加高效地訓(xùn)練和使用 diffusion 模型。以訓(xùn)練 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 等熱門模型和社區(qū)提供了良好支持及優(yōu)化。

低成本微調(diào)

Colossal-AI 為了滿足用戶通過較少資源短時(shí)間訓(xùn)練出可以生成有自己風(fēng)格的模型的需求，提供了基于 HuggingFace 上開源的 Stable Diffusion 模型權(quán)重進(jìn)行微調(diào)的功能。用戶只需簡單修改 Dataloader 載入自己的微調(diào)數(shù)據(jù)集并讀取預(yù)訓(xùn)練權(quán)重，簡單修改參數(shù)配置 yaml 文件并運(yùn)行訓(xùn)練腳本，便可在個人電腦微調(diào)屬于自己的個性化模型。

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)行其他改動，方便新用戶熟悉推理流程并可以讓習(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

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

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

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

開源地址：

https://github.com/hpcaitech/ColossalAI