天天看點

NLP界“威震天”襲來!英偉達1小時成功訓練BERT,83億參數打造史上最大語言模型

NLP界“威震天”襲來!英偉達1小時成功訓練BERT,83億參數打造史上最大語言模型

新智元報道

來源:VB、TechCrunch、GitHub

編輯:金磊、小芹

【新智元導讀】英偉達一舉創造了2個壯舉!訓練出了世界上最大的語言模型——MegatronLM,包含83億參數,比BERT大24倍,比GPT-2大5.6倍;還打破了實時對話AI的記錄,僅耗時53分鐘即可訓練出行業标準BERT模型、2毫秒就能對答案做出推斷!

世界上最大的語言模型來了,順便還破了個記錄!

英偉達宣布,目前已經訓練出了世界上最大的語言模型——MegatronLM。

這個模型有多大?83億個參數!比谷歌的 BERT 大24倍,比 OpenAI 的 GPT-2 大5.6倍!

不僅如此,英偉達還宣布打破了實時對話 AI 的記錄——耗時53分鐘就可以訓練出行業标準的BERT模型、2毫秒左右就能對答案做出推斷。

為了實作這一壯舉,英偉達利用模型的并行性,将一個神經網絡分割成多個部分,建立了因資料太大無法容納在單個GPU的訓練模型。

最重要的是,代碼已開源!

NLP界“威震天”襲來!英偉達1小時成功訓練BERT,83億參數打造史上最大語言模型

GitHub項目位址:

https://github.com/NVIDIA/Megatron-LM

MegatronLM,堪稱 NLP 界的“威震天”!

有錢任性:訓練史上最大語言模型需要多少GPU?

更大的語言模型對于諸如文章完成、問題回答和對話系統等NLP任務非常有用。最近,訓練最大的神經語言模型已經成為提高NLP應用水準的最佳方法。

最近的兩篇論文,BERT和GPT-2,展示了大規模語言模組化的好處。這兩篇論文都利用了計算機和可用文本語料庫的進步,在自然語言了解、模組化和生成方面顯著超越了目前的最優水準。

訓練這些模型需要數以百計exaflops級的計算力和巧妙的記憶體管理,以換取減少記憶體占用的重新計算。然而,對于超過10億參數的超大型的模型,單個GPU上的記憶體不足以比對模型以及訓練所需的參數,需要利用模型并行性來将參數分割到多個GPU上。有幾種模組化并行性的方法,但是它們很難使用,因為它們依賴于自定義編譯器,或者擴充性很差,或者需要對優化器進行更改。

在這項工作中,我們通過對現有PyTorch transformer實作進行少量有針對性的修改,實作了一種簡單而有效的模型并行方法。我們的代碼是用原生Python編寫的,利用混合精度訓練,并利用NCCL庫在GPU之間進行通信。

我們通過在512個GPU上訓練一個transformer語言模型證明了這種方法的有效性,該模型具有8路模型并行性和64路資料并行性,83億參數,使其成為有史以來規模最大的基于transformer的語言模型,其大小為BERT的24倍,GPT-2的5.6倍。我們已經在GitHub存儲庫中釋出了實作此方法的代碼。

我們的實驗是在英偉達的DGX SuperPOD上進行的。在沒有模型并行性的情況下,我們可以在單個V100 32GB GPU上訓練一個12億參數的基線模型,并在整個訓練過程中保持39 TeraFLOPS,這是DGX2-H伺服器上單個GPU理論峰值的30%。

我們将模型參數擴充到83億,使用512個GPU,通過8路模型并行化,在整個應用程式中我們實作了高達15.1 PetaFLOPS的持續性能,與單GPU相比,擴充效率達到76%。圖1顯示了擴充的結果。

NLP界“威震天”襲來!英偉達1小時成功訓練BERT,83億參數打造史上最大語言模型

圖1:模型并行(藍色):多達8路模型并行弱擴充,每個GPU大約有10億個參數(例如2個GPU有20億參數,4個GPU有40億參數)。模型+資料并行(綠色):類似于模型并行的64路資料并行的配置。

多GPU并行性

訓練模型的典型範例是利用 weak scaling 方法和分布式資料并行性,根據GPU的數量來擴充訓練批大小。這種方法允許模型在更大的資料集上進行訓練,但有一個限制,即所有參數必須适合一個GPU。

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

我們利用 transformer網絡的結構,通過添加一些同步primitives來建立一個簡單的模型并行實作。

transformer層由一個self attention block和一個2層的多層感覺器(MLP)組成。我們分别在這兩個子產品中引入模型并行性。

如圖2a所示,這是MLP的結構,由兩個GEMM組成,中間有一個GeLU非線性,後面有一個dropout層。我們以列并行方式劃分第一個GEMM。這使得GeLU 非線性可以獨立地應用于每個分塊GEMM的輸出。子產品中的第二個GEMM沿着行并行化,直接擷取GeLU層的輸出,不需要任何通信。然後,在将輸出傳遞到dropout層之前,跨GPU減少第二個GEMM的輸出。這種方法将MLP block中的GEMM跨GPU分割了,隻需要在正向傳遞(g算子)中執行一個all-reduce操作,在反向傳遞(f算子)中執行一個all-reduce操作。

NLP界“威震天”襲來!英偉達1小時成功訓練BERT,83億參數打造史上最大語言模型

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

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

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

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

NLP界“威震天”襲來!英偉達1小時成功訓練BERT,83億參數打造史上最大語言模型

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

這種方法實作起來很簡單,因為它隻需要在向前和向後傳遞中添加一些額外的all-reduce操作。它不需要編譯器,并且與gPipe等方法提倡的那種pipeline模型并行性是正交的。

性能

為了測試我們的實作的計算性能,我們考慮了表1中四組參數的GPT-2模型。

NLP界“威震天”襲來!英偉達1小時成功訓練BERT,83億參數打造史上最大語言模型

表1:用于scaling 研究的參數。

所有的實驗都是在NVIDIA的DGX SuperPOD上進行的,我們使用了多達32台DGX- 2h伺服器(總共512個Tesla V100 SXM3 32GB GPU)。該系統針對多節點深度學習應用程式進行了優化,伺服器内部GPU之間的帶寬為300 GB/s,伺服器之間的互連帶寬為100 GB/s。

圖4顯示了模型和模型+資料并行性的擴充值。我們在這兩種設定中都觀察到了出色的擴充數字。例如,8路(8 GPU)模型并行的83億參數模型實作了77%的線性擴充。模型+資料并行性要求在反向傳播步驟之後進一步通信梯度,是以擴充數略有下降。然而,即使是運作在512個GPU上的最大配置(83億參數),相對于強大的基準單GPU配置(12億個參數),我們仍然可以實作74%的擴充性。

NLP界“威震天”襲來!英偉達1小時成功訓練BERT,83億參數打造史上最大語言模型

圖4:模型(左)和模型+資料(右)随着GPU的數量并行地進行weak scaling。

最後,我們研究了attention heads對模型并行擴充的影響。為此,我們考慮了83億參數、具有8路模型并行性的參數配置,并将attention heads的數目從16個改為32個。結果如表2所示。随着attention heads數量的增加,self attention層中的一些GEMM變小,同時softmax中的元素數量增加。這導緻了輕微的scaling decrease。未來的研究在設計大型transformer模型時應該警惕這種超參數,平衡模型性能和模型效率。

NLP界“威震天”襲來!英偉達1小時成功訓練BERT,83億參數打造史上最大語言模型

表2:attention heads 數量對scaling的影響。

GPT-2訓練

為了訓練GPT-2模型,我們建立了一個從_Reddit_下載下傳的37 GB _WebText_ dataset,它類似于原始GPT-2論文中描述的webtext資料集。資料集最終有810萬個url。我們将WebText資料集随機分割為95:5的比例,分别得到訓練集和驗證集。我們考慮了4種參數規模的模型:3.45億、7.75億、25億和83億。

NLP界“威震天”襲來!英偉達1小時成功訓練BERT,83億參數打造史上最大語言模型

圖5:訓練子集的驗證困惑度。在對37GB資料集過拟合之後,8.3B模型提前停止了。

圖5顯示了驗證的困惑度(perplexity)。我們發現。最大的83億參數的語言模型在~6epoch之後開始overfit,一種1 epoch被定義為15200次疊代。我們認為這可以通過使用更大規模的資料集來緩解,類似于XLNet和RoBERTa等最近論文中使用的資料集。

GPT-2評估

為了分析大型語言模型的訓練性能,我們在wikitext-103資料集上計算了perplexity,在Lambada資料集上計算了closize風格的預測精度。

正如預期的一樣,wikitext perplexity随着模型尺寸的增大而減小,lambada準确率随着模型尺寸的增大而增加(表3)。

NLP界“威震天”襲來!英偉達1小時成功訓練BERT,83億參數打造史上最大語言模型

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

結論

在這項工作中,我們在現有的深度學習硬體、軟體和模型的基礎上,建構了世界上最大的基于transformer的語言模型。

在此過程中,我們成功地突破了傳統的單GPU訓練的限制,實作了一種簡單而高效的模型并行方法,隻需對現有PyTorch transformer實作進行少量有針對性的修改。

我們在512台NVIDIA V100 GPU上高效地訓練了83億參數的語言模型(分别比BERT和GPT-2大24倍和5.6倍),具有8路模型并行性,并在整個應用程式中實作了高達15.1千萬億次浮點運算(PetaFLOPS)。

我們發現,與較小的transformer模型相比,更大的transformer模型可以在相同的時間内進行訓練,并且可以顯著提高性能。

然而,正如我們在工作中所展示的,NLP仍然需要合适的資料集、問題和技術來正确地訓練這些大型語言模型,否則會出現過拟合。

我們将我們的工作開源,以便社群就可以複制并擴充它們。

英偉達官方GitHub項目已開源!

英偉達在官方GitHub上對MegatronLM開源了代碼,也提供了相應的教程。

項目位址:https://github.com/NVIDIA/Megatron-LM

安裝

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

此外,代碼庫的一部分利用tensorflow-cpu(可選)執行TFRecords的資料加載以進行BERT訓練。

建議要麼使用./docker/中提供的Dockerfile,要麼建立一個虛拟環境(以避免破壞現有的tf安裝)并安裝requirements.txt。

1python -m pip install virtualenv
2virtualenv bert_env
3source bert_env/bin/activate
4pip install -r requirements.txt           

複制

用法

提供了5個預訓練BERT的腳本和3個預訓練GPT2的腳本。使用 --save 和 --load 儲存并加載模型檢查點(checkpoint)。

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

BERT預訓練

1bash scripts/pretrain_bert.sh           

複制

此腳本運作單個gpu BERT預訓練,主要用于調試目的。優化參數設定為64路分布式訓練。

要使用此腳本,請 --train-data以loose json格式放置,每行一個json。json字典的文本字段應該對應于 --text-key。

1python pretrain_bert.py \
 2       --num-layers 24 \
 3       --hidden-size 1024 \
 4       --num-attention-heads 16 \
 5       --batch-size 4 \
 6       --seq-length 512 \
 7       --max-preds-per-seq 80 \
 8       --max-position-embeddings 512 \
 9       --train-iters 1000000 \
10       --save checkpoints/bert_345m \
11       --load checkpoints/bert_345m \
12       --resume-dataloader \
13       --train-data wikipedia \
14       --lazy-loader \
15       --tokenizer-type BertWordPieceTokenizer \
16       --tokenizer-model-type bert-large-uncased \
17       --presplit-sentences \
18       --cache-dir cache \
19       --split 949,50,1 \
20       --distributed-backend nccl \
21       --lr 0.0001 \
22       --lr-decay-style linear \
23       --lr-decay-iters 990000 \
24       --weight-decay 1e-2 \
25       --clip-grad 1.0 \
26       --warmup .01 \
27       --fp16 \
28       --fp32-embedding           

複制

GPT2 預訓練

1bash scripts/pretrain_gpt2.sh           

複制

此腳本運作單gpu gpt2預訓練,主要用于調試目的。優化參數設定為64路分布式訓練。

它與前一個腳本格式大緻相同,但有一些值得注意的差異:

  • --tokenizer-type已切換為GPT2BPETokenizer;
  • --lr-decay-style已切換為cosine decay等等。

另外,GPT2使用來自BERT的不同參數初始化,用于訓練深度殘差網絡。要使用此初始化來訓練BERT,請使用--deep-init。

1python pretrain_gpt2.py \
 2       --num-layers 24 \
 3       --hidden-size 1024 \
 4       --num-attention-heads 16 \
 5       --batch-size 8 \
 6       --seq-length 1024 \
 7       --max-position-embeddings 1024 \
 8       --train-iters 320000 \
 9       --save checkpoints/gpt2_345m \
10       --load checkpoints/gpt2_345m \
11       --resume-dataloader \
12       --train-data wikipedia \
13       --lazy-loader \
14       --tokenizer-type GPT2BPETokenizer \
15       --cache-dir cache \
16       --split 949,50,1 \
17       --distributed-backend nccl \
18       --lr 0.00015 \
19       --lr-decay-style cosine \
20       --weight-decay 1e-2 \
21       --clip-grad 1.0 \
22       --warmup .01 \
23       --checkpoint-activations \
24       --fp16           

複制

更多細節内容,讀者可前往官方GitHub浏覽:

https://github.com/NVIDIA/Megatron-LM

參考連結:

VB:https://venturebeat.com/2019/08/13/nvidia-trains-worlds-largest-transformer-based-language-model/

TechCrunch:https://techcrunch.com/2019/08/13/nvidia-breaks-records-in-training-and-inference-for-real-time-conversational-ai/

GitHub:https://github.com/NVIDIA/Megatron-LM