LLM-分布式训练工具（一）：DeepSpeed【微软】【大模型分布式训练工具，实现ZeRO并行训练算法】【zero3配置将模型参数切分后分配到不同的显卡中，突破单张显卡容量不足以加载模型参数的限制】

偏执的太偏执、 2023-10-12 13:44 19阅读 0赞

DeepSpeed是微软推出的大规模模型分布式训练的工具，主要实现了ZeRO并行训练算法。

原始文档链接：

[DeepSpeed][]

### 一、DeepSpeed目前支持的功能 ###

1.  Optimizer state partitioning (ZeRO stage 1)
2.  Gradient partitioning (ZeRO stage 2)
3.  Parameter partitioning (ZeRO stage 3)
4.  Custom mixed precision training handling
5.  A range of fast CUDA-extension-based optimizers
6.  ZeRO-Offload to CPU and NVMe

### 二、DeepSpeed的使用 ###

#### 2.1 调用方法 ####

# 单卡的使用方法
    deepspeed --num_gpus=1 examples/pytorch/translation/run_translation.py ...
    # 单卡，并指定对应的GPU
    deepspeed --include localhost:1 examples/pytorch/translation/run_translation.py ...
    
    # 多GPU的使用方法1
    torch.distributed.run --nproc_per_node=2 your_program.py <normal cl args> --deepspeed ds_config.json
    # 多GPU的使用方法2
    deepspeed --num_gpus=2 your_program.py <normal cl args> --deepspeed ds_config.json
    
    # 多节点多卡方法1，需要在多个节点上手动启动
    python -m torch.distributed.run --nproc_per_node=8 --nnode=2 --node_rank=0 --master_addr=hostname1 --master_port=9901 your_program.py <normal cl args> --deepspeed ds_config.json
    # 多节点多卡方法2，需要创建一个 hostfile 文件，只需在一个节点上启动
    hostname1 slots=8
    hostname2 slots=8
    # 然后运行
    deepspeed --num_gpus 8 --num_nodes 2 --hostfile hostfile --master_addr hostname1 --master_port=9901 your_program.py <normal cl args> --deepspeed ds_config.json
    
    # 在SLURM上运行，略，参见原始文档
    # 在jupyter中运行，略，参见原始文档

#### 2.2 为什么单卡的情况，也可以使用deepspeed？ ####

1.  使用ZeRO-offload，将部分数据offload到CPU，降低对显存的需求
2.  提供了对显存的管理，减少显存中的碎片

#### 2.3 传递参数 ####

TrainingArguments(..., deepspeed="/path/to/ds_config.json")
    
    # or
    ds_config_dict = dict(scheduler=scheduler_params, optimizer=optimizer_params)
    TrainingArguments(..., deepspeed=ds_config_dict)

### 三、ZeRO中的配置 ###

#### ZeRO-2 ####

*  一个配置示例

{
        "fp16": {
            "enabled": "auto",
            "loss_scale": 0,
            "loss_scale_window": 1000,
            "initial_scale_power": 16,
            "hysteresis": 2,
            "min_loss_scale": 1
        },
    
        "optimizer": {
            "type": "AdamW",
            "params": {
                "lr": "auto",
                "betas": "auto",
                "eps": "auto",
                "weight_decay": "auto"
            }
        },
    
        "scheduler": {
            "type": "WarmupLR",
            "params": {
                "warmup_min_lr": "auto",
                "warmup_max_lr": "auto",
                "warmup_num_steps": "auto"
            }
        },
    
        "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
        },
    
        "gradient_accumulation_steps": "auto",
        "gradient_clipping": "auto",
        "steps_per_print": 2000,
        "train_batch_size": "auto",
        "train_micro_batch_size_per_gpu": "auto",
        "wall_clock_breakdown": false
    }

*  `overlap_comm`：控制是否使用通信与计算的重叠。当设置为`True`时，DeepSpeed将在梯度计算时尝试并行执行梯度通信。可以有效地减少通信时间，从而加速整个训练过程。
 *  `allgather_bucket_size`：用于控制Allgather操作的分桶大小。Allgather操作是指在分布式训练中，每个进程收集其他所有进程的张量，并将这些张量按顺序拼接起来。通过将张量划分为较小的桶（buckets），可以在通信过程中更高效地传输数据。`allgather_bucket_size`值越大，每个桶的大小越大，通信操作可能会变得更快，但也需要更多的内存来存储中间结果。合适的桶大小要根据实际情况调整。
 *  `reduce_bucket_size`：类似于`allgather_bucket_size`，用于控制Allreduce操作的分桶大小。Allreduce操作是将所有进程的某个张量进行规约（例如求和），并将结果广播回所有进程。通过将张量划分为较小的桶，可以更高效地传输数据。`reduce_bucket_size`值越大，每个桶的大小越大，通信操作可能会变得更快，但同时也需要更多的内存来存储中间结果。合适的桶大小需要根据实际情况进行调整。
 *  `overlap_comm`使用的是`allgather_bucket_size`和`reduce_bucket_size`值的4.5倍。如果它们被设置为5e8，需要9GB显存（5e8 x 2Bytes x 2 x 4.5）。如果内存大小是8GB或更小，需要将这些参数减少到约2e8，从而避免OOM，这需要3.6GB显存。如果在大容量GPU上也出现OOM，也需要做同样的调整。
 *  在deepspeed==0.4.4中新增了 `round_robin_gradients` 选项，可以并行化CPU的offload。当梯度累积的步数增加，或者GPU数量增加时，会有更好的性能优势。

#### ZeRO-3 ####

*  一个配置示例

{
        "fp16": {
            "enabled": "auto",
            "loss_scale": 0,
            "loss_scale_window": 1000,
            "initial_scale_power": 16,
            "hysteresis": 2,
            "min_loss_scale": 1
        },
    
        "optimizer": {
            "type": "AdamW",
            "params": {
                "lr": "auto",
                "betas": "auto",
                "eps": "auto",
                "weight_decay": "auto"
            }
        },
    
        "scheduler": {
            "type": "WarmupLR",
            "params": {
                "warmup_min_lr": "auto",
                "warmup_max_lr": "auto",
                "warmup_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,
            "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
        },
    
        "gradient_accumulation_steps": "auto",
        "gradient_clipping": "auto",
        "steps_per_print": 2000,
        "train_batch_size": "auto",
        "train_micro_batch_size_per_gpu": "auto",
        "wall_clock_breakdown": false
    }

*  `stage3_max_live_parameters` 是保留在 GPU 上的完整参数数量的上限。
 *  `stage3_max_reuse_distance` 是指将来何时再次使用参数的指标，从而决定是丢弃参数还是保留参数。 如果一个参数在不久的将来要再次使用（小于 stage3\_max\_reuse\_distance），可以保留以减少通信开销。 使用activation checkpointing时，这一点非常有用。
 *  如果遇到 OOM，可以减少 `stage3_max_live_parameters` 和 `stage3_max_reuse_distance`。 除非正在使用activation checkpointing，否则它们对性能的影响应该很小。 1e9 会消耗 ~2GB。 内存由 `stage3_max_live_parameters` 和 `stage3_max_reuse_distance` 共享，所以不是相加的，一共 2GB。
 *  `stage3_gather_16bit_weights_on_model_save` 在保存模型时启用模型 fp16 权重合并。 对大型模型和多GPU，在内存和速度方面都是一项昂贵的操作。 如果打算恢复训练，目前需要使用它。 未来的更新将消除此限制。
 *  `sub_group_size` 控制在optimizer steps中更新参数的粒度。 参数被分组到 `sub_group_size` 的桶中，每个桶一次更新一个。 当与 ZeRO-Infinity 中的 NVMe offload一起使用时，`sub_group_size` 控制模型状态在optimizer steps期间从 NVMe 移入和移出 CPU 内存的粒度。 防止超大模型耗尽 CPU 内存。不使用NVMe offload时，使其保持默认值。出现OOM时，减小`sub_group_size`。当优化器迭代很慢时，可以增大`sub_group_size` 。
 *  ZeRO-3 中未使用 `allgather_partitions`、`allgather_bucket_size` 和 `reduce_scatter` 配置参数

#### ZeRO-stage-0 ####

*  stage 0会禁用所有的分片，然后把DeepSpeed当作时DDP来使用。

{
        "zero_optimization": {
            "stage": 0
        }
    }

#### ZeRO-stage-1 ####

*  只对优化器参数进行分片，可以加速一丢丢

{
        "zero_optimization": {
            "stage": 1
        }
    }

#### NVMe Support ####

*  ZeRO-Infinity 需要使用 ZeRO-3
 *  ZeRO-3 会比 ZeRO-2 慢很多。使用以下策略，可以使得ZeRO-3 的速度更接近ZeRO-2
    
     *  将`stage3_param_persistence_threshold`参数设置的很大，比如`6 * hidden_size * hidden_size`
     *  将`offload_params`参数关闭（可以极大改善性能）

#### 如何选择不同的Zero stage和offload ####

*  从左到右，越来越慢  
    Stage 0 (DDP) > Stage 1 > Stage 2 > Stage 2 + offload > Stage 3 > Stage 3 + offloads
 *  从左到右，所需GPU显存越来越少  
    Stage 0 (DDP) < Stage 1 < Stage 2 < Stage 2 + offload < Stage 3 < Stage 3 + offloads

#### 四、调参步骤 ####

1.  将`batch_size`设置为1，通过梯度累积实现任意的有效`batch_size`
2.  如果OOM则，设置`--gradient_checkpointing 1` (HF Trainer)，或者 `model.gradient_checkpointing_enable()`
3.  如果OOM则，尝试ZeRO stage 2
4.  如果OOM则，尝试ZeRO stage 2 + `offload_optimizer`
5.  如果OOM则，尝试ZeRO stage 3
6.  如果OOM则，尝试offload\_param到CPU
7.  如果OOM则，尝试offload\_optimizer到CPU
8.  如果OOM则，尝试降低一些默认参数。比如使用generate时，减小beam search的搜索范围
9.  如果OOM则，使用混合精度训练，在Ampere的GPU上使用bf16，在旧版本GPU上使用fp16
10. 如果仍然OOM，则使用ZeRO-Infinity ，使用offload\_param和offload\_optimizer到NVME
11. 一旦使用batch\_size=1时，没有导致OOM，测量此时的有效吞吐量，然后尽可能增大batch\_size
12. 开始优化参数，可以关闭offload参数，或者降低ZeRO stage，然后调整batch\_size，然后继续测量吞吐量，直到性能比较满意（调参可以增加66%的性能）

#### 一些其他建议 ####

*   *  如果训模型from scratch，hidden size最好可以被16整除
     *  batch size最好可以被2整除

### 五、优化器和调度器 ###

*  当不使用`offload_optimizer` 时，可以按照下表，混合使用HF和DS的优化器和迭代器，除了HF Scheduler和DS Optimizer这一种情况。

<table> 
 <tbody> 
  <tr> 
   <th>Combos</th> 
   <th>HF Scheduler</th> 
   <th>DS Scheduler</th> 
  </tr> 
  <tr> 
   <td>HF Optimizer</td> 
   <td>Yes</td> 
   <td>Yes</td> 
  </tr> 
  <tr> 
   <td>DS Optimizer</td> 
   <td>No</td> 
   <td>Yes</td> 
  </tr> 
 </tbody> 
</table>

#### 5.1 优化器 ####

*  启用 offload\_optimizer 时可以使用非 DeepSpeed 的优化器，只要它同时具有 CPU 和 GPU 的实现（LAMB 除外）。
 *  DeepSpeed 的主要优化器是 Adam、AdamW、OneBitAdam 和 Lamb。 这些已通过 ZeRO 进行了彻底测试，建议使用。
 *  如果没有在配置文件中配置优化器参数，Trainer 将自动将其设置为 AdamW，并将使用命令行参数的默认值：--learning\_rate、--adam\_beta1、--adam\_beta2、 --adam\_epsilon 和 --weight\_decay。
 *  与 AdamW 类似，可以配置其他官方支持的优化器。 请记住，它们可能具有不同的配置值。 例如 对于 Adam，需要将 weight\_decay 设置为 0.01 左右。
 *  此外，offload在与 Deepspeed 的 CPU Adam 优化器一起使用时效果最佳。 如果想对offload使用不同的优化器，deepspeed==0.8.3 以后的版本，还需要添加：

{
        "zero_force_ds_cpu_optimizer": false
    }

#### 5.2 调度器 ####

*  DeepSpeed 支持 LRRangeTest、OneCycle、WarmupLR 和 WarmupDecayLR 学习率调度器。
 *  Transformers和DeepSpeed中调度器的overlap

WarmupLR 使用 --lr_scheduler_type constant_with_warmup
    WarmupDecayLR 使用 --lr_scheduler_type linear

### 六、训练精度 ###

*  由于 fp16 混合精度大大减少了内存需求，并可以实现更快的速度，因此只有在在此训练模式下表现不佳时，才考虑不使用混合精度训练。 通常，当模型未在 fp16 混合精度中进行预训练时，会出现这种情况（例如，使用 bf16 预训练的模型）。 这样的模型可能会溢出，导致loss为NaN。 如果是这种情况，使用完整的 fp32 模式。
 *  如果是基于 Ampere 架构的 GPU，pytorch 1.7 及更高版本将自动切换为使用更高效的 tf32 格式进行某些操作，但结果仍将采用 fp32。
 *  使用 Trainer，可以使用 --tf32 启用它，或使用 --tf32 0 或 --no\_tf32 禁用它。 PyTorch 默认值是使用tf32。

#### 自动混合精度 ####

*  fp16
    
     *  可以使用 pytorch-like AMP 方式或者 apex-like 方式
     *  使用 `--fp16``--fp16_backend amp` 或 `--fp16_full_eval` 命令行参数时启用此模式
 *  bf16
    
     *  使用`--bf16` or `--bf16_full_eval` 命令行参数时启用此模式

#### NCCL ####

*  通讯会采用一种单独的数据类型
 *  默认情况下，半精度训练使用 fp16 作为reduction操作的默认值
 *  可以增加一个小的开销并确保reduction将使用 fp32 作为累积数据类型

{
        "communication_data_type": "fp32"
    }

#### apex ####

*  Apex 是一个在 PyTorch 深度学习框架下用于加速训练和提高性能的库。Apex 提供了混合精度训练、分布式训练和内存优化等功能，帮助用户提高训练速度、扩展训练规模以及优化 GPU 资源利用率。
 *  使用`--fp16`、 `--fp16_backend apex`、 `--fp16_opt_level 01` 命令行参数时启用此模式

"amp": {
         "enabled": "auto",
         "opt_level": "auto"
    }

### 七、获取模型参数 ###

*  deepspeed会在优化器参数中存储模型的主参数，存储在`global_step*/*optim_states.pt` 文件中，数据类型为fp32。因此，想要从checkpoint中恢复训练，则保持默认即可
 *  如果模型是在ZeRO-2模式下保存的，模型参数会以fp16的形式存储在`pytorch_model.bin`中
 *  如果模型是在ZeRO-3模式下保存的，需要如下所示设置参数，否则pytorch\_model.bin将不会被创建

{
      "zero_optimization": {
             "stage3_gather_16bit_weights_on_model_save": true
        }
    }

*  在线fp32权重恢复（需要很多的RAM）略
 *  离线获取fp32权重

python zero_to_fp32.py . pytorch_model.bin

#### ZeRO-3 and Infinity Nuances ####

*  构造超大模型（略）
 *  搜集参数（略）

### 八、ZeRO inference ###

只有ZeRO-3是有意义的，因为可以将参数分片：

deepspeed --num_gpus=2 your_program.py <normal cl args> --do_eval --deepspeed ds_config.json

### 九、估算需要的显存 ###

可以通过下面的代码，先估算不同配置需要的显存数量，从而决定开始尝试的ZeRO stage。

python -c 'from transformers import AutoModel; \
    from deepspeed.runtime.zero.stage3 import estimate_zero3_model_states_mem_needs_all_live; \
    model = AutoModel.from_pretrained("bigscience/T0_3B"); \
    estimate_zero3_model_states_mem_needs_all_live(model, num_gpus_per_node=2, num_nodes=1)'
    [...]
    Estimated memory needed for params, optim states and gradients for a:
    HW: Setup with 1 node, 2 GPUs per node.
    SW: Model with 2783M total params, 65M largest layer params.
      per CPU  |  per GPU |   Options
     70.00GB |   0.25GB | offload_param=cpu , offload_optimizer=cpu , zero_init=1
     70.00GB |   0.25GB | offload_param=cpu , offload_optimizer=cpu , zero_init=0
     62.23GB |   2.84GB | offload_param=none, offload_optimizer=cpu , zero_init=1
     62.23GB |   2.84GB | offload_param=none, offload_optimizer=cpu , zero_init=0
     0.74GB |  23.58GB | offload_param=none, offload_optimizer=none, zero_init=1
     31.11GB |  23.58GB | offload_param=none, offload_optimizer=none, zero_init=0

### 十、可能遇到的问题 ###

*  启动时，进程被杀死，并且没有打印出traceback：CPU显存不够
 *  loss是NaN：训练时用的是bf16，使用时是fp16。常常发生于google在TPU上train的模型，如T5。此时需要使用fp32或者bf16。

如果觉得本文有帮助，欢迎点赞收藏评论转发。也欢迎关注同名公众号【JOYWIN】，获取第一手更新文章。

--------------------

[deepspeed入门教程 - 知乎][deepspeed_ -]

[微软DeepSpeed Chat，人人可快速训练百亿、千亿级ChatGPT大模型][DeepSpeed Chat_ChatGPT]

[【自然语言处理】【大模型】DeepSpeed+Transformers：简单快捷上手百亿参数模型微调 - 知乎][DeepSpeed_Transformers_ -]

[DeepSpeed]: https://link.zhihu.com/?target=https%3A//huggingface.co/docs/transformers/main/main_classes/deepspeed
[deepspeed_ -]: https://zhuanlan.zhihu.com/p/630734624
[DeepSpeed Chat_ChatGPT]: https://baijiahao.baidu.com/s?id=1763048327731793454&wfr=spider&for=pc
[DeepSpeed_Transformers_ -]: https://zhuanlan.zhihu.com/p/615235322