火爆全球的 Stable Diffusion，預訓練和微調(diào)成本，正式被打下來了！
現(xiàn)在，預訓練的速度被提升到 6.5 倍，成本直降 85%；
微調(diào)只需 RTX 2070/3050 即可搞定，硬件成本直降近九成。
這就是國產(chǎn)項目 Colossal-AI 最新開源的加速方案。
今年，Stable Diffusion 可謂火爆圈內(nèi)圈外，而且最近憑借免費開源還效果好的特點，進一步出圈，讓人看到了 AIGC 的巨大潛力和前景。
藝術圈和 AI 圈大神都爭相拿它創(chuàng)作，GitHub star 數(shù)已達 33.5k，背后公司估值甚至飆升至 10 億美金。
Stable Diffusion 由此也成為擴散模型中最火出圈的項目。
BUT，它唯一的缺點是，訓練成本真的很高！
僅 Stable Diffusion v1 版本的模型單次訓練便需要 150000 個 A100 GPU Hour；
微調(diào)至少需要 RTX 3090/4090 這類高端消費級顯卡。
其背后公司 Stability AI 為了維護運營和算力需求，擁有一個由 4000 塊 A100 組成的 GPU 集群，光成本就超過 5000 萬美元（折合人民幣約 3.6 億）。
實際上，由于擴散模型本身原理機制特殊，導致它本身對硬件顯存、計算速度的要求都很高，而且使用的數(shù)據(jù)庫也是千億量級的。
所以想要提升訓練速度、降低成本，還真不是件容易事。
那么，Colossal-AI 到底是如何實現(xiàn)的？
從顯存入手
首先來解釋，為什么擴散模型對硬件的要求高。
擴散模型主要包括兩個過程：前向擴散和反向生成。
前向擴散是對一張圖像逐漸添加高斯噪聲，直至變成隨機噪聲。
反向生成是去噪聲過程，將一個隨機噪聲使用多個 U-Net 進行逐漸去噪，直至生成一張圖像，這也是擴散模型的訓練過程。
相比于傳統(tǒng)端到端深度學習模型，擴散模型的訓練過程顯然更為復雜。
其中 Stable Diffusion 除了擴散模型本身，還包含一個 Frozen CLIP Textcoder 來輸入文本提示詞（text prompts）、一個自動編碼器（Autoencoder）將高分辨率圖像壓縮到潛在空間（Latent Space），并且還要在每個 time step 計算 loss。
這樣一來，訓練方案對顯存開銷、計算速度都提出了極大挑戰(zhàn)。
針對如上問題，Colossal-AI 分別從預訓練、微調(diào)兩個步驟上進行了優(yōu)化。
預訓練階段，主要做的就是降低顯存開銷。
一般來說 batch size 越大，訓練速度越快，擴散模型亦是如此。
通過 ZeRO、Gemini、Chunk-based 內(nèi)存管理等策略以及 Flash Attention 模塊優(yōu)化 Cross-attention 計算，Colossal-AI 可以讓用戶在 8GB 顯存的普通消費級顯卡（如 RTX 2070/3050）上訓練擴散模型。
如果在 A100 上，則能將單卡 Batch Size 直接擴大到 256。
對比 Stable-diffusion-v1 FP32 的 Distributed Data Parallel (DDP) ，訓練可以提速 6.5 倍。
由此可以大幅縮短訓練時長，從而降低高達數(shù)百萬美元的訓練成本。
在微調(diào)上，對于大多數(shù) AIGC 玩家而言，都傾向于選擇使用開源的預訓練模型權重來進行微調(diào)個性化下游任務。
一方面是由于擴散模型本身復雜，另一方面是 Stable Diffusion 預訓練采用的是 LAION-5B 數(shù)據(jù)集，包含 5850 億個圖片文本對，需要 240TB 儲存空間。
但現(xiàn)有的很多開源微調(diào)方案中，使用的訓練并行方式主要為 DDP，這導致訓練過程中占用的顯存很多。
即使微調(diào)也至少需要 RTX 3090/4090 這類最高端的消費級顯卡。
加之，現(xiàn)階段很多開源的訓練框架都沒有給出完整的訓練配置與腳本，還需要自己來做一些補全和調(diào)試工作。
基于現(xiàn)狀，Colossal-AI 提出了首個同時開源完整訓練配置參數(shù)和訓練腳本的方案。
這種方案一方面可以讓用戶不再局限于現(xiàn)有的權重推理，更加靈活地訓練出針對新下游任務的最新版細分模型，擴大應用范圍；
另一方面，由于引入了顯存優(yōu)化等技術，方案支持在普通個人電腦的單張顯卡（如 GeForce RTX 2070/3050 8GB）上快速完成微調(diào)任務。
要知道，RTX 2070 現(xiàn)在的價格可比 3090/4090 便宜了將近 9 成。
此外，對于速度不敏感的任務，還可以進一步使用 Colossal-AI NVMe，即利用低成本的硬盤空間降低顯存消耗。
優(yōu)化秘籍是啥？
所以顯存開銷是通過哪些技術降下去的？
主要基于三個方面：
ZeRO+Gemini
Flash Attention
FP16、acitvation checkpoint 等常見優(yōu)化技術
先來看 ZeRO+Gemini。
ZeRO 是零冗余優(yōu)化器。
這種方法主要通過切分優(yōu)化器狀態(tài)、梯度、模型參數(shù)，使 GPU 僅保存當前計算所需的部分，從而來消除數(shù)據(jù)并行、模型并行中存在的內(nèi)存冗余。
和經(jīng)典的數(shù)據(jù)并行策略相比，它能極大提高內(nèi)存使用效率，同時不犧牲計算粒度和通信效率。
與此同時，Colossal-AI 還引入了 Chunk 機制進一步提升 ZeRO 性能。
將運算順序上連續(xù)的一組參數(shù)存入一個 Chunk 中（Chunk 即一段連續(xù)的內(nèi)存空間），每個 Chunk 的大小相同。
Chunk 方式組織內(nèi)存可以保證 PCI-e 和 GPU-GPU 之間網(wǎng)絡帶寬的高效利用，減小了通信次數(shù)，同時避免潛在的內(nèi)存碎片。
Gemini 是 Colossal-AI 的異構內(nèi)存空間管理器。
它支持將優(yōu)化器從 GPU 卸載到 CPU，以節(jié)省 GPU 內(nèi)存占用。
打破 GPU 內(nèi)存墻的限制，可以實現(xiàn)同時利用 GPU 內(nèi)存和 CPU 內(nèi)存（由 CPU DRAM 或 NVMe SSD 內(nèi)存組成），從而進一步擴展可訓練模型規(guī)模。
第二方面是 Flash Attention 機制。
LDM（Latent Diffusion Models）是改善擴散模型效率的一種方法，它通過在模型框架中引入交叉注意力機制（Cross-Attention）來實現(xiàn)多模態(tài)訓練，使得擴散模型可以更靈活地實現(xiàn)對 class-condition, text-to-image, layout-to-image 的支持。
但是 Cross-Attention 層對比原始擴散模型的 CNN 層增加了額外的計算開銷，極大增加了訓練成本。
相比之下，Flash Attention 是針對長序列 Attention 的加速版本。
它使用 Flatten 來減少 GPU 高帶寬內(nèi)存（HBM）之間的內(nèi)存讀 / 寫次數(shù)，F(xiàn)lash Attention 同時針對塊狀稀疏的 Attention，設計了一個近似的注意力算法，比現(xiàn)有任何 attention 方法都要快。
它能將 Attention 的速度提升 104%，將端到端訓練的峰值顯存減少 23%。
除了以上兩方面，Colossal-AI 還集成了 FP16、activation checkpoint 等常見優(yōu)化技術。
例如，activate checkpoint 通過用計算換取內(nèi)存來工作。
它避免存儲整個計算圖的所有中間激活用于反向計算，在檢查點部分不保存中間激活，而是在反向傳遞中重新計算它們，進一步降低了顯存。
而 FP16 在基本不影響精度前提下，將原本的 32 位浮點數(shù)運算轉(zhuǎn)為 16 位，降低顯存使用，提升計算效率。
上手也超簡單
到了實操環(huán)節(jié)，Colossal-AI 只需少量代碼，即可快捷啟動擴散模型訓練。
不同于常見 PyTorch 開源項目，Stable Diffusion 是基于 PyTorch Lighting 搭建的。
PyTorch Lightning 為流行的深度學習框架 PyTorch 提供了簡潔易用、靈活高效的高級接口，為廣大 AI 研究人員提供了簡潔易用的高層次抽象，從而使深度學習實驗更易于閱讀和再現(xiàn)，
受 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, precision=16, strategy="colossalai") trainer.fit(model)
除此之外，Colossal-AI 和 PyTorch Lightning 也對 OPT、HuggingFace 等熱門模型和社區(qū)提供了良好支持及優(yōu)化。
低成本微調(diào)
在微調(diào)方面，Colossal-AI 提供了基于 HuggingFace 上開源的 Stable Diffusion 模型權重進行微調(diào)的功能。
用戶只需簡單修改 Dataloader 載入自己的微調(diào)數(shù)據(jù)集并讀取預訓練權重，再修改參數(shù)配置 yaml 文件并運行訓練腳本，便可在個人電腦微調(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 同時支持原生 Stable Diffusion 推理管道。
在完成訓練或精調(diào)后，只需直接調(diào)用 diffuser 庫并加載自己保存的模型參數(shù)，即可直接進行推理，無需進行其他改動。
這樣可以方便新用戶熟悉推理流程，并讓習慣使用原版框架的用戶快速上手。
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 不只能用來加速訓練擴散模型，它面向各種大模型。
通過多維自動并行、異構內(nèi)存管理、大規(guī)模優(yōu)化庫、自適應任務調(diào)度等方法，Colossal-AI 可以高效快速部署 AI 大模型訓練和推理，降低應用成本。
自開源以來，Colossal-AI 多次在 GitHub 及 Paper With Code 熱榜上位列世界第一。
經(jīng)國際專家的嚴格評審，Colossal-AI 已成功入選為 SC、AAAI、PPoPP 等國際 AI 與 HPC 頂級會議的官方教程。
目前，其相關解決方案已經(jīng)成功在自動駕駛、云計算、零售、醫(yī)藥、芯片等行業(yè)知名廠商落地應用。
例如針對生物醫(yī)藥行業(yè)的蛋白質(zhì)結構預測模型 AlphaFold，基于 Colossal-AI 的優(yōu)化方案，F(xiàn)astFold 成功將單張 GPU 可推理氨基酸序列長度擴增到一萬，覆蓋 99.9999% 的蛋白質(zhì)。
只用筆記本上的消費級顯卡就能解析 90% 蛋白質(zhì)。
而且還能對訓練、推理進行全流程并行加速，目前已助力多家新型藥物研發(fā)企業(yè)縮短開發(fā)流程，降低開發(fā)成本。
傳送門
開源地址：
https://github.com/hpcaitech/ColossalAI
參考鏈接：
[1]https://medium.com/@yangyou_berkeley/diffusion-pretraining-and-hardware-fine-tuning-can-be-almost-7x-cheaper-85e970fe207b
[2]https://github.com/CompVis/stable-diffusion
[3]https://arxiv.org/abs/2205.14135
[4]https://arxiv.org/abs/2112.10752
[5]https://openai.com/blog/triton/
本文來自微信公眾號：量子位 （ID：QbitAI），作者：明敏
廣告聲明：文內(nèi)含有的對外跳轉(zhuǎn)鏈接（包括不限于超鏈接、二維碼、口令等形式），用于傳遞更多信息，節(jié)省甄選時間，結果僅供參考，IT之家所有文章均包含本聲明。