鑽石舞台

公眾號關注「GitHubDaily」

設為 「星標」，每天帶你逛 GitHub！

大家好，我是小 G。

最近兩年，大模型已成為 AI 圈的一種潮流，不僅橫掃各大性能榜單，更產生了諸多有趣應用。

例如，微軟和 OpenAI 開發的自動代碼建議補全神器 Copilot，化身程序員最佳助手，提升工作效率。

OpenAI 剛剛發布能以假亂真的文本生成圖像模型 DALL-E 2，Google 便緊接着發布了 Imagen，在大模型上，大公司也是相當的卷，絲毫不比 CV 刷榜差。

文本到圖像生成樣例 「一個被貓絆倒的希臘人雕像」（左側兩列為 Imagen，右側兩列為 DALL・E 2）

模型增大帶來的神奇表現，使得近幾年預訓練模型規模呈現爆炸式增長。然而，訓練甚至微調大模型都需要非常高的硬件成本，動輒幾十、上百張 GPU。

此外，PyTorch、TensorFlow 等現有深度學習框架也難以有效處理超大模型，通常需要專業的 AI 系統工程師做針對具體模型做適配和優化。

更重要的是，不是每一個實驗室以及研發團隊都具備 「鈔」 能力，能夠隨時調用大規模 GPU 集群來使用大模型，更不用提僅有一張顯卡的個人開發者。因此，儘管大模型已經吸引了大量關注，高昂的上手門檻卻令大眾 「望塵莫及」。

導致大模型使用成本增高的核心原因是顯存限制。GPU 計算雖快，但顯存容量有限，無法容納大模型。

在前不久的文章《國人主導的開源項目，斬獲 3700 個 GitHub Star！》中，我們介紹了 Colossal-AI 最新研發的異構內存系統，它能夠高效地同時使用 GPU 顯存以及價格低廉的 CPU 內存，可提升模型容量十餘倍，以低成本高效訓練 AI 大模型，兼容目前主流的訓練任務和優化加速技術，還能便捷擴展至大規模分布式。

Hugging Face 為深度學習社區提供了超過 5 萬個 AI 模型的實現，最其中也不乏像 GPT, OPT 這樣的大模型，現已成為最流行的 AI 庫之一。

Colossal-AI 無縫支持 Hugging Face 社區模型，讓大模型對每一位開發者都變得觸手可及。接下來，我們將以 Meta 發布的大模型 OPT 為例，展現如何使用 Colossal-AI，僅需添加幾行代碼，便可實現大模型的低成本訓練和微調。

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

低成本加速大模型 OPT

OPT 模型

OPT 的全稱為 Open Pretrained Transformer，是 Meta (Facebook) AI 實驗室發布的對標 GPT-3 的大規模 Transformer 模型。與 OpenAI 尚未公開模型權重的 GPT-3 相比，Meta AI 慷慨地開源了所有的代碼以及模型權重，極大推動了 AI 大模型的民主化，每一位開發者都能以此為基礎開發個性化的下游任務。

接下來，我們將用 Hugging Face 提供的 OPT 模型的預訓練權重進行 Casual Language Modelling 的微調。

添加配置文件

想要使用 Colossal-AI 中各個強大功能，用戶無需更改代碼訓練邏輯，只用添加一個簡單的配置文件，即可賦予模型所期望的功能，比如混合精度、梯度累積、多維並行訓練、冗餘內存優化等。

在一張 GPU 上，以異構訓練為例，我們只需在配置文件里加上相關配置項。其中 tensor_placement_policy 決定了我們異構訓練的策略，這個參數可以為 cuda、cpu 以及 auto。各個策略有不同的優點：

cuda: 將全部模型參數都放置於 GPU 上，適合不 offload 時仍然能進行訓練的傳統場景；

cpu 則會將模型參數都放置在 CPU 內存中，僅在 GPU 顯存中保留當前參與計算的權重，適合超大模型的訓練；

auto 則會根據實時的內存信息，自動決定保留在 GPU 顯存中的參數量，這樣能最大化利用 GPU 顯存，同時減少 CPU-GPU 之間的數據傳輸。

對於一般用戶而言，僅需選擇 auto 策略，由 Colossal-AI 自動化地實時動態選擇最佳異構策略，最大化計算效率。

from colossalai.zero.shard_utils import TensorShardStrategyzero = dict(model_config=dict(shard_strategy=TensorShardStrategy(), tensor_placement_policy="auto"), optimizer_config=dict(gpu_margin_mem_ratio=0.8))

運行啟動

在配置文件準備好之後，我們只需插入幾行代碼即可啟動聲明的新功能。

首先，通過一行代碼，使用配置文件啟動 Colossal-AI，Colossal-AI 會自動初始化分布式環境，並讀取相關配置，之後將配置里的功能自動注入到模型以及優化器等組件中。

colossalai.launch_from_torch(config='./configs/colossalai_zero.py')

接下來，用戶可以照常定義數據集、模型、優化器、損失函數等，例如直接使用原生 PyTorch 代碼。在定義模型時，只需將模型放置於 ZeroInitContext 下初始化即可。

在例子裡，我們使用 Hugging Face 提供的 OPTForCausalLM 模型以及預訓練權重，在 Wikitext 數據集上進行微調。

with ZeroInitContext(target_device=torch.cuda.current_device(), shard_strategy=shard_strategy, shard_param=True): model = OPTForCausalLM.from_pretrained( 'facebook/opt-1.3b' config=config )

接着，只需要調用 colossalai.initialize，便可將配置文件里定義的異構內存功能統一注入到訓練引擎中，即可啟動相應功能。

engine, train_dataloader, eval_dataloader, lr_scheduler = colossalai.initialize(model=model, optimizer=optimizer, criterion=criterion, train_dataloader=train_dataloader, test_dataloader=eval_dataloader, lr_scheduler=lr_scheduler)優勢顯著

在單張 GPU，與微軟 DeepSpeed 相比，Colossal-AI 的使用自動化的 auto 策略，在不同的模型規模上相比 DeepSpeed 的 ZeRO Offloading 策略，均體現出顯著優勢，最快可實現 40% 的加速。而 PyTorch 等傳統深度學習框架，在單張 GPU 上已經無法運行如此大的模型。

對於使用 8 張 GPU 的並行訓練，Colossal-AI 僅需在啟動命令中添加 - nprocs 8 即可實現！

背後秘訣

如此顯著的提升來自於 Colossal-AI 的高效異構內存管理子系統 Gemini。簡單的來說，在模型訓練時，Gemini 在前面的幾個 step 進行預熱，收集 PyTorch 動態計算圖中的內存消耗信息；在預熱結束後，計算一個算子前，利用收集的內存使用記錄，Gemini 將預留出這個算子在計算設備上所需的峰值內存，並同時從 GPU 顯存里移動一些模型張量到 CPU 內存。

Gemini 內置的內存管理器給每個張量都標記一個狀態信息，包括 HOLD、COMPUTE、FREE 等。然後，根據動態查詢到的內存使用情況，不斷動態轉換張量狀態、調整張量位置，相比起 DeepSpeed 的 ZeRO Offload 的靜態劃分，Colossal-AI Gemini 能更高效利用 GPU 顯存和 CPU 內存，實現在硬件極其有限的情況下，最大化模型容量和平衡訓練速度。

對於大模型的代表 GPT，使用 Colossal-AI 在搭載 RTX 2060 6GB 的普通遊戲筆記本上，也足以訓練高達 15 億參數模型；對於搭載 RTX3090 24GB 的個人電腦，更是可以直接訓練 180 億參數的模型；對於 Tesla V100 等專業計算卡，Colossal-AI 也能顯示出顯著改善。

便捷高效並行擴展

想要以最短時間訓練當今世界最大最前沿的 AI 模型，仍離不開高效的分布式並行擴展。對於同時使用數據並行、流水並行、2.5 維張量並行等複雜並行策略，僅需簡單聲明，即可自動實現，Colossal-AI 無需像其他系統和框架侵入代碼，手動處理複雜的底層邏輯。

parallel = dict( pipeline=2, tensor=dict(mode='2.5d', depth = 1, size=4))

在面對擴展至數十甚至數百張 GPU 的大規模並行場景時，相比英偉達 Megatron-LM 等現有系統，Colossal-AI 在性能上仍顯示出顯著的加速性能和資源節約。這意味着對於預訓練 GPT-3 等超大 AI 模型，可節省數百萬元的訓練費用。

Colossal-AI 相關解決方案已成功在自動駕駛、雲計算、零售、醫藥、芯片等行業知名廠商落地應用，廣受好評。

Colossal-AI 注重開源社區建設，提供中文教程，開放用戶社群及論壇，對於用戶反饋進行高效交流與迭代更新，不斷添加 PaLM、AlphaFold 等前沿應用。

自開源以來，Colossal-AI 已經多次在 GitHub 及 Papers With Code 熱榜位列世界第一，與眾多已有數萬 star 的明星開源項目一起受到海內外關注！

傳送門

GitHub：https://github.com/hpcaitech/ColossalAI

參考文獻：

https://arxiv.org/abs/2202.05924v2

https://arxiv.org/abs/2205.11487

https://github.com/features/copilot

https://github.com/huggingface/transformers

https://medium.com/@yangyou_berkeley/colossal-ai-seamlessly-accelerates-large-models-at-low-costs-with-hugging-face-4d1a887e500d