大模型分布式训练并行技术
三种方式:数据并行、模型并行、3D并行
参考:https://juejin.cn/post/7254001262646738981
一、数据并行DP (Data Parallel)
数据并行是最常见的并行形式,因为它很简单。在数据并行训练中,数据集被分割成几个碎片,每个碎片被分配到一个设备上。这相当于沿批次(Batch)维度对训练过程进行并行化。每个设备将持有一个完整的模型副本,并在分配的数据集碎片上进行训练。在反向传播之后,模型的梯度将被全部减少,以便在不同设备上的模型参数能够保持同步。典型的数据并行实现:PyTorch DDP。
二、模型并行
通常有两种类型的模型并行:张量并行和流水线并行
- 流水线并行是在各层之间进行并行计算。
- 张量并行是在一个操作中进行并行计算,如:矩阵-矩阵乘法。
1、流水线并行 PP(Pipeline Parallelism)
流水线并行的核心思想是,模型按层分割成若干块,每块都交给一个设备。
在前向传播过程中,每个设备将中间的激活传递给下一个阶段。
在后向传播过程中,每个设备将输入张量的梯度传回给前一个流水线阶段。
这允许设备同时进行计算,从而增加训练的吞吐量。
流水线并行训练的一个明显缺点是训练设备容易出现空闲状态(因为后一个阶段需要等待前一个阶段执行完毕),导致计算资源的浪费,加速效率没有数据并行高。
典型的流水线并行实现:GPipe、PipeDream、PipeDream-2BW、PipeDream Flush(1F1B)
2、张量并行TP(Tensor Parallelism)
3、优化器并行 ???
目前随着模型越来越大,单个GPU的显存目前通常无法装下那么大的模型了。那么就要想办法对占显存的地方进行优化。
通常来说,模型训练的过程中,GPU上需要进行存储的参数包括了模型本身的参数、优化器状态、激活函数的输出值、梯度以及一些零时的Buffer。各种数据的占比如下图所示:
可以看到模型参数仅占模型训练过程中所有数据的一部分,当进行混合精度运算时,其中模型状态参数(优化器状态 + 梯度+ 模型参数)占到了一大半以上。因此,我们需要想办法去除模型训练过程中的冗余数据。
而优化器相关的并行就是一种去除冗余数据的并行方案,目前这种并行最流行的方法是 ZeRO(即零冗余优化器)。针对模型状态的存储优化(去除冗余),ZeRO使用的方法是分片,即每张卡只存 1/N 的模型状态量,这样系统内只维护一份模型状态。ZeRO有三个不同级别,对模型状态进行不同程度的分片:
ZeRO-1 : 对优化器状态分片(Optimizer States Sharding)
ZeRO-2 : 对优化器状态和梯度分片(Optimizer States & Gradients Sharding)
ZeRO-3 : 对优化器状态、梯度分片以及模型权重参数分片(Optimizer States & Gradients & Parameters Sharding)
三、3D并行
总的来说,3D并行是由数据并行(DP)、张量并行(TP)和流水线并行(PP)组成