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