如何更好、更快和更便宜地實(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)注。
?

由于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ú)角獸行列。

但高昂的硬件需求和訓(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)練的部分。

?

對(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)門檻！

?

• 個(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，即利用低成本的硬盤空間降低顯存消耗。

背后優(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)存碎片。

?

此外，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ī)模。

• 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等常見優(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ì)算效率。

快速上手使用

不同于常見的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相關(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

參考鏈接:

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

https://arxiv.org/abs/2205.14135

https://arxiv.org/abs/2112.10752

https://openai.com/blog/triton/文章來(lái)源地址http://www.zghlxwxcb.cn/news/detail-487291.html

到了這里，關(guān)于Diffusion預(yù)訓(xùn)練成本降低6.5倍，微調(diào)硬件成本降低7倍！Colossal-AI完整開源方案低成本加速AIGC產(chǎn)業(yè)落地的文章就介紹完了。如果您還想了解更多內(nèi)容，請(qǐng)?jiān)谟疑辖撬阉鱐OY模板網(wǎng)以前的文章或繼續(xù)瀏覽下面的相關(guān)文章，希望大家以后多多支持TOY模板網(wǎng)！