大模型分布式训练方法FDSP和DeepSpeed

备注： 本文部分内容参考自其他作者的内容，如有不妥，请联系，立即删除。

pytorch单精度、半精度、混合精度、单卡、多卡（DP / DDP）、FSDP、DeepSpeed模型训练
相关代码：pytorch-model-train-template

1. Pytorch2.0新特性 FSDP（Fully Sharded Data Parallel）

Fully Sharded Data Parallel

Introducing PyTorch Fully Sharded Data Parallel (FSDP) API

PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel

详解PyTorch FSDP数据并行(Fully Sharded Data Parallel)

2023 年了，大模型训练还要不要用 PyTorch 的 FSDP ？

FSDP 的实现借鉴了 FairScale，对优化器状态、梯度、模型参数进行分区，实现在更大规模的数据集上训练参数量更大的模型。

FairScale（你真的需要 FSDP 么）
ZeRO（Zero Redundancy Optimizer）

模型训练的时候，显存占用大体可以分成三部分，即激活值（根据输入数据得到的中间结果）、模型权重、模型梯度和优化器状态。对于视觉模型而言，显存占比最大的是激活值，因此使用混合精度训练能够大幅度的降低激活值的显存占用（fp16）。然而对于大语言模型或者多模态模型而言，优化后三者的显存占用则显得更重要。

以 PyTorch 为例，当你使用 DistributedDataParallel 时，其实会在每个进程为模型参数、模型梯度、优化器状态分配内存，并在训练过程中同步地更新这些数据。这样的做法虽然能够通过数据并行以达到加速训练的目的，但是它在显存分配上的策略，显然是非常糟糕的。既然每个进行的参数都是一样的，为什么每个进程还需要保存完整的参数呢？所以 ZeRO 就主张每个进程只保存参数的一部分，用到的时候再 all gather 到各个进程。ZeRO 有三个阶段的优化策略，即：

ZeRO1：只把优化器状态进行分片
ZeRO2：对优化器状态 + 梯度进行分片
ZeRO3：对优化器状态 + 梯度 + 模型参数进行分片

FSDP的工作原理

FSDP是一种数据并行训练，但与传统的数据并行训练不同，传统的数据并行训练在每一片GPU上独立维护模型参数、梯度和优化器状态，FSDP可以在多个worker之间将所有这些状态分片，并可以有选择的将分片的模型参数卸载到cpu上。

下图显示了FSDP如何为2个数据并行进程工作:
在这里插入图片描述

all-gather：与All-Reduce的作用类似，All-Reduce对数据做规约，All Gather会收集所有数据到所有节点上。从最基础的角度来看，All Gather相当于一个Gather操作之后跟着一个Broadcast操作。

在这里插入图片描述

reduce-scatter：Reduce Scatter操作将各个节点的输入先进行求和，然后在第0维度按卡数切分，将数据分发到对应的卡上。例如上图所示，每卡的输入均为4x1的Tensor。Reduce Scatter先对输入求和得到[0, 4, 8, 12]的Tensor，然后进行分发，每卡获得1x1大小的Tensor。例如卡0对应的输出结果为[[0.0]]，卡1对应的输出结果为[[4.0]]。

在这里插入图片描述

通常，模型层以嵌套的方式用FSDP封装，在向前或向后计算期间，只有单个FSDP实例中的层需要收集单个设备的全部参数。收集到的完整参数将在计算后立即释放，释放的内存可用于下一层的计算。通过这种方式，可以节省峰值GPU内存，从而可以训练更大的模型或使用更大的batch size。为了进一步提高内存效率，FSDP可以在实例运行时将参数、梯度和优化器状态卸载给CPU

在Pytorch中使用FSDP

有两种方法可以用PyTorch FSDP对模型进行封装。自动封装是DDP的直接替代品；手动封装需要对模型定义代码进行最小的更改，并且能够探索复杂的分片策略。

自动封装

fsdp_auto_wrap_policy参数允许指定一个可调用的函数来递归地用FSDP对模型的Layer进行封装
PyTorch FSDP提供的default_auto_wrap_policy函数递归地封装参数量大于100M的层。
自定义自动封装策略详见 FSDP API doc
此外，可以选择性的配置cpu_offload，当参数在计算过程中暂时用不到时，将这些参数卸载到CPU，可以进一步提高内存效率，但代价是主机和设备之间的数据传输开销。

Pytorch FSDP自动封装的样例代码如下：

from torch.distributed.fsdp import (
   FullyShardedDataParallel,
   CPUOffload,
)
from torch.distributed.fsdp.wrap import (
   default_auto_wrap_policy,
)
import torch.nn as nn
 
class model(nn.Module):
   def __init__(self):
       super().__init__()
       self.layer1 = nn.Linear(8, 4)
       self.layer2 = nn.Linear(4, 16)
       self.layer3 = nn.Linear(16, 4)
 
model = DistributedDataParallel(model())
fsdp_model = FullyShardedDataParallel(
   model(),
   fsdp_auto_wrap_policy=default_auto_wrap_policy,
   cpu_offload=CPUOffload(offload_params=True),
)

手动封装

通过有选择地对模型的某些部分应用wrap，手动包装对于探索复杂的分片策略非常有用。总体设置可以传递给enable_wrap()上下文管理器。

Pytorch FSDP手动封装的样例代码如下：

from torch.distributed.fsdp import (
   FullyShardedDataParallel,
   CPUOffload,
)
from torch.distributed.fsdp.wrap import (
   enable_wrap,
   wrap,
)
import torch.nn as nn
from typing import Dict
 
 
class model(nn.Module):
   def __init__(self):
       super().__init__()
       self.layer1 = wrap(nn.Linear(8, 4))
       self.layer2 = nn.Linear(4, 16)
       self.layer3 = wrap(nn.Linear(16, 4))
 
wrapper_kwargs = Dict(cpu_offload=CPUOffload(offload_params=True))
with enable_wrap(wrapper_cls=FullyShardedDataParallel, **wrapper_kwargs):
   fsdp_model = wrap(model())

FSDP在模型训练阶段的使用

在使用上述两种方法将模型用FSDP封装后，可以用通常训练的方式训练模型，如下所示：

optim = torch.optim.Adam(fsdp_model.parameters(), lr=0.0001)
for sample, label in next_batch():
  out = fsdp_model(input)
  loss = criterion(out, label)
  loss.backward()
  optim.step()

FSDP Benchmark

在单节点具有8片NVIDIA A100-SXM4-40G的集群环境，使用随机生成的数据分别训练1750亿和10000亿的GPT模型。在训练过程中，除了使用CPU Offload之外，还是用了checkpoint技术来节省GPU内存的占用
对于GPT 1750亿参数的模型，每个GPU的最大吞吐量为159 teraFLOP/s，达到了NVIDIA A100峰值理论性能312 teraFLOP/s/GPU的51%，批处理大小为20，序列长度为512，在128个GPU上实现。进一步增加GPU的数量会导致每个gpu的吞吐量下降，因为节点之间的通信越来越多。
对于GPT 10000亿参数的模型，每个GPU的最大吞吐量为84 teraFLOP/s，达到了NVIDIA A100峰值理论性能312 teraFLOP/s/GPU的27%，批处理大小为4，序列长度为2048，在128个GPU上实现。然而，GPU数量的进一步增加并不会对每个GPU的吞吐量产生太大的影响，因为我们观察到，10000亿模型的训练中最大的瓶颈不是来自通信，而是来自GPU峰值内存达到极限时缓慢的CUDA缓存分配器。使用具有更大内存容量的A100 80G GPU将在很大程度上解决这个问题，并且还有助于扩展批处理大小以实现更大的吞吐量。

2. Accelerate中FSDP的用法

使用accelerate config命令配置FSDP参数

accelerate config --config_file ./config/default_config.yml

compute_environment: LOCAL_MACHINE
deepspeed_config: {}
distributed_type: FSDP
downcast_bf16: 'no'
fsdp_config:
  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
  fsdp_backward_prefetch_policy: BACKWARD_PRE
  fsdp_offload_params: false
  fsdp_sharding_strategy: 1
  fsdp_state_dict_type: FULL_STATE_DICT
  fsdp_transformer_layer_cls_to_wrap: BertLayer
machine_rank: 0
main_process_ip: null
main_process_port: null
main_training_function: main
mixed_precision: 'no'
num_machines: 1
num_processes: 2
use_cpu: false

accelerate中支持的FSDP的命令行参数

Sharding Strategy（分区策略）

FULL_SHARD (shards optimizer states, gradients and parameters)。完全分区
SHARD_GRAD_OP (shards optimizer states and gradients)。只对优化器和参数梯度分区
NO_SHARD (DDP)。不做分区
HYBRID_SHARD (shards optimizer states, gradients and parameters within each node while each node has full copy)。混合分区
HYBRID_SHARD_ZERO2 (shards optimizer states and gradients within each node while each node has full copy)。混合分区

Offload Params（卸载参数）：Decides Whether to offload parameters and gradients to CPU。是否将模型参数和梯度卸载到CPU上

Auto Wrap Policy（自动封装策略）

TRANSFORMER_BASED_WRAP。
SIZE_BASED_WRAP
NO_WRAP

Transformer Layer Class to Wrap: When using TRANSFORMER_BASED_WRAP, user specifies comma-separated string of transformer layer class names (case-sensitive) to wrap ,e.g, BertLayer, GPTJBlock, T5Block, BertLayer,BertEmbeddings,BertSelfOutput... This is important because submodules that share weights (e.g., embedding layer) should not end up in different FSDP wrapped units. Using this policy, wrapping happens for each block containing Multi-Head Attention followed by couple of MLP layers. Remaining layers including the shared embeddings are conveniently wrapped in same outermost FSDP unit. Therefore, use this for transformer based models. You can use the model._no_split_modulesfor 🤗 Transformer models by answeringyesto Do you want to use the model’s _no_split_modules to wrap. Only applicable for 🤗 Transformers. It will try to use model._no_split_modules` when available.

Min Num Params: minimum number of parameters when using SIZE_BASED_WRAP

Backward Prefetch（反向预取）

BACKWARD_PRE
BACKWARD_POST
NO_PREFETCH

State Dict Type

FULL_STATE_DICT
LOCAL_STATE_DICT
SHARDED_STATE_DICT

Forward Prefetch: if True, then FSDP explicitly prefetches the next upcoming all-gather while executing in the forward pass. only use with Static graphs.

Use Orig Params: If True, allows non-uniform requires_grad during init, which means support for interspersed frozen and trainable paramteres. Useful in cases such as parameter-efficient fine-tuning.
Please refer this blog

CPU RAM Efficient Model loading: If True, only the first process loads the pretrained model checkoint while all other processes have empty weights. Only applicable for 🤗 Transformers models. This should be set to False if you experience errors when loading the pretrained 🤗 Transformers model via from_pretrained method. When using this, Sync Module States needs to be True else all the processes expect the main process would have random empty weights leading to unexpected behaviour during training.

Sync Module States`: If True, each individually wrapped FSDP unit will broadcast module parameters from rank 0

3. DeepSpeed

DeepSpeed使用指南(简略版)

DeepSpeed

ZeRO & DeepSpeed: New system optimizations enable training models with over 100 billion parameters

eepSpeed之ZeRO系列：将显存优化进行到底

模型显存占用分析

假设模型采用混合精度训练，采用Adam优化器，Adam在SGD基础上，为每个参数梯度增加了一阶动量（momentum）和二阶动量（variance）。混合精度训练，字如其名，同时存在fp16和fp32两种格式的数值，其中模型参数、模型梯度都是fp16，此外还有fp32的模型参数，如果优化器是Adam，则还有fp32的momentum和variance。

模型状态：模型参数（fp16）、模型梯度（fp16）和Adam状态（fp32的模型参数备份，fp32的momentum和fp32的variance）。假设模型参数量 $\Phi$ ，则共需要 $2\Phi + 2\Phi + (4\Phi + 4\Phi + 4\Phi) = 16\Phi$ 字节存储，其中Adam状态占比达到75%
剩余状态: 除了模型状态之外的显存占用，包括激活值（activation）、各种临时缓冲区（buffer）以及无法使用的显存碎片（fragmentation）。其中激活值占用的显存可以采用activation checkpoint方法以计算换空间的方式来处理

模型参数划分

参数占用逻辑上可以分成三种类型。将这些类型的参数划分：

optimizer states：即优化器的参数状态。例如，Adam的动量参数。
gradients：梯度缓存，对应于optimizer。
parameters：模型参数。

对应的，DeepSpeed的ZeRO configp配置文件就可以分为如下几类：

ZeRO Stage 1: 划分optimizer states。优化器参数被划分到多个memory上，每个momoey上的进程只负责更新它自己那部分参数。
ZeRO Stage 2: 划分gradient。每个memory，只保留它分配到的optimizer state所对应的梯度。这很合理，因为梯度和optimizer是紧密联系在一起的。只知道梯度，不知道optimizer state，是没有办法优化模型参数的。主要用于模型训练阶段
ZeRO Stage 3: 划分parameter模型参数，或者说，不同的layer. ZeRO-3会在forward和backward的时候，自动将模型参数分配到多个memory。既能用于模型训练阶段，也能用在推理阶段，可以在推理阶段将一个大模型拆分到多个GPU上部署

DeepSpeed功能

Stage 1：将优化器状态在多个workers/GPUs之间进行分区，显存占用可以降低4倍，同时保持和采用DDP一样的通信开销。对优化器Adam的状态进行分片，也就是图中的 $P_{os}$ ，这里os指的是optimizer states。模型参数（parameters）和梯度（gradients）仍旧是每张卡保持一份，此时，每张卡的模型状态所需显存是 $4\Phi + \frac{12\Phi}{N_{gpus}}$ 字节，当 $N_{gpus}$ 比较大时，趋向于 $4\Phi$ ，也就是原来 $16\Phi$ 的 $\frac{1}{4}$ 。
Stage 2：将优化器状态 + 参数梯度在多个workers/GPUs之间进行分区，显存占用可以降低8倍，同时保持和采用DDP一样的通信开销。继续对模型梯度进行分片，也就是图中的 $p_{os+g}$ ，模型参数仍旧是每张卡保持一份，此时，每张卡的模型状态所需显存是 $2\Phi + \frac{2\Phi + 12\Phi}{N_{gpus}}$ 字节，当 $N_{gpus}$ 比较大时，趋向于 $2\Phi$ ，也即是原来 $16\Phi$ 的 $\frac{1}{8}$ 。
Stage 3：将优化器状态 + 参数梯度 + 模型参数在多个workers/GPUs之间进行分区，显存占用的降低和GPU的数量成正比，但是会带来50%的额外通信开销。继续对模型参数进行分片，也就是图中的 $P_{os+g+p}$ ，此时每张卡的模型状态所需显存是 $\frac{16\Phi}{N_{gpus}}$ 字节，当 $N_{gpus}$ 比较大时，趋向于 0 。
Optimizer Offload：在Stage 2的基础上，将优化器状态 + 参数梯度卸载到CPU或者硬盘上，进一步降低GPU内存的占用
Param Offload：在Stage 3的基础上，将模型参数卸载到CPU或者硬盘上，进一步降低GPU内存的占用

DeepSpeed通信开销
基于accelerate配置DeepSpeed参数

accelerate config --config_file default_ds.yml

基于accelerate和DeepSpeed参数执行代码

accelerate launch my_script.py --default_ds.yml

ZeRO Stage-2 参数示例

compute_environment: LOCAL_MACHINE
deepspeed_config:
 deepspeed_config_file: /home/ubuntu/accelerate/examples/configs/deepspeed_config_templates/zero_stage2_config.json
 zero3_init_flag: true
distributed_type: DEEPSPEED
fsdp_config: {}
machine_rank: 0
main_process_ip: null
main_process_port: null
main_training_function: main
mixed_precision: fp16
num_machines: 1
num_processes: 2
use_cpu: false

zero_stage2_config.json文件如下：

{
    "fp16": {
        "enabled": true,
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 16,
        "hysteresis": 2,
        "min_loss_scale": 1
    },
    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": "auto",
            "weight_decay": "auto",
            "torch_adam": true,
            "adam_w_mode": true
        }
    },
    "scheduler": {
        "type": "WarmupDecayLR",
        "params": {
            "warmup_min_lr": "auto",
            "warmup_max_lr": "auto",
            "warmup_num_steps": "auto",
            "total_num_steps": "auto"
        }
    },
    "zero_optimization": {
        "stage": 2,
        "allgather_partitions": true,
        "allgather_bucket_size": 2e8,
        "overlap_comm": true,
        "reduce_scatter": true,
        "reduce_bucket_size": "auto",
        "contiguous_gradients": true
    },
    "gradient_accumulation_steps": 1,
    "gradient_clipping": "auto",
    "steps_per_print": 2000,
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "wall_clock_breakdown": false
}

ZeRO Stage-3 with CPU Offload 参数示例

compute_environment: LOCAL_MACHINE
deepspeed_config:
 deepspeed_config_file: /home/ubuntu/accelerate/examples/configs/deepspeed_config_templates/zero_stage3_offload_config.json
 zero3_init_flag: true
distributed_type: DEEPSPEED
fsdp_config: {}
machine_rank: 0
main_process_ip: null
main_process_port: null
main_training_function: main
mixed_precision: fp16
num_machines: 1
num_processes: 2
use_cpu: false

zero_stage3_offload_config.json文件如下：

{
    "fp16": {
        "enabled": true,
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 16,
        "hysteresis": 2,
        "min_loss_scale": 1
    },
    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": "auto",
            "weight_decay": "auto"
        }
    },
    "scheduler": {
        "type": "WarmupDecayLR",
        "params": {
            "warmup_min_lr": "auto",
            "warmup_max_lr": "auto",
            "warmup_num_steps": "auto",
            "total_num_steps": "auto"
        }
    },
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": true
        },
        "offload_param": {
            "device": "cpu",
            "pin_memory": true
        },
        "overlap_comm": true,
        "contiguous_gradients": true,
        "reduce_bucket_size": "auto",
        "stage3_prefetch_bucket_size": "auto",
        "stage3_param_persistence_threshold": "auto",
        "sub_group_size": 1e9,
        "stage3_max_live_parameters": 1e9,
        "stage3_max_reuse_distance": 1e9,
        "stage3_gather_16bit_weights_on_model_save": "auto"
    },
    "gradient_accumulation_steps": 1,
    "gradient_clipping": "auto",
    "steps_per_print": 2000,
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "wall_clock_breakdown": false
}

accelerate支持的DeepSpeed参数列表

zero_stage:

Disabled 不使用ZeRO
optimizer state partitioning 对优化器状态进行分区
optimizer+gradient state partitioning 对优化器状态 + 参数梯度进行分区
optimizer+gradient+parameter partitioning 对优化器状态 + 参数梯度 + 模型参数进行分区

gradient_accumulation_steps: 梯度累加的部署

gradient_clipping: 梯度裁剪的数值，默认1.0

offload_optimizer_device: 优化器状态 + 参数梯度卸载

[none] Disable optimizer offloading 禁用优化器状态 + 参数梯度卸载
[cpu] offload optimizer to CPU 优化器状态 + 参数梯度卸载到CPU
[nvme] offload optimizer to NVMe SSD. Only applicable with ZeRO >= Stage-2. 优化器状态 + 参数梯度卸载到SSD硬盘

offload_param_device: 模型参数卸载

[none] Disable parameter offloading 禁用参数卸载
[cpu] offload parameters to CPU 模型参数卸载到CPU
[nvme] offload parameters to NVMe SSD. Only applicable with ZeRO Stage-3. 模型参数卸载到SSD硬盘

zero3_init_flag: Decides whether to enable deepspeed.zero.Init for constructing massive models. Only applicable with ZeRO Stage-3.

zero3_save_16bit_model: Decides whether to save 16-bit model weights when using ZeRO Stage-3.

mixed_precision: 混合精度训练

no for FP32 training 使用单精度FP32
fp16 for FP16 mixed-precision training 使用半精度fp16
bf16 for BF16 mixed-precision training. 使用半精度bf16，需要硬件支持

4. Accelerate中的DeepSpeed用法

accelerate集成DeepSpeed的两种方式：

通过accelerate config --config_file XXX.yal的方式配置DeepSpeed参数
通过DeepSpeedPlugin的方式

5. GPU分布式训练通讯原语

分布式训练硬核技术——通信原语

Overview of NCCL

AI框架基础技术之深度学习中的通信优化

Collective Operations

MPI和NCCL是最常用的通讯库，MPI专注于CPU的并行通讯，NCCL则专注于GPU的通讯。

目前机器学习中主要由两种分布式架构：

参数服务器架构（Parameter Server，PS）
去中心化架构（Decentralized Network）

其中，分布式训练通常在计算集群上进行，集群的每个节点分别执行一部分计算。不同节点的计算之间有数据依赖和共享，需要将数据在不同节点间传输，这就是通信。

分布式的通信一般有两大类：

集合通信（Collective communication，CC）：在一组节点内进行通信
点对点通信（Point to point communication，P2P）：在两个节点之间进行通信

深度学习训练过程中因为需要传输大量的网络模型权重参数和训练过程中产生的大量临时变量等，因此主要使用集合通信的方式。可以理解为，机器学习/深度学习的分布式训练，主要是采用在PS架构下的集合通讯模式；而在大模型的分布式训练中，因为减少跟单点参数服务器统一更新，更多直接采用纯集合通讯模式。

集合通讯中包含多个sender和多个receiver，一般的通信原语包括broadcast、gather、all-gather、scatter、reduce、all-reduce、reduce-scatter、all-to-all等通信操作进行数据传输，下面将会分别介绍其具体含义。

broadcast：在集合通信中，如果某个节点想把自身的数据发送到集群中的其他节点，那么就可以使用广播Broadcast的操作。

在这里插入图片描述

scatter：Scatter操作表示一种散播行为，将主节点的数据进行划分并散布至其他指定的节点。

Scatter与Broadcast非常相似，都是一对多的通信方式，不同的是Broadcast的0号节点将相同的信息发送给所有的节点，而Scatter则是将数据的不同部分，按需发送给所有的节点。
在这里插入图片描述

reduce：Reduce称为规约运算，是一系列简单运算操作的统称，细分可以包括：SUM、MIN、MAX、PROD、LOR等类型的规约操作。Reduce意为减少/精简，因为其操作在每个节点上获取一个输入元素数组，通过执行操作后，将得到精简的更少的元素。

在NCCL中的Reduce，从多个sender那里接收数据，最终combine到一个节点上。
在这里插入图片描述

all-reduce：Reduce是一系列简单运算操作的统称，All Reduce则是在所有的节点上都应用同样的Reduce操作。All Reduce操作可通过单节点上Reduce + Broadcast操作完成。

在这里插入图片描述

gather：Gather操作将多个sender上的数据收集到单个节点上，Gather可以理解为反向的Scatter。Gather操作会从多个节点里面收集数据到一个节点上面，而不是从一个节点分发数据到多个节点。

在这里插入图片描述

Gather和Reduce的区别：Gather对数据进行收集，组成一个大的集合，Reduce对数据进行规约，规约后的几何大小与各个子节点保持一致。

all-gather：很多时候发送多个元素到多个节点也很有用，即在多对多通信模式的场景。与All-Reduce的作用类似，All-Reduce对数据做规约，All Gather会收集所有数据到所有节点上。从最基础的角度来看，All Gather相当于一个Gather操作之后跟着一个Bcast操作。

在这里插入图片描述

reduce-scatter：Reduce Scatter操作将各个节点的输入先进行求和，然后在第0维度按卡数切分，将数据分发到对应的卡上。例如上图所示，每卡的输入均为4x1的Tensor。Reduce Scatter先对输入求和得到[0, 4, 8, 12]的Tensor，然后进行分发，每卡获得1x1大小的Tensor。例如卡0对应的输出结果为[[0.0]]，卡1对应的输出结果为[[4.0]]。

在这里插入图片描述