無縫支持Hugging Face社區(qū)，Colossal-AI低成本輕松加速大模型
?大模型已成為 AI 圈的一種潮流，不僅橫掃各大性能榜單，更產(chǎn)生了諸多有趣應(yīng)用。例如，微軟和 OpenAI 開發(fā)的自動(dòng)代碼建議補(bǔ)全神器 Copilot，化身程序員最佳助手，提升工作效率。
OpenAI 剛剛發(fā)布能以假亂真的文本生成圖像模型 DALL-E 2，Google 便緊接著發(fā)布了 Imagen，在大模型上，大公司也是相當(dāng)?shù)木?，絲毫不比 CV 刷榜差。
文本到圖像生成樣例“一個(gè)被貓絆倒的希臘人雕像”（左側(cè)兩列為 Imagen，右側(cè)兩列為 DALL·E 2）模型增大帶來的神奇表現(xiàn)，使得近幾年預(yù)訓(xùn)練模型規(guī)模呈現(xiàn)爆炸式增長。然而，訓(xùn)練甚至微調(diào)大模型都需要非常高的硬件成本，動(dòng)輒幾十、上百張 GPU。此外，PyTorch、TensorFlow 等現(xiàn)有深度學(xué)習(xí)框架也難以有效處理超大模型，通常需要專業(yè)的 AI 系統(tǒng)工程師做針對具體模型做適配和優(yōu)化。更重要的是，不是每一個(gè)實(shí)驗(yàn)室以及研發(fā)團(tuán)隊(duì)都具備 “鈔” 能力，能夠隨時(shí)調(diào)用大規(guī)模 GPU 集群來使用大模型，更不用提僅有一張顯卡的個(gè)人開發(fā)者。因此，盡管大模型已經(jīng)吸引了大量關(guān)注，高昂的上手門檻卻令大眾“望塵莫及”。
導(dǎo)致大模型使用成本增高的核心原因是顯存限制。GPU 計(jì)算雖快，但顯存容量有限，無法容納大模型。Colossal-AI 針對這一痛點(diǎn)，通過異構(gòu)內(nèi)存系統(tǒng)，高效地同時(shí)使用 GPU 顯存以及價(jià)格低廉的 CPU 內(nèi)存，在僅有一塊 GPU 的個(gè)人 PC 上便能訓(xùn)練高達(dá) 180 億參數(shù) GPT，可提升模型容量十余倍，大幅度降低了 AI 大模型微調(diào)和推理等下游任務(wù)和應(yīng)用部署的門檻，還能便捷擴(kuò)展至大規(guī)模分布式。Hugging Face 為深度學(xué)習(xí)社區(qū)提供了超過 5 萬個(gè) AI 模型的實(shí)現(xiàn)，最其中也不乏像 GPT, OPT 這樣的大模型，現(xiàn)已成為最流行的 AI 庫之一。
Colossal-AI 無縫支持 Hugging Face 社區(qū)模型，讓大模型對每一位開發(fā)者都變得觸手可及。接下來，我們將以 Meta 發(fā)布的大模型 OPT 為例，展現(xiàn)如何使用 Colossal-AI，僅需添加幾行代碼，便可實(shí)現(xiàn)大模型的低成本訓(xùn)練和微調(diào)。
開源地址：https://github.com/hpcaitech/ColossalAI低成本加速大模型 OPTOPT 模型OPT 的全稱為 Open Pretrained Transformer，是 Meta(Facebook) AI 實(shí)驗(yàn)室發(fā)布的對標(biāo) GPT-3 的大規(guī)模 Transformer 模型。與 OpenAI 尚未公開模型權(quán)重的 GPT-3 相比，Meta AI 慷慨地開源了所有的代碼以及模型權(quán)重，極大推動(dòng)了 AI 大模型的民主化，每一位開發(fā)者都能以此為基礎(chǔ)開發(fā)個(gè)性化的下游任務(wù)。接下來，我們將用 Hugging Face 提供的 OPT 模型的預(yù)訓(xùn)練權(quán)重進(jìn)行 Casual Language Modelling 的微調(diào)。添加配置文件 想要使用 Colossal-AI 中各個(gè)強(qiáng)大功能，用戶無需更改代碼訓(xùn)練邏輯，只用添加一個(gè)簡單的配置文件，即可賦予模型所期望的功能，比如混合精度、梯度累積、多維并行訓(xùn)練、冗余內(nèi)存優(yōu)化等。在一張 GPU 上，以異構(gòu)訓(xùn)練為例，我們只需在配置文件里加上相關(guān)配置項(xiàng)。其中 tensor_placement_policy 決定了我們異構(gòu)訓(xùn)練的策略，這個(gè)參數(shù)可以為 cuda、cpu 以及 auto。各個(gè)策略有不同的優(yōu)點(diǎn)：
- cuda: 將全部模型參數(shù)都放置于 GPU 上，適合不 offload 時(shí)仍然能進(jìn)行訓(xùn)練的傳統(tǒ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ù)傳輸。
對于一般用戶而言，僅需選擇 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)
運(yùn)行啟動(dòng)在配置文件準(zhǔn)備好之后，我們只需插入幾行代碼即可啟動(dòng)聲明的新功能。首先，通過一行代碼，使用配置文件啟動(dòng) Colossal-AI，Colossal-AI 會自動(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
)
接著，只需要調(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）
優(yōu)勢顯著在單張 GPU，與微軟 DeepSpeed 相比，Colossal-AI 的使用自動(dòng)化的 auto 策略，在不同的模型規(guī)模上相比 DeepSpeed 的 ZeRO Offloading 策略，均體現(xiàn)出顯著優(yōu)勢，最快可實(shí)現(xiàn) 40% 的加速。而 PyTorch 等傳統(tǒng)深度學(xué)習(xí)框架，在單張 GPU 上已經(jīng)無法運(yùn)行如此大的模型。
對于使用 8 張 GPU 的并行訓(xùn)練，Colossal-AI 僅需在啟動(dòng)命令中添加 - nprocs 8 即可實(shí)現(xiàn)！背后秘訣如此顯著的提升來自于 Colossal-AI 的高效異構(gòu)內(nèi)存管理子系統(tǒng) Gemini。簡單的來說，在模型訓(xùn)練時(shí)，Gemini 在前面的幾個(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)整張量位置，相比起 DeepSpeed 的 ZeRO Offload 的靜態(tài)劃分，Colossal-AI Gemini 能更高效利用 GPU 顯存和 CPU 內(nèi)存，實(shí)現(xiàn)在硬件極其有限的情況下，最大化模型容量和平衡訓(xùn)練速度。
對于大模型的代表 GPT，使用 Colossal-AI 在搭載 RTX 2060 6GB 的普通游戲筆記本上，也足以訓(xùn)練高達(dá) 15 億參數(shù)模型；對于搭載 RTX3090 24GB 的個(gè)人電腦，更是可以直接訓(xùn)練 180 億參數(shù)的模型；對于 Tesla V100 等專業(yè)計(jì)算卡，Colossal-AI 也能顯示出顯著改善。更進(jìn)一步：便捷高效并行擴(kuò)展并行分布式技術(shù)是進(jìn)一步加速模型訓(xùn)練的重要手段，想要以最短時(shí)間訓(xùn)練當(dāng)今世界最大最前沿的 AI 模型，仍離不開高效的分布式并行擴(kuò)展。針對現(xiàn)有方案并行維度有限、效率不高、通用性差、部署困難、缺乏維護(hù)等痛點(diǎn)，Colossal-AI 通過高效多維并行和異構(gòu)并行等技術(shù)，讓用戶僅需極少量修改，即可高效快速部署 AI 大模型訓(xùn)練。例如，對于同時(shí)使用數(shù)據(jù)并行、流水并行、2.5 維張量并行等復(fù)雜并行策略，僅需簡單聲明，即可自動(dòng)實(shí)現(xiàn)，Colossal-AI 無需像其他系統(tǒng)和框架侵入代碼，手動(dòng)處理復(fù)雜的底層邏輯。
Python
parallel = dict(
pipeline=2,
tensor=dict(mode='2.5d', depth = 1, size=4)
)
具體來說，對于 GPT-3 這樣的超大 AI 模型，相比英偉達(dá)方案，Colossal-AI 僅需一半的計(jì)算資源，即可啟動(dòng)訓(xùn)練；若使用相同計(jì)算資源，則能提速 11%，可降低 GPT-3 訓(xùn)練成本超百萬美元。Colossal-AI 相關(guān)解決方案已成功自動(dòng)駕駛、云計(jì)算、零售、醫(yī)藥、芯片等行業(yè)知名廠商落地應(yīng)用，廣受好評。
例如，對于蛋白質(zhì)結(jié)構(gòu)預(yù)測應(yīng)用 AlphaFold2，基于 Colossal-AI 的加速方案的 FastFold，成功超越谷歌和哥倫比亞大學(xué)的方案，將 AlphaFold2 訓(xùn)練時(shí)間從 11 天減少到 67 小時(shí)，且總成本更低，在長序列推理中也實(shí)現(xiàn) 9.3～11.6 倍的速度提升。
Colossal-AI 注重開源社區(qū)建設(shè)，提供中文教程，開放用戶社群及論壇，對于用戶反饋進(jìn)行高效交流與迭代更新，不斷添加 PaLM、AlphaFold 等前沿應(yīng)用。自然開源以來，Colossal-AI 已經(jīng)多次在 GitHub 及 Papers With Code 熱榜位列世界第一，與眾多已有數(shù)萬 star 的明星開源項(xiàng)目一起受到海內(nèi)外關(guān)注！
傳送門項(xiàng)目地址：https://github.com/hpcaitech/ColossalAI?
