分布式训练
AITraining 通过 Accelerate 支持多 GPU 训练，并为大型模型提供可选的 DeepSpeed Zero-3 优化。
| 组件 | 必需 | 安装 |
|---|
| Accelerate | 是（已包含） | 随 AITraining 包含 |
| DeepSpeed | 可选 | pip install deepspeed |
| 多个 GPU | 是 | NVIDIA CUDA GPU |
分布式后端
| 后端 | 值 | 描述 |
|---|
| DDP | ddp 或 None | PyTorch 分布式数据并行（默认） |
| DeepSpeed | deepspeed | 带自动分片的 DeepSpeed Zero-3 |
快速开始
DDP（默认）
使用多个 GPU 时，自动使用 DDP：
aitraining llm --train \
--model meta-llama/Llama-3.2-1B \
--data-path ./data \
--project-name my-model \
--trainer sft \
--peft
DeepSpeed
对于大型模型，使用 DeepSpeed Zero-3：
aitraining llm --train \
--model meta-llama/Llama-3.2-3B \
--data-path ./data \
--project-name my-model \
--trainer sft \
--distributed-backend deepspeed \
--peft
Python API
from autotrain.trainers.clm.params import LLMTrainingParams
from autotrain.project import AutoTrainProject
params = LLMTrainingParams(
model="meta-llama/Llama-3.2-3B",
data_path="./data",
project_name="distributed-model",
trainer="sft",
# Distribution
distributed_backend="deepspeed", # or None for DDP
# Training
epochs=3,
batch_size=2,
gradient_accumulation=4,
mixed_precision="bf16",
peft=True,
lora_r=16,
)
project = AutoTrainProject(params=params, backend="local", process=True)
project.create()
YAML 配置
task: llm-sft
backend: local
base_model: meta-llama/Llama-3.2-3B
project_name: distributed-model
data:
path: ./data
train_split: train
valid_split: null
chat_template: tokenizer
column_mapping:
text_column: text
log: wandb
params:
distributed_backend: deepspeed # or null for DDP
epochs: 3
batch_size: 2
gradient_accumulation: 4
mixed_precision: bf16
peft: true
lora_r: 16
工作原理
Accelerate 启动
训练通过 Accelerate 启动：
- AITraining 检测可用 GPU
- 通过
accelerate launch 启动训练
- 对于 DeepSpeed，添加
--use_deepspeed 和 Zero-3 标志
- 记录
accelerate env 用于调试
DDP 设置
使用 DDP 时：
- 为性能设置
ddp_find_unused_parameters=False
- 每个 GPU 处理批次的一部分
- 梯度在 GPU 之间同步
DeepSpeed Zero-3
使用 DeepSpeed 时：
- 模型参数在 GPU 之间分片
- 多节点使用
--deepspeed_multinode_launcher standard
- 自动应用 Zero-3 配置
- 模型保存使用
accelerator.get_state_dict() 并解包
多节点训练
对于多节点 DeepSpeed 训练：
# 在每个节点上
aitraining llm --train \
--model meta-llama/Llama-3.2-3B \
--data-path ./data \
--project-name my-model \
--distributed-backend deepspeed \
--peft
--deepspeed_multinode_launcher standard 标志会自动传递。
任务特定行为
LLM 训练
- 默认：检测到多个 GPU 时使用 DDP
- DeepSpeed：显式设置
--distributed-backend deepspeed
Seq2Seq 和 VLM
- 多 GPU 情况自动选择 DeepSpeed
- 对于 PEFT + 量化 + bf16 组合使用多 GPU DDP
使用 DeepSpeed 的检查点
使用 DeepSpeed 时，PEFT 适配器保存的处理方式不同。不使用 SavePeftModelCallback；而是使用 accelerator.get_state_dict(trainer.deepspeed) 保存并解包模型。
GPU 选择
控制使用哪些 GPU：
# 使用特定 GPU
CUDA_VISIBLE_DEVICES=0,1 aitraining llm --train ...
# 使用所有可用 GPU（默认）
aitraining llm --train ...
故障排除
检查 Accelerate 环境
常见问题
| 问题 | 解决方案 |
|---|
| 找不到 DeepSpeed | pip install deepspeed |
| NCCL 错误 | 检查 GPU 连接和 CUDA 版本 |
| OOM 错误 | 减少批量大小或使用 DeepSpeed |
| 训练缓慢 | 确保 GPU 在同一 PCIe 总线上 |
下一步
