DeepSpeed是微软推出的大规模模型分布式训练的工具,主要实现了ZeRO并行训练算法。
原始文档链接:
一、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】,获取第一手更新文章。