1. AI技术聚合首页
  2. IT

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

青葱年少 • • IT • 阅读 55

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这一种情况。
CombosHF SchedulerDS Scheduler
HF OptimizerYesYes
DS OptimizerNoYes

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：简单快捷上手百亿参数模型微调 – 知乎

文章出处登录后可见！

已经登录？立即刷新

共计人评分，平均分

到目前为止还没有投票！成为第一位评论此文章。

赞 (0)
青葱年少的头像青葱年少普通用户
0
上一篇 2023年12月27日
下一篇 2023年12月27日

相关推荐