探究超大Transformer語言模型的分布式訓(xùn)練框架

NVIDIA Megatron 是一個基于 PyTorch 的框架，用于訓(xùn)練基于 Transformer 架構(gòu)的巨型語言模型。本系列文章將詳細(xì)介紹Megatron的設(shè)計和實(shí)踐，探索這一框架如何助力大模型的預(yù)訓(xùn)練計算。 上篇主要介紹了大模型訓(xùn)練的發(fā)展趨勢、NVIDIA Megatron的模型并行設(shè)計，本篇將承接上篇的內(nèi)容，解析Megatron 在NVIDIA DGX SuperPOD 上的實(shí)踐。

優(yōu)化的分布式集群架構(gòu)：NVIDIA DGX SuperPOD

有了高效的分布式訓(xùn)練框架，自然也需要優(yōu)化的分布式訓(xùn)練集群。

NVIDIA DGX SuperPOD 便是 NVIDIA 推出的一個分布式集群的參考架構(gòu)，最新一代是基于NVIDIA DGX A100 和NVIDIA HDR 200Gb/s ConnectX-6 HCAs組建，旨在幫助AI 研究人員快速搭建一套強(qiáng)大、靈活、高效的系統(tǒng)，以滿足工業(yè)界日益復(fù)雜、多變的模型對計算資源不同程度的需求。尤其對于超大語言模型預(yù)訓(xùn)練這一復(fù)雜場景，DGX SuperPOD 架構(gòu)尤為重要。

DGX SuperPOD 采用模塊化的設(shè)計，支持不同規(guī)模大小的設(shè)計。一個標(biāo)準(zhǔn)的SuperPOD 由140 臺DGX A100和三層Infiniband 交換機(jī)通過胖樹結(jié)構(gòu)全互聯(lián)起來。每臺DGX A100 配有8個200Gb/s 的高速計算網(wǎng)，并配有2個200Gb/s的高速存儲網(wǎng)，采用計算和存儲網(wǎng)絡(luò)分離的方案。

多個POD之間可以通過核心層交換機(jī)直連起來，可以支持多達(dá)560 臺DGX A100的互聯(lián)規(guī)模。

更多關(guān)于NVIDIA DGX SuperPOD 架構(gòu)的詳細(xì)設(shè)計，請參閱下列連接中的白皮書：https://images.nvidia.com/aem-dam/Solutions/Data-Center/gated-resources/nvidia-dgx-superpod-a100.pdf

NVIDIA Megatron 在 DGX SuperPOD 上的實(shí)踐

基于DGX SuperPOD 的Megatron實(shí)踐在不同大小的模型上，都表現(xiàn)出了很好的計算效率。

模型從1.7B 到1T ，訓(xùn)練規(guī)模從32 卡到3072 卡。

基于GPT-3 175B 模型的訓(xùn)練，采用如下的配置：

128 臺 DGX A100，總共 1024張 A100

Tensor 并行度：8；Pipeline 并行度：16； 數(shù)據(jù)并行度：8

全局Batch size ： 1536；Micro-batch size： 1

在如此大的訓(xùn)練規(guī)模下，GPU 仍可達(dá)到44% 左右的計算效率，在規(guī)模和效率上，都遠(yuǎn)遠(yuǎn)超過已有的公開結(jié)果。

詳細(xì)內(nèi)容請參考以下鏈接：

Megatron repro： https://github.com/nvidia/megatron-lm

GPT3-175B training scripts： https://github.com/NVIDIA/Megatron-LM/blob/main/examples/pretrain_gpt3_175B.sh

總結(jié)

1. 大模型是大勢所趨。

2. 大規(guī)模分布式訓(xùn)練是訓(xùn)練大模型的必須。

3. NVIDIA Megatron 是開源的、軟硬協(xié)同設(shè)計的訓(xùn)練框架，專為Transformer-based的超大語言模型設(shè)計。

4. NVIDIA DGX SuperPOD 是開放的集群參考設(shè)計架構(gòu)，專為大規(guī)模分布式訓(xùn)練而準(zhǔn)備。

5. Megatron 優(yōu)化的Tensor模型并行：用于intra-transformer 層，可以高效地執(zhí)行在HGX based的系統(tǒng)上。

6. Megatron優(yōu)化的 Pipeline 模型并行：用于inter-transformer 層，可以有效地利用集群中多網(wǎng)卡的設(shè)計。

7. 數(shù)據(jù)并行的加持，可以擴(kuò)展到更大規(guī)模、訓(xùn)練更快。

8. GPT-3 175B 的大模型，在1024 張 A100上可達(dá)到44%左右的計算效率。

9. NVIDIA Megatron 的設(shè)計和實(shí)踐，已廣泛用于學(xué)術(shù)界和工業(yè)界。

編輯：jq