Hugging Face Transformers 分布式训练实战：用 🤗 Accelerate 将原生 PyTorch 训练循环改造为多卡/多机训练

本文基于 Transformers 仓库的官方教程 Accelerate 分布式训练指南 展开，讲解如何用仅 4 行代码改动，把一条普通的 PyTorch 训练循环改造成支持单机多卡乃至多机多卡的分布式训练；并结合仓库中 examples/pytorch/ 下的无 Trainer 训练脚本与 依赖声明，补充参数默认值、梯度累积、断点续训和实验跟踪等实战细节。读完后可独立完成 accelerate config、accelerate launch 全流程，并理解 Accelerator.prepare 与 accelerator.backward 在底层承担的职责。

一、背景：为什么选择 🤗 Accelerate

随着模型规模持续增大，并行化已成为在有限硬件上训练更大模型、并将训练速度提升数个数量级的核心策略。Hugging Face 为此开发了 🤗 Accelerate 库，帮助用户在任意分布式环境（单机的多块 GPU，或多机的多块 GPU）下训练 Transformers 模型。官方教程的核心主张是：无需重写训练逻辑，只需把原生 PyTorch 训练循环中的设备管理代码交给 Accelerator 对象处理。

从仓库依赖看，Accelerate 已深度内建于 Transformers 生态：setup.py 中声明了 accelerate>=1.1.0 这一硬性依赖，并且 torch、deepspeed 等安装组均自动包含 Accelerate（extras["torch"] = deps_list("torch", "accelerate")、extras["deepspeed"] = deps_list("deepspeed", "accelerate")）。此外，Trainer 本身的实现也依赖 Accelerate（参见 integrations/accelerate.py 与 trainer.py），这意味着掌握本文的无 Trainer 工作流，对理解 Trainer 的分布式行为同样有帮助。

二、环境配置：安装并初始化 Accelerator

2.1 安装

先安装 🤗 Accelerate：

pip install accelerate

若通过 Transformers 的 extras 安装，则无需单独安装：

pip install transformers[torch]   # torch 安装组已包含 accelerate>=1.1.0
pip install transformers[deepspeed]  # DeepSpeed 集成同样依赖 accelerate

2.2 创建 Accelerator 对象

接下来导入并创建一个 Accelerator 对象。Accelerator 会自动探测当前运行环境的分布式形态（单进程 CPU、单卡、DDP 多卡、多机 DDP、FSDP、DeepSpeed、TPU 等），并完成训练所需的全部初始化。关键在于：你不需要再手动把模型放到指定设备上，设备放置由 Accelerator 统一管理。

>>> from accelerate import Accelerator

>>> accelerator = Accelerator()

仓库中的真实脚本在此基础上还展示了两个常用构造参数（见 run_clm_no_trainer.py）：

accelerator_log_kwargs = {}

if args.with_tracking:
    accelerator_log_kwargs["log_with"] = args.report_to      # 指定跟踪后端
    accelerator_log_kwargs["project_dir"] = args.output_dir  # 日志输出目录

accelerator = Accelerator(
    gradient_accumulation_steps=args.gradient_accumulation_steps,  # 梯度累积步数
    **accelerator_log_kwargs,
)

• gradient_accumulation_steps：把多次前向的小批量梯度累积起来再统一更新，等效放大 batch size，是显存受限时扩大有效批量（effective batch size）的标准手段；
• log_with / project_dir：声明实验跟踪后端（如 TensorBoard、WandB 等）及其输出目录，跟踪器默认只在主进程上初始化。

脚本中还会用 logger.info(accelerator.state, main_process_only=False) 打印 Accelerator 的运行时状态（分布类型、进程数等），便于在调试阶段确认当前究竟以何种并行模式运行。

三、准备加速：把训练对象交给 accelerator.prepare

改造的第二步是把所有参与训练的组件——训练/验证 DataLoader、模型、优化器——一次性交给 Accelerator.prepare 方法：

>>> train_dataloader, eval_dataloader, model, optimizer = accelerator.prepare(
...     train_dataloader, eval_dataloader, model, optimizer
... )

prepare 返回的对象顺序与传入顺序一致。它内部完成的事情包括：把模型移动到正确的设备并用对应的分布式包装器（如 DDP）包裹、为 DataLoader 注入分布式采样器使数据按进程切分、将优化器替换为感知梯度累积/混合精度语义的版本等。

仓库的完整示例（run_clm_no_trainer.py）额外展示了两个细节：

# 官方示例把 lr_scheduler 也一并 prepare
model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
)

# TPU 场景：权重绑定关系在 TPU 上会断开，需要手动恢复
if accelerator.distributed_type == DistributedType.TPU:
    model.tie_weights()

也就是说，prepare 的入参并不限于文档列出的四类对象，学习率调度器同样推荐纳入，以保证其 step 语义与分布式/梯度累积行为一致；而在 TPU 这种特殊后端上，prepare 后还需要处理权重共享（weight tying）的细节。

四、反向传播：用 accelerator.backward 替换 loss.backward

改造的最后一步，是把训练循环中惯用的 loss.backward() 替换为 Accelerate 的 Accelerator.backward 方法：

>>> for epoch in range(num_epochs):
...     for batch in train_dataloader:
...         outputs = model(**batch)
...         loss = outputs.loss
...         accelerator.backward(loss)

...         optimizer.step()
...         lr_scheduler.step()
...         optimizer.zero_grad()
...         progress_bar.update(1)

accelerator.backward 会按当前精度配置对损失做缩放后再执行反向传播（混合精度下自动调用 GradScaler），并在需要时处理梯度累积期间的同步时机。下面是官方教程给出的完整 diff——只需 4 行新增代码即可开启分布式训练：

+ from accelerate import Accelerator
  from transformers import AdamW, AutoModelForSequenceClassification, get_scheduler

+ accelerator = Accelerator()

  model = AutoModelForSequenceClassification.from_pretrained(checkpoint, num_labels=2)
  optimizer = AdamW(model.parameters(), lr=3e-5)

- device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
- model.to(device)

+ train_dataloader, eval_dataloader, model, optimizer = accelerator.prepare(
+     train_dataloader, eval_dataloader, model, optimizer
+ )

  num_epochs = 3
  num_training_steps = num_epochs * len(train_dataloader)
  lr_scheduler = get_scheduler(
      "linear",
      optimizer=optimizer,
      num_warmup_steps=0,
      num_training_steps=num_training_steps
  )

  progress_bar = tqdm(range(num_training_steps))

  model.train()
  for epoch in range(num_epochs):
      for batch in train_dataloader:
-         batch = {k: v.to(device) for k, v in batch.items()}
          outputs = model(**batch)
          loss = outputs.loss
-         loss.backward()
+         accelerator.backward(loss)

          optimizer.step()
          lr_scheduler.step()
          optimizer.zero_grad()
          progress_bar.update(1)

注意被删除的三行设备管理代码：torch.device 选择、model.to(device)、逐 batch 的 batch.to(device)。这些正是 Accelerator + prepare 接管后的职责——prepare 过的 DataLoader 产出的 batch 已经位于正确设备上。

从源码结构看，仓库的 no_trainer 系列脚本在生产级用法上比教程更进一步（以 run_clm_no_trainer.py 为例）：

1. 调度器步数乘以进程数：num_warmup_steps=args.num_warmup_steps * accelerator.num_processes，因为 prepare 后的 DataLoader 每个进程只遍历自己的数据分片，全局步数需要换算；
2. 梯度累积上下文管理器：训练步进整体包在 with accelerator.accumulate(model): 中，accelerator.backward(loss) 放在其内部（第 636-642 行），由 accumulate 自动管理“何时同步梯度、何时真正执行优化器更新”；
3. 有效批量计算公式：total_batch_size = per_device_train_batch_size * accelerator.num_processes * gradient_accumulation_steps，其中 accelerator.num_processes 即当前参与训练的进程总数；
4. 只在主进程显示进度条：tqdm(..., disable=not accelerator.is_local_main_process)，避免多进程下进度条互相覆盖。

五、启动训练

5.1 以脚本方式训练

若训练代码位于一个 Python 脚本中，先运行配置命令生成本机（或集群）的分布式配置文件：

accelerate config

accelerate config 会以交互式问答收集 GPU 数量、机器数、分布式后端、混合精度等选项，并落盘为配置文件。随后一条命令即可按配置启动训练，进程编排（torchrun/deepspeed 等）全部由 Accelerate 处理：

accelerate launch train.py

5.2 以 Notebook 方式训练

如果计划在 Colaboratory 等环境中使用 TPU，🤗 Accelerate 也支持 Notebook 运行方式：把所有训练相关代码封装进一个函数，再把它交给 notebook_launcher：

>>> from accelerate import notebook_launcher

>>> notebook_launcher(training_function)

notebook_launcher 会在 Notebook 内核内部拉起正确的多进程训练流程（TPU 上通常为 XLA 编译后的多核执行），training_function 内部同样遵循「Accelerator() → prepare → backward」三步范式即可。

六、进阶实战：从仓库 no_trainer 示例看完整分布式工作流

Transformers 在 examples/pytorch/ 下为每个任务都提供了 run_*_no_trainer.py 参考实现，它们就是官方教程这套范式的完整落地版本，可整体对照阅读：

• run_glue_no_trainer.py：文本分类（即教程 diff 中的场景）
• run_clm_no_trainer.py：语言建模
• run_qa_no_trainer.py：问答
• run_ner_no_trainer.py：NER
• run_translation_no_trainer.py、run_summarization_no_trainer.py、run_image_classification_no_trainer.py、run_object_detection_no_trainer.py 等

这些脚本中值得直接借鉴的分布式工程实践包括：

断点续训（checkpointing）。周期性用 accelerator.save_state(output_dir) 把模型、优化器、调度器与数据加载器的完整状态写入 step_{n} 目录；恢复时用 accelerator.load_state(checkpoint_path) 还原，并配合 accelerator.skip_first_batches(train_dataloader, resume_step) 跳过已完成的数据批次（run_clm_no_trainer.py）。

实验跟踪。accelerator.init_trackers("clm_no_trainer", experiment_config) 一次性初始化所有后端，训练循环中通过 accelerator.log(..., step=...) 上报指标，日志默认仅在主进程写入。

多进程输出治理。用 accelerator.print(...) 替代 print 保证只从主进程输出；用 accelerator.is_local_main_process 控制「每台机器只打印一次」的日志（如 datasets/transformers 的日志级别设置，参见 run_clm_no_trainer.py）。

DeepSpeed 等后端零改动切换。同一份 prepare/backward 代码，仅靠 accelerate config / accelerate launch 切换 DeepSpeed ZeRO 配置即可运行，这正是 setup.py 将 accelerate 列入 deepspeed 安装组的意义所在。

七、小结

步骤	代码	职责
初始化	accelerator = Accelerator(gradient_accumulation_steps=...)	自动探测并行环境、接管设备放置
准备	train_dataloader, eval_dataloader, model, optimizer = accelerator.prepare(...)	分布式包装模型、切分数据、适配优化器/调度器
反向	accelerator.backward(loss)	精度感知的反向传播与梯度同步时机管理
启动	accelerate config + accelerate launch train.py / notebook_launcher(training_function)	进程编排、多机/TPU 环境适配

以上四步即是官方教程给出的全部改造内容：删去手动设备管理、引入 Accelerator、prepare 与 accelerator.backward，即可让一条原生 PyTorch 训练循环在从单机单卡到多机多卡的任意拓扑上运行。更多 Accelerate 的功能（FSDP、DeepSpeed ZeRO、混合精度、多后端跟踪等）可查阅 🤗 Accelerate 官方文档获取。