PyTorch分布式训练神器：Accelerate 简洁实用指南

Accelerate 是 Hugging Face 的轻量级分布式训练工具，3 行代码实现单卡到多卡的切换。

核心优势

特性	说明
非侵入式	仅修改 3 行代码
全硬件适配	自动检测 CPU/GPU/TPU 环境
零性能损耗	基于 PyTorch 原生分布式 API
大模型友好	支持 DeepSpeed、FSDP、量化

快速上手

from accelerate import Accelerator
 
# 1. 初始化
accelerator = Accelerator(mixed_precision="fp16")
 
# 2. 一键包装
model, optimizer, dataloader = accelerator.prepare(
    model, optimizer, dataloader
)
 
# 3. 训练循环
for batch in dataloader:
    optimizer.zero_grad()
    loss = model(batch)
    accelerator.backward(loss)  # 替换 loss.backward()
    optimizer.step()

启动方式

# 单卡
python train.py
 
# 多卡
accelerate launch --num_processes=2 train.py

常见场景

大模型训练

accelerate config  # 交互式配置 DeepSpeed/FSDP
accelerate launch --config_file config.yaml train.py

模型保存

accelerator.wait_for_everyone()
unwrapped_model = accelerator.unwrap_model(model)
if accelerator.is_main_process:
    torch.save(unwrapped_model.state_dict(), "model.pth")

分布式评估

preds, labels = accelerator.gather_for_metrics((preds, batch["labels"]))
if accelerator.is_main_process:
    metric.add_batch(predictions=preds, references=labels)

避坑要点

• 保存模型必须用 unwrap_model()
• 移除手动 to(device)
• 打印/保存操作加 is_main_process 判断
• 学习率调度器需传入 prepare()