探究超大Transformer语言模型的分布式训练框架

描述

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

总结

1. 大模型是大势所趋。

2. 大规模分布式训练是训练大模型的必须。

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 的设计和实践,已广泛用于学术界和工业界。

编辑:jq

打开APP阅读更多精彩内容
声明:本文内容及配图由入驻作者撰写或者入驻合作网站授权转载。文章观点仅代表作者本人,不代表电子发烧友网立场。文章及其配图仅供工程师学习之用,如有内容侵权或者其他违规问题,请联系本站处理。 举报投诉

全部0条评论

快来发表一下你的评论吧 !

×
20
完善资料,
赚取积分