当前位置:数码通 > 科技

大模型分布式训练并行技术(一)

来源于 数码通 2023-10-04 16:22

近年来,随着 Transformer 和 MOE 架构的引入,深度学习模型已经轻松突破万亿参数的规模,传统的单机单卡模型已经不能满足训练超大型模型的要求。因此,我们需要在单机甚至多机多卡上训练基于多卡的分布式大模型。

使用AI集群使深度学习算法能够从大量数据中更高效地训练具有卓越性能的大型模型,这是分布式机器学习的主要焦点目标。为了实现这一目标,一般需要根据硬件资源和数据/模型规模的匹配来考虑划分计算任务、训练数据和模型,从而进行分布式存储和分布式训练。因此,分布式训练相关技术值得我们深入分析其背后的机制。

下面主要讲解大模型分布式训练的并行技术。本系列大致分为九篇文章。

大模型分布式训练并行技术(一)——概述

大模型分布式训练并行技术(二)——数据并行

大模型分布式训练并行技术(三)——流水线并行

大模型分布式训练并行技术(四)——张量并行

大模型分布式训练并行技术(五)——序列并行

大模型分布式训练并行技术(六)——多维混合并行

大模型分布式训练并行技术(七)-自动并行

大模型分布式训练并行技术(八)-MOE并行

大模型分布式训练并行技术(九)-总结

本文是分布式训练并行技术的第一篇文章。简要介绍了大型模型分布式训练的常用并行技术。

数据并行

数据并行是最常见的并行形式,因为它很简单。在数据并行训练中,数据集被分成多个分片,每个分片分配给一个设备。这相当于沿着批量维度并行化训练过程。每个设备将保存模型的完整副本,并在分配的数据集分片上进行训练。经过反向传播后,模型的梯度将完全减小,从而使不同设备上的模型参数保持同步。典型的数据并行实现:PyTorch DDP。

图片.png

模型并行性

在数据并行训练中,一个明显的特征是每个GPU都持有整个模型权重的副本。这会产生冗余问题。另一种并行模式是模型并行,其中模型被分割并分布在一系列设备上。

模型并行通常有两种类型:张量并行和管道并行。

张量并行是单个操作中的并行计算,例如矩阵-矩阵乘法。

管道并行是层与层之间的并行计算。

因此,从另一个角度来看,张量并行可以看作是层内并行,而管道并行可以看作是层间并行。

张量平行

张量并行训练就是将张量沿特定维度划分为N个块。每个设备只保存整个张量的1/N,不影响计算图的正确性。这需要额外的沟通以确保结果的正确性。

以一般矩阵乘法为例,假设我们有 C = AB。我们可以沿着列将 B 分成 [B0 B1 B2...Bn],每个设备持有一列。然后,我们在每个设备上将 A 乘以 B 中的每一列,我们将得到 [AB0 AB1 AB2 ... ABn]。此时,每个设备仍然保存着一部分结果,例如设备(rank=0)保存着AB0。为了确保结果的正确性,我们需要收集所有结果并沿列连接张量。这样我们就能够在设备上分配张量,同时确保计算流程保持正确。

图片.png

典型的张量并行实现:Megatron-LM (1D)、Colossal-AI (2D、2.5D、3D)。

管道并行度

流水线并行的核心思想是,将模型逐层划分为若干块,每个块交付给一个设备。

在前向传播期间,每个设备将中间激活传递到下一个阶段。

在反向传播期间,每个设备将输入张量的梯度传递回前一个管道阶段。

这允许设备同时执行计算,从而提高训练吞吐量。

图片.png

流水线并行训练的一个明显缺点是训练设备容易出现空闲状态(因为后一个阶段需要等待前一个阶段执行),导致计算资源浪费,且加速效率不高作为数据并行性。

图片.png

典型的管道并行实现:GPipe、PipeDream、PipeDream-2BW、PipeDream Flush (1F1B)。

优化器相关的并行性

随着模型变得越来越大,单个GPU的显存通常无法容纳如此大的模型。那么你就得想办法优化显存占用的空间。

一般来说,在模型训练的过程中,需要存储在GPU上的参数包括模型本身的参数、优化器状态、激活函数的输出值、梯度以及一些零时间缓冲区。各种数据占比如下图所示:

图片.png

可以看出,模型训练过程中模型参数只占全部数据的一部分。进行混合精度运算时,模型状态参数(优化器状态+梯度+模型参数)占一半以上。因此,我们需要想办法在模型训练过程中去除冗余数据。

优化器相关的并行性是一种去除冗余数据的并行方案。目前,最流行的并行方法是ZeRO(零冗余优化器)。对于模型状态存储优化(去除冗余),ZeRO 采用分片的方式,即每张卡只存储模型状态量的 1/N,这样系统中只维护一份模型状态的副本。 ZeRO 具有三个不同级别,可不同程度地分割模型状态:

ZeRO-1:Optimizer 状态分片)

ZeRO-2:优化器状态和梯度分片

ZeRO-3:优化器状态和梯度和 Parameters 分片)

图片.png

异构系统并行性

在上述方法中,通常需要大量的GPU来训练大型模型。然而,人们经常忽视的是,CPU 与 GPU 相比,内存要大得多。在典型的服务器上,CPU 可以轻松拥有数百 GB 甚至 TB 的内存,而每个 GPU 卡通常只有 48 或 80 GB 的内存。这促使人们思考为什么不使用CPU内存来进行分布式训练。

最近的进展依赖 CPU 甚至 NVMe 磁盘来训练大型模型。主要思想是在不使用张量时将其卸载回 CPU 内存或 NVMe 磁盘。

通过使用异构系统架构,可以在单台机器上容纳巨大的模型。

图片.png

多维混合并行

多维混合并行是指结合数据并行、模型并行、管道并行等多种并行技术进行分布式训练。

图片.png

通常,在超大规模模型的预训练和全参数微调时,需要使用多维混合并行。

图片.png

为了充分利用带宽,张量并行通常需要最大的通信量,而数据并行和管道并行则需要相对较小的通信量。因此,在同一服务器内使用张量并行,而在服务器之间使用数据并行和管道并行。

图片.png

自动并行化

前面提到的数据并行、张量并行、流水线并行等多维混合并行,需要将模型拆分为多个AI加速卡。如果让用户手动实现,对于开发者来说会非常困难,需要考虑性能、内存、沟通、训练效果等问题。如果模型能够自动按算子或按层划分为不同的加速卡,则可以大大降低开发者的使用难度。因此,自动并行应运而生。

图片.png

MOE并行/专家并行

一般来说,模型规模的扩大会导致训练成本的大幅增加,计算资源的限制成为大规模集约化模型训练的瓶颈。为了解决这个问题,提出了一种基于稀疏MoE层的深度学习模型架构,即将大模型拆分为多个小模型(专家),每次迭代确定激活一部分专家进行计算基于样本,达到节省计算资源的效果;引入可训练且保证稀疏性的门机制,保证算力的优化。

利用MoE结构,可以实现超大规模的模型训练,同时计算成本呈次线性增长,为恒定的计算资源预算带来巨大收益。 MOE并行本质上是一种模型并行方法。下图显示了一个模型,其中六个专家网络由两个专家通道并行训练。其中,专家1-3放置在第一计算单元上,专家4-6放置在第二计算单元上。

图片.png

结论

本文简要介绍大型模型分布式训练常用的并行技术。后续章节将对常见并行技术的不同解决方案进行详细说明。

-->
登录后参与评论