如何更好、更快和更便宜地實現訓練、微調 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