不得不說，為了讓更多人能用上大模型，技術圈真是各出奇招！
模型不夠開放？有人自己上手搞免費開源版。
比如最近風靡全網的
DALL·E Mini
，Meta開放的
OPT-175B
（Open Pretrained Transformer）。
都是透過復刻的方式，讓原本不夠open的大模型，變成
人人可用
。
還有人覺得模型太大，個人玩家很難承受起天價成本。
所以提出異構記憶體、平行計算等方法，讓大模型訓練加速又降本。
比如開源專案
Colossal-AI
，前不久剛實現了讓一塊英偉達3090就能單挑180億引數大模型。
而在這兩天，他們又來了一波上新：
無縫支援Hugging Face社群模型，
只需新增幾行程式碼
，就能實現大模型的低成本訓練和微調。
要知道，Hugging Face作為當下最流行的AI庫之一，提供了
超過5萬個
AI模型的實現，是許多AI玩家訓練大模型的首選。
而Colossal-AI這波操作，是讓公開模型的訓練微調變得更加切實可行。
並且在訓練效果上也有提升。
單張GPU
上，相比於微軟的DeepSpeed，使用Colossal-AI的自動最佳化策略，最快能實現
40%
的加速。
而PyTorch等傳統深度學習框架，在單張GPU上已經
無法執行
如此大的模型。
對於使用8張GPU的並行訓練，僅需在啟動命令中新增-nprocs 8就能實現。
這波下來，可以說是把個人AI玩家需要考慮的成本、效率、實操問題，都拿捏住了~
無需修改程式碼邏輯
光說不練假把式。
下面就以OPT為例，詳細展開看看Colossal-AI的新功能到底怎麼用。
OPT，全稱為Open Pretrained Transformer。
它由Meta AI釋出，
對標GPT-3
，最大引數量可達1750億。
最大特點就是，GPT-3沒有公開模型權重，而OPT開源了所有程式碼及權重。
因此，每一位開發者都能在此基礎上開發個性化的下游任務。
下面的舉例，就是根據OPT提供的預訓練權重，進行因果語言模型（Casual Language Modelling）的微調。
主要分為兩個步驟：
新增配置檔案執行啟動
第一步，是根據想進行的任務新增配置檔案。
比如在一張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 TensorShardStrategy zero = dict（model_config=dict（shard_strategy=TensorShardStrategy（）， ）， 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+CPU異構
而能夠讓使用者實現如上“傻瓜式”操作的關鍵，還是AI系統本身要足夠聰明。
發揮核心作用的是Colossal-AI系統的高效異構記憶體管理子系統
Gemini
。
它就像是系統內的一個總管，在收集好計算所需的資訊後，動態分配CPU、GPU的記憶體使用。
具體工作原理，就是在前面幾個step進行預熱，收集PyTorch動態計算圖中的記憶體消耗資訊。
在預熱結束後，計算一個運算元前，利用收集的記憶體使用記錄，Gemini將預留出這個運算元在計算裝置上
所需的峰值記憶體
，並同時從GPU視訊記憶體移動一些模型張量到CPU記憶體。
Gemini內建的記憶體管理器給每個張量都標記一個狀態資訊，包括HOLD、COMPUTE、FREE等。
然後，根據
動態查詢
到的記憶體使用情況，不斷動態轉換張量狀態、調整張量位置。
帶來的直接好處，就是能在硬體非常有限的情況下，
最大化
模型容量和平衡訓練速度。
要知道，業界主流方法ZeRO （Zero Reduency Optimizer），儘管也利用CPU+GPU異構記憶體的方法，但是由於是靜態劃分，還是會引起系統崩潰、不必要通訊量等問題。
而且，使用動態異構CPU+GPU記憶體的辦法，還能用
加記憶體條
的辦法來擴充記憶體。
怎麼也比買高階顯示卡划算多了。
目前，使用Colossal-AI的方法，RTX 2060 6GB普通遊戲本能訓練
15億
引數模型；RTX 3090 24GB主機直接單挑
180億
引數大模型；Tesla V100 32GB連
240億
引數都能拿下。
除了最大化利用記憶體外，Colossal-AI還使用分散式並行的方法，讓訓練速度不斷提升。
它提出同時使用資料並行、流水並行、2。5維張量並行等
複雜並行策略
。
方法雖複雜，但上手卻還是非常“傻瓜操作”，只需簡單宣告，就能自動實現。
無需像其他系統和框架侵入程式碼，手動處理複雜的底層邏輯。
parallel = dict（ pipeline=2， ， depth = 1， size=4） ）
Colossal-AI還能做什麼？
實際上，自開源以來，Colossal-AI已經多次在GitHub及Papers With Code熱榜位列
世界第一
，在技術圈小有名氣。
除了如上提到的用單張GPU訓練大模型外，Colossal-AI在擴充套件至數十張甚至數百張GPU的
大規模並行場景
時，相比於英偉達Megatron-LM等現有系統，效能可以翻倍，使用資源可以降低至其十分之一之下。
換算一下，在預訓練GPT-3等超大AI模型上，節省的費用可以達到數百萬元。
據透露，Colossal-AI相關的解決方案已經被自動駕駛、雲計算、零售、醫藥、晶片等行業的知名廠商用上了。
與此同時，他們也非常注重開源社群建設，提供中文教程、開放使用者社群論壇，根據大家的需求反饋不斷更新迭代。
比如我們發現，之前有粉絲留言詢問，Colossal-AI能否直接載入Hugging Face上的一些模型？
好嘛，這次更新就來了。
所以，對於大模型訓練，你覺得現在還有哪些難點亟需解決呢？
歡迎評論區留言討論~
傳送門
專案地址：https：//github。com/hpcaitech/ColossalAI
參考連結：
［1］https：//medium。com/@yangyou_berkeley/colossal-ai-seamlessly-accelerates-large-models-at-low-costs-with-hugging-face-4d1a887e500d
［2］https：//arxiv。org/abs/2202。05924v2
［3］https：//arxiv。org/abs/2205。11487
［4］https：//github。com/features/copilot
［5］https：//github。com/huggingface/transformers