DeepSpeed原理

DeepSpeed 是由微软开发的一种深度学习优化库，专为高性能训练和推理而设计，尤其适用于大规模深度学习模型（如 GPT 系列、BERT 等）。它通过一系列技术和优化策略，帮助研究者和开发者高效利用硬件资源，实现快速训练、降低内存使用以及提升推理速度。
正如其官方描述那样：

Image From: https://github.com/microsoft/DeepSpeed

Deepspeed作为一种显存优化技术，那么就会有一个问题：模型训练显存都被谁占用了？
参考论文（https://arxiv.org/pdf/1910.02054）中的描述在一个1.5B的GPT-2模型参数量为3G（半精度）但是一块32G的显卡可能无法训练下来，这是因为显存都被 模型状态 以及 剩余状态（Residual Memory Consumption）

模型状态显存占用
主要指的是：优化器状态，梯度，模型参数。比如说在训练过程中一般都会选择使用Adam作为一种优化器进行使用，而在Adam计算过程中就会存储两部分内容：1、动量（上一轮梯度累计）；2、二阶动量（存储梯度平方的滑动平均值）。如何去避免这部分结果对显存占用的影响，就提出了 混合精度训练（用FP16存储和计算梯度及优化器状态）
比如说：用Adam作为优化器在混合精度下训练参数量为\(\Phi\)的模型显存占用：1、一部分用来存储FP16的参数以及梯度：\(2\Phi, 2\Phi\)；2、另外一部分需要存储优化器状态（FP32存储：模型参数，动量，二阶动量）：\(4\Phi, 4\Phi, 4\Phi\)。那么显存占用上就有：\(2+ 2+ 4+ 4+ 4=16\Phi\)。那么回到上面提到的1.5B的GPT-2至少需要：\(1.5 \times 16=24G\)

剩余状态显存占用
这部分主要指的是： 除了模型状态之外的显存占用，包括激活值（activation）（可以通过Activation checkpointing减少）、各种临时缓冲区（buffer）以及无法使用的显存碎片（fragmentation）

了解模型训练过程中显存占用之后再去了解DeepSpeed中核心内容ZeRO（按照论文中表述作者是分了两部分介绍：ZeRO-DP和ZeRO-R分别去优化上面两部分显存占用）

ZeRO-DP原理

主要是通过切分（partitioning）的方式来减少 模型状态显存占用

第一种方式为\(P_{OS}\)：对优化器的状态进行切分，将\(N\)块GPU上每块只存储\(\frac{1}{N}\)，那么最后显存占用（按上面的显存分析为例）就为：\(4\Phi+ \frac{12\times \Phi}{N}\)

第二种方式为\(P_{OS+g}\)也就是在对优化器切分的基础上补充一个对梯度的切分，那么显存占用上就变成为：\(2\Phi+ \frac{(2+ 12)\times \Phi}{N}\)

第三种方式为\(P_{OS+g+p}\)也就是对模型状态三个都进行切分，显存占用为：\(\frac{4\Phi+ 12\Phi}{N}\)

对于上面3种方式显存减少上分别为：\(4\text{x}, 8\text{x}, N\)（其中N表示的为设备数量）

进一步理解上面3个操作

Image From: https://zhuanlan.zhihu.com/p/618865052

第一种方式\(P_{OS}\)

因为会将优化器状态切分，那么在3个不同设备上分别存储3分优化器状态（o1, o2, o3）,对于这3部分优化器（因为优化器最后还是去“作用”到梯度上），分别对各自的梯度进行优化，但是会有一个问题：每块GPU上存储的是 一部分优化器状态，那么对于每份优化器也只能去优化各自的参数，每次更新需要通过 All-Gather 操作合并梯度，完成优化器状态更新

第二种方式\(P_{OS+g}\)

在进行前向+反向传播之后，得到完整的梯度，因为要实现梯度拆分，那么就对梯度进行reduce-scatter对于不同的GPU就会存储不同的梯度（g1, g2, g3白色的就会剔除掉）前向和反向传播需要通过 All-Gather 和 All-Reduce 操作同步梯度和参数

第三种方式为\(P_{OS+g+p}\)

通过 All-Gather和 Reduce-Scatter 高效完成参数同步和更新。

总的来说：ZeRO-DP是一种 用完就丢 的套路，计算时候是完整内容，但是使用完之后就丢掉

补充1：All-Gather, All-Reduce, reduce-scatter什么意思？

All-Gather：将每个设备上的数据片段收集起来并广播到所有设备上。最终，每个设备都会拥有所有设备的数据片段

比如说4个GPU分别存储不同的值：\(GPU_i: i(i=1,2,3,4)\)通过 all-gather那么不同GPU值为\(GPU_i: [1,2,3,4]\)

reduce-scatter：将数据分片并执行 聚合，然后将结果分发给每个设备。每个设备最终只保留聚合后的部分结果

比如说4个GPU分别存储不同的值：\(GPU_i: [i_1, i_2, i_3, i_4](i=1,2,3,4)\)通过 reduce-scatter（假设为按照加法聚合）那么不同GPU值为\(GPU_1: [1_1, 2_1, 3_1, 4_1]...\)

All-Reduce：用于在所有设备之间对数据进行 聚合（Reduce） 和 广播（Broadcast）。每个设备都会执行相同的聚合操作，并最终持有相同的聚合结果

比如说4个GPU分别存储不同的值：\(GPU_i: i(i=1,2,3,4)\)通过 all-reduce（假设为sum）那么不同GPU值为\(GPU_i: 10\)
Ring-ALLReduce操作：
第一阶段，通过reduce-sactter传递参数

通过3次参数更新之后，这样就会出现不同设备上都会有一个都具有参数\(a_i+ b_i+ c_i+ d_i\)那么下一阶段就是通过all-gather将不同设备上参数广播到不同设备最后实现参数都实现更新。

补充2：通信量和传统的数据并行之间有无区别？
这部分描述来自论文（https://arxiv.org/pdf/1910.02054）中的描述：
传统的数据并行方式：传统的DDP主要使用的是Ring AllReduce在通信量上为：\(2\Phi\)（主要来自两部分：）
\(P_{OS} \text{和} P_{OS+g}\)通信量：\(2\Phi\)，以后者为例：因为每部分设备只保留了部分梯度信息，因此首先需要通过reduce-scatter操作（\(\Phi\)）在梯度都通一之后需要对所有的参数进行更新（\(\Phi\)）
\(P_{OS+g+p}\)：\(3\Phi\)。前向传播过程中每个设备都只保存了部分参数，因此需要对设备之间进行一次参数广播，在前向操作结束之后，将其他参数删除掉（比如\(GPU_i\)接受了\(i+1,...,n\)的参数，那么就将这部分参数删除）此部分通信量为：\(\frac{\Phi \times N}{N}=\Phi\)，类似的反向传播还需要再来一次，梯度还需要进行reduce-scatter

ZeRO-R原理

1、对于激活值的占用。通过\(P_a\)：Partitioned Activation Checkpointing通过分区+checkpointing方式
2、对于临时缓冲区。模型训练过程中经常会创建一些大小不等的临时缓冲区，比如对梯度进行AllReduce，解决办法就是预先创建一个固定的缓冲区，训练过程中不再动态创建，如果要传输的数据较小，则多组数据bucket后再一次性传输，提高效率
3、对于显存碎片。显存出现碎片的一大原因是时候gradient checkpointing后，不断地创建和销毁那些不保存的激活值，解决方法是预先分配一块连续的显存，将常驻显存的模型状态和checkpointed activation存在里面，剩余显存用于动态创建和销毁discarded activation

DeepSpeed代码

Deepspeed代码也比较简单，首先安装deepspeed:pip install deepspeed。使用deepspeed之前一般先去初始化，代码如下：

def initialize(args=None,
               model: torch.nn.Module = None,
               optimizer: Optional[Union[Optimizer, DeepSpeedOptimizerCallable]] = None,
               model_parameters: Optional[torch.nn.Module] = None,
               training_data: Optional[torch.utils.data.Dataset] = None,
               lr_scheduler: Optional[Union[_LRScheduler, DeepSpeedSchedulerCallable]] = None,
               distributed_port: int = TORCH_DISTRIBUTED_DEFAULT_PORT,
               mpu=None,
               dist_init_required: Optional[bool] = None,
               collate_fn=None,
               config=None,
               mesh_param=None,
               config_params=None):
    """初始化 DeepSpeed 引擎。

    参数:
        args: 一个包含 `local_rank` 和 `deepspeed_config` 字段的对象。
            如果提供了 `config`，此参数是可选的。

        model: 必填项：在应用任何包装器之前的 nn.Module 类。

        optimizer: 可选：用户定义的 Optimizer 或返回 Optimizer 对象的 Callable。
            如果提供，将覆盖 DeepSpeed JSON 配置中的任何优化器定义。

        model_parameters: 可选：torch.Tensors 或字典的可迭代对象。
            指定需要优化的张量。

        training_data: 可选：torch.utils.data.Dataset 类型的数据集。

        lr_scheduler: 可选：学习率调度器对象或一个 Callable，接收一个 Optimizer 并返回调度器对象。
            调度器对象应定义 `get_lr()`、`step()`、`state_dict()` 和 `load_state_dict()` 方法。

        distributed_port: 可选：主节点（rank 0）用于分布式训练期间通信的空闲端口。

        mpu: 可选：模型并行单元对象，需实现以下方法：
            `get_{model,data}_parallel_{rank,group,world_size}()`。

        dist_init_required: 可选：如果为 None，将根据需要自动初始化 torch 分布式；
            否则用户可以通过布尔值强制初始化或不初始化。

        collate_fn: 可选：合并样本列表以形成一个小批量的张量。
            在从 map-style 数据集中使用批量加载时使用。

        config: 可选：可以作为路径或字典传递的 DeepSpeed 配置，
            用于替代 `args.deepspeed_config`。

        config_params: 可选：与 `config` 相同，为了向后兼容保留。

    返回值:
        返回一个包含 `engine`, `optimizer`, `training_dataloader`, `lr_scheduler` 的元组。

        * `engine`: DeepSpeed 运行时引擎，用于包装客户端模型以进行分布式训练。

        * `optimizer`: 如果提供了用户定义的 `optimizer`，返回包装后的优化器；
          如果在 JSON 配置中指定了优化器也会返回；否则为 `None`。

        * `training_dataloader`: 如果提供了 `training_data`，则返回 DeepSpeed 数据加载器；
          否则为 `None`。

        * `lr_scheduler`: 如果提供了用户定义的 `lr_scheduler`，或在 JSON 配置中指定了调度器，
          返回包装后的学习率调度器；否则为 `None`。
    """

deepspeed具体案例可以查看其官方示例：https://github.com/microsoft/DeepSpeedExamples
具体使用也很简单,因为Deepspeed将各种功能都封装好了，可以直接使用，一个建议Demo如下：

# 首先初始化
model_engine, optimizer, train_loader, _ = deepspeed.initialize(
        model=model,
        model_parameters=model.parameters(),
        training_data=train_dataset,
        config=config['deepspeed_config'] # 这里的话是直接将deepspeed的设置都存储到一个json文件里面了
    )
def train(model_engine, optimizer, train_loader, ...):
    ...
    image = image.to(model_engine.local_rank)
    out = model_engine(..)
    ...
    model_engine.backward()
    model_engine.step()
    ...

值得注意的是：

• 1、如果需要访问设备，可以直接用：model_engine.local_ranl()进行访问即可
• 2、如果再deepspeed参数（更加多的参数可以参考官方文档：1，2）中设置了 半精度 训练，在数据里面要设定：images.to(model.local_rank).half()

{
  "train_batch_size": 512,
  "gradient_accumulation_steps": 1,
  "fp16": {
    "enabled": true,
    "loss_scale": 0,
    "loss_scale_window": 1000,
    "hysteresis": 2,
    "min_loss_scale": 1
  }, //开启半精度训练
  "optimizer": {
    "type": "Adam",
    "params": {
      "lr": 0.001,
      "betas": [0.9, 0.999],
      "eps": 1e-8,
      "weight_decay": 0.01
    }
  }, // 设置优化器
  "zero_optimization": {
    "stage": 2
  } // 指定zero的方式：1，2，3
}

• 3、理论上分析，在显存占用上是 \(P_{OS}<P_{OS+g}<P_{OS+g+p}\) 但是实验过程中会出现相反的情况，参考这部分讨论：1、在使用deepspeed中的zero设定时，需要保证模型的大小足够大（大小>1B的参数）。于此同时在使用stage=2或者stage=3的时候可以分别指定下面参数：1、reduce_bucket_size，allgather_bucket_size；2、stage3_max_live_parameters， stage3_max_reuse_distance
• 4、对于zero中stage设定，通过结合github上的讨论：

1、stage=2时：

{
"zero_optimization": {
    "stage": 2,
    "offload_optimizer": {
        "device": "cpu",
        "pin_memory": true
    },
    "allgather_partitions": true,
    "allgather_bucket_size": 2e8,
    "overlap_comm": true,
    "reduce_scatter": true,
    "reduce_bucket_size": 2e8,
    "contiguous_gradients": true
},
}

2、stage=3时：

{
"zero_optimization": {
      "stage": 3,
      "offload_optimizer": {
          "device": "cpu",
          "pin_memory": true
      }, //是否将优化器状态迁移到CPU
      "offload_param": {
          "device": "cpu",
          "pin_memory": true
      },
      "overlap_comm": true,
      "contiguous_gradients": true,
      "sub_group_size": 1e9,
      "reduce_bucket_size": "auto",
      "stage3_prefetch_bucket_size": "auto",
      "stage3_param_persistence_threshold": "auto",
      "stage3_max_live_parameters": 1e9,
      "stage3_max_reuse_distance": 1e9,
      "stage3_gather_16bit_weights_on_model_save": true
  },
}

3、其他
实际参数过程中，可能还需要设置train_batch_size，gradient_accumulation_steps（梯度累计次数），optimizer（优化器选择）

代码

https://gitee.com/a-ha-a/deep-learning-note/tree/master/DeepLearning-Summary/Computer-Vision/deepspeed

参考

1、https://arxiv.org/pdf/1910.02054
2、https://zhuanlan.zhihu.com/p/513571706
3、https://zhuanlan.zhihu.com/p/618865052
4、https://zhuanlan.zhihu.com/p/504957661
5、https://deepspeed.readthedocs.io/en/latest/initialize.html#
6、https://www.deepspeed.ai/docs/config-json/#batch-size-related-parameters
7、https://zhuanlan.zhihu.com/p/630734624