背景介紹

2018 年，預(yù)訓(xùn)練語(yǔ)言模型技術(shù)橫空出世并引發(fā)了人工智能領(lǐng)域的性能革命。研究表明，增大參數(shù)量與數(shù)據(jù)規(guī)模是進(jìn)一步提升語(yǔ)言模型性能的有效手段，對(duì)十億、百億乃至千億級(jí)大模型的探索成為業(yè)界的熱門話題。這引發(fā)了國(guó)內(nèi)外研究機(jī)構(gòu)與互聯(lián)網(wǎng)企業(yè)的激烈競(jìng)爭(zhēng)，將模型規(guī)模與性能不斷推向新的高度。除 Google、OpenAI 等國(guó)外知名機(jī)構(gòu)外，近年來(lái)國(guó)內(nèi)相關(guān)研究機(jī)構(gòu)與公司也異軍突起，形成了大模型的研究與應(yīng)用熱潮，人工智能由此進(jìn)入“大模型時(shí)代”。

然而在“大模型時(shí)代”，模型巨大的參數(shù)規(guī)模和算力需求，帶來(lái)了很大的?訓(xùn)練和微調(diào)難題：

01??算力成本高

模型巨大的參數(shù)量無(wú)法在單張顯卡中完成存儲(chǔ)與計(jì)算。OpenAI 訓(xùn)練 GPT-3 （約?1750?億參數(shù)）使用了上千張 GPU，Google 訓(xùn)練 PaLM （約 5000?億參數(shù)）使用了六千片 TPU?，帶來(lái)了高昂的訓(xùn)練成本，配置規(guī)模龐大的計(jì)算集群也十分困難。

02??編程難度大

為了利用分布式算力加速大模型的訓(xùn)練和微調(diào)，程序員需要編寫復(fù)雜的分布式程序來(lái)驅(qū)動(dòng)大模型?，F(xiàn)有的框架（如?DeepSpeed ）已經(jīng)能夠較好地支持模型的分布式訓(xùn)練，但依然需要用戶進(jìn)行較為復(fù)雜的編程與配置。

03??計(jì)算速度低

在分布式訓(xùn)練框架中，不同計(jì)算節(jié)點(diǎn)之間需要頻繁地進(jìn)行通信，此時(shí)通信帶寬往往成為模型的訓(xùn)練瓶頸，若不能合理設(shè)計(jì)通信策略并進(jìn)行有效且正確的代碼實(shí)現(xiàn)，訓(xùn)練模型所需的時(shí)間將被大大延長(zhǎng)，不能充分發(fā)揮計(jì)算設(shè)備的潛力。

針對(duì)上述難題，我們推出了大模型高效訓(xùn)練工具包 BMTrain 與模型倉(cāng)庫(kù) ModelCenter?。

BMTrain & ModelCenter

在 OpenBMB 全流程加速工具框架中，致力于解決大模型訓(xùn)練難題的 BMTrain 是極為重要的一環(huán)。在設(shè)計(jì)中，我們追求以下特點(diǎn)：

01??高效率

作為大模型訓(xùn)練的 “發(fā)動(dòng)機(jī)”，BMTrain 能夠在任意數(shù)量的 GPU 上進(jìn)行高效的大模型預(yù)訓(xùn)練與微調(diào)，最優(yōu)化分布式框架的通信開銷，在超大規(guī)模模型訓(xùn)練場(chǎng)景下與 DeepSpeed 等框架相比可以節(jié)省 90% 的算力成本。

02??低資源

不同于?OpenAI、Google 等機(jī)構(gòu)的算力條件，為了讓更多實(shí)驗(yàn)室和企業(yè)也能夠訓(xùn)練大模型，我們致力于在保持高效計(jì)算前提下最小化大模型的算力需求，實(shí)現(xiàn)單張消費(fèi)級(jí)顯卡全參數(shù)微調(diào) BERT-Large，8 臺(tái) A100 小集群訓(xùn)練 GPT-3，進(jìn)一步降低大模型的算力門檻。

03??可擴(kuò)展

在編程難度方面，我們致力做最簡(jiǎn)潔最有效的封裝，僅使用少量的代碼替換，即可達(dá)到與原生 PyTorch 一致的編程體驗(yàn)，一鍵安裝工具包降低配置難度，讓大模型真正飛入千家萬(wàn)戶。

與此同時(shí)，我們還基于 BMTrain 實(shí)現(xiàn)了一系列大規(guī)模預(yù)訓(xùn)練語(yǔ)言模型，集成于?ModelCenter?倉(cāng)庫(kù)。BMTrain 與 ModelCenter 共同組成了高效分布式預(yù)訓(xùn)練框架，適用于任意的 Transformer 結(jié)構(gòu)，可以在較少數(shù)量的 GPU 上訓(xùn)練，且高度兼容 PyTorch、Transformers 庫(kù)，學(xué)習(xí)成本極低。目前我們已經(jīng)支持了常用的英文模型如 BERT、GPT、T5、RoBERTa 以及中文模型如 CPM-1、CPM-2 等。

簡(jiǎn)單易用

BMTrain 基于簡(jiǎn)單易用的設(shè)計(jì)原則，致力做最簡(jiǎn)潔最有效的封裝，僅使用少量的代碼替換，即可達(dá)到與原生 PyTorch 一致的編程體驗(yàn)，快速實(shí)現(xiàn)大模型的訓(xùn)練加速。我們還實(shí)現(xiàn)了模型倉(cāng)庫(kù) ModelCenter，方便使用者快速使用主流架構(gòu)大模型。

01??無(wú)縫銜接

貼合 PyTorch?使用習(xí)慣，上手門檻更低，僅需?簡(jiǎn)單替換?即可完成訓(xùn)練提速：?

下面將提供一個(gè)簡(jiǎn)單的對(duì)比圖來(lái)直觀展示 BMTrain?的便捷易用（左邊是原始代碼，右邊是替換后的代碼）。

02??一鍵轉(zhuǎn)換

對(duì)于?PyTorch?實(shí)現(xiàn)的原始模型，直接使用?BMTrain 提供的?BMTrainModelWrapper?自動(dòng)包裝模型，實(shí)現(xiàn)分布式加速。

03??模型倉(cāng)庫(kù)

我們基于 BMTrain 構(gòu)建了模型倉(cāng)庫(kù) ModelCenter，可以直接加載進(jìn)行使用。ModelCenter ?的使用方式與 Transformers 類似，且兼容 Transformers 中各類數(shù)據(jù)處理接口。

效果拔群

BMTrain在大模型訓(xùn)練上效果出色，在不同規(guī)模的算力條件下均有較好的性能表現(xiàn)。

01??消費(fèi)級(jí)算力（單卡2080Ti）

在消費(fèi)級(jí)顯卡 2080Ti 上，BMTrain?可以實(shí)現(xiàn) BERT-Large?的微調(diào)（3 億參數(shù)，樣本長(zhǎng)度 512）。

02??入門級(jí)算力（單卡V100）

在入門級(jí)算力條件下（顯卡為 V100 32GB），BMTrain?可以實(shí)現(xiàn)?BERT-Large?的高效訓(xùn)練（3 億參數(shù)，樣本長(zhǎng)度 512）。

03??中等算力（單機(jī)8卡A100）

在中等規(guī)模算力條件下（顯卡為?A100 40GB，NVLink），BMTrain 可以訓(xùn)練大規(guī)模的 GPT-13B (130億參數(shù)，樣本長(zhǎng)度 512）。

04??高級(jí)算力（多機(jī)8卡A100）

在較高規(guī)模算力條件下（顯卡為A100 40GB，NVLink，400Gbps IB），BMTrain 可以訓(xùn)練超大規(guī)模的 GPT-3（1750億參數(shù)）。

使用 BMTrain 或者 ColossalAI ，64?卡?A100?跑完 GPT-3?的 300B token大概需要2年，服務(wù)器與顯卡租金大約 900 萬(wàn)左右。根據(jù)我們的實(shí)驗(yàn)估算，使用 128 張?A100?時(shí)，單卡吞吐量可以提升 2.5 倍以上，6?個(gè)月可以跑完 GPT-3，服務(wù)器租金大約 500?萬(wàn)左右。雖然訓(xùn)練出 GPT-3 的成本依然高昂，但與 GPT-3?的?1200?萬(wàn)美元相比，成本仍然節(jié)約了 90%?以上。

上手教程

BMTrain 和 ModelCenter 提供了豐富的文檔，方便使用者快速上手，便捷地體驗(yàn)大模型的魅力。

STEP 01 ·?快速安裝?

01 ?BMTrain安裝

· ?可以通過(guò) pip?直接安裝（推薦）

·??也可以選擇從 GitHub?進(jìn)行源碼安裝

02 ?ModelCenter安裝

· ?可以通過(guò) pip?直接安裝（推薦）

·??也可以選擇從 GitHub?進(jìn)行源碼安裝

STEP 02?· 啟動(dòng) BMTrain

首先，你需要在代碼開頭初始化 BMTrain。正如在使用 PyTorch 的分布式訓(xùn)練模塊需要在代碼開頭使用?init_process_group?一樣，使用 BMTrain 需要在代碼開頭使用?init_distributed。

注意：?使用 BMTrain 時(shí)請(qǐng)不要使用 PyTorch 自帶的?distributed?模塊包括?torch.distributed.init_process_group?以及相關(guān)通信函數(shù)。

STEP?03?· 性能優(yōu)化

01 ?ZeRO-3優(yōu)化

使用ZeRO-3優(yōu)化只需要對(duì)模型代碼進(jìn)行簡(jiǎn)單替換：·??torch.nn.Module?替換為?bmtrain.DistributedModule?·??torch.nn.Parameter?替換為?bmtrain.DistributedParameter?

02 ?Checkpointing優(yōu)化

在模塊上套用?bmtrain.CheckpointBlock?即可

03 ?通信優(yōu)化

為了進(jìn)一步縮短通信額外開銷，將通信與運(yùn)算時(shí)間重疊，可以使用?TransformerBlockList ?來(lái)進(jìn)一步優(yōu)化。在使用時(shí)需要對(duì)代碼進(jìn)行簡(jiǎn)單替換：·??torch.nn.ModuleList?替換為?bmtrain.TransformerBlockList?·??for module in self.module_list:??x= module(x, ...)?替換為?x = self.module_list(x, ...)?

STEP?04?·?運(yùn)行分布式訓(xùn)練代碼

BMTrain 使用 PyTorch 原生分布式訓(xùn)練啟動(dòng)器，你可以根據(jù) PyTorch 版本選擇下列命令中的一個(gè)。

· ?torch.distributed.launch

· ?torchrun

其中：

${MASTER_ADDR}?為主節(jié)點(diǎn)的 IP 地址，只有一個(gè)節(jié)點(diǎn)可以寫?localhost?或 127.0.0.1

${MASTER_PORT}?為主節(jié)點(diǎn)的端口?

${NNODES}?為節(jié)點(diǎn)數(shù)量（一般為機(jī)器數(shù)量）?

${GPU_PER_NODE}?為每個(gè)節(jié)點(diǎn)的 GPU 數(shù)量?

${NODE_RANK}?為本節(jié)點(diǎn)的 RANK

${ARGS}?為代碼輸入的其他參數(shù)

STEP?05?·?使用ModelCenter

本節(jié)將首先以如何在一個(gè)分類數(shù)據(jù)集上微調(diào) BERT 模型為例，介紹如何使用?ModelCenter 的現(xiàn)有模型。其次將簡(jiǎn)單介紹如何使用 ModelCenter 實(shí)現(xiàn)一個(gè)新的模型。

01 ?準(zhǔn)備模型

接下來(lái)，你可以從??model_center??中獲取預(yù)訓(xùn)練好的 BERT 模型，例如?bert-base-uncased。由于我們是在一個(gè)分類任務(wù)上微調(diào) BERT 模型，所以需要在最后一層后添加一個(gè)全連接層。

如果只需要 config 來(lái)構(gòu)建模型，而不需要現(xiàn)成的預(yù)訓(xùn)練參數(shù)，可以參考下面的方法：

02 ?準(zhǔn)備數(shù)據(jù)集

下一步是準(zhǔn)備數(shù)據(jù)集，用于訓(xùn)練和驗(yàn)證模型。這里，我們使用 SuperGLUE benchmark 中的 BoolQ 數(shù)據(jù)集。你需要下載該數(shù)據(jù)集，并將解壓后的文件夾放在?your_path_to_dataset?路徑下。

03 ?模型訓(xùn)練

現(xiàn)在，在設(shè)置?優(yōu)化器、學(xué)習(xí)率調(diào)整策略和損失函數(shù)?后，就可以開始訓(xùn)練模型了！在?示例代碼（https://modelcenter.readthedocs.io/en/latest/notes/quickstart.html）中，我們訓(xùn)練 BERT 模型 5 輪，并且在每輪訓(xùn)練結(jié)束后驗(yàn)證模型的性能。

04 ?運(yùn)行代碼

運(yùn)行代碼的方式和上節(jié)所給出的一樣，你可以根據(jù) PyTorch 版本選擇其中的一個(gè)，具體命令可以參照上節(jié)：運(yùn)行分布式訓(xùn)練代碼?中所給出的示例。

05 ?如何創(chuàng)建新的模型

• 在 ModelCenter 的?model_center/layer?中，我們給出了包括像Linear、LayerNorm、Embedding?等等?常用模塊?的實(shí)現(xiàn)。且這些模塊都是基于bmtrain.DistributedParameter?和?bmtrain.DistributedModule?實(shí)現(xiàn)的，支持分布式訓(xùn)練。因此在創(chuàng)建新的模型時(shí)，我們可以直接調(diào)用這些模塊，完成網(wǎng)絡(luò)的搭建。

• 在 ModelCenter 的?model_center/layer?中，我們同樣實(shí)現(xiàn)了一些?常用的模塊的組合，被稱作?block，如組合了?Attention?和?Add&Norm?的?SelfAttentionBlock?。每一個(gè) block?都有不同的選擇， 且支持?pre-layernorm?和?post-layernorm。

• 我們使用?bmtrain.CheckpointBlock?和?bmtrain.TransformerBlockList?來(lái)包裝我們的?tranformer block。這些在不增加大量計(jì)算時(shí)間的情況下大量減少了?GPU?內(nèi)存使用量。

• 在 ModelCenter 所提供的?常用模型?的基礎(chǔ)上，通過(guò)?加入相應(yīng)的模塊，你可以很容易地構(gòu)建一個(gè)新的模型。你只需要將模型特定功能添加到通用的模型結(jié)構(gòu)中，就可以輕而易舉地實(shí)現(xiàn)一個(gè)新模型。

• 可以參考示例代碼（https://modelcenter.readthedocs.io/en/latest/notes/write_model.html）中給出的實(shí)現(xiàn)方式。

未來(lái)展望

在未來(lái)，我們將在 BMTrain 中進(jìn)一步完善訓(xùn)練的工具鏈，并達(dá)到更好的 PyTorch 兼容性；在 ModelCenter 中支持更多的預(yù)訓(xùn)練大模型，并完善模型對(duì)推理的支持。我們誠(chéng)邀更多的開發(fā)者加入我們的開發(fā)隊(duì)伍，共同維護(hù) OpenBMB 開源社區(qū)。也歡迎大家使用 BMTrain 實(shí)現(xiàn)更多類型的大模型，共同維護(hù)大模型構(gòu)建與應(yīng)用生態(tài)，為大模型時(shí)代貢獻(xiàn)一份力量。

??BMTrain：https://github.com/OpenBMB/BMTrain

關(guān)注我們

微信搜索關(guān)注?“OpenBMB開源社區(qū)”

獲取更多大模型干貨知識(shí)和前沿資訊！

??傳送門｜相關(guān)鏈接

?? ?官方網(wǎng)站：https://www.openbmb.org

????GitHub：https://github.com/OpenBMB

????交流QQ群：735930538

????啟智社區(qū)：https://git.openi.org.cn/OpenBMB

????微博：http://weibo.cn/OpenBMB

????知乎：https://www.zhihu.com/people/OpenBMB

????Twitter：https://twitter.com/OpenBMB