1塊GPU+幾行代碼，大模型訓(xùn)練提速40%！無(wú)縫支持HuggingFace，來(lái)自國(guó)產(chǎn)開源項(xiàng)目

明敏 2022-07-13 13:38:32 來(lái)源：量子位

明敏 發(fā)自 凹非寺

量子位 | 公眾號(hào) QbitAI

不得不說(shuō)，為了讓更多人能用上大模型，技術(shù)圈真是各出奇招！

模型不夠開放？有人自己上手搞免費(fèi)開源版。

比如最近風(fēng)靡全網(wǎng)的DALL·E Mini，Meta開放的OPT-175B（Open Pretrained Transformer）。

都是通過(guò)復(fù)刻的方式，讓原本不夠open的大模型，變成人人可用。

還有人覺得模型太大，個(gè)人玩家很難承受起天價(jià)成本。

所以提出異構(gòu)內(nèi)存、并行計(jì)算等方法，讓大模型訓(xùn)練加速又降本。

比如開源項(xiàng)目Colossal-AI，前不久剛實(shí)現(xiàn)了讓一塊英偉達(dá)3090就能單挑180億參數(shù)大模型。

而在這兩天，他們又來(lái)了一波上新：

無(wú)縫支持Hugging Face社區(qū)模型，只需添加幾行代碼，就能實(shí)現(xiàn)大模型的低成本訓(xùn)練和微調(diào)。

要知道，Hugging Face作為當(dāng)下最流行的AI庫(kù)之一，提供了超過(guò)5萬(wàn)個(gè)AI模型的實(shí)現(xiàn)，是許多AI玩家訓(xùn)練大模型的首選。

而Colossal-AI這波操作，是讓公開模型的訓(xùn)練微調(diào)變得更加切實(shí)可行。

并且在訓(xùn)練效果上也有提升。

單張GPU上，相比于微軟的DeepSpeed，使用Colossal-AI的自動(dòng)優(yōu)化策略，最快能實(shí)現(xiàn)40%的加速。

而PyTorch等傳統(tǒng)深度學(xué)習(xí)框架，在單張GPU上已經(jīng)無(wú)法運(yùn)行如此大的模型。

對(duì)于使用8張GPU的并行訓(xùn)練，僅需在啟動(dòng)命令中添加-nprocs 8就能實(shí)現(xiàn)。

這波下來(lái)，可以說(shuō)是把個(gè)人AI玩家需要考慮的成本、效率、實(shí)操問題，都拿捏住了~

無(wú)需修改代碼邏輯

光說(shuō)不練假把式。

下面就以O(shè)PT為例，詳細(xì)展開看看Colossal-AI的新功能到底怎么用。

OPT，全稱為Open Pretrained Transformer。

它由Meta AI發(fā)布，對(duì)標(biāo)GPT-3，最大參數(shù)量可達(dá)1750億。

最大特點(diǎn)就是，GPT-3沒有公開模型權(quán)重，而OPT開源了所有代碼及權(quán)重。

因此，每一位開發(fā)者都能在此基礎(chǔ)上開發(fā)個(gè)性化的下游任務(wù)。

下面的舉例，就是根據(jù)OPT提供的預(yù)訓(xùn)練權(quán)重，進(jìn)行因果語(yǔ)言模型（Casual Language Modelling）的微調(diào)。

主要分為兩個(gè)步驟：

• 添加配置文件
• 運(yùn)行啟動(dòng)

第一步，是根據(jù)想進(jìn)行的任務(wù)添加配置文件。

比如在一張GPU上，以異構(gòu)訓(xùn)練為例，只需在配置文件里加上相關(guān)配置項(xiàng)，并不需要更改代碼的訓(xùn)練邏輯。

比如，tensor_placement_policy決定了異構(gòu)訓(xùn)練的策略，參數(shù)可以為CUDA、CPU及auto。

每個(gè)策略的優(yōu)點(diǎn)不同、適應(yīng)的情況也不一樣。

CUDA：將全部模型參數(shù)都放置于GPU上，適合不offload時(shí)仍然能進(jìn)行訓(xùn)練的傳統(tǒng)場(chǎng)景。

CPU：將模型參數(shù)都放置在CPU內(nèi)存中，僅在GPU顯存中保留當(dāng)前參與計(jì)算的權(quán)重，適合超大模型的訓(xùn)練。

auto：根據(jù)實(shí)時(shí)的內(nèi)存信息，自動(dòng)決定保留在GPU顯存中的參數(shù)量，這樣能最大化利用GPU顯存，同時(shí)減少CPU-GPU之間的數(shù)據(jù)傳輸。

對(duì)于普通用戶來(lái)說(shuō)，使用auto策略是最便捷的。

這樣可以由Colossal-AI自動(dòng)化地實(shí)時(shí)動(dòng)態(tài)選擇最佳異構(gòu)策略，最大化計(jì)算效率。

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))
            

第二步，是在配置文件準(zhǔn)備好后，插入幾行代碼來(lái)啟動(dòng)新功能。

首先，通過(guò)一行代碼，使用配置文件來(lái)啟動(dòng)Colossal-AI。

Colossal-AI會(huì)自動(dòng)初始化分布式環(huán)境，讀取相關(guān)配置，然后將配置里的功能自動(dòng)注入到模型及優(yōu)化器等組件中。

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

然后，還是像往常一樣定義數(shù)據(jù)集、模型、優(yōu)化器、損失函數(shù)等。

比如直接使用原生PyTorch代碼，在定義模型時(shí)，只需將模型放置于ZeroInitContext下初始化即可。

在這里，使用的是Hugging Face提供的OPTForCausalLM模型以及預(yù)訓(xùn)練權(quán)重，在Wikitext數(shù)據(jù)集上進(jìn)行微調(diào)。

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
            )

接下來(lái)，只需要調(diào)用colossalai.initialize，便可將配置文件里定義的異構(gòu)內(nèi)存功能統(tǒng)一注入到訓(xùn)練引擎中，即可啟動(dòng)相應(yīng)功能。

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異構(gòu)

而能夠讓用戶實(shí)現(xiàn)如上“傻瓜式”操作的關(guān)鍵，還是AI系統(tǒng)本身要足夠聰明。

發(fā)揮核心作用的是Colossal-AI系統(tǒng)的高效異構(gòu)內(nèi)存管理子系統(tǒng)Gemini。

它就像是系統(tǒng)內(nèi)的一個(gè)總管，在收集好計(jì)算所需的信息后，動(dòng)態(tài)分配CPU、GPU的內(nèi)存使用。

具體工作原理，就是在前面幾個(gè)step進(jìn)行預(yù)熱，收集PyTorch動(dòng)態(tài)計(jì)算圖中的內(nèi)存消耗信息。

在預(yù)熱結(jié)束后，計(jì)算一個(gè)算子前，利用收集的內(nèi)存使用記錄，Gemini將預(yù)留出這個(gè)算子在計(jì)算設(shè)備上所需的峰值內(nèi)存，并同時(shí)從GPU顯存移動(dòng)一些模型張量到CPU內(nèi)存。

Gemini內(nèi)置的內(nèi)存管理器給每個(gè)張量都標(biāo)記一個(gè)狀態(tài)信息，包括HOLD、COMPUTE、FREE等。

然后，根據(jù)動(dòng)態(tài)查詢到的內(nèi)存使用情況，不斷動(dòng)態(tài)轉(zhuǎn)換張量狀態(tài)、調(diào)整張量位置。

帶來(lái)的直接好處，就是能在硬件非常有限的情況下，最大化模型容量和平衡訓(xùn)練速度。

要知道，業(yè)界主流方法ZeRO （Zero Reduency Optimizer），盡管也利用CPU+GPU異構(gòu)內(nèi)存的方法，但是由于是靜態(tài)劃分，還是會(huì)引起系統(tǒng)崩潰、不必要通信量等問題。

而且，使用動(dòng)態(tài)異構(gòu)CPU+GPU內(nèi)存的辦法，還能用加內(nèi)存條的辦法來(lái)擴(kuò)充內(nèi)存。

怎么也比買高端顯卡劃算多了。

目前，使用Colossal-AI的方法，RTX 2060 6GB普通游戲本能訓(xùn)練15億參數(shù)模型；RTX 3090 24GB主機(jī)直接單挑180億參數(shù)大模型；Tesla V100 32GB連240億參數(shù)都能拿下。

除了最大化利用內(nèi)存外，Colossal-AI還使用分布式并行的方法，讓訓(xùn)練速度不斷提升。

它提出同時(shí)使用數(shù)據(jù)并行、流水并行、2.5維張量并行等復(fù)雜并行策略。

方法雖復(fù)雜，但上手卻還是非?！吧倒喜僮鳌保恍韬?jiǎn)單聲明，就能自動(dòng)實(shí)現(xiàn)。

無(wú)需像其他系統(tǒng)和框架侵入代碼，手動(dòng)處理復(fù)雜的底層邏輯。

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

Colossal-AI還能做什么？

實(shí)際上，自開源以來(lái)，Colossal-AI已經(jīng)多次在GitHub及Papers With Code熱榜位列世界第一，在技術(shù)圈小有名氣。

除了如上提到的用單張GPU訓(xùn)練大模型外，Colossal-AI在擴(kuò)展至數(shù)十張甚至數(shù)百?gòu)圙PU的大規(guī)模并行場(chǎng)景時(shí)，相比于英偉達(dá)Megatron-LM等現(xiàn)有系統(tǒng)，性能可以翻倍，使用資源可以降低至其十分之一之下。

換算一下，在預(yù)訓(xùn)練GPT-3等超大AI模型上，節(jié)省的費(fèi)用可以達(dá)到數(shù)百萬(wàn)元。

據(jù)透露，Colossal-AI相關(guān)的解決方案已經(jīng)被自動(dòng)駕駛、云計(jì)算、零售、醫(yī)藥、芯片等行業(yè)的知名廠商用上了。

與此同時(shí)，他們也非常注重開源社區(qū)建設(shè)，提供中文教程、開放用戶社群論壇，根據(jù)大家的需求反饋不斷更新迭代。

比如我們發(fā)現(xiàn)，之前有粉絲留言詢問，Colossal-AI能否直接加載Hugging Face上的一些模型？

好嘛，這次更新就來(lái)了。

所以，對(duì)于大模型訓(xùn)練，你覺得現(xiàn)在還有哪些難點(diǎn)亟需解決呢？

歡迎評(píng)論區(qū)留言討論~

傳送門

項(xiàng)目地址：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