DeepSpeed ​

DeepSpeed是一个专为PyTorch设计的加速库，它在分布式训练时能够显著提升显存使用效率和加快训练速度。其核心是Zero Redundancy Optimizer（ZeRO），ZeRO的工作分为几个阶段：

• ZeRO-1：支持优化器切分。
• ZeRO-2：支持优化器切分和梯度切分。
• ZeRO-3：支持优化器切分，梯度切分和参数切分。

以下指南将提供具体训练脚本，提供不同配置的ds_config.json，以便学习在分布式训练时如何使用DeepSpeed。

目前与openMind Library联动使用时，该生态库只支持PyTorch框架。

环境及配置 ​

pip install deepspeed==0.13.1

权重、数据集和微调脚本 ​

实际运行脚本可参考如下命令：

if [ -d ./test/output ];then
    rm -rf ./test/output
    mkdir -p ./test/output
else
    mkdir -p ./test/output
fi

# master_port参数需用户根据实际情况进行配置
torchrun --nproc_per_node=8 --master_port=xxx finetune.py \
    --model_name_or_path "PyTorch-NPU/qwen1.5_7b_chat"\
    --data_path alpaca_data.json \
    --deepspeed ds_config.json \
    --bf16 True \
    --output_dir ./test/output \
    --max_steps 2000 \
    --per_device_train_batch_size 1 \
    --evaluation_strategy "no" \
    --save_strategy "steps" \
    --save_steps 3000 \
    --save_total_limit 1 \
    --learning_rate 1e-6 \
    --weight_decay 0. \
    --warmup_ratio 0.03 \
    --lr_scheduler_type "cosine" \
    --seed 1234 \
    --logging_steps 1  > ./test/output/train.log 2>&1 &

其中ds_config.json 可参考下文DeepSpeed基础配置和DeepSpeed进阶配置。

DeepSpeed基础配置 ​

ZeRO-1 ​

{
    "zero_optimization": {
        "stage": 1
    }
}

ZeRO-2 ​

• allgather_partitions：在每个步骤结束时，从所有计算卡收集更新后的参数，可以选择使用allgather集体通信操作或者一系列broadcast集体通信操作。
• allgather_bucket_size：每次allgather操作的元素数量。限制了大模型的allgather所需的内存。
• overlap_comm：控制是否使用通信与计算的重叠。
• reduce_scatter：使用reduce或reduce scatter代替allreduce来平均梯度。
• reduce_bucket_size：每次reduce或allreduce操作的元素数量。限制了大模型的allgather所需的内存。
• contiguous_gradients：在产生梯度时将其复制到连续的缓冲区中。避免了反向传播过程中的内存碎片化。
• round_robin_gradients：优化CPU offload性能。当梯度累积的步数增加，或者NPU数量增加时，会有更好的性能优势。

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

ZeRO-3 ​

• overlap_comm：控制是否使用通信与计算的重叠。
• contiguous_gradients：在产生梯度时将其复制到连续的缓冲区中。避免了反向传播过程中的内存碎片化。
• sub_group_size：控制在optimizer steps中更新参数的粒度。
• reduce_bucket_size：控制Allreduce操作的分桶大小。
• stage3_prefetch_bucket_size：预取参数的固定缓冲区大小。较小的值使用较少的内存，但可能会增加由于通信导致的停滞。
• stage3_param_persistence_threshold：不要对小于这个阈值的参数进行分区。较小的值使用较少的内存，但可能会大幅增加通信量。
• stage3_max_live_parameters：每个计算卡上驻留的最大参数数量，在释放之前。较小的值使用较少的内存，但需要进行更多的通信。
• stage3_max_reuse_distance：如果参数在接下来的这个参数阈值内会被再次使用，则不要释放该参数。较小的值可以减少内存使用，但会增加通信次数。

{
    "zero_optimization": {
        "stage": 3,
        "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
    }
}

DeepSpeed进阶配置 ​

Optimizer和scheduler ​

DeepSpeed提供了多种优化器选项，包括Adam、AdamW、OneBitAdam和LAMB，同时它也支持从PyTorch中导入其他优化器以满足更广泛的需求。若在配置文件中未指定优化器，DeepSpeed将默认采用AdamW作为优化器。

{
   "optimizer": {
       "type": "AdamW",
       "params": {
         "lr": "auto",
         "betas": "auto",
         "eps": "auto",
         "weight_decay": "auto"
       }
   }
}

DeepSpeed不仅支持多种优化器，还兼容LRRangeTest、OneCycle、WarmupLR以及WarmupDecayLR等学习率调度器。若在配置文件中未指定调度器，DeepSpeed将自动选择WarmupDecayLR作为默认调度器。

{
   "scheduler": {
         "type": "WarmupDecayLR",
         "params": {
             "total_num_steps": "auto",
             "warmup_min_lr": "auto",
             "warmup_max_lr": "auto",
             "warmup_num_steps": "auto"
         }
     }
}

Precision ​

DeepSpeed支持fp32、fp16和bf16。

{
    "bf16": {
        "enabled": "auto",
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 16,
        "hysteresis": 2,
        "min_loss_scale": 1
    }
}

Batch size ​

数据批次大小可以自动配置或者显式设置。如果选择使用auto选项，Trainer会将train_micro_batch_size_per_gpu设置为args.per_device_train_batch_size的值，将train_batch_size设置为args.world_size * args.per_device_train_batch_size * args.gradient_accumulation_steps的值。

{
    "train_micro_batch_size_per_gpu": "auto",
    "train_batch_size": "auto"
}

Gradient accumulation ​

梯度累积可以自动配置或显式设置。如果选择使用auto选项，Trainer会将其设置为args.gradient_accumulation_steps的值。

{
    "gradient_accumulation_steps": "auto"
}

Gradient clipping ​

梯度裁剪可以自动配置或显式设置。如果选择使用auto选项，Trainer会将其设置为args.max_grad_norm的值。这是一种常用的技术，用于防止在训练过程中出现梯度爆炸，即梯度值过大导致数值计算问题。通过设置一个最大梯度范数，可以确保梯度在反向传播过程中保持在一个合理的范围内，有助于模型的稳定训练。

{
    "gradient_clipping": "auto"
}

Communication data type ​

在进行诸如reduction、gathering和scattering等通信操作时，会使用一种单独的数据类型。可以通过在配置文件中设置communication_data_type参数来选择通信数据类型。

{
    "communication_data_type": "fp32"
}

配置示例 ​

以下将提供完整集合以上功能的ZeRO-2和ZeRO-3配置示例，可以根据需要选择使用。

{   
    "train_micro_batch_size_per_gpu": "auto",
    "train_batch_size": "auto",
    "gradient_accumulation_steps": "auto",
    "gradient_clipping": "auto",
    "communication_data_type": "fp32",
    
    "bf16": {
        "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"
       }
   }, 
  
    "zero_optimization": {
        "stage": 2,
        "allgather_partitions": true,
        "allgather_bucket_size": 2e8,
        "overlap_comm": false,
        "reduce_scatter": true,
        "reduce_bucket_size": 2e8,
        "contiguous_gradients": true
    }
}

{   
    "train_micro_batch_size_per_gpu": "auto",
    "train_batch_size": "auto",
    "gradient_accumulation_steps": "auto",
    "gradient_clipping": "auto",
    "communication_data_type": "fp32",
    
    "bf16": {
        "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"
       }
   }, 
  
    "zero_optimization": {
        "stage": 3,
        "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
    }
}