無縫支持Hugging Face社區,Colossal-AI低成本輕鬆加速大模型
大模型已成為AI圈的一種潮流,不僅橫掃各大性能榜單,更產生了諸多有趣應用。
例如,微軟和OpenAI開發的自動代碼建議補全神器Copilot,化身程序員最佳助手,提升工作效率。
![](http://oscimg.oschina.net/oscnet/up-b7087ab0b667c1ddd7238bbc60597779d5e.gif)
OpenAI 剛剛發佈能以假亂真的文本生成圖像模型 DALL-E 2,Google 便緊接着發佈了 Imagen,在大模型上,大公司也是相當的卷,絲毫不比 CV 刷榜差。
![](http://oscimg.oschina.net/oscnet/up-db13b0004531179928e62770aefb930b5b2.png)
文本到圖像生成樣例“一個被貓絆倒的希臘人雕像”(左側兩列為Imagen,右側兩列為DALL·E 2)
模型增大帶來的神奇表現,使得近幾年預訓練模型規模呈現爆炸式增長。然而,訓練甚至微調大模型都需要非常高的硬件成本,動輒幾十、上百張GPU。此外,PyTorch、TensorFlow等現有深度學習框架也難以有效處理超大模型,通常需要專業的AI系統工程師做針對具體模型做適配和優化。
更重要的是,不是每一個實驗室以及研發團隊都具備“鈔”能力,能夠隨時調用大規模GPU集羣來使用大模型,更不用提僅有一張顯卡的個人開發者。因此,儘管大模型已經吸引了大量關注,高昂的上手門檻卻令大眾“望塵莫及”。
![](http://oscimg.oschina.net/oscnet/up-3a9a03121ad9f2922ad2eec66b128203040.png)
導致大模型使用成本增高的核心原因是顯存限制。GPU計算雖快,但顯存容量有限,無法容納大模型。Colossal-AI針對這一痛點,通過異構內存系統,高效地同時使用GPU顯存以及價格低廉的CPU內存,在僅有一塊GPU的個人PC上便能訓練高達180億參數GPT,
可提升模型容量十餘倍,大幅度降低了AI大模型微調和推理等下游任務和應用部署的門檻,還能便捷擴展至大規模分佈式。
Hugging Face為深度學習社區提供了超過5萬個AI模型的實現,最其中也不乏像GPT, OPT這樣的大模型,現已成為最流行的AI庫之一。
![](http://oscimg.oschina.net/oscnet/up-6022e2aaf5fec78bdcb9030d565bed13b81.png)
Colossal-AI無縫支持Hugging Face社區模型,讓大模型對每一位開發者都變得觸手可及。接下來,我們將以Meta發佈的大模型OPT為例,展現如何使用Colossal-AI,
僅需添加幾行代碼,便可實現大模型的低成本訓練和微調。
低成本加速大模型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 TensorShardStrategy
zero = 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上已經無法運行如此大的模型。
![](http://oscimg.oschina.net/oscnet/up-3e4ee188f4cec16541bbe8dbc442a84656a.jpg)
對於使用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內存,實現在硬件極其有限的情況下,最大化模型容量和平衡訓練速度。
![](http://oscimg.oschina.net/oscnet/up-28605316fa1ead9871de91888ce07039cf3.png)
對於大模型的代表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模型,可節省數百萬元的訓練費用。
![](http://oscimg.oschina.net/oscnet/up-a6ee9fa0351f52f37799a7ed4ffd06c853b.png)
Colossal-AI相關解決方案已成功在
自動駕駛、雲計算、零售、醫藥、芯片等行業知名廠商落地應用,廣受好評。
Colossal-AI注重開源社區建設,提供中文教程,開放用户社羣及論壇,對於用户反饋進行高效交流與迭代更新,不斷添加PaLM、AlphaFold等前沿應用。
自然開源以來,Colossal-AI已經
多次在GitHub及Papers With Code熱榜位列世界第一,與眾多已有數萬star的明星開源項目一起受到海內外關注!
傳送門
項目地址:
招聘
潞晨科技還在廣納英才,招聘全職/實習AI分佈式系統、架構、編譯器、網絡、CUDA、SaaS、k8s等核心系統研發人員,開源社區運營、銷售人員。
潞晨科技提供有競爭力的薪資回報,特別優秀的,還可以申請遠程工作。也歡迎各位向潞晨科技引薦優秀人才,如果您推薦優秀人才成功簽約潞晨科技,我們將為您提供數千元至數萬元的推薦費。
工作地點:北京,新加坡,美國 (可相互轉崗)。
簡歷投遞郵箱:
[email protected]
參考文獻:
「其他文章」
- 開源方案復現ChatGPT流程!1.62GB顯存即可體驗,單機訓練提速7.73倍
- 硬件預算最高直降46倍!低成本上手AIGC和千億大模型,一行代碼自動並行,Colossal-AI再升級
- 潞晨科技完成600萬美元種子及天使輪融資,藍馳領投天使輪
- 無縫支持Hugging Face社區,Colossal-AI低成本輕鬆加速大模型
- 推理加速性能超越英偉達FasterTransformer 50%,開源方案打通大模型落地關鍵路徑
- 在個人電腦用單塊GPU帶動180億參數GPT!熱門開源項目再添新特性
- 使用Colossal-AI分佈式訓練BERT模型
- 使用Colossal-AI復現Pathways Language Model
- 霸榜GitHub熱門第一多日後,Colossal-AI正式版發佈
- 助力藥物研發,低成本加速AlphaFold訓練從11天到67小時,11倍推理加速——開源解決方案FastFold