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

这篇具有很好参考价值的文章主要介绍了LLM-分布式训练工具（一）：DeepSpeed【微软】【大模型分布式训练工具，实现ZeRO并行训练算法】【zero3配置将模型参数切分后分配到不同的显卡中，突破单张显卡容量不足以加载模型参数的限制】。希望对大家有所帮助。如果存在错误或未考虑完全的地方，请大家不吝赐教，您也可以点击"举报违法"按钮提交疑问。

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

原始文档链接：

DeepSpeed

一、DeepSpeed目前支持的功能

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

二、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)

三、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

四、调参步骤

将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%的性能）

一些其他建议

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

五、优化器和调度器

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

Combos	HF Scheduler	DS Scheduler
HF Optimizer	Yes	Yes
DS Optimizer	No	Yes

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 Chat，人人可快速训练百亿、千亿级ChatGPT大模型

【自然语言处理】【大模型】DeepSpeed+Transformers：简单快捷上手百亿参数模型微调 - 知乎文章来源地址https://www.toymoban.com/news/detail-589181.html

到了这里，关于LLM-分布式训练工具（一）：DeepSpeed【微软】【大模型分布式训练工具，实现ZeRO并行训练算法】【zero3配置将模型参数切分后分配到不同的显卡中，突破单张显卡容量不足以加载模型参数的限制】的文章就介绍完了。如果您还想了解更多内容，请在右上角搜索TOY模板网以前的文章或继续浏览下面的相关文章，希望大家以后多多支持TOY模板网！