张芷铭的个人博客

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

核心优势

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

快速上手

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16

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()

启动方式

1
2
3
4
5

# 单卡
python train.py

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

常见场景

大模型训练

1
2

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

模型保存

1
2
3
4

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

分布式评估

1
2
3

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()

Comments