跳转到内容
Jason's Study

Accelerate — HuggingFace 设备/分布式抽象

是什么

Accelerate 是 HuggingFace 出的一个薄壳库，目标只有一句话：让你的单机 PyTorch 训练脚本，加几行就能在 8 卡 / TPU / DeepSpeed / FSDP 上跑。日常类比：像变压器——你家电器是 220V 的，进了工厂要接 380V 三相电；Accelerate 就是中间那个不让你换电器的变压器。

你写：

from accelerate import Accelerator


accelerator = Accelerator()
model, optimizer, dataloader = accelerator.prepare(model, optimizer, dataloader)


for batch in dataloader:
    loss = model(batch).loss
    accelerator.backward(loss)   # 替代 loss.backward()
    optimizer.step()

只多了 3 处改动。然后 accelerate config 选一次后端，accelerate launch train.py 启动——同一份代码在单卡 / 8 卡 DDP / FSDP / DeepSpeed ZeRO-3 / TPU 全跑得起来。

为什么重要

不理解 Accelerate，下面这些事会卡住你：

  • 为什么 transformers 的 Trainer / peft / trl / diffusers 训练脚本不写 torch.distributed.init_process_group、不写 DistributedSampler，却能多卡训——它们底层全是 Accelerate
  • 为什么 from_pretrained("llama-70b", device_map="auto") 能把 70B 模型自动切到 2×A100 + CPU offload——这是 Accelerate 的 Big Model Inference
  • 为什么社区从原生 DDP 转向 Accelerate / pytorch-lightning——样板代码（rank / world_size / .to(device)）让人疯掉
  • 为什么 PyTorch 官方微调库 torchtune 和 Lightning 都在解同一个问题，但风格完全不同

核心要点

记 3 个抽象 + 1 个 launcher：

  1. Accelerator 总管对象：你只持有这一个实例。它替你判断”现在是单机还是 8 卡？是 fp16 还是 bf16？是 DDP 还是 FSDP？“——所有”看后端脸色”的逻辑都收敛到它身上。

  2. prepare() 是变压器接口：把 model / optimizer / dataloader / scheduler 一起交给它，它返回包好的版本——model 自动 wrap 成 DDP/FSDP，dataloader 自动加 DistributedSampler 切片，optimizer 自动接梯度同步。你的训练循环本身不变。

  3. backward() / gather() / print() 等替代品：在分布式下，loss.backward() 要走梯度同步；print 会被 N 个进程刷 N 次；多卡 metric 要 all-reduce。Accelerate 各给一个等价方法，单卡跑时它们就退化成原生行为。

  4. accelerate launch：替代 python script.py。accelerate config 写一份 yaml（比如”4 卡 DDP + bf16”），launch 命令读它，自动 spawn 进程、设环境变量。业务代码不知道自己跑在哪个后端上。

合起来就是一句话：Accelerator 把”我跑在哪”和”我训什么”解耦。

实践案例

案例 1：从单卡到 8 卡 DDP，只改 3 行

原来的训练循环：

model.to("cuda")
for batch in dataloader:
    loss = model(batch).loss
    loss.backward()
    optimizer.step()

改成：

from accelerate import Accelerator
accelerator = Accelerator()
model, optimizer, dataloader = accelerator.prepare(model, optimizer, dataloader)
for batch in dataloader:                # batch 已经在正确 device 上
    loss = model(batch).loss
    accelerator.backward(loss)          # 自动梯度同步
    optimizer.step()

然后 accelerate launch train.py。没有 init_process_group、没有 LOCAL_RANK、没有 DistributedSampler。

案例 2：切到 DeepSpeed ZeRO-3 不改一行业务代码

只动 accelerate config：选 “DeepSpeed → ZeRO stage 3 → CPU offload”。生成的 yaml：

compute_environment: LOCAL_MACHINE
distributed_type: DEEPSPEED
deepspeed_config:
  zero_stage: 3
  offload_optimizer_device: cpu
mixed_precision: bf16

业务代码一字不改，再 accelerate launch train.py 就跑在 ZeRO-3 上了。把 yaml 换成 FSDP 同理。这就是”解耦”的现金价值。

案例 3：Big Model Inference 跑 70B 模型

from accelerate import init_empty_weights, load_checkpoint_and_dispatch


with init_empty_weights():               # 只造结构，不分配权重
    model = AutoModelForCausalLM.from_config(config)


model = load_checkpoint_and_dispatch(
    model, "llama-70b/", device_map="auto",   # 自动切到 GPU+CPU+disk
)

device_map="auto" 会量一遍每张卡剩多少显存，把 transformer block 一层一层放进去；放不下的层走 CPU offload，再不够就 disk。transformers 的 from_pretrained(device_map="auto") 就是这套机制的封装。

踩过的坑

  1. prepare() 之后访问原模型要 unwrap_model：DDP/FSDP 会把 model 包一层，model.config 这种属性访问会变成 model.module.config。直接写会在单卡能跑、多卡崩。统一写 accelerator.unwrap_model(model).config。

  2. 保存模型不能直接 torch.save(model.state_dict())：state_dict 的 key 会带 module. / _orig_mod. 前缀，加载到推理代码会全部 mismatch。要么 unwrap_model 后再 save，要么用 accelerator.save_state 走 Accelerate 自己的协议。

  3. 每个 rank 都在下数据集：dataset 下载 / tokenize 这种昂贵操作要包在 with accelerator.main_process_first(): 里，否则 8 个进程同时下，会把 HuggingFace cache 撞崩。

  4. print 会被刷 8 次：分布式下每个进程都执行一次 print。用 accelerator.print(...)，它只在 main process 输出。同理 accelerator.is_main_process 守 wandb / logging 初始化。

  5. 梯度累积别手动除：传统写法 loss = loss / accum_steps; loss.backward()。Accelerate 要写成 Accelerator(gradient_accumulation_steps=N) + with accelerator.accumulate(model):——它会处理”中间步骤不同步梯度只在累积满时同步”的优化，手动除会让多卡通信白白发生。

适用 vs 不适用场景

适用：

  • 已经有一份能跑的单机 PyTorch 训练脚本，想扩到多卡或换后端
  • 用 transformers / peft / trl / diffusers 做 LLM / 扩散模型微调（它们已经基于 Accelerate）
  • 想在 DDP / FSDP / DeepSpeed 之间快速切换做 benchmark
  • 推理时要加载比单卡显存大的模型（Big Model Inference）

不适用：

  • 想完全不写训练循环、只填 training_step → 用 pytorch-lightning，Accelerate 不替你写循环
  • 训练任务高度定制化（自定义 RL rollout / 多模型异步更新）→ 直接写原生 pytorch DDP，更可控
  • 极端规模（千卡、流水线并行 PP + TP + DP 三维切）→ 上 Megatron-LM / 自研，Accelerate 在那个量级抽象不够
  • 不是 PyTorch 生态（JAX / TF）→ 不适用

历史小故事（可跳过）

  • 2021：Sylvain Gugger 在 HuggingFace 写 transformers.Trainer，发现”设备 + 分布式”逻辑反复重复，抽出独立库 accelerate。
  • 2022：DeepSpeed 和 FSDP 集成进来，统一 launcher。社区开始从手写 DDP 迁移过来。
  • 2023：Big Model Inference 落地——device_map="auto" 让单机推 70B 成为日常操作，间接推动了开源 LLM 浪潮。
  • 2024：fp8（TransformerEngine）+ Megatron-LM 接入，覆盖到企业级训练。

短短 3 年从内部工具变成 PyTorch 生态分布式训练事实标准。

学到什么

  1. 抽象的价值不在功能多，在解耦——Accelerate 不发明新算法，只是把”我跑在哪”和”我训什么”切开
  2. 薄壳 vs 厚壳：pytorch-lightning 把训练循环吃掉是厚壳；Accelerate 只接管设备/分布式是薄壳——薄壳侵入小、迁移成本低
  3. launcher + config 模式：业务代码无关后端，后端切换只动 yaml——这一招在 torchtune / pytorch-lightning 都在用
  4. unwrap_model 是抽象漏出来的代价：再薄的包装也会泄漏，DDP 的 module. 前缀就是经典案例

延伸阅读

关联

  • pytorch —— Accelerate 是 PyTorch 分布式 API（DDP / FSDP）的薄壳封装
  • pytorch-lightning —— 同领域厚壳方案：吃掉训练循环，只让你写 training_step
  • torchtune —— PyTorch 官方 LLM 微调库，自己造抽象不用 Accelerate，对比设计取舍
  • hindley-milner —— 不直接相关，但同样体现”抽象的价值在解耦”