devbox
程序质量控制师
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

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

作者：程序质量控制师 | 2024-02-01 11:49:44

赞

踩

deepspeed zero3

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 调用方法

  1. # 单卡的使用方法
  2. deepspeed --num_gpus=1 examples/pytorch/translation/run_translation.py ...
  3. # 单卡，并指定对应的GPU
  4. deepspeed --include localhost:1 examples/pytorch/translation/run_translation.py ...
  5. ​
  6. # 多GPU的使用方法1
  7. torch.distributed.run --nproc_per_node=2 your_program.py <normal cl args> --deepspeed ds_config.json
  8. # 多GPU的使用方法2
  9. deepspeed --num_gpus=2 your_program.py <normal cl args> --deepspeed ds_config.json
  10. ​
  11. # 多节点多卡方法1，需要在多个节点上手动启动
  12. 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
  13. # 多节点多卡方法2，需要创建一个 hostfile 文件，只需在一个节点上启动
  14. hostname1 slots=8
  15. hostname2 slots=8
  16. # 然后运行
  17. 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
  18. ​
  19. # 在SLURM上运行，略，参见原始文档
  20. # 在jupyter中运行，略，参见原始文档

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

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

2.3 传递参数

  1. TrainingArguments(..., deepspeed="/path/to/ds_config.json")
  2. ​
  3. # or
  4. ds_config_dict = dict(scheduler=scheduler_params, optimizer=optimizer_params)
  5. TrainingArguments(..., deepspeed=ds_config_dict)

三、ZeRO中的配置

ZeRO-2

  • 一个配置示例
  1. {
  2.     "fp16": {
  3.         "enabled": "auto",
  4.         "loss_scale": 0,
  5.         "loss_scale_window": 1000,
  6.         "initial_scale_power": 16,
  7.         "hysteresis": 2,
  8.         "min_loss_scale": 1
  9.     },
  10.  
  11.     "optimizer": {
  12.         "type": "AdamW",
  13.         "params": {
  14.             "lr": "auto",
  15.             "betas": "auto",
  16.             "eps": "auto",
  17.             "weight_decay": "auto"
  18.         }
  19.     },
  20.  
  21.     "scheduler": {
  22.         "type": "WarmupLR",
  23.         "params": {
  24.             "warmup_min_lr": "auto",
  25.             "warmup_max_lr": "auto",
  26.             "warmup_num_steps": "auto"
  27.         }
  28.     },
  29.  
  30.     "zero_optimization": {
  31.         "stage": 2,
  32.         "offload_optimizer": {
  33.             "device": "cpu",
  34.             "pin_memory": true
  35.         },
  36.         "allgather_partitions": true,
  37.         "allgather_bucket_size": 2e8,
  38.         "overlap_comm": true,
  39.         "reduce_scatter": true,
  40.         "reduce_bucket_size": 2e8,
  41.         "contiguous_gradients": true
  42.     },
  43.  
  44.     "gradient_accumulation_steps": "auto",
  45.     "gradient_clipping": "auto",
  46.     "steps_per_print": 2000,
  47.     "train_batch_size": "auto",
  48.     "train_micro_batch_size_per_gpu": "auto",
  49.     "wall_clock_breakdown": false
  50. }
  • 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

  • 一个配置示例
  1. {
  2.     "fp16": {
  3.         "enabled": "auto",
  4.         "loss_scale": 0,
  5.         "loss_scale_window": 1000,
  6.         "initial_scale_power": 16,
  7.         "hysteresis": 2,
  8.         "min_loss_scale": 1
  9.     },
  10.  
  11.     "optimizer": {
  12.         "type": "AdamW",
  13.         "params": {
  14.             "lr": "auto",
  15.             "betas": "auto",
  16.             "eps": "auto",
  17.             "weight_decay": "auto"
  18.         }
  19.     },
  20.  
  21.     "scheduler": {
  22.         "type": "WarmupLR",
  23.         "params": {
  24.             "warmup_min_lr": "auto",
  25.             "warmup_max_lr": "auto",
  26.             "warmup_num_steps": "auto"
  27.         }
  28.     },
  29.  
  30.     "zero_optimization": {
  31.         "stage": 3,
  32.         "offload_optimizer": {
  33.             "device": "cpu",
  34.             "pin_memory": true
  35.         },
  36.         "offload_param": {
  37.             "device": "cpu",
  38.             "pin_memory": true
  39.         },
  40.         "overlap_comm": true,
  41.         "contiguous_gradients": true,
  42.         "sub_group_size": 1e9,
  43.         "reduce_bucket_size": "auto",
  44.         "stage3_prefetch_bucket_size": "auto",
  45.         "stage3_param_persistence_threshold": "auto",
  46.         "stage3_max_live_parameters": 1e9,
  47.         "stage3_max_reuse_distance": 1e9,
  48.         "stage3_gather_16bit_weights_on_model_save": true
  49.     },
  50.  
  51.     "gradient_accumulation_steps": "auto",
  52.     "gradient_clipping": "auto",
  53.     "steps_per_print": 2000,
  54.     "train_batch_size": "auto",
  55.     "train_micro_batch_size_per_gpu": "auto",
  56.     "wall_clock_breakdown": false
  57. }
  • 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来使用。
  1. {
  2.     "zero_optimization": {
  3.         "stage": 0
  4.     }
  5. }

ZeRO-stage-1

  • 只对优化器参数进行分片，可以加速一丢丢
  1. {
  2.     "zero_optimization": {
  3.         "stage": 1
  4.     }
  5. }

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 以后的版本，还需要添加：
  1. {
  2.     "zero_force_ds_cpu_optimizer": false
  3. }

5.2 调度器

  • DeepSpeed 支持 LRRangeTest、OneCycle、WarmupLR 和 WarmupDecayLR 学习率调度器。
  • Transformers和DeepSpeed中调度器的overlap
  1. WarmupLR 使用 --lr_scheduler_type constant_with_warmup
  2. 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 作为累积数据类型
  1. {
  2.     "communication_data_type": "fp32"
  3. }

apex

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

七、获取模型参数

  • deepspeed会在优化器参数中存储模型的主参数，存储在global_step*/*optim_states.pt 文件中，数据类型为fp32。因此，想要从checkpoint中恢复训练，则保持默认即可
  • 如果模型是在ZeRO-2模式下保存的，模型参数会以fp16的形式存储在pytorch_model.bin中
  • 如果模型是在ZeRO-3模式下保存的，需要如下所示设置参数，否则pytorch_model.bin将不会被创建
  1. {
  2.   "zero_optimization": {
  3.          "stage3_gather_16bit_weights_on_model_save": true
  4.     }
  5. }
  • 在线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。

  1. python -c 'from transformers import AutoModel; \
  2. from deepspeed.runtime.zero.stage3 import estimate_zero3_model_states_mem_needs_all_live; \
  3. model = AutoModel.from_pretrained("bigscience/T0_3B"); \
  4. estimate_zero3_model_states_mem_needs_all_live(model, num_gpus_per_node=2, num_nodes=1)'
  5. [...]
  6. Estimated memory needed for params, optim states and gradients for a:
  7. HW: Setup with 1 node, 2 GPUs per node.
  8. SW: Model with 2783M total params, 65M largest layer params.
  9.   per CPU  |  per GPU |   Options
  10.  70.00GB |   0.25GB | offload_param=cpu , offload_optimizer=cpu , zero_init=1
  11.  70.00GB |   0.25GB | offload_param=cpu , offload_optimizer=cpu , zero_init=0
  12.  62.23GB |   2.84GB | offload_param=none, offload_optimizer=cpu , zero_init=1
  13.  62.23GB |   2.84GB | offload_param=none, offload_optimizer=cpu , zero_init=0
  14.  0.74GB |  23.58GB | offload_param=none, offload_optimizer=none, zero_init=1
  15.  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：简单快捷上手百亿参数模型微调 - 知乎

声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/article/detail/53662
推荐阅读
相关标签

Copyright © 2003-2013 www.wpsshop.cn 版权所有，并保留所有权利。

  

闽ICP备14008679号