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`参数决定
|