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

NVIDIA Megatron 是一個(gè)基于 PyTorch 的框架，用于訓(xùn)練基于 Transformer 架構(gòu)的巨型語(yǔ)言模型。本系列文章將詳細(xì)介紹Megatron的設(shè)計(jì)和實(shí)踐，探索這一框架如何助力大模型的預(yù)訓(xùn)練計(jì)算。 上篇主要介紹了大模型訓(xùn)練的發(fā)展趨勢(shì)、NVIDIA Megatron的模型并行設(shè)計(jì)，本篇將承接上篇的內(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è)分布式集群的參考架構(gòu)，最新一代是基于NVIDIA DGX A100 和NVIDIA HDR 200Gb/s ConnectX-6 HCAs組建，旨在幫助AI 研究人員快速搭建一套強(qiáng)大、靈活、高效的系統(tǒng)，以滿(mǎn)足工業(yè)界日益復(fù)雜、多變的模型對(duì)計(jì)算資源不同程度的需求。尤其對(duì)于超大語(yǔ)言模型預(yù)訓(xùn)練這一復(fù)雜場(chǎng)景，DGX SuperPOD 架構(gòu)尤為重要。

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

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

更多關(guān)于NVIDIA DGX SuperPOD 架構(gòu)的詳細(xì)設(shè)計(jì)，請(qǐng)參閱下列連接中的白皮書(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)出了很好的計(jì)算效率。

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

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

128 臺(tái) DGX A100，總共 1024張 A100

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

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

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

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

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. 大模型是大勢(shì)所趨。

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

3. NVIDIA Megatron 是開(kāi)源的、軟硬協(xié)同設(shè)計(jì)的訓(xùn)練框架，專(zhuān)為T(mén)ransformer-based的超大語(yǔ)言模型設(shè)計(jì)。

4. NVIDIA DGX SuperPOD 是開(kāi)放的集群參考設(shè)計(jì)架構(gòu)，專(zhuān)為大規(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è)計(jì)。

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

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

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

編輯：jq