700 億參數(shù) LLaMA2 訓(xùn)練加速 195%,基礎(chǔ)大模型最佳實(shí)踐再升級
ChatGPT 引發(fā)的大模型熱潮愈演愈烈,全球科技巨頭和明星初創(chuàng)爭相入局,打造以 AI 大模型為核心的競爭力和多樣化商業(yè)使用需求。其中 LLaMA 系列模型,因良好的基礎(chǔ)能力和開放生態(tài),已積累了海量的用戶和實(shí)際應(yīng)用案例,成為無數(shù)開源模型后來者的模仿和競爭的標(biāo)桿對象。
但如何降低類 LLaMA2 大模型預(yù)訓(xùn)練成本,如何基于 LLaMA2 通過繼續(xù)預(yù)訓(xùn)練和微調(diào),低成本構(gòu)建 AI 大模型實(shí)際應(yīng)用,仍是 AIGC 相關(guān)企業(yè)面臨的關(guān)鍵瓶頸。
作為全球規(guī)模最大、最活躍的大模型開發(fā)工具與社區(qū),Colossal-AI 再次迭代,提供開箱即用的 8 到 512 卡 LLaMA2 訓(xùn)練、微調(diào)、推理方案,對 700 億參數(shù)訓(xùn)練加速 195%,并提供一站式云平臺解決方案,極大降低大模型開發(fā)和落地應(yīng)用成本。
開源地址:https://github.com/hpcaitech/ColossalAI
LLaMA2 訓(xùn)練加速 195%
Meta 開源的 LLaMA 系列大模型進(jìn)一步激發(fā)了打造類 ChatGPT 的熱情,并由此衍生出了諸多項(xiàng)目和應(yīng)用。

最新的 7B~70B LLaMA2 大模型,則進(jìn)一步提高了語言模型的基礎(chǔ)能力。但由于 LLaMA2 的預(yù)訓(xùn)練預(yù)料大部分來自英文通用知識,而僅用微調(diào)能夠提升和注入的領(lǐng)域知識和多語言能力也相對有限。此外,高質(zhì)量的專業(yè)知識和數(shù)據(jù)集通常被視為各個行業(yè)和公司的核心資產(chǎn),僅能以私有化形式保存。因此,以低成本預(yù)訓(xùn)練 / 繼續(xù)預(yù)訓(xùn)練 / 微調(diào) LLaMA2 系列大模型,結(jié)合高質(zhì)量私有化業(yè)務(wù)數(shù)據(jù)積累,幫助業(yè)務(wù)降本增效是眾多行業(yè)與企業(yè)的迫切需求與瓶頸。但 LLaMA2 大模型僅發(fā)布了原始模型權(quán)重與推理腳本,不支持訓(xùn)練 / 微調(diào),也未提供數(shù)據(jù)集。
針對上述空白與需求,Colossal-AI 開源了針對 LLaMA2 的全流程方案,并具備高可擴(kuò)展性,支持從 70 億到 700 億參數(shù)的模型,從 8 卡到 512 卡都可保持良好的性能。
在使用 8 卡訓(xùn)練 / 微調(diào) LLaMA2-7B 時,Colossal-AI 能達(dá)到約 54% 的硬件利用率(MFU),處于業(yè)界領(lǐng)先水平。而對于預(yù)訓(xùn)練任務(wù),以使用 512 張 A100 40GB 預(yù)訓(xùn)練 LLaMA2-70B 為例,DeepSpeed ZeRO3 策略因顯存不足而無法啟動,僅能通過速度衰減較大的 ZeRO3-offload 策略啟動。而 Colossal-AI 則因卓越的系統(tǒng)優(yōu)化和擴(kuò)展性,仍能保持良好性能,訓(xùn)練提速 195%。

Colossal-AI LLaMA-2 訓(xùn)練 / 微調(diào)方案的高性能來源于新的異構(gòu)內(nèi)存管理系統(tǒng) Gemini 和高性能算子(包括 Flash attention 2)等系統(tǒng)優(yōu)化。新 Gemini 提供了高可擴(kuò)展性,高魯棒性,高易用性的接口。其 Checkpoint 格式與 HuggingFace 完全兼容,減小了使用和轉(zhuǎn)換成本。其對于切分、offload 等的設(shè)置更加靈活且易用,能夠覆蓋更多硬件配置下的 LLaMA-2 訓(xùn)練 / 微調(diào)任務(wù)。僅需數(shù)行代碼即可使用:
from colossalai.booster import Booster
from colossalai.booster.plugin import GeminiPlugin
plugin = GeminiPlugin ()
booster = Booster (plugin=plugin)
model, optimizer, train_dataloader, criterion = booster.boost (model, optimizer, train_dataloader, criterion)ShardFormer 多維細(xì)粒度并行
雖然對于主流硬件條件和大多數(shù)模型,Colossal-AI 的新 Gemini 已經(jīng)能夠提供良好的性能。但是對于一些極端硬件條件,或者是特殊模型,可能仍然需要多維并行的細(xì)粒度優(yōu)化?,F(xiàn)有其他方案通常需要分布式系統(tǒng)資深專家,手動對代碼進(jìn)行大規(guī)模重構(gòu)和調(diào)優(yōu),Colossal-AI 的 ShardFormer 提供了開箱即用的多維并行和算子優(yōu)化的能力,僅需數(shù)行代碼即可使用,在單機(jī) / 大規(guī)模集群上都能提供良好的性能。
from colossalai.booster import Booster
from colossalai.booster.plugin import HybridParallelPlugin
from transformers.models.llama import LlamaForCausalLM, LlamaConfig
plugin = HybridParallelPlugin (tp_size=2, pp_size=2, num_microbatches=4, zero_stage=1)
booster = Booster (plugin=plugin)
model = LlamaForCausalLM (LlamaConfig ())
model, optimizer, train_dataloader, criterion = booster.boost (model, optimizer, train_dataloader, criterion)Colossal-AI ShardFormer 支持包括 LLaMA1/2、BLOOM、OPT、T5、GPT-2、BERT、GLM 在內(nèi)的主流開源模型,也可以直接使用 Huggingface/transformers 模型導(dǎo)入,Checkpoint 格式也與 HuggingFace 完全兼容,對比 Megatron-LM 等需重寫大量代碼的方案,大大提升了易用性。

對于并行策略,已支持以下多種并行方式:張量并行、流水線并行、序列并行、數(shù)據(jù)并行、Zero 數(shù)據(jù)并行等,并可將多種并行方式組合使用,只需通過簡單的配置命令,即可適配各種復(fù)雜的硬件環(huán)境 / 模型。同時,其內(nèi)置了各種高性能算子,免去了繁瑣的兼容 / 配置過程。其中包括:
- Flash attention 2
- Memory efficient attention (xformers)
- Fused Normalization Layer
- JIT kernels
云平臺大模型一站式解決
為了進(jìn)一步提升開發(fā)和部署效率,Colossal-AI 團(tuán)隊還將上述系統(tǒng)優(yōu)勢與算力結(jié)合,提供 Colossal-AI 云平臺,提供廉價算力和開箱即用的 AI 主流應(yīng)用,包括對話大模型,多模態(tài)模型,生物醫(yī)藥等,現(xiàn)已開啟內(nèi)測。

通過屏蔽大模型底層的分布式并行計算、內(nèi)存、通信管理與優(yōu)化等,AI 開發(fā)者可以繼續(xù)專注于 AI 模型與算法設(shè)計,以更低成本更快速度完成 AI 大模型助力業(yè)務(wù)降本增效。

用戶只需要上傳相關(guān)數(shù)據(jù),即可無代碼訓(xùn)練個性化私有模型,并將訓(xùn)練好的模型一鍵部署。相關(guān)的應(yīng)用都經(jīng)過 Colossal-AI 團(tuán)隊精心優(yōu)化,得益于算法和系統(tǒng)的雙面優(yōu)化,能大大降低模型訓(xùn)練以及部署的成本。

- Colossal-AI 云平臺:platform.luchentech.com
- Colossal-AI 開源地址:https://github.com/hpcaitech/ColossalAI


































