GPT-3难以复现,为什么说PyTorch走上了一条“大弯路”?

来源: 网络 日期:2022-03-25

2020 年,最轰动的 AI 新闻莫过于 OpenAI 发布的 GPT-3 了。它的1750亿参数量及其在众多NLP任务上超过人类的出众表现让人们开始坚信:大模型才是未来。但与之带来的问题是,训练超大模型所需的算力、存储已不再是单机就能搞定。

据 NVIDIA 估算,如果要训练GPT-3 ,即使单个机器的显存/内存能装得下,用 8 张 V100 的显卡,训练时长预计要 36 年;即使用 512 张 V100 ,训练也需要将近 7 个月;如果拥有 1024 张 80GB A100, 那么完整训练 GPT-3 的时长可以缩减到 1 个月。

除去硬件资源这个经济问题,在技术层面,意味着训练大模型一定是一个分布式问题。因为算力需求还是一个相对容易解决的问题,毕竟拥有大集群的组织并不只 OpenAI 一家,而如何解决上千块 GPU 的分布式训练问题才是关键。

根据目前业界已有的分布式训练方案,即便你是一位非常优秀的数据科学家,知晓并能解决 Transformer 相关的所有算法问题,但如果你不知道如何解决分布式训练时上百台服务器之间的通信、拓扑、模型并行、流水并行等问题,你甚至都无法启动这次训练。一定程度上,这解释了GPT-3发布时隔一年,却只有 NVIDIA 、微软等大企业可以复现 GPT-3 。

目前,开源的 GPT 模型库主要是 NVIDIA开发的 Megatron-LM 和经过微软深度定制开发的 DeepSpeed,其中,DeepSpeed 的模型并行等内核取自 Megatron,它们都是专门为支持 PyTorch 分布式训练 GPT 而设计。

不过在实际训练中,PyTorch 、 Megatron、DeepSpeed 都走了一条非常长的弯路。不仅是弯路,你会发现 Megatron 的代码只能被 NVIDIA 的分布式训练专家所复用,它对于 PyTorch 的算法工程师而言门槛极高,以至于任何想要用 PyTorch 复现一个分布式大模型的算法工程师,都得先等 NVIDIA 开发完才能再使用 Megatron 提供的模型。

作为新一代深度学习开源框架,致力于“大模型分布式”高效开发的 OneFlow 框架用一套通用设计非常简单清晰地解决了GPT模型的分布式训练难题,同时还在已有的测试规模上性能超过 NVIDIA 的 Megatron,这为大规模分布式训练框架提出了更优的设计理念和路径。

一、PyTorch 分布式训练GPT的痛点是什么?

此前,NVIDIA 放出了一篇重量级的论文:Efficient Large-Scale Language Model Training on GPU Clusters ,用 3072 张 80 GB A100 训练 GPT,最大规模的模型参数量达到了 1T,这是 GPT-3 原版规模的 5 倍。

NVIDIA 训练 GPT-3 最大到 1T 参数规模

论文里,NVIDIA 介绍了分布式训练超大规模模型的三种必须的并行技术:

· 数据并行(Data Parallelism)

· 模型并行(Tensor Model Parallelism)

· 流水并行(Pipeline Model Parallelism)

其中,数据并行是最常见的并行方式。而模型并行是对某一层(如 Linear/Dense Layer 里的 Variable )的模型 Tensor 切分,从而将大的模型 Tensor 分成多个相对较小的 Tensor 进行并行计算;流水并行,是将整个网络分段(stage),不同段在不同的设备上,前后阶段流水分批工作,通过一种“接力”的方式并行。

对于 1T 规模的模型,NVIDIA 一共使用了 384 台 DGX-A100 机器(每台装有 8 张 80GB A100 GPU),机器内部各 GPU 间使用超高速 NVLink 和 NVSwitch 互联,每台机器装有 8 个 200Gbps 的 InfiniBand (IB) 网卡,可以说是硬件集群顶配中的顶配。

那么,这些机器是如何协同工作的?GPT 网络是由很多层 Transformer Layer 组成,每一层内部是一个由多层 MLP 和 attenTIon 机制组成的子图,对于参数规模 1T 的 GPT 而言就有 128 层的 Transformer Layer,这个超大超深的网络被分割成了 64 个 stage (阶段),每个 stage 跑在 6 台 DGX-A100 上,其中 6 台机器之间进行数据并行,每台机器内部的 8 张卡之间做模型并行,整个集群的 3072 张 A100 按照机器拓扑被划分成了 [6 x 8 x 64] 的矩阵,同时使用数据并行 & 模型并行 & 流水并行进行训练。

3072 张 A100 集群拓扑

GPipe、梯度累加、重计算(CheckpoinTIng)和 1F1B(One Forward pass followed by One Backward pass)是分布式训练 GPT 的流水并行的核心技术。无论是 NVIDIA 的Megatron(PyTorch),还是 OneFlow、PaddlePaddle、MindSpore ,都是通过不同的设计实现了上述相同的功能。

基于 PyTorch 开发的 Megatron,本质上是一个专用于 GPT 的模型库,所有的代码都是 Python 脚本,NVIDIA 为 GPT 专门定制了分布式训练所需的算子、流水并行调度器、模型并行所需的通信原语等功能,在 GPU 上的性能表现上,Megatron已经非常优异。可以说,NVIDIA 在使用 PyTorch 做分布式训练上已经做到极致了。

 1/4    1 2 3 4 下一页 尾页
    A+
声明:本文转载自其它媒体,转载目的在于传递更多信息,并不代表赞同其观点和对其真实性负责。