127 lines
4.0 KiB
Markdown
127 lines
4.0 KiB
Markdown
|
# 使用Accelerate+DeepSpeed进行分布式训练
|
|||
|
|
|
|||
|
|
本文档介绍如何使用Accelerate和DeepSpeed进行MiniMind模型的分布式训练。
|
|||
|
|
|
|||
|
|
## 环境准备
|
|||
|
|
|
|||
|
|
首先，确保安装了必要的依赖：
|
|||
|
|
|
|||
|
|
```bash
|
|||
|
|
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`修改而来，主要变化包括：
|
|||
|
|
|
|||
|
|
1. 使用Accelerator替代了PyTorch原生的分布式功能
|
|||
|
|
2. 移除了torchrun相关的分布式初始化代码
|
|||
|
|
3. 使用Accelerator的API进行模型、优化器和数据加载器的准备
|
|||
|
|
4. 使用Accelerator的API进行反向传播和梯度裁剪
|
|||
|
|
5. 处理了位置编码和未使用参数的问题
|
|||
|
|
|
|||
|
|
## 启动训练
|
|||
|
|
|
|||
|
|
有两种方式启动训练：
|
|||
|
|
|
|||
|
|
### 方法1：使用预先配置的accelerate配置文件
|
|||
|
|
|
|||
|
|
```bash
|
|||
|
|
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
|
|||
|
|
|
|||
|
|
```bash
|
|||
|
|
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
|
|||
|
|
bash run_accelerate.sh
|
|||
|
|
```
|
|||
|
|
|
|||
|
|
## Accelerate与DeepSpeed配置的关系
|
|||
|
|
|
|||
|
|
1. **Accelerate**是一个高级API，用于简化分布式训练的设置和启动，它可以与多种分布式训练后端（如DeepSpeed、FSDP等）一起使用。
|
|||
|
|
|
|||
|
|
2. **DeepSpeed**是一个优化库，专注于大规模模型训练的内存优化和性能提升，提供了ZeRO优化等功能。
|
|||
|
|
|
|||
|
|
3. **配置关系**：
|
|||
|
|
- Accelerate配置文件（YAML）定义了使用哪种分布式后端以及基本的分布式设置
|
|||
|
|
- DeepSpeed配置文件（JSON）定义了DeepSpeed特有的优化参数
|
|||
|
|
- Accelerate通过`deepspeed_config_file`参数引用DeepSpeed配置文件
|
|||
|
|
|
|||
|
|
## 注意事项
|
|||
|
|
|
|||
|
|
1. **位置编码处理**：
|
|||
|
|
- 在模型中，`pos_cis`是一个复数张量，在分布式训练中需要特别处理
|
|||
|
|
- 在新的训练脚本中，我们使用Accelerator的API来处理这个问题，不再需要`_ddp_params_and_buffers_to_ignore`
|
|||
|
|
|
|||
|
|
2. **未使用参数处理**：
|
|||
|
|
- 原代码中使用`find_unused_parameters=True`来处理未使用的参数
|
|||
|
|
- 在新的训练脚本中，我们直接使用Accelerator的API，它会自动处理这个问题
|
|||
|
|
|
|||
|
|
3. **混合精度训练**：
|
|||
|
|
- DeepSpeed配置文件中的`fp16`和`bf16`设置为`"auto"`
|
|||
|
|
- 实际使用的精度由Accelerate的`--mixed_precision`参数决定
|
|||
|
|
|
|||
|
|
4. **梯度累积**：
|
|||
|
|
- DeepSpeed配置文件中的`gradient_accumulation_steps`设置为`"auto"`
|
|||
|
|
- 实际的梯度累积步数由训练脚本的`--accumulation_steps`参数决定
|