本文經AI新媒體量子位（公眾號ID:QbitAI）授權轉載，轉載請聯系出處。

火爆全球的Stable Diffusion，預訓練和微調成本，正式被打下來了！

現在，預訓練的速度被提升到6.5倍，成本直降85%；

微調只需RTX 2070/3050即可搞定，硬件成本直降近九成。

這就是國產項目Colossal-AI最新開源的加速方案。

開源地址：https://github.com/hpcaitech/ColossalAI

?

今年，Stable Diffusion可謂火爆圈內圈外，而且最近憑借免費開源還效果好的特點，進一步出圈，讓人看到了AIGC的巨大潛力和前景。

藝術圈和AI圈大神都爭相拿它創作，GitHub star數已達33.5k，背后公司估值甚至飆升至10億美金。

Stable Diffusion由此也成為擴散模型中最火出圈的項目。

BUT，它唯一的缺點是，訓練成本真的很高！

僅Stable Diffusion v1 版本的模型單次訓練便需要150000個A100 GPU Hour；

微調至少需要RTX 3090/4090這類高端消費級顯卡。

其背后公司Stability AI為了維護運營和算力需求，擁有一個由4000塊A100組成的GPU集群，光成本就超過5000萬美元（折合人民幣約3.6億）。

實際上，由于擴散模型本身原理機制特殊，導致它本身對硬件顯存、計算速度的要求都很高，而且使用的數據庫也是千億量級的。

所以想要提升訓練速度、降低成本，還真不是件容易事。

那么，Colossal-AI到底是如何實現的？

從顯存入手

首先來解釋，為什么擴散模型對硬件的要求高。

擴散模型主要包括兩個過程：前向擴散和反向生成。

前向擴散是對一張圖像逐漸添加高斯噪聲，直至變成隨機噪聲。

反向生成是去噪聲過程，將一個隨機噪聲使用多個U-Net進行逐漸去噪，直至生成一張圖像，這也是擴散模型的訓練過程。

相比于傳統端到端深度學習模型，擴散模型的訓練過程顯然更為復雜。

其中Stable Diffusion除了擴散模型本身，還包含一個Frozen CLIP Textcoder 來輸入文本提示詞（text prompts）、一個自動編碼器（Autoencoder）將高分辨率圖像壓縮到潛在空間（Latent Space），并且還要在每個time step計算loss。

這樣一來，訓練方案對顯存開銷、計算速度都提出了極大挑戰。

針對如上問題，Colossal-AI分別從預訓練、微調兩個步驟上進行了優化。

預訓練階段，主要做的就是降低顯存開銷。

一般來說batch size越大，訓練速度越快，擴散模型亦是如此。

通過ZeRO、Gemini、Chunk-based內存管理等策略以及Flash Attention模塊優化Cross-attention計算，Colossal-AI可以讓用戶在8GB顯存的普通消費級顯卡（如RTX 2070/3050）上訓練擴散模型。

如果在A100上，則能將單卡Batch Size直接擴大到256。

對比Stable-diffusion-v1 FP32的Distributed Data Parallel (DDP) ，訓練可以提速6.5倍。

由此可以大幅縮短訓練時長，從而降低高達數百萬美元的訓練成本。

在微調上，對于大多數AIGC玩家而言，都傾向于選擇使用開源的預訓練模型權重來進行微調個性化下游任務。

一方面是由于擴散模型本身復雜，另一方面是Stable Diffusion預訓練采用的是LAION-5B數據集，包含5850億個圖片文本對，需要240TB儲存空間。

但現有的很多開源微調方案中，使用的訓練并行方式主要為DDP，這導致訓練過程中占用的顯存很多。

即使微調也至少需要RTX 3090/4090這類最高端的消費級顯卡。

加之，現階段很多開源的訓練框架都沒有給出完整的訓練配置與腳本，還需要自己來做一些補全和調試工作。

基于現狀，Colossal-AI提出了首個同時開源完整訓練配置參數和訓練腳本的方案。

這種方案一方面可以讓用戶不再局限于現有的權重推理，更加靈活地訓練出針對新下游任務的最新版細分模型，擴大應用范圍；

另一方面，由于引入了顯存優化等技術，方案支持在普通個人電腦的單張顯卡（如GeForce RTX 2070/3050 8GB）上快速完成微調任務。

要知道，RTX 2070現在的價格可比3090/4090便宜了將近9成。

此外，對于速度不敏感的任務，還可以進一步使用Colossal-AI NVMe，即利用低成本的硬盤空間降低顯存消耗。

優化秘籍是啥？

所以顯存開銷是通過哪些技術降下去的？

主要基于三個方面：

• ZeRO+Gemini
• Flash Attention
• FP16、acitvation checkpoint等常見優化技術

先來看ZeRO+Gemini。

ZeRO是零冗余優化器。

這種方法主要通過切分優化器狀態、梯度、模型參數，使GPU僅保存當前計算所需的部分，從而來消除數據并行、模型并行中存在的內存冗余。

和經典的數據并行策略相比，它能極大提高內存使用效率，同時不犧牲計算粒度和通信效率。

與此同時，Colossal-AI還引入了Chunk機制進一步提升ZeRO性能。

將運算順序上連續的一組參數存入一個Chunk中（Chunk即一段連續的內存空間），每個Chunk的大小相同。

Chunk方式組織內存可以保證PCI-e和GPU-GPU之間網絡帶寬的高效利用，減小了通信次數，同時避免潛在的內存碎片。

△Chunk機制

Gemini是Colossal-AI的異構內存空間管理器。

它支持將優化器從GPU卸載到CPU，以節省GPU內存占用。

打破GPU內存墻的限制，可以實現同時利用GPU內存和CPU內存（由 CPU DRAM 或 NVMe SSD內存組成），從而進一步擴展可訓練模型規模。

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

第二方面是Flash Attention機制。

LDM（Latent Diffusion Models）是改善擴散模型效率的一種方法，它通過在模型框架中引入交叉注意力機制（Cross-Attention）來實現多模態訓練，使得擴散模型可以更靈活地實現對class-condition, text-to-image, layout-to-image的支持。

但是Cross-Attention層對比原始擴散模型的CNN層增加了額外的計算開銷，極大增加了訓練成本。

相比之下，Flash Attention是針對長序列Attention的加速版本。

它使用Flatten來減少GPU高帶寬內存（HBM）之間的內存讀/寫次數，Flash Attention同時針對塊狀稀疏的Attention，設計了一個近似的注意力算法，比現有任何attention方法都要快。

它能將Attention的速度提升104%，將端到端訓練的峰值顯存減少23%。

△ Flash Attention

除了以上兩方面，Colossal-AI還集成了FP16、activation checkpoint等常見優化技術。

例如，activate checkpoint通過用計算換取內存來工作。

它避免存儲整個計算圖的所有中間激活用于反向計算，在檢查點部分不保存中間激活，而是在反向傳遞中重新計算它們，進一步降低了顯存。

而FP16在基本不影響精度前提下，將原本的32位浮點數運算轉為16位，降低顯存使用，提升計算效率。

上手也超簡單

到了實操環節，Colossal-AI只需少量代碼，即可快捷啟動擴散模型訓練。

不同于常見PyTorch開源項目，Stable Diffusion是基于PyTorch Lighting搭建的。

PyTorch Lightning為流行的深度學習框架PyTorch提供了簡潔易用、靈活高效的高級接口，為廣大AI研究人員提供了簡潔易用的高層次抽象，從而使深度學習實驗更易于閱讀和再現，

受PyTorch Lightning的邀請，Colossal-AI已集成作為PyTorch Lightning的官方大模型解決方案。

以Stable Diffusion為例，啟動代碼如下：

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, precisinotallow=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

實際上，如上提到的Colossal-AI不只能用來加速訓練擴散模型，它面向各種大模型。

通過多維自動并行、異構內存管理、大規模優化庫、自適應任務調度等方法，Colossal-AI可以高效快速部署AI大模型訓練和推理，降低應用成本。

自開源以來，Colossal-AI多次在GitHub及Paper With Code熱榜上位列世界第一。

經國際專家的嚴格評審，Colossal-AI已成功入選為SC、AAAI、PPoPP等國際AI與HPC頂級會議的官方教程。

目前，其相關解決方案已經成功在自動駕駛、云計算、零售、醫藥、芯片等行業知名廠商落地應用。

例如針對生物醫藥行業的蛋白質結構預測模型AlphaFold，基于Colossal-AI的優化方案，FastFold成功將單張GPU可推理氨基酸序列長度擴增到一萬，覆蓋99.9999%的蛋白質。

只用筆記本上的消費級顯卡就能解析90%蛋白質。

而且還能對訓練、推理進行全流程并行加速，目前已助力多家新型藥物研發企業縮短開發流程，降低開發成本。

傳送門

開源地址：??https://github.com/hpcaitech/ColossalAI??