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

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

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

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

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

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

更多關于NVIDIA DGX SuperPOD 架構的詳細設計，請參閱下列連接中的白皮書：https://images.nvidia.com/aem-dam/Solutions/Data-Center/gated-resources/nvidia-dgx-superpod-a100.pdf

NVIDIA Megatron 在 DGX SuperPOD 上的實踐

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

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

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

128 臺 DGX A100，總共 1024張 A100

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

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

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

詳細內容請參考以下鏈接：

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

總結

1. 大模型是大勢所趨。

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

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

4. NVIDIA DGX SuperPOD 是開放的集群參考設計架構，專為大規(guī)模分布式訓練而準備。

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

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

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

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

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

編輯：jq