使用Accelerate+DeepSpeed进行分布式训练

本文档介绍如何使用Accelerate和DeepSpeed进行MiniMind模型的分布式训练。

环境准备

首先，确保安装了必要的依赖：

pip install accelerate deepspeed

配置文件说明

1. DeepSpeed配置文件 (ds_config.json)

DeepSpeed配置文件定义了优化器、学习率调度器和ZeRO优化等参数。主要配置包括：

ZeRO优化：使用ZeRO-2进行优化，可以减少GPU内存使用
优化器设置：使用AdamW优化器
混合精度训练：支持FP16和BF16
梯度累积：通过"auto"自动设置，与训练脚本参数保持一致

2. Accelerate配置文件 (accelerate_config.yaml)

Accelerate配置文件定义了分布式训练的基本设置，包括：

分布式类型：使用DeepSpeed
混合精度：使用BF16
进程数量：设置为4（可根据GPU数量调整）
DeepSpeed配置：指向ds_config.json文件

训练脚本说明

新的训练脚本train_pretrain_accelerate.py基于原有的train_pretrain.py修改而来，主要变化包括：

使用Accelerator替代了PyTorch原生的分布式功能
移除了torchrun相关的分布式初始化代码
使用Accelerator的API进行模型、优化器和数据加载器的准备
使用Accelerator的API进行反向传播和梯度裁剪
处理了位置编码和未使用参数的问题

启动训练

有两种方式启动训练：

方法1：使用预先配置的accelerate配置文件

accelerate launch --config_file accelerate_config.yaml train_pretrain_accelerate.py \
    --epochs 3 \
    --batch_size 24 \
    --learning_rate 2e-4 \
    --dtype bfloat16 \
    --accumulation_steps 32 \
    --grad_clip 1.0 \
    --log_interval 100 \
    --save_interval 10000 \
    --dim 1024 \
    --n_layers 32 \
    --max_seq_len 1024 \
    --use_flash_attn \
    --profile \
    --profile_interval 10

方法2：使用命令行参数直接配置accelerate

CUDA_VISIBLE_DEVICES=0,1,2,3 accelerate launch \
    --multi_gpu \
    --num_processes=4 \
    --mixed_precision=bf16 \
    --main_process_port=29500 \
    --deepspeed_config_file ds_config.json \
    train_pretrain_accelerate.py \
    --epochs 3 \
    --batch_size 24 \
    --learning_rate 2e-4 \
    --dtype bfloat16 \
    --accumulation_steps 32 \
    --grad_clip 1.0 \
    --log_interval 100 \
    --save_interval 10000 \
    --dim 1024 \
    --n_layers 32 \
    --max_seq_len 1024 \
    --use_flash_attn \
    --profile \
    --profile_interval 10

也可以直接使用提供的脚本：

bash run_accelerate.sh

Accelerate与DeepSpeed配置的关系

Accelerate是一个高级API，用于简化分布式训练的设置和启动，它可以与多种分布式训练后端（如DeepSpeed、FSDP等）一起使用。
DeepSpeed是一个优化库，专注于大规模模型训练的内存优化和性能提升，提供了ZeRO优化等功能。
配置关系：
- Accelerate配置文件（YAML）定义了使用哪种分布式后端以及基本的分布式设置
- DeepSpeed配置文件（JSON）定义了DeepSpeed特有的优化参数
- Accelerate通过deepspeed_config_file参数引用DeepSpeed配置文件

注意事项

位置编码处理：
- 在模型中，pos_cis是一个复数张量，在分布式训练中需要特别处理
- 在新的训练脚本中，我们使用Accelerator的API来处理这个问题，不再需要_ddp_params_and_buffers_to_ignore
未使用参数处理：
- 原代码中使用find_unused_parameters=True来处理未使用的参数
- 在新的训练脚本中，我们直接使用Accelerator的API，它会自动处理这个问题
混合精度训练：
- DeepSpeed配置文件中的fp16和bf16设置为"auto"
- 实际使用的精度由Accelerate的--mixed_precision参数决定
梯度累积：
- DeepSpeed配置文件中的gradient_accumulation_steps设置为"auto"
- 实际的梯度累积步数由训练脚本的--accumulation_steps参数决定

4.0 KiB Raw Permalink Blame History Unescape Escape