NVIDIA Megatron 是一个基于 PyTorch 的框架，用于训练基于 Transformer 架构的巨型语言模型。本系列文章将详细介绍Megatron的设计和实践，探索这一框架如何助力大模型的预训练计算。 上篇主要介绍了大模型训练的发展趋势、NVIDIA Megatron的模型并行设计，本篇将承接上篇的内容，解析Megatron 在NVIDIA DGX SuperPOD 上的实践。
优化的分布式集群架构：NVIDIA DGX SuperPOD
有了高效的分布式训练框架，自然也需要优化的分布式训练集群。
NVIDIA DGX SuperPOD 便是 NVIDIA 推出的一个分布式集群的参考架构，最新一代是基于NVIDIA DGX A100 和NVIDIA HDR 200Gb/s ConnectX-6 HCAs组建，旨在帮助AI 研究人员快速搭建一套强大、灵活、高效的系统，以满足工业界日益复杂、多变的模型对计算资源不同程度的需求。尤其对于超大语言模型预训练这一复杂场景，DGX SuperPOD 架构尤为重要。
DGX SuperPOD 采用模块化的设计，支持不同规模大小的设计。一个标准的SuperPOD 由140 台DGX A100和三层Infiniband 交换机通过胖树结构全互联起来。每台DGX A100 配有8个200Gb/s 的高速计算网，并配有2个200Gb/s的高速存储网，采用计算和存储网络分离的方案。
多个POD之间可以通过核心层交换机直连起来，可以支持多达560 台DGX A100的互联规模。
更多关于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实践在不同大小的模型上，都表现出了很好的计算效率。
如下表，模型从1.7B 到1T ，训练规模从32 卡到3072 卡。
基于GPT-3 175B 模型的训练，采用如下的配置：
128 台 DGX A100，总共 1024张 A100
Tensor 并行度：8；Pipeline 并行度：16； 数据并行度：8
全局Batch size : 1536；Micro-batch size: 1
在如此大的训练规模下，GPU 仍可达到44% 左右的计算效率，在规模和效率上，都远远超过已有的公开结果。
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
3. NVIDIA Megatron 是开源的、软硬协同设计的训练框架，专为Transformer-based的超大语言模型设计。
4. NVIDIA DGX SuperPOD 是开放的集群参考设计架构，专为大规模分布式训练而准备。
5. Megatron 优化的Tensor模型并行：用于intra-transformer 层，可以高效地执行在HGX based的系统上。
6. Megatron优化的 Pipeline 模型并行：用于inter-transformer 层，可以有效地利用集群中多网卡的设计。
7. 数据并行的加持，可以扩展到更大规模、训练更快。
8. GPT-3 175B 的大模型，在1024 张 A100上可达到44%左右的计算效率。
9. NVIDIA Megatron 的设计和实践，已广泛用于学术界和工业界。
分享到微信
分享到微博
