DeepSpeed是微软推出的大规模模型分布式训练的工具,主要实现了ZeRO并行训练算法。
本文是huggingface的DeepSpeed文档的笔记,做查询和备忘,初次学习建议结合原始文档食用。原始文档链接: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的使用
# 单卡的使用方法
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,降低对显存的需求
- 提供了对显存的管理,减少显存中的碎片
参数解释:
ZeRO-2
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
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
配置参数
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
如何选择stage
python -c 'import transformers; \
from deepspeed.runtime.zero.stage3 import estimate_zero3_model_states_mem_needs_all_live; \
model = transformers.LlamaForCausalLM.from_pretrained("model_path"); \
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 6738M total params, 131M largest layer params.
per CPU | per GPU | Options
169.44GB | 0.49GB | offload_param=cpu , offload_optimizer=cpu , zero_init=1
169.44GB | 0.49GB | offload_param=cpu , offload_optimizer=cpu , zero_init=0
150.62GB | 6.76GB | offload_param=none, offload_optimizer=cpu , zero_init=1
150.62GB | 6.76GB | offload_param=none, offload_optimizer=cpu , zero_init=0
1.46GB | 56.97GB | offload_param=none, offload_optimizer=none, zero_init=1
75.31GB | 56.97GB | offload_param=none, offload_optimizer=none, zero_init=0