Accelerate 是 Hugging Face 的轻量级分布式训练包装器:同一份训练脚本不改逻辑,3 行代码接入,即可在单卡、多卡 DDP、DeepSpeed、FSDP 之间无缝切换。

定位与核心优势

Accelerate 不发明新的并行策略,只把 torch.distributed 通信原语 之上的样板代码——进程组初始化、设备选择、DDP 包装、数据切分、梯度同步——收敛为统一接口,是 分布式训练总览 中数据并行路线的工程化封装。

特性说明
非侵入式仅改 3 行代码,训练循环保持 PyTorch 原生写法
全硬件适配自动检测 CPU/GPU/TPU/NPU,单卡零开销直跑
零性能损耗底层就是 PyTorch 原生分布式 API,无额外抽象税
大模型友好一份配置切换 DeepSpeed ZeRO、FSDP、CPU offload、混合精度

快速上手:三步改造

from accelerate import Accelerator
 
# 1. 初始化:声明混合精度与梯度累积,替代手写 AMP 与 GradScaler
accelerator = Accelerator(mixed_precision="bf16")
 
# 2. 一键包装:设备迁移、DDP 包装、DataLoader 切分全部在此完成
model, optimizer, dataloader, scheduler = accelerator.prepare(
    model, optimizer, dataloader, scheduler
)
 
# 3. 训练循环:没有任何分布式样板代码
for batch in dataloader:
    with accelerator.accumulate(model):  # 梯度累积上下文,自动缩放 loss
        loss = model(**batch).loss
        accelerator.backward(loss)  # 替换 loss.backward()
        optimizer.step()
        scheduler.step()
        optimizer.zero_grad()

prepare 做了什么

prepare() 是 Accelerate 的核心分发点,按当前后端对传入对象逐一改造:

对象改造内容
model迁移到正确设备;包装为 DDP / FSDP / DeepSpeed engine
dataloader按 rank 切分 batch,各进程拿到互不重叠的 shard,无需手写 DistributedSampler
optimizer接管混合精度 loss scaling;DeepSpeed 下并入引擎统一管理
scheduler与梯度累积步对齐,step 时机由包装器控制

与手写 torch.distributed 通信原语 的对照:

手写 torch.distributedAccelerate 等价物
init_process_group("nccl") + 解析 envAccelerator() 内部完成
model.to(local_rank) + DDP(model)prepare(model)
DistributedSampler + set_epoch()prepare(dataloader) 自动切分
GradScaler + scaler.scale(loss).backward()accelerator.backward(loss)
手写 all_gather + 不等长 batch 拼接gather_for_metrics()

backward() 必须替换 loss.backward() 的原因:混合精度下要插入 unscale,DeepSpeed 引擎的梯度规约(内部仍是 AllReduce 集合通信)挂在自定义 backward 上,梯度累积时还要按比例缩放——三者都被这一个调用收敛。Dataset 与 DataLoader 与 torch.cuda 与混合精度 中的手动细节同理被吸收。

启动方式:accelerate launch

python train.py                            # 单卡直接跑,无需改动
accelerate launch train.py                 # 自动用满本机所有 GPU
accelerate launch --num_processes=2 \
  --multi_gpu --num_machines=2 train.py    # 多机多卡

accelerate launch 与 torchrun 等价:同样通过环境变量(RANK、WORLD_SIZE、LOCAL_RANK)拉起 worker 进程。区别在于 launch 会先读取 Accelerate 配置,把进程数、分布式后端、混合精度一并注入,省去命令行传参。脚本内通过 accelerator.device、accelerator.process_index 拿到解析结果,永远不需要手动 os.environ["LOCAL_RANK"]。

大模型训练:accelerate config

accelerate config 交互式问答生成 default_config.yaml,之后 accelerate launch --config_file config.yaml train.py 即可切换后端。DeepSpeed ZeRO 配置示例:

compute_environment: LOCAL_MACHINE
distributed_type: DEEPSPEED
num_processes: 4
mixed_precision: bf16
deepspeed_config:
  zero_optimization_stage: 2      # ZeRO-2 切分 optimizer state 与梯度
  gradient_accumulation_steps: 4
  offload_optimizer_device: cpu   # 显存不足时卸载到内存

FSDP 配置示例:

compute_environment: LOCAL_MACHINE
distributed_type: FSDP
num_processes: 4
fsdp_config:
  fsdp_auto_wrap_policy: TRANSFORMER_BASED
  fsdp_transformer_layer_cls_to_wrap: LlamaDecoderLayer  # 按 Transformer 层切分
  fsdp_state_dict_type: FULL_STATE_DICT

训练代码完全不变——这正是 prepare 分发点的价值:DDP/FSDP/DeepSpeed 的差异被配置文件隔离。

模型保存与检查点

accelerator.wait_for_everyone()          # 同步所有进程后再保存
unwrapped = accelerator.unwrap_model(model)
accelerator.save(unwrapped.state_dict(), "model.pt")  # 单文件权重,续见 [[模型权重保存]]
 
# 断点续训:连同 optimizer、scheduler、RNG 状态一起保存
accelerator.save_state("checkpoint/")
# DeepSpeed/FSDP 下各 rank 分片落盘,load_state 原样恢复
accelerator.load_state("checkpoint/")

save_state 在 FSDP/DeepSpeed 后端下自动按 rank 分片写盘,恢复时聚合,与 大模型检查点 的分片格式对应;需要导出单文件全量权重时用 accelerator.get_state_dict(model) 先聚合再 save。

分布式评估

for batch in eval_dataloader:
    preds = model(**batch).logits.argmax(-1)
    # 各 rank 张量收集到主进程,自动丢弃最后不完整 batch 的重复样本
    preds, labels = accelerator.gather_for_metrics((preds, batch["labels"]))
if accelerator.is_main_process:
    metric.compute()

避坑要点

  • 保存必须 unwrap_model():包装后的模型 state_dict 键带 module. 前缀,且 DeepSpeed 引擎没有标准 state_dict 接口
  • 删除手动 to(device):prepare 已完成迁移,dataloader 吐出的 batch 也在正确设备上;手动指定 cuda:0 在多卡下必错
  • 打印 / 落盘 / wandb 初始化包上 is_main_process:否则每 rank 重复执行;跨进程有依赖的同步点前调 wait_for_everyone()
  • scheduler 必须传入 prepare():否则 step 时机与梯度累积、混合精度不对齐,学习率曲线错位
  • 梯度累积勿手动 loss / N:accelerator.accumulate(model) 已按配置缩放,手动再除一次等于双重缩放,DeepSpeed 下直接报错
  • 收集指标用 gather_for_metrics 而非裸 gather:后者会因各 rank batch 数不均在指标里重复计入尾部样本