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

機器之心發表於2022-11-09
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


參考連結

https://github.com/CompVis/stable-diffusion

https://arxiv.org/abs/2205.14135

https://arxiv.org/abs/2112.10752

https://openai.com/blog/triton/

https://medium.com/@yangyou_berkeley/diffusion-pretraining-and-hardware-fine-tuning-can-be-almost-7x-cheaper-85e970fe207b

相關文章