英偉達(dá)一舉創(chuàng)造了2個(gè)壯舉！訓(xùn)練出了世界上最大的語言模型——MegatronLM，包含83億參數(shù)，比BERT大24倍，比GPT-2大5.6倍；還打破了實(shí)時(shí)對話AI的記錄，僅耗時(shí)53分鐘即可訓(xùn)練出行業(yè)標(biāo)準(zhǔn)BERT模型、2毫秒就能對答案做出推斷！

世界上最大的語言模型來了，順便還破了個(gè)記錄！

英偉達(dá)宣布，目前已經(jīng)訓(xùn)練出了世界上最大的語言模型——MegatronLM。

這個(gè)模型有多大？83億個(gè)參數(shù)！比谷歌的 BERT 大24倍，比 OpenAI 的 GPT-2 大5.6倍！

不僅如此，英偉達(dá)還宣布打破了實(shí)時(shí)對話 AI 的記錄——耗時(shí)53分鐘就可以訓(xùn)練出行業(yè)標(biāo)準(zhǔn)的BERT模型、2毫秒左右就能對答案做出推斷。

為了實(shí)現(xiàn)這一壯舉，英偉達(dá)利用模型的并行性，將一個(gè)神經(jīng)網(wǎng)絡(luò)分割成多個(gè)部分，創(chuàng)建了因數(shù)據(jù)太大無法容納在單個(gè)GPU的訓(xùn)練模型。

最重要的是，代碼已開源！

GitHub項(xiàng)目地址：
https://github.com/NVIDIA/Megatron-LM

MegatronLM，堪稱NLP 界的“威震天”！

有錢任性：訓(xùn)練史上最大語言模型需要多少GPU？

更大的語言模型對于諸如文章完成、問題回答和對話系統(tǒng)等NLP任務(wù)非常有用。最近，訓(xùn)練最大的神經(jīng)語言模型已經(jīng)成為提高NLP應(yīng)用水平的最佳方法。

最近的兩篇論文，BERT和GPT-2，展示了大規(guī)模語言建模的好處。這兩篇論文都利用了計(jì)算機(jī)和可用文本語料庫的進(jìn)步，在自然語言理解、建模和生成方面顯著超越了當(dāng)前的最優(yōu)水平。

訓(xùn)練這些模型需要數(shù)以百計(jì)exaflops級的計(jì)算力和巧妙的內(nèi)存管理，以換取減少內(nèi)存占用的重新計(jì)算。然而，對于超過10億參數(shù)的超大型的模型，單個(gè)GPU上的內(nèi)存不足以匹配模型以及訓(xùn)練所需的參數(shù)，需要利用模型并行性來將參數(shù)分割到多個(gè)GPU上。有幾種建模并行性的方法，但是它們很難使用，因?yàn)樗鼈円蕾囉谧远x編譯器，或者擴(kuò)展性很差，或者需要對優(yōu)化器進(jìn)行更改。

在這項(xiàng)工作中，我們通過對現(xiàn)有PyTorch transformer實(shí)現(xiàn)進(jìn)行少量有針對性的修改，實(shí)現(xiàn)了一種簡單而有效的模型并行方法。我們的代碼是用原生Python編寫的，利用混合精度訓(xùn)練，并利用NCCL庫在GPU之間進(jìn)行通信。 我們通過在512個(gè)GPU上訓(xùn)練一個(gè)transformer語言模型證明了這種方法的有效性，該模型具有8路模型并行性和64路數(shù)據(jù)并行性，83億參數(shù)，使其成為有史以來規(guī)模最大的基于transformer的語言模型，其大小為BERT的24倍，GPT-2的5.6倍。我們已經(jīng)在GitHub存儲庫中發(fā)布了實(shí)現(xiàn)此方法的代碼。

我們的實(shí)驗(yàn)是在英偉達(dá)的DGX SuperPOD上進(jìn)行的。在沒有模型并行性的情況下，我們可以在單個(gè)V100 32GB GPU上訓(xùn)練一個(gè)12億參數(shù)的基線模型，并在整個(gè)訓(xùn)練過程中保持39 TeraFLOPS，這是DGX2-H服務(wù)器上單個(gè)GPU理論峰值的30%。

我們將模型參數(shù)擴(kuò)展到83億，使用512個(gè)GPU，通過8路模型并行化，在整個(gè)應(yīng)用程序中我們實(shí)現(xiàn)了高達(dá)15.1 PetaFLOPS的持續(xù)性能，與單GPU相比，擴(kuò)展效率達(dá)到76%。圖1顯示了擴(kuò)展的結(jié)果。

圖1：模型并行(藍(lán)色)：多達(dá)8路模型并行弱擴(kuò)展，每個(gè)GPU大約有10億個(gè)參數(shù)(例如2個(gè)GPU有20億參數(shù)，4個(gè)GPU有40億參數(shù))。模型+數(shù)據(jù)并行(綠色)：類似于模型并行的64路數(shù)據(jù)并行的配置。

多GPU并行性

訓(xùn)練模型的典型范例是利用 weak scaling 方法和分布式數(shù)據(jù)并行性，根據(jù)GPU的數(shù)量來擴(kuò)展訓(xùn)練批大小。這種方法允許模型在更大的數(shù)據(jù)集上進(jìn)行訓(xùn)練，但有一個(gè)約束，即所有參數(shù)必須適合一個(gè)GPU。

模型并行訓(xùn)練可以通過跨多個(gè)GPU劃分模型來克服這一限制。近年來出現(xiàn)了幾個(gè)通用模型并行框架，如GPipe和Mesh-TensorFlow。gPipe在不同的處理器上劃分層組，而Mesh-TensorFlow使用層內(nèi)模型并行性。我們的方法在概念上類似于Mesh-TensorFlow，我們關(guān)注層內(nèi)并行性并融合GEMM以減少同步。然而，我們只對現(xiàn)有PyTorch transformer實(shí)現(xiàn)進(jìn)行了一些有針對性的修改，以便使用模型并行性來訓(xùn)練大型transformers。我們的方法很簡單，不需要任何新的編譯器或代碼重新連接來實(shí)現(xiàn)模型并行性，并且可以通過插入一些簡單的primitives(圖2中的f和g 算子)完全實(shí)現(xiàn)。

我們利用 transformer網(wǎng)絡(luò)的結(jié)構(gòu)，通過添加一些同步primitives來創(chuàng)建一個(gè)簡單的模型并行實(shí)現(xiàn)。

transformer層由一個(gè)self attention block和一個(gè)2層的多層感知器(MLP)組成。我們分別在這兩個(gè)模塊中引入模型并行性。

如圖2a所示，這是MLP的結(jié)構(gòu)，由兩個(gè)GEMM組成，中間有一個(gè)GeLU非線性，后面有一個(gè)dropout層。我們以列并行方式劃分第一個(gè)GEMM。這使得GeLU 非線性可以獨(dú)立地應(yīng)用于每個(gè)分塊GEMM的輸出。模塊中的第二個(gè)GEMM沿著行并行化，直接獲取GeLU層的輸出，不需要任何通信。然后，在將輸出傳遞到dropout層之前，跨GPU減少第二個(gè)GEMM的輸出。這種方法將MLP block中的GEMM跨GPU分割了，只需要在正向傳遞(g算子)中執(zhí)行一個(gè)all-reduce操作，在反向傳遞(f算子)中執(zhí)行一個(gè)all-reduce操作。

圖2：(a): MLP， (b)：transformer的self attention block。

如圖2(b)所示，在self attention block上，我們利用multihead attention操作中的固有并行性，以列并行方式劃分與鍵（K），查詢（Q）和值（V）相關(guān)聯(lián)的 GEMM。

這使得我們可以在GPU之間分割每個(gè)attention head參數(shù)和工作負(fù)載，并且不需要任何即時(shí)通信來完成self attention。

這種方法對于MLP和self-attention層都融合了兩個(gè)GEMM的組，消除了中間的同步點(diǎn)，并獲得了更好的scaling性能。這使我們能夠在一個(gè)簡單的transformer層中執(zhí)行所有GEMM，只使用前向路徑的2個(gè)all reduce和后向路徑的2個(gè)all reduce，如圖3所示。

圖3：GPT-2 transformer層的模型并行性。

這種方法實(shí)現(xiàn)起來很簡單，因?yàn)樗恍枰谙蚯昂拖蚝髠鬟f中添加一些額外的all-reduce操作。它不需要編譯器，并且與gPipe等方法提倡的那種pipeline模型并行性是正交的。

性能

為了測試我們的實(shí)現(xiàn)的計(jì)算性能，我們考慮了表1中四組參數(shù)的GPT-2模型。

表1：用于scaling 研究的參數(shù)。

所有的實(shí)驗(yàn)都是在NVIDIA的DGX SuperPOD上進(jìn)行的，我們使用了多達(dá)32臺DGX- 2h服務(wù)器(總共512個(gè)Tesla V100 SXM3 32GB GPU)。該系統(tǒng)針對多節(jié)點(diǎn)深度學(xué)習(xí)應(yīng)用程序進(jìn)行了優(yōu)化，服務(wù)器內(nèi)部GPU之間的帶寬為300 GB/s，服務(wù)器之間的互連帶寬為100 GB/s。

圖4顯示了模型和模型+數(shù)據(jù)并行性的擴(kuò)展值。我們在這兩種設(shè)置中都觀察到了出色的擴(kuò)展數(shù)字。例如，8路(8 GPU)模型并行的83億參數(shù)模型實(shí)現(xiàn)了77%的線性擴(kuò)展。模型+數(shù)據(jù)并行性要求在反向傳播步驟之后進(jìn)一步通信梯度，因此擴(kuò)展數(shù)略有下降。然而，即使是運(yùn)行在512個(gè)GPU上的最大配置(83億參數(shù))，相對于強(qiáng)大的基準(zhǔn)單GPU配置(12億個(gè)參數(shù))，我們?nèi)匀豢梢詫?shí)現(xiàn)74%的擴(kuò)展性。

圖4：模型(左)和模型+數(shù)據(jù)(右)隨著GPU的數(shù)量并行地進(jìn)行weak scaling。

最后，我們研究了attention heads對模型并行擴(kuò)展的影響。為此，我們考慮了83億參數(shù)、具有8路模型并行性的參數(shù)配置，并將attention heads的數(shù)目從16個(gè)改為32個(gè)。結(jié)果如表2所示。隨著attention heads數(shù)量的增加，self attention層中的一些GEMM變小，同時(shí)softmax中的元素?cái)?shù)量增加。這導(dǎo)致了輕微的scaling decrease。未來的研究在設(shè)計(jì)大型transformer模型時(shí)應(yīng)該警惕這種超參數(shù)，平衡模型性能和模型效率。

表2：attention heads 數(shù)量對scaling的影響。

GPT-2訓(xùn)練

為了訓(xùn)練GPT-2模型，我們創(chuàng)建了一個(gè)從_Reddit_下載的37 GB _WebText_ dataset，它類似于原始GPT-2論文中描述的webtext數(shù)據(jù)集。數(shù)據(jù)集最終有810萬個(gè)url。我們將WebText數(shù)據(jù)集隨機(jī)分割為95:5的比例，分別得到訓(xùn)練集和驗(yàn)證集。我們考慮了4種參數(shù)規(guī)模的模型：3.45億、7.75億、25億和83億。

圖5：訓(xùn)練子集的驗(yàn)證困惑度。在對37GB數(shù)據(jù)集過擬合之后，8.3B模型提前停止了。

圖5顯示了驗(yàn)證的困惑度(perplexity)。我們發(fā)現(xiàn)。最大的83億參數(shù)的語言模型在~6epoch之后開始o(jì)verfit，一種1 epoch被定義為15200次迭代。我們認(rèn)為這可以通過使用更大規(guī)模的數(shù)據(jù)集來緩解，類似于XLNet和RoBERTa等最近論文中使用的數(shù)據(jù)集。

GPT-2評估

為了分析大型語言模型的訓(xùn)練性能，我們在wikitext-103數(shù)據(jù)集上計(jì)算了perplexity，在Lambada數(shù)據(jù)集上計(jì)算了closize風(fēng)格的預(yù)測精度。

正如預(yù)期的一樣，wikitext perplexity隨著模型尺寸的增大而減小，lambada準(zhǔn)確率隨著模型尺寸的增大而增加(表3)。

表3：wikitext perplexity(越低越好)和Lambada完形精度(越高越好)的評估結(jié)果。

結(jié)論

在這項(xiàng)工作中，我們在現(xiàn)有的深度學(xué)習(xí)硬件、軟件和模型的基礎(chǔ)上，構(gòu)建了世界上最大的基于transformer的語言模型。

在此過程中，我們成功地突破了傳統(tǒng)的單GPU訓(xùn)練的限制，實(shí)現(xiàn)了一種簡單而高效的模型并行方法，只需對現(xiàn)有PyTorch transformer實(shí)現(xiàn)進(jìn)行少量有針對性的修改。

我們在512臺NVIDIA V100 GPU上高效地訓(xùn)練了83億參數(shù)的語言模型(分別比BERT和GPT-2大24倍和5.6倍)，具有8路模型并行性，并在整個(gè)應(yīng)用程序中實(shí)現(xiàn)了高達(dá)15.1千萬億次浮點(diǎn)運(yùn)算(PetaFLOPS)。

我們發(fā)現(xiàn)，與較小的transformer模型相比，更大的transformer模型可以在相同的時(shí)間內(nèi)進(jìn)行訓(xùn)練，并且可以顯著提高性能。

然而，正如我們在工作中所展示的，NLP仍然需要合適的數(shù)據(jù)集、問題和技術(shù)來正確地訓(xùn)練這些大型語言模型，否則會出現(xiàn)過擬合。

我們將我們的工作開源，以便社區(qū)就可以復(fù)制并擴(kuò)展它們。

英偉達(dá)官方GitHub項(xiàng)目已開源！

英偉達(dá)在官方GitHub上對MegatronLM開源了代碼，也提供了相應(yīng)的教程。

項(xiàng)目地址：https://github.com/NVIDIA/Megatron-LM

安裝

官方只支持 Python 3.6。請安裝支持GPU的最新版本PyTorch。

此外，代碼庫的一部分利用tensorflow-cpu（可選）執(zhí)行TFRecords的數(shù)據(jù)加載以進(jìn)行BERT訓(xùn)練。

建議要么使用./docker/中提供的Dockerfile，要么創(chuàng)建一個(gè)虛擬環(huán)境(以避免破壞現(xiàn)有的tf安裝)并安裝requirements.txt。

1python-mpipinstallvirtualenv 2virtualenvbert_env 3sourcebert_env/bin/activate 4pipinstall-rrequirements.txt

用法

提供了5個(gè)預(yù)訓(xùn)練BERT的腳本和3個(gè)預(yù)訓(xùn)練GPT2的腳本。使用 --save 和 --load 保存并加載模型檢查點(diǎn)(checkpoint)。

此外，還提供 GPT2 腳本，用于在wiki文本和LAMBADA上生成GPT2的交互式文本生成和零樣本(zero shot)評估。

BERT預(yù)訓(xùn)練

1bashscripts/pretrain_bert.sh

此腳本運(yùn)行單個(gè)gpu BERT預(yù)訓(xùn)練，主要用于調(diào)試目的。優(yōu)化參數(shù)設(shè)置為64路分布式訓(xùn)練。

要使用此腳本，請 --train-data以loose json格式放置，每行一個(gè)json。json字典的文本字段應(yīng)該對應(yīng)于 --text-key。

1pythonpretrain_bert.py 2--num-layers24 3--hidden-size1024 4--num-attention-heads16 5--batch-size4 6--seq-length512 7--max-preds-per-seq80 8--max-position-embeddings512 9--train-iters1000000 10--savecheckpoints/bert_345m 11--loadcheckpoints/bert_345m 12--resume-dataloader 13--train-datawikipedia 14--lazy-loader 15--tokenizer-typeBertWordPieceTokenizer 16--tokenizer-model-typebert-large-uncased 17--presplit-sentences 18--cache-dircache 19--split949,50,1 20--distributed-backendnccl 21--lr0.0001 22--lr-decay-stylelinear 23--lr-decay-iters990000 24--weight-decay1e-2 25--clip-grad1.0 26--warmup.01 27--fp16 28--fp32-embedding

GPT2 預(yù)訓(xùn)練

1bashscripts/pretrain_gpt2.sh

此腳本運(yùn)行單gpu gpt2預(yù)訓(xùn)練，主要用于調(diào)試目的。優(yōu)化參數(shù)設(shè)置為64路分布式訓(xùn)練。

它與前一個(gè)腳本格式大致相同，但有一些值得注意的差異：

--tokenizer-type已切換為GPT2BPETokenizer；

--lr-decay-style已切換為cosine decay等等。

另外，GPT2使用來自BERT的不同參數(shù)初始化，用于訓(xùn)練深度殘差網(wǎng)絡(luò)。要使用此初始化來訓(xùn)練BERT，請使用--deep-init。

1pythonpretrain_gpt2.py 2--num-layers24 3--hidden-size1024 4--num-attention-heads16 5--batch-size8 6--seq-length1024 7--max-position-embeddings1024 8--train-iters320000 9--savecheckpoints/gpt2_345m 10--loadcheckpoints/gpt2_345m 11--resume-dataloader 12--train-datawikipedia 13--lazy-loader 14--tokenizer-typeGPT2BPETokenizer 15--cache-dircache 16--split949,50,1 17--distributed-backendnccl 18--lr0.00015 19--lr-decay-stylecosine 20--weight-decay1e-2 21--clip-grad1.0 22--warmup.01 23--checkpoint-activations 24--fp16