Deepspeed应用中涉及的知识点和tricks本文记录了我在学习deepspeed应用中涉及到的相关知识点。包括c

本文记录了我在学习deepspeed应用中涉及到的相关知识点。

deepspeed提前预估显存占用

zhuanlan.zhihu.com/p/630734624

from transformers import AutoModel
from deepspeed.runtime.zero.stage3 import estimate_zero3_model_states_mem_needs_all_live

## specify the model you want to train on your device
model_path = "/lpai/volumes/zeus-perception/weidafeng/pretrained_models/huggingface_converted/opt-350m"
model = AutoModel.from_pretrained(model_path) 

## estimate the memory cost (both CPU and GPU)
estimate_zero3_model_states_mem_needs_all_live(model, num_gpus_per_node= 1 , num_nodes= 1 )

bash脚本：

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

deepspeed保存模型参数

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

注：deepspeed的三个stage，分别优化的是：

Optimizer status 优化器状态
Gradient 梯度
Parameters 模型参数

所以只有stage3的参数是分散的

Deepspeed config参数详解

zero2

"zero_optimization": {
        "stage": 2,
        "allgather_partitions": true,
        "allgather_bucket_size": 3e8,
        "overlap_comm": true,
        "reduce_scatter": true,
        "reduce_bucket_size": 3e8,
        "contiguous_gradients": true
    }

allgather_partitions：在每个步骤结束时，从所有GPU中选择使用allgather集体操作或一系列广播集体操作之间的方式，以收集更新后的参数。 (默认值：true)
allgather_bucket_size：用于调节Allgather操作的分桶大小。将张量分成较小的桶有助于在通信过程中更高效地传输数据。较大的allgather_bucket_size值会导致每个桶的尺寸增大，可能加速通信操作，但也需要更多内存来存储中间结果。选择合适的桶大小需要根据实际情况进行调整。(默认值：5e8)
overlap_comm：控制通信与计算是否交叠执行。当设置为True时，DeepSpeed将尝试在梯度计算期间并行进行梯度通信。这有效地缩短通信时间，从而加速整个训练过程。(默认值：false)
reduce_scatter：使用reduce或reduce scatter来替代allreduce以平均梯度。(默认值：true)
reduce_bucket_size：用于控制Allreduce操作的分桶大小。将张量分为较小的桶有助于数据在通信过程中的更高效传输。随着reduce_bucket_size值的增大，每个桶的尺寸也随之增大，这或许能加速通信操作，但同时也需要更多内存来存储中间结果。合适的桶大小应根据实际情况进行适当调整。(默认值：5e8)
contiguous_gradients：在梯度产生时将其复制到一个连续的缓冲区中。在反向传播过程中避免了内存碎片化问题。(默认值：true)

zero3

ZeRO-DP 会将模型拆分为 N 个 GPU，在某种程度上，这与张量并行性的水平切片相同，与垂直切片相反（垂直切片将整个层组放在不同的 GPU 上）

blog.csdn.net/qq_56591814…

"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": 1e6,
        "stage3_prefetch_bucket_size": 4e6,
        "stage3_param_persistence_threshold": 1e4,
        "stage3_max_live_parameters": 1e9,
        "stage3_max_reuse_distance": 1e9,
        "stage3_gather_16bit_weights_on_model_save": true
    },

sub_group_size：控制在优化器步骤中参数更新的粒度。参数被分组到大小为sub_group_size的桶中，每个桶依次进行一次更新。当与ZeRO-Infinity中的NVMe offload同时使用时，sub_group_size决定了在优化器步骤期间从NVMe迁移到CPU内存的模型状态的粒度。这有助于避免超大模型对CPU内存的过度占用。在不使用NVMe offload时，请保持其默认值。若遇到内存不足（OOM）情况，可以考虑减小sub_group_size。当优化器迭代较缓慢时，也可以考虑增大sub_group_size。(默认值：1e9)
stage3_prefetch_bucket_size：预取参数的固定缓冲区大小。较小的值使用的内存较少，但可能会因通信而增加停顿。(默认值：5e8)
stage3_max_live_parameters：保留在GPU上的完整参数数量的上限。(默认值：1e9)
stage3_max_reuse_distance：根据参数在未来何时再次使用的指标来决定是舍弃还是保留参数。如果一个参数在不久的将来会再次被使用（小于stage3_max_reuse_distance），则会保留该参数以减少通信开销。在遇到内存不足（OOM）的情况下，可以降低stage3_max_live_parameters和stage3_max_reuse_distance的值。(默认值：1e9)
stage3_gather_16bit_weights_on_model_save：在保存模型时启用模型FP16权重合并。对于大型模型和多 GPU 环境，这是一项在内存和速度方面代价较高的操作。(默认值：false)
ZeRO-3 中不使用 allgather_partitions、allgather_bucket_size 和 reduce_scatter 配置参数
（其他参数如grad_hooks、round_robin_gradients本文未提及）

混合精度

"fp16": {
    "enabled": true,
    "auto_cast": false,
    "loss_scale": 0,
    "initial_scale_power": 16,
    "loss_scale_window": 1000,
    "hysteresis": 2,
    "consecutive_hysteresis": false,
    "min_loss_scale": 1
}


"bf16": {
   "enabled": true
 }

auto_cast：是否将输入强制转换为fp16数据类型 (默认值：false)
loss_scale：表示FP16训练的损失缩放值。默认值0.0启用动态损失缩放，否则该值将用于静态固定损失缩放 (默认值：0.0)
initial_scale_power：表示初始动态损失比例值的功率，实际损失规模计算为 $2^{initial_scale_power$ (默认值：16)
loss_scale_window：代表动态损失缩放值上升/下降的窗口范围。(默认值：1000)
hysteresis：表示动态损耗缩放中的延迟偏移 (默认值：2)
consecutive_hysteresis：表示是否在达到不会溢出的迭代时重新填充滞后。(默认值：false)
min_loss_scale：表示最小动态损失比例值 (默认值：1)

注意：开启fp16后可能出现如上图所示overflow情况

BF16：配置以bfloat16浮点格式作为FP16的替代方式。bfloat16需要硬件支持（例如，NVIDIA A100）。使用bfloat16进行训练不需要损失缩放。(默认值：false)

Batch size 相关参数

"train_batch_size": "auto",
"gradient_accumulation_steps": "auto",
"train_micro_batch_size_per_gpu": "auto",

train_batch_size：有效的训练批量大小。这指的是每次模型更新所涉及的数据样本数量。train_batch_size是单个 GPU 在一次前向/后向传递中处理的批量大小，也称为train_micro_batch_size_per_gpu，以及梯度累积步骤（也称为gradient_accumulation_steps），还有GPU数量，这些因素共同决定。
- 如果同时提供train_micro_batch_size_per_gpu和gradient_accumulation_steps，可以忽略train_batch_size。(默认值：32)
gradient_accumulation_steps：在计算平均并应用梯度之前累积梯度的训练步骤数。这个功能有时候对于提高可扩展性非常有用，因为它减少了步骤之间梯度通信的频率。另一个影响是，可以在每个GPU上使用更大的批量大小进行训练。
- 如果同时提供train_batch_size和train_micro_batch_size_per_gpu，可以忽略gradient_accumulation_steps。(默认值：1)
train_micro_batch_size_per_gpu：单个GPU在一个步骤中处理的微批量大小（不进行梯度累积）。
- 如果同时提供train_batch_size和gradient_accumulation_steps，可以忽略train_micro_batch_size_per_gpu。(默认值：train_batch_size的值)

train_batch_size = train_micro_batch_size_per_gpu * gradient_accumulation * number of GPUs.（即训练批次的大小 = 每个 GPU 上的微批次大小 * 几个微批次 * 几个GPU）

关于auto

可以发现，配置示例中有参数被设置为auto。由于DeepSpeed目前已经被集成到了HuggingFace Transformer框架。而DeepSpeed的很多参数，和Transformer的Trainer参数设置是一模一样的。因此，官方推荐将很多常用的模型训练参数，设置为auto，在使用Trainer进行训练的时候，这些值都会自动更新为Trainer中的设置，或者帮你自动计算。

大多数情况下只需要注意DeepSpeed-specific参数(如:ZeRO,offload)。其他和Trainner重复的参数项，强烈建议设置成auto。

therefore set these values to auto and the Trainer will automatically assign the recommended values. But, of course, feel free to set these explicitly as well.

如何选择最佳性能的ZeRO阶段和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

因此，当想要在适合最少数量的GPU的情况下获得最快的执行速度时，可以遵循以下过程。我们从最快的方法开始，如果遇到GPU内存不足，然后转到下一个速度较慢但使用更少GPU内存的方法，依此类推。

具体方法

首先，将批量大小设置为1
1. 启用 --gradient_checkpointing 1（HF Trainer）或直接使用 model.gradient_checkpointing_enable() - 如果出现内存不足（OOM），则
尝试使用ZeRO stage2。如果出现内存不足（OOM），则
1. 尝试使用ZeRO stage2+ offload optimizer - 如果出现内存不足（OOM），则
切换到ZeRO stage3 - 如果出现内存不足（OOM），则
1. 将 offload_param 启用到CPU - 如果出现内存不足（OOM），则
2. 将 offload_optimizer 启用到CPU - 如果出现内存不足（OOM），则
使用混合精度进行训练而不是fp32
如果仍然出现内存不足（OOM），可以添加更多硬件或启用ZeRO-Infinity - 即将卸载 offload_param 和 offload_optimizer 切换到nvme。

torchrun启动和deepspeed启动的区别

deepspeed需要所有实例之间ssh免密，并且远程的话要使用pdsh启动进程；
torchrun的话是需要在所有实例上输入命令。

效果应该没区别

You can use a launcher of your choice here. You can continue using the pytorch launcher:

torch.distributed.run --nproc_per_node=2 your_program.py <normal cl args> \
        --deepspeed ds_config.json

or use the launcher provided by deepspeed:

deepspeed --num_gpus=2 your_program.py <normal cl args> \
--deepspeed ds_config.json

As you can see the arguments aren’t the same, but for most needs either of them works.

HF预训练模型启用deepspseed方法

如果你想使用Hugging Face Transformers库中的预训练模型来进行DeepSpeed训练，要确保在创建模型之前设置了DeepSpeed配置。

具体来说，需要先创建一个TrainingArguments对象（training_args），并在其中指定包括DeepSpeed配置（ds_config）在内的训练参数。然后再加载与训练模型，最后创建Trainer对象，加载模型传递训练参数，开始训练。整个顺序如下所示（官方的示例脚本就是这种代码顺序）：

from transformers import AutoModel, Trainer, TrainingArguments

# 包含Deepspeed参数的args，需要比模型更先定义
training_args = TrainingArguments(..., deepspeed=ds_config)

model = AutoModel.from_pretrained("t5-small")
trainer = Trainer(model=model, args=training_args, ...)

恢复fp32权重

FP32 权重：

虽然 fp16 权重在继续训练时很好，但是如果完成了微调模型并希望将其上传到 models hub 或传递给其他人，则最有可能需要获取 fp32 权重。这最好不要在训练期间执行，因为这是一个需要大量内存的过程，因此最好在训练完成后离线执行。但是，如果希望并且您有足够的空闲 CPU 内存，则可以在同一个训练脚本中执行。

在线恢复：

from transformers.trainer_utils import get_last_checkpoint
from deepspeed.utils.zero_to_fp32 import load_state_dict_from_zero_checkpoint

checkpoint_dir = get_last_checkpoint(trainer.args.output_dir)
fp32_model = load_state_dict_from_zero_checkpoint(trainer.model, checkpoint_dir)

离线恢复：

DeepSpeed创建了一个特殊的转换脚本zero_to_fp32.py，它将其放置在检查点文件夹的顶层。使用此脚本，您可以在任何时候提取权重。该脚本是独立的，您不再需要配置文件或Trainer即可执行提取。

假设您的检查点文件夹如下所示：

$ ls -l output_dir/checkpoint-1/
-rw-rw-r-- 1 stas stas 1.4K Mar 27 20:42 config.json
drwxrwxr-x 2 stas stas 4.0K Mar 25 19:52 global_step1/
-rw-rw-r-- 1 stas stas   12 Mar 27 13:16 latest
-rw-rw-r-- 1 stas stas 827K Mar 27 20:42 optimizer.pt
-rw-rw-r-- 1 stas stas 231M Mar 27 20:42 pytorch_model.bin
-rw-rw-r-- 1 stas stas  623 Mar 27 20:42 scheduler.pt
-rw-rw-r-- 1 stas stas 1.8K Mar 27 20:42 special_tokens_map.json
-rw-rw-r-- 1 stas stas 774K Mar 27 20:42 spiece.model
-rw-rw-r-- 1 stas stas 1.9K Mar 27 20:42 tokenizer_config.json
-rw-rw-r-- 1 stas stas  339 Mar 27 20:42 trainer_state.json
-rw-rw-r-- 1 stas stas 2.3K Mar 27 20:42 training_args.bin
-rwxrw-r-- 1 stas stas 5.5K Mar 27 13:16 zero_to_fp32.py*

在此示例中，只有一个DeepSpeed检查点子文件夹global_step1。因此，要重建fp32权重，只需运行：

python zero_to_fp32.py . pytorch_model.bin

pytorch_model.bin现在将包含从多个GPU中合并的完整fp32模型权重。

该脚本将自动能够处理ZeRO-2或ZeRO-3检查点。