近年来, 大多数出现在顶级人工智能会议上的模型都是在多个 GPU 上训练的, 特别是随着基于 Transformer 的语言模型的提出。当研究人员和工程师开发人工智能模型时, 分布式训练无疑是一种常见的做法。传统的单 机单卡模式已经无法满足超大模型进行训练的要求,这一趋势背后有几个原因。
- 模型规模迅速增加。2018 年的 BERT-Large 有 3.45 亿的参数, 2018 年的 GPT-2 有 15 亿的参数, 而 2020 年 的 GPT-3 有 1750 亿个参数。很明显, 模型规模随着时间的推移呈指数级增长。目前最大的模型已经超过 了万亿个参数。而与较小的模型相比,超大型模型通常能提供更优越的性能。
- 数据集规模迅速增加。从 GPT- 1 ,GPT-2 再到 GPT-3,训练数据分别为 5GB ,40GB ,45TB (清洗后约为 570GB),数量量爆炸性增长。
- 计算能力越来越强。随着半导体行业的进步, 显卡变得越来越强大。由于核的数量增多, GPU 是深度学习 最常见的算力资源。从 2012 年的 K10 GPU 到 2020 年的 A100 GPU,计算能力已经增加了几百倍。这使我 们能够更快地执行计算密集型任务,而深度学习正是这样一项任务。
如今, 我们接触到的模型可能太大, 以致于无法装入一个 GPU,而数据集也可能大到足以在一个 GPU 上训 练上百天。这时, 只有用不同的并行化技术在多个 GPU 上训练我们的模型, 我们才能完成并加快模型训练, 以 追求在合理的时间内获得想要的结果。因此,我们需要进行单机多卡、甚至是多机多卡进行大模型的训练。
而利用 AI 集群, 使深度学习算法更好地从大量数据中训练出性能优良的大模型是分布式机器学习的首要目 标。为了实现该目标, 一般需要根据硬件资源与数据/模型规模的匹配情况, 考虑对计算任务、训练数据和模型 进行划分,进行分布式存储和分布式训练。
分布式系统
分布式系统由多个软件组件组成, 在多台机器上运行。例如, 传统的数据库运行在一台机器上。随着数据量 的爆发式增长, 单台机器已经不能为企业提供理想的性能。特别是在双十一这样的网络狂欢节, 网络流量会出 乎意料的大。为了应对这种压力, 现代高性能数据库被设计成在多台机器上运行, 它们共同为用户提供高吞吐 量和低延迟。
分布式系统示例
分布式系统的一个重要评价指标是可扩展性。例如, 当我们在 4 台机器上运行一个应用程序时, 我们自然希望该应用程序的运行速度能提高 4 倍。然而, 由于通信开销和硬件性能的差异, 很难实现线性提速。因此, 当我 们实现应用程序时, 必须考虑如何使其更快。良好的设计和系统优化的算法可以帮助我们提供良好的性能。有 时,甚至有可能实现线性和超线性提速。总训练速度可以用如下公式简略估计:
总训练速度 ∝ 单设备计算速度 × 计算设备总量 × 多设备加速比
其中, 单设备计算速度主要由单块计算加速芯片的运算速度和数据 I/O 能力来决定, 对单设备训练效率进行优 化, 主要的技术手段有混合精度训练、算子融合、梯度累加等; 分布式训练系统中计算设备数量越多, 其理论峰 值计算速度就会越高, 但是受到通讯效率的影响, 计算设备数量增大则会造成加速比急速降低; 多设备加速比 则是由计算和通讯效率决定, 需要结合算法和网络拓扑结构进行优化, 分布式训练并行策略主要目标就是提升 分布式训练系统中的多设备加速比。
分布式训练需要多台机器/GPU。在训练期间, 这些设备之间会有通信。为了更好地理解分布式训练, 有几个重要的术语需要我们了解清楚。
- host: 主机 (host) 是通信网络中的主要设备。在初始化分布式环境时,经常需要它作为一个参数。 port: 这里的端口 (port) 主要是指主机上用于通信的主端口。
- rank: 在网络中赋予设备的唯一 ID。
- world size: 网络中设备的数量。
- process group: 进程组 (process group) 是一个通信网络, 包括设备的一个子集。总是有一个默认的进程组, 它包含所有的设备。一个子集的设备可以形成一个进程组,以便它们只在组内的设备之间进行通信。
让我们假设我们有 2 台机器(也称为节点),每台机器有 4 个 GPU,参考上图。当我们在这两台机器上初始 化分布式环境时,我们基本上启动了 8 个进程(每台机器上有 4 个进程),每个进程被绑定到一个 GPU 上。在初始化分布式环境之前, 我们需要指定主机(主地址) 和端口(主端口)。在这个例子中, 我们可以让主 机为节点 0,端口为一个数字, 如 29500。所有的 8 个进程将寻找地址和端口并相互连接, 默认的进程组将被创 建。默认进程组的 world size 为 8,细节如下表。
进程组示例
我们还可以创建一个新的进程组。这个新的进程组可以包含任何进程的子集。例如, 我们可以创建一个只 包含偶数进程的组,如下表所示:
在进程组中,各进程可以通过两种方式进行通信,下图中展示了部分进程组中进程通信示例。
- collective: 一组进程一起执行分散、聚集、 all-reduce 、广播等操作。
进程组中进程通信示例
目前的大语言模型都采用了分布式训练架构完成训练。GPT-3 的训练全部使用 NVIDIA V100 GPU 。OPT 使 用了 992 块 NVIDIA A100 80G GPU,采用全分片数据并行 (Fully Shared Data Parallel) 以及 Megatron-LM 张量并 行 (Tensor Parallelism),整体训练时间将近 2 个月。BLOOM 模型的研究人员则公开了更多在硬件和所采用的系 统架构方面的细节。该模型的训练一共花费 3.5 个月, 使用 48 个计算节点, 每个节点包含 8 块 NVIDIA A100 80G GPU (总计 384 个 GPU),并且使用 4*NVLink 用于节点内部 GPU 之间通信, 节点之间采用四个 Omni-Path 100 Gbps 网卡构建的增强 8 维超立方体全局拓扑网络进行通信。LLaMA 模型训练给出了不同参数规模的总 GPU 小 时数。LLaMA 模型训练采用 A100-80GB GPU,LLaMA-7B 模型训练需要 82432 GPU 小时, LLaMA- 13B 模型训 练需要 135168 GPU 小时, LLaMA-33B 模型训练花费了 530432 GPU 小时, 而 LLaMA-65B 模型训练花费则高达 1022362 GPU 小时。由于 LLaMA 所使用的训练数据量远超 OPT 和 BLOOM 模型, 因此, 虽然模型参数量远小 于上述两个模型,但是其所需计算量仍然非常惊人。
通过使用分布式训练系统, 大语言模型训练周期可以从单计算设备花费几十年, 缩短到使用数千个计算设 备花费几十天就可以完成。然而, 分布式训练系统仍然需要克服计算墙、显存墙、通信墙等多种挑战, 以确保集 群内的所有资源得到充分利用,从而加速训练过程并缩短训练周期。
- 计算墙:单个计算设备所能提供的计算能力与大语言模型所需的总计算量之间存在巨大差异。2022 年 3 年 发布的 NVIDIA H100 SXM 的单卡 FP16 算力也只有 2000 TFLOPs,而 GPT-3 则需要 314 ZFLOPs 的总算 力,两者相差了 8 个数量级。
- 显存墙:单个计算设备无法完整存储一个大语言模型的参数。GPT-3 包含 1750 亿参数, 如果采用FP16 格 式进行存储,需要 700GB 的计算设备内存空间,而 NVIDIA H100 GPU 只有 80 GB 显存。
- 通信墙:分布式训练系统中各计算设备之间需要频繁地进行参数传输和同步。由于通信的延迟和带宽限制, 这可能成为训练过程的瓶颈。GPT-3 训练过程中, 如果分布式系统中存在 128 个模型副本, 那么在每次迭 代过程中至少需要传输 89.6TB 的梯度数据。而截止 2023 年 8 月, 单个 InfiniBand 链路仅能够提供不超过800Gb/s 带宽。
计算墙和显存墙源于单计算设备的计算和存储能力有限, 与模型对庞大计算和存储需求之间存在矛盾。这个问 题可以通过采用分布式训练方法来解决, 但分布式训练又会面临通信墙的挑战。在多机多卡的训练中, 这些问 题逐渐显现。随着大模型参数的增大, 对应的集群规模也随之增加, 这些问题变得更加突出。同时, 在大型集群 进行长时间训练时,设备故障可能会影响或中断训练过程,对分布式系统的问题性也提出了很高要求。