最近，AI大模型連續(xù)火爆出圈，人工智能生成模型（AIGC）的熱度尚未褪去，聊天機(jī)器人ChatGPT便引發(fā)全網(wǎng)熱議，兩周吸引百萬(wàn)用戶。還有卷趴一半程序員的AlphaCode，生成全新蛋白質(zhì)的ESM2等，不斷探索AI大模型落地的新領(lǐng)域。面對(duì)大模型帶來(lái)的技術(shù)革命，連谷歌都拉響“紅色警報(bào)”，擔(dān)心ChatGPT砸掉自己搜索引擎的飯碗。

作為當(dāng)下最火熱的開源AI大模型解決方案，Colossal-AI已收獲Github Star七千多顆，此前在Stable Diffusion、GPT-3、AlphaFold等大模型上展現(xiàn)卓越性能優(yōu)勢(shì)。針對(duì)AI大模型落地成本高昂這一痛點(diǎn)，Colossal-AI本次更新聚焦于降低大模型應(yīng)用成本以及增強(qiáng)易用性，包括:

• Stable Diffusion 2.0低成本訓(xùn)練/微調(diào)/推理，最多可降低顯存消耗5.6倍，使硬件成本直降46倍，一行代碼即可使用;

• 1750億參數(shù)BLOOM模型單機(jī)推理，顯存消耗降低4倍，使硬件成本降低10余倍;

• 一行代碼實(shí)現(xiàn)自動(dòng)搜索最佳并行策略，顯著降低分布式訓(xùn)練上手門檻，原生支持Hugging Face，Timm等熱門AI模型庫(kù)。

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

Stable Diffusion 2.0低成本訓(xùn)練/微調(diào)/推理

AIGC（人工智能生成內(nèi)容）目前是AI領(lǐng)域最熱門的話題之一，還被Science評(píng)為2022年度十大科學(xué)突破。短短數(shù)月，爆紅AI圈的Stable Diffusion已升級(jí)到2.0版本，正如部分用戶所言：“我從未見過(guò)任何技術(shù)發(fā)展如此之快。V1還沒整透徹，V2就來(lái)了”。

然而，高昂的成本也一定程度上限制了AIGC產(chǎn)業(yè)的大規(guī)模普及。以Stable Diffusion為例，其背后的Stability AI維護(hù)超過(guò) 4000 個(gè)英偉達(dá) A100 的 GPU 集群，并已為此支出超過(guò) 5000 萬(wàn)美元的運(yùn)營(yíng)成本。面對(duì)快速迭代的模型、算法和下游任務(wù)，如何降低應(yīng)用成本成為AIGC真正走向落地的核心問(wèn)題。

Stable Diffusion 2.0基于簡(jiǎn)單易用的PyTorch Lightning框架搭建。作為PyTorch Lightning的官方大模型解決方案，Colossal-AI在第一時(shí)間跟進(jìn)，開源了更為高效，硬件門檻更低的全套訓(xùn)練/微調(diào)/推理方案:

• 可減少訓(xùn)練顯存消耗5.6倍，硬件成本最高降低46倍；

• 支持DreamBooth單GPU快速個(gè)性化微調(diào)；

• 推理顯存消耗降低2.5倍。

該方案也將于近期合并進(jìn)入當(dāng)下最火熱的AI模型社區(qū)Hugging Face，進(jìn)一步方便用戶使用。

訓(xùn)練

為了加快訓(xùn)練速度，降低訓(xùn)練成本，使用更大的batch size已成為被廣泛使用的有效手段。但GPU有限的顯存容量，嚴(yán)重限制了batch size大小，推高了訓(xùn)練硬件門檻。

Colossal-AI通過(guò)一系列顯存優(yōu)化技術(shù)和支持Stable Diffusion 2.0，使Stable Diffusion平均在每個(gè)GPU上使用大batch size 16訓(xùn)練的顯存需求從64.5GB降低5.6倍至11.6GB，還可擴(kuò)展至單GPU或多GPU并行。相比使用最先進(jìn)的A100 80GB，目前僅需3060等消費(fèi)級(jí)顯卡即可滿足需求，硬件成本最高直降46倍。更多用戶可以在消費(fèi)級(jí)GPU上，低成本地開展Stable Diffusion的相關(guān)研究與應(yīng)用落地。

背后顯存優(yōu)化

Flash Attention：早在Stable Diffusion 1.0版本，Colossal-AI就率先引入了Flash Attention技術(shù)，成功將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 方法都要快。在Stable Diffusion 1.0版本，整個(gè)Diffusion Model只有少量attention層，F(xiàn)lash attention還沒有體現(xiàn)出其性能優(yōu)勢(shì)。在Stable Diffusion 2.0中，由于將大量卷積層替換為attention層，進(jìn)一步發(fā)揮了Flash Attention的顯存優(yōu)化潛力。

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

一行代碼快速上手

作為PyTorch Lightning的官方合作伙伴，僅需一行代碼即可調(diào)用Colossal-AI的上述顯存優(yōu)化。

from lightning.pytorch import trainer, LightningModule
from lightning.pytorch.strategies import ColossalAIStrategy

Mystrategy = ColossalAIStrategy(use_chunk=True, enable_distributed_storage=True, placement_policy=auto)

trainer = Trainer(accelerator="gpu", devices=4, precision=16, strategy=Mystrategy)
trainer.fit(model)

DreamBooth微調(diào)

DreamBooth是一種用于個(gè)性化文本到圖像模型（例如Stable Diffusion）的方法，只需要幾張（3~5張）期望主題的圖像，即可生成指定主題的一系列圖像。用戶只需直接運(yùn)行文件train_dreambooth_colossalai.py，即可在該微調(diào)任務(wù)上充分發(fā)揮Colossal-AI的顯存優(yōu)化，個(gè)性化快速微調(diào)自己的圖文模型，極大地降低了使用門檻。

推理

由于模型推理對(duì)數(shù)值精度不敏感，這為實(shí)現(xiàn)低精度的低成本推理提供了可能。對(duì)于Stable Diffusion 2.0模型，可以通過(guò)添加一行代碼，支持模型的Int8量化推理，顯存消耗降低2.5倍，僅需3.1GB，且不造成顯著性能損失。

model = replace_module(model)

1750億BLOOM模型低成本推理

隨著模型規(guī)模的不斷增大，模型推理的顯存及內(nèi)存占用也變成了不可忽視的關(guān)鍵因素。以Hugging Face發(fā)布的1750億參數(shù)開源模型BLOOM為例，如果直接使用常見的FP32/FP16進(jìn)行推理，在單節(jié)點(diǎn)8張GPU使用模型并行，每張GPU需要消耗至少87.5GB/43.8GB的顯存。如此大的顯存占用，即使是最先進(jìn)的8卡A100（80GB/40GB）服務(wù)器，也無(wú)法直接部署推理服務(wù)，而多節(jié)點(diǎn)推理又會(huì)帶來(lái)沉重的額外成本和通信開銷。

Colossal-AI實(shí)現(xiàn)了高效的Int8量化和模型并行推理，可以將1750億參數(shù)的BLOOM等大模型的推理服務(wù)，部署到3090/4090等消費(fèi)級(jí)顯卡的8卡服務(wù)器，同時(shí)不產(chǎn)生顯著的CPU內(nèi)存占用提升及性能損耗。相比原有的A100方案，可降低硬件部署成本十余倍。

在Colossal-AI中，通過(guò)對(duì)模型進(jìn)行了Int8量化，可將模型總體顯存占用從352.3GB（FP16）降低到185.6GB， 同時(shí)使用Colossal-AI的模型并行技術(shù)，將每張顯卡的占用減少到了23.2GB。在模型并行中，為了不增加CPU內(nèi)存占用，Colossal-AI在主進(jìn)程中進(jìn)行模型的量化和切分工作，其余的每個(gè)進(jìn)程中分別使用lazy_init獲得幾乎不占顯存和內(nèi)存的meta model，再通過(guò)gloo backend在進(jìn)程之間傳遞模型參數(shù)。通過(guò)上述方案，在沒有分段加載模型參數(shù)的情況下，便可以使得CPU內(nèi)存占用峰值，達(dá)到理論上的較優(yōu)水平。相較于將模型按層切分的“類流水線”分布方式，模型并行可以提高非密集請(qǐng)求下的顯存使用效率。

一行代碼自動(dòng)并行

大模型的分布式混合部署是一個(gè)非常復(fù)雜的問(wèn)題，目前常見的分布式大模型訓(xùn)練方案，都依賴用戶人工反復(fù)嘗試以及系統(tǒng)專家的經(jīng)驗(yàn)來(lái)進(jìn)行配置部署。然而這對(duì)絕大多數(shù)AI開發(fā)者來(lái)說(shuō)十分不友好，因?yàn)椴幌Ｍ麑r(shí)間精力花費(fèi)在研究分布式系統(tǒng)和試錯(cuò)上。

Colossal-AI提供了業(yè)界急需的一套高效易用自動(dòng)并行系統(tǒng)。相比現(xiàn)有其他手動(dòng)配置復(fù)雜并行策略和修改模型的解決方案，Colossal-AI僅需增加一行代碼，提供cluster信息以及單機(jī)訓(xùn)練模型即可獲得分布式訓(xùn)練能力，并且原生支持包括Hugging Face，Timm等熱門AI模型庫(kù)。

# wrap the model using auto_engine
model, optimizer = auto_engine(model, optimizer, cluster_info)
# normal training loop
...

因此，Colossal-AI可以極大地降低AI開發(fā)者的使用分布式技術(shù)訓(xùn)練和微調(diào)大模型門檻。同時(shí)，自動(dòng)并行系統(tǒng)可以從更細(xì)粒度搜索并行策略，找到更加高效的并行方案。

Graph Tracing

Colossal-AI是首個(gè)基于PyTorch框架使用靜態(tài)圖分析的自動(dòng)并行系統(tǒng)。PyTorch作為一個(gè)動(dòng)態(tài)圖框架，獲取其靜態(tài)的執(zhí)行計(jì)劃是機(jī)器學(xué)習(xí)系統(tǒng)領(lǐng)域被長(zhǎng)期研究的問(wèn)題。Colossal-AI使用基于torch.FX Tracer的ColoTracer，在tracing過(guò)程中推導(dǎo)并記錄了每個(gè)tensor的元信息，例如tensor shape，dims，dtype等，可以為后續(xù)的自動(dòng)并行策略搜索提供幫助。因此Colossal-AI具有更好的模型泛化能力，而不是依靠模型名或手動(dòng)修改來(lái)適配并行策略。

細(xì)粒度分布式訓(xùn)練策略搜索

Colossal-AI會(huì)在滿足內(nèi)存預(yù)算的限制下，以最快運(yùn)行時(shí)間為目標(biāo)，為每個(gè)op進(jìn)行策略搜索，最終得到真實(shí)訓(xùn)練時(shí)的策略，包括每個(gè)tensor的切分策略，不同計(jì)算節(jié)點(diǎn)間需要插入的通信算子類型，是否要進(jìn)行算子替換等。現(xiàn)有系統(tǒng)中的張量并行，數(shù)據(jù)并行，NVIDIA在Megatron-LM等并行系統(tǒng)中使用的column切分和row切分并行等混合并行，都是自動(dòng)并行可以搜索到的策略的子集。除了這些可以手動(dòng)指定的并行方式外，Colossal-AI的自動(dòng)并行系統(tǒng)有能力為每個(gè)op指定獨(dú)特的并行方式，因此有可能找到比依賴專家經(jīng)驗(yàn)和試錯(cuò)配置的手動(dòng)切分更好的并行策略。

分布式tensor與shape consistency系統(tǒng)

與PyTorch最新發(fā)布的DTensor類似，Colossal-AI也使用了device mesh對(duì)集群進(jìn)行了抽象管理。具體來(lái)說(shuō)，Colossal-AI使用sharding spec對(duì)tensor的分布式存儲(chǔ)狀態(tài)進(jìn)行標(biāo)注，使用shape consistency manager自動(dòng)地對(duì)同一tensor在不同sharding spec間進(jìn)行轉(zhuǎn)換。這讓Colossal-AI的通用性和易用性極大地提升，借助shape consistency manager可以沒有負(fù)擔(dān)地切分tensor，而不用擔(dān)心上游op的output與下游的input在集群中的存儲(chǔ)方式不同。

相較于PyTorch DTensor，Colossal-AI有以下3個(gè)優(yōu)勢(shì)：

1. Colossal-AI的device mesh可以profiling到集群性能指標(biāo)，對(duì)不同的通信算子進(jìn)行耗時(shí)估算。

2. Colossal-AI的shape consistency會(huì)貪心地搜索sharding spec間的轉(zhuǎn)換方式，而不是樸素地逐dimension進(jìn)行轉(zhuǎn)換，這樣能找到更高效的轉(zhuǎn)換路徑，進(jìn)而使得sharding spec間的轉(zhuǎn)換通信開銷更小。

3. 加入了all_to_all操作，使得Colossal-AI的擴(kuò)展性更強(qiáng)，這在大規(guī)模集群上進(jìn)行訓(xùn)練時(shí)，可以展現(xiàn)出很大的優(yōu)勢(shì)。

與activation checkpoint結(jié)合

作為大模型訓(xùn)練中必不可少的顯存壓縮技術(shù)，Colossal-AI也提供了對(duì)于activation checkpoint的自動(dòng)搜索功能。相比于大部分將最大顯存壓縮作為目標(biāo)的技術(shù)方案，Colossal-AI的搜索目標(biāo)是在顯存預(yù)算以內(nèi)，找到最快的activation checkpoint方案。同時(shí)，為了避免將activation checkpoint的搜索一起建模到SPMD solver中導(dǎo)致搜索時(shí)間爆炸，Colossal-AI做了2-stage search的設(shè)計(jì)，因此可以在合理的時(shí)間內(nèi)搜索到有效可行的分布式訓(xùn)練方案。

關(guān)于Colossal-AI

面向大模型時(shí)代的通用深度學(xué)習(xí)系統(tǒng)Colossal-AI，可實(shí)現(xiàn)高效快速部署AI大模型訓(xùn)練和推理，降低AI大模型應(yīng)用成本。自開源以來(lái)，Colossal-AI已經(jīng)多次在GitHub熱榜位列世界第一，獲得GitHub Star超七千顆，并成功入選SC、AAAI、PPoPP等國(guó)際AI與HPC頂級(jí)會(huì)議的官方教程。

Colossal-AI相關(guān)解決方案已成功在自動(dòng)駕駛、云計(jì)算、零售、醫(yī)藥、芯片等行業(yè)知名廠商落地應(yīng)用，廣受好評(píng)。例如，近期火爆的ChatGPT尚未開源，且不具備聯(lián)網(wǎng)功能。Colossal-AI已成功幫助某世界500強(qiáng)企業(yè)，開發(fā)具備在線搜索引擎能力增強(qiáng)的聊天機(jī)器人模型。

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

參考鏈接：https://www.hpc-ai.tech/blog/colossal-ai-0-2-0文章來(lái)源地址http://www.zghlxwxcb.cn/news/detail-496562.html