從 AI 畫畫到 NLP 大模型，AIGC 的落地成本被一次性打下來了！
話不多說，直接看結果：
Stable Diffusion 2.0 訓練 / 微調 / 推理，顯存消耗最多可節省 5.6 倍，使硬件成本直降至1/46，一行代碼即可啟用；
1750 億參數大模型 BLOOM 單機推理，顯存消耗節省 4 倍，硬件成本直降至十幾分之一。
一行代碼實現自動搜索最佳并行策略，顯著降低分布式訓練上手門檻，原生支持 Hugging Face、Timm 等熱門 AI 模型庫。
要知道，在 AIGC 爆火的另一面，居高不下的成本困擾着整個行業。
上周，首批 AI 畫畫公司之一 StockAI 被迫宣布關閉平台。原因無他，創始人表示：
公司驅動成本太高了，目前的收入難以為繼。
即便如 ChatGPT 身後有 OpenAI 和微軟撐腰，也在平台開放幾周後發出公告，開始限制用戶每日使用次數。
言下之意無非四個字：燒不起了。
總而言之，降低 AI 大模型落地成本，是目前行業内亟需解決的問題。
與此同時，開源 AI 大模型解決方案Colossal-AI在過去一年迅速蹿紅，GitHub 上已經收獲 7k+Star。
如上降本方案，便出自它之手。
具體是怎麼實現的？往下看，
開源地址：
https://github.com/hpcaitech/ColossalAI
Stable Diffusion2.0 低成本訓練 / 微調 / 推理
相較于 1.0 版本，Stable Diffusion 2.0 不僅提升了模型生成圖像像素，還引入了 Depth2img 模型、text-guided 修複模型等，功能更加完善。
這波上新其實讓用戶們既驚喜又措手不及。
（畢竟 1.0 都還沒玩明白呢）
但話說回來，還是老問題，AIGC 模型落地的成本高啊。
以 Stable Diffusion 為例，其背後的 Stability AI 維護超過 4000 個英偉達 A100 的 GPU 集群，并已為此支出超過 5000 萬美元的運營成本。
面對快速叠代的模型、算法和下遊任務，如何降低應用成本成為 AIGC 真正走向落地的核心問題。
Stable Diffusion 2.0 基于簡單易用的 PyTorch Lightning 框架搭建。
作為 PyTorch Lightning 的官方大模型解決方案，Colossal-AI 第一時間進行跟進。
具體内容有以下幾點：
顯存消耗可節省 5.6 倍，硬件成本最多降至 1/46
支持 DreamBooth 單 GPU 快速個性化微調
推理顯存消耗節省 2.5 倍
而且該方案也将于近期合并進入 Hugging Face，進一步方便用戶使用。
訓練
為了加快訓練速度，降低訓練成本，使用更大的 batch size已成為被廣泛使用的有效手段。但 GPU 有限的顯存容量，嚴重限制了 batch size 大小，推高了訓練硬件門檻。
通過一系列顯存優化技術，Colossal-AI 使得 Stable Diffusion 平均在每個 GPU 上使用大 batch size 16 訓練的顯存需求，從 64.5GB 降低到了 11.6GB、節省 5.6 倍，還可擴展至單 GPU 或多 GPU 并行。
相比使用最先進的 A100 80GB，目前僅需3060 等消費級顯卡即可滿足需求，硬件成本最高直降至 1/46。
由此可以讓更多用戶在消費級 GPU 上，低成本地開展 Stable Diffusion 的相關研究與應用落地。
背後顯存優化 Flash Attention
早在 Stable Diffusion 1.0 版本，Colossal-AI 就率先引入了 Flash Attention 技術，成功将 attention 的速度提升 104%，将端到端訓練的峰值顯存減少 23%。
Flash Attention 是針對長序列 attention 的加速版本，使用 Flatten 來減少 GPU 高帶寬内存（HBM）之間的内存讀 / 寫次數，Flash attention 同時針對塊狀稀疏的 attention，設計了一個近似的注意力算法，比任何現有的近似 attention 方法都要快。
在 Stable Diffusion 1.0 版本，整個 Diffusion Model 隻有少量 attention 層，Flash attention 還沒有體現出其性能優勢。
在 Stable Diffusion 2.0 中，由于将大量卷積層替換為 attention 層，進一步發揮了 Flash Attention 的顯存優化潛力。
ZeRO + Gemini
Colossal-AI 支持使用零冗餘優化器（ZeRO）的方法來消除内存冗餘，與經典的數據并行性策略相比，可極大提高内存使用效率，同時不犧牲計算粒度和通信效率。
此外，Colossal-AI 還引入了Chunk 機制進一步提升 ZeRO 的性能。
運算順序上連續的一組參數存入一個 Chunk 中（Chunk 即一段連續的内存空間），每個 Chunk 的大小相同。
Chunk 方式組織内存可以保證 PCI-e 和 GPU-GPU 之間網絡帶寬的高效利用，減小了通信次數，同時避免潛在的内存碎片。
而 Colossal-AI 的異構内存空間管理器 Gemini 支持将優化器狀态從 GPU 卸載到 CPU，以節省 GPU 内存占用。
可以同時利用 GPU 内存、CPU 内存（由 CPU DRAM 或 NVMe SSD 内存組成）來突破單 GPU 内存牆的限制，進一步擴展了可訓練模型規模。
一行代碼快速上手
作為 PyTorch Lightning 的官方合作夥伴，僅需一行代碼即可調用 Colossal-AI 的上述顯存優化。
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 微調
在推出 Stable Diffusion 2.0 加速方案的同時，Colossal-AI 還 " 順手 " 發布了 DreamBooth 模型的微調方案。
這是谷歌在今年 8 月發布的模型。它隻需3-5 張圖片，再加上文字表述，就能讓指定物體遷移到其他場景或風格中去。
和 Dall-E 2、Imagen 等最大的不同是，DreamBooth 能對選定對象忠實還原。
方案中，用戶隻需直接運行文件 train_dreambooth_colossalai.py，即可在該微調任務上充分發揮 Colossal-AI 的顯存優化，個性化快速微調自己的圖文模型，極大降低使用門檻。
推理
由于模型推理對數值精度不敏感，這為實現低精度的低成本推理提供了可能。
對于 Stable Diffusion 2.0 模型，可以通過添加一行代碼，支持模型的 Int8 量化推理，顯存消耗節省 2.5 倍，僅需 3.1GB，且不造成顯著性能損失。
model = replace_module ( model )
用 RTX3090 推理 1750 億 BLOOM 模型
AI 畫畫爆火的另一邊，NLP 大模型的趨勢也還在延續。
今年 7 月，Hugging Face 發布了 1750 億參數開源模型BLOOM，它動用 384 塊 A100 煉成。
如果直接使用常見的 FP32/FP16 進行推理，在單節點 8 張 GPU 使用模型并行，每張 GPU 需要消耗至少 87.5GB/43.8GB 的顯存。
如此大的顯存占用，即使是最先進的 8 卡 A100（80GB/40GB）服務器，也無法直接部署推理服務，而多節點推理又會帶來沉重的額外成本和通信開銷。
基于這一現狀，Colossal-AI 實現了高效的 Int8 量化和模型并行推理，可以将 1750 億參數的 BLOOM 等大模型的推理服務，部署到 3090/4090 等消費級顯卡的 8 卡服務器，同時不産生顯著的 CPU 内存占用提升及性能損耗。
相比原有的 A100 方案，可将硬件部署成本降低到原有的十幾分之一。
通過對模型進行 Int8 量化，Colossal-AI 可将模型總體顯存占用從 352.3GB（FP16）降低到 185.6GB， 同時使用 Colossal-AI 的模型并行技術，将每張顯卡的占用減少到了 23.2GB。
在模型并行中，為了不增加 CPU 内存占用，Colossal-AI 在主進程中進行模型的量化和切分工作，其餘的每個進程中分别使用 lazy_init 獲得幾乎不占顯存和内存的 meta model，再通過 gloo backend 在進程之間傳遞模型參數。
通過上述方案，在沒有分段加載模型參數的情況下，便可以使得 CPU 内存占用峰值，達到理論上的較優水平。相較于将模型按層切分的 " 類流水線 " 分布方式，模型并行可以提高非密集請求下的顯存使用效率。
一行代碼自動并行
大模型的分布式混合部署是一個非常複雜的問題。
目前常見的分布式大模型訓練方案，都依賴于用戶人工反複嘗試以及系統專家基于經驗進行配置部署。
然而，這對于絕大多數 AI 開發者來說很不友好，因為大家都不希望把過多時間精力花費在研究分布式系統和試錯上。
由此，Colossal-AI 的高效易用自動并行系統，可以說是解大家燃眉之急了。
僅需增加一行代碼，它就能提供 cluster 信息以及單機訓練模型即可獲得分布式訓練能力，并且原生支持包括 Hugging Face，Timm 等熱門 AI 模型庫。
# wrap the model using auto_engine
model, optimizer = auto_engine ( model, optimizer, cluster_info )
# normal training loop
...
因此，Colossal-AI 可以極大地降低 AI 開發者的使用分布式技術訓練和微調大模型門檻。同時，自動并行系統可以從更細粒度搜索并行策略，找到更加高效的并行方案。
Graph Tracing
Colossal-AI 是首個基于 PyTorch 框架使用靜态圖分析的自動并行系統。
PyTorch 作為一個動态圖框架，獲取其靜态的執行計劃是機器學習系統領域被長期研究的問題。
Colossal-AI 使用基于 torch.FX Tracer 的 ColoTracer，在 tracing 過程中推導并記錄了每個 tensor 的元信息，例如 tensor shape，dims，dtype 等，可以為後續的自動并行策略搜索提供幫助。
因此 Colossal-AI 具有更好的模型泛化能力，而不是依靠模型名或手動修改來适配并行策略。
細粒度分布式訓練策略搜索
Colossal-AI 會在滿足内存預算的限制下，以最快運行時間為目标，為每個 op 進行策略搜索，最終得到真實訓練時的策略，包括每個 tensor 的切分策略，不同計算節點間需要插入的通信算子類型，是否要進行算子替換等。
現有系統中的張量并行、數據并行，NVIDIA 在 Megatron-LM 等并行系統中使用的 column 切分和 row 切分并行等混合并行，都是自動并行可以搜索到的策略的子集。
除了這些可以手動指定的并行方式外，Colossal-AI 的自動并行系統有能力為每個 op 指定獨特的并行方式，因此有可能找到比依賴專家經驗和試錯配置的手動切分更好的并行策略。
分布式 tensor 與 shape consistency 系統
與 PyTorch 最新發布的 DTensor 類似，Colossal-AI 也使用了 device mesh 對集群進行了抽象管理。
具體來說，Colossal-AI 使用 sharding spec 對 tensor 的分布式存儲狀态進行标注，使用 shape consistency manager 自動地對同一 tensor 在不同 sharding spec 間進行轉換。
這讓 Colossal-AI 的通用性和易用性極大地提升，借助 shape consistency manager 可以沒有負擔地切分 tensor，而不用擔心上遊 op 的 output 與下遊的 input 在集群中的存儲方式不同。
相較于 PyTorch DTensor，Colossal-AI 有以下 3 個優勢：
Colossal-AI 的 device mesh 可以 profiling 到集群性能指标，對不同的通信算子進行耗時估算。
Colossal-AI 的 shape consistency 會貪心地搜索 sharding spec 間的轉換方式，而不是樸素地逐 dimension 進行轉換，這樣能找到更高效的轉換路徑，進而使得 sharding spec 間的轉換通信開銷更小。
加入了 all_to_all 操作，使得 Colossal-AI 的擴展性更強，這在大規模集群上進行訓練時，可以展現出很大的優勢
與 activation checkpoint 結合
作為大模型訓練中必不可少的顯存壓縮技術，Colossal-AI 也提供了對于 activation checkpoint 的自動搜索功能。
相比于大部分将最大顯存壓縮作為目标的技術方案，Colossal-AI 的搜索目标是在顯存預算以内，找到最快的 activation checkpoint 方案。
同時，為了避免将 activation checkpoint 的搜索一起建模到 SPMD solver 中導緻搜索時間爆炸，Colossal-AI 做了 2-stage search 的設計，因此可以在合理的時間内搜索到有效可行的分布式訓練方案。
關于 Colossal-AI
通用深度學習系統 Colossal-AI 面向大模型時代，它可實現高效快速部署 AI 大模型訓練和推理，降低 AI 大模型應用成本。
自開源以來，Colossal-AI 已經多次在 GitHub 熱榜位列世界第一，獲得 GitHub Star 超七千顆，并成功入選 SC、AAAI、PPoPP 等國際 AI 與 HPC 頂級會議的官方教程。
Colossal-AI 相關解決方案已成功在自動駕駛、雲計算、零售、醫藥、芯片等行業知名廠商落地應用，廣受好評。
例如近期爆火的ChatGPT，尚未開源且不具備聯網功能。Colossal-AI 已成功幫助某世界 500 強企業，開發具備在線搜索引擎能力增強的聊天機器人模型。
傳送門
參考鍊接：
https://www.hpc-ai.tech/blog/colossal-ai-0-2-0
— 完 —
點這裡關注我，記得标星哦～
一鍵三連「分享」、「點贊」和「在看」
科技前沿進展日日相見 ~
