DeepSpeed是微软推出的大规模模型分布式训练的工具, 主要实现了ZeRO并行训练算法。
DeepSpeed是一个开源深度学习训练优化库, 其中包含的一个新的显存优化技术—— ZeRO(零冗余优化器), 通过扩大规模, 提升速度, 控制成本, 提升可用性, 极大地推进了大模型训练能力。DeepSpeed 已经帮助研究人员开发了图灵自然语言生成模型( Turing-NLG), 其在发表时为世界上最大的语言模型(拥有 170 亿参数), 并有着最佳的精度。我们在 5 月份发布了 ZeRO-2——支持有着 2000 亿参数的模型训练, 与最新技术相比, 训练速度可达 10 倍——以及一系列计算、IO 和收敛优化功能, 从而助力最快速的 BERT 训练。自那时起, 我们持续高速地进行创新, 不断突破深度学习模型训练的速度和规模的边界。

1. DeepSpeed 概述

DeepSpeed的核心就在于, GPU显存不够, CPU内存来凑。

1.1 目前支持的功能

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

1.2 安装

pip install deepspeed

# windows环境下, 没有async_io库, 需要设置环境变量屏蔽async_io库
set DS_BUILD_AIO=0
set DS_BUILD_SPARSE_ATTN=0
python setup.py bdist_wheel

# 默认使用通信后端mpi（conda安装比较方便, 也可以是用nccl）
conda install mpi4y mpich -c conda-forges

1.3 DeepSpeed 与 Transformer

HuggingFace提供了对DeepSpeed的友好集成, DeepSpeed使用所需要的很多参数, 都可以由Transformer的Trainer来自动指定。可以说, DeepSpeed在HuggingFace Transformer上的使用, 会更为便捷(当然, DeepSpeed也可以独立使用, 并不依赖于Transformer)。

pip install transformers
# pip install transformers[deepspeed]

由于DeepSpeed目前已经被集成到了HuggingFace Transformer框架。而DeepSpeed的很多参数, 和Transformer的Trainer参数设置是一模一样的, 例如 “optimizer”, “scheduler”。因此, 官方推荐将很多常用的模型训练参数, 设置为auto, 在使用Trainer进行训练的时候, 这些值都会自动更新为Trainer中的设置, 或者帮你自动计算。

当然, 你也可以自己设置, 但一定要确保和Trainer中的设置一样。

1.4 基础组件

分布式训练需要掌握分布式环境中的基础配置，包括节点变化、全局进程编号、局部进程编号、全局总进程数、主节点等。这些组件都跟分布式训练紧密相关，同时组件之间也有非常大的联系，例如通信联系等。

1.5 通信策略

既然是分布式训练，那机器之间必须要保持通信，这样才可以传输模型参数，梯度参数等信息。
DeepSpeed提供了mpi、gioo、nccl等通信策略

我们在使用DeepSpeed进行分布式训练的时候，可以根据自身的情况选择合适的通信库，通常情况下，如果是GPU进行分布式训练，可以选择nccl。

2. 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?

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

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)

2.4 使用非共享的文件系统

DeepSpeed默认在多节点环境下使用一个共享的存储。如果不使用共享存储并且每个节点只能看到它自己的本地文件系统, 你需要在deepspeed的配置文件中加入checkpoint部分:

{
  "checkpoint": {
    "use_node_local_storage": true
  }
}

你也可以使用Trainer中的--save_on_each_node参数作为替代, 这样做上面的配置会被自动的添加。

3. ZeRO中的配置

Zero Redundancy Optimizer(ZeRO)是DeepSpeed的workhorse. 用户可以提供不同的ZeRO config文件, 来实现DeepSpeed的不同功能特性。ZeRO的实现方法, 就是把参数占用, 逻辑上分成三种类型。将这些类型的参数划分:

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

3.1 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数量增加时, 会有更好的性能优势。

3.2 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 配置参数

3.3 ZeRO-stage-0

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

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

3.4 ZeRO-stage-1

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

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

3.5 NVMe Support

ZeRO-Infinity可以看成是stage-3的进阶版本, 需要依赖于NVMe的支持。他可以offload所有模型参数状态到CPU以及NVMe上。得益于NMVe协议, 除了使用CPU内存之外, ZeRO可以额外利用SSD(固态), 从而极大地节约了memory开销, 加速了通信速度。

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

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

原则是: 能直接多卡训练, 就不要用ZeRO; 能用ZeRO-2就不要用ZeRO-3.
从左到右, 越来越慢:

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

4. 调参步骤

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

4.1 一些其他建议

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

5. 优化器和调度器

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

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

6. 训练精度

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

6.1 自动混合精度

6.1.1 fp16

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

6.1.2 bf16

使用–bf16 or –bf16_full_eval 命令行参数时启用此模式

6.2 NCCL

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

{
    "communication_data_type": "fp32"
}

6.3 apex

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

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

7. 获取模型参数

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

7.1 ZeRO-3 and Infinity Nuances

构造超大模型
搜集参数

8. ZeRO inference

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

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

9. 估算需要的显存

由于DeepSpeed会通过占用CPU内存来减缓GPU的开销, 当系统CPU不够的时候, DeepSpeed进程就会自动被系统停止, 造成没有任何报错, DeepSpeed无法启动的现象。建议先用上文介绍的estimation估计一下CPU内存占用, 然后用free -h查看一下机器的CPU内存空余量, 来判断能否使用DeepSpeed。

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

from transformers import AutoModel
from deepspeed.runtime.zero.stage3 import estimate_zero3_model_states_mem_needs_all_live

## specify the model you want to train on your device
model = AutoModel.from_pretrained("t5-large") 
## estimate the memory cost (both CPU and GPU)
estimate_zero3_model_states_mem_needs_all_live(model, num_gpus_per_node=1, 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

10. 可能遇到的问题

• 启动时, 进程被杀死, 并且没有打印出traceback: CPU显存不够
  • 由于DeepSpeed会通过占用CPU内存来减缓GPU的开销, 当系统CPU不够的时候, DeepSpeed进程就会自动被系统停止, 造成没有任何报错, DeepSpeed无法启动的现象。建议先用上文介绍的estimation估计一下CPU内存占用, 然后用free -h查看一下机器的CPU内存空余量, 来判断能否使用DeepSpeed。
• loss是NaN: 训练时用的是bf16, 使用时是fp16。常常发生于google在TPU上train的模型, 如T5。此时需要使用fp32或者bf16。
  • 有可能因为训练精度问题, 出现loss为NAN的情况。
• 使用DeepSpeed stage2之后, 就不能灵活地更改optimizer了
• DeepSpeed进程目前还不支持在Vscode进行debug, 因为缺少相应的VScode编译插件的支持

参考:
DeepSpeed Docs
deepspeed入门教程
DeepSpeed Integration
DeepSpeed: Extreme-scale model training for everyone