Transformers 分布式训练实战：用 🤗 Accelerate 让 PyTorch 训练脚本跑在多卡与多机环境

本文围绕 Transformers 仓库中的 Accelerate 分布式训练指南展开，讲清楚如何用极少代码改动把一条原生 PyTorch 训练循环改造成支持多 GPU、多机训练的形式：包括 Accelerator 的创建、prepare 与 backward 两个关键 API、accelerate config / accelerate launch 的启动流程，以及 Colab 笔记本中的 notebook_launcher 用法。读完后，你可以直接按文中的 diff 改造自己的训练脚本，并在仓库示例 run_glue_no_trainer.py 与 Trainer 源码中验证这些 API 的真实调用方式。

为什么需要 Accelerate

随着大语言模型规模的增长，并行化成为在有限硬件上训练更大模型、并显著加速训练的核心策略之一。Transformers 与 🤗 Accelerate 的配合目标，是让你无需手写任何设备管理（model.to(device)）、分布式包装（DDP/FSDP）或混合精度缩放逻辑，就能在多种分布式环境——单机的多块 GPU、跨多台机器的多 GPU 集群——训练任意 Transformers 模型。

从仓库源码可以印证 Accelerate 在 Transformers 中的地位。Trainer 并不直接操作分布式细节，而是在 create_accelerator_and_postprocess 中统一创建 Accelerator 对象，训练步中的反向传播也是委托给 self.accelerator.backward(loss) 完成（见 training_step）。也就是说，本指南中手动编写的四行改动，正是 Trainer 内部为你所做事情的"手工版"。

一、安装与创建 Accelerator

第一步是安装 Accelerate：

pip install accelerate

然后导入并创建 [Accelerator] 对象。它会自动探测你当前的分布式环境（GPU 数量、分布式后端、混合精度等）并初始化训练所需的全部组件，你不再需要手动把模型放到某个特定设备上：

>>> from accelerate import Accelerator

>>> accelerator = Accelerator()

二、用 accelerator.prepare 准备训练对象

下一步是把所有与训练相关的对象传给 [accelerator.prepare] 方法。这通常包括训练/评估 DataLoader、模型和优化器：

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

prepare 做的事情可以归纳为三类：

• 模型：包装进合适的分布式封装（如 DDP），并移动到正确设备；
• DataLoader：按进程数切分批次（分片采样器），使每个 GPU 拿到互不重叠的数据；
• 优化器/调度器：创建与混合精度策略兼容的封装版本。

在仓库的真实示例 run_glue_no_trainer.py 中可以看到完整用法，lr_scheduler 也一并传入：

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

注意一个细节：prepare 之后 DataLoader 的长度会随进程数变化（每个进程只迭代自己的分片）。因此 示例代码 在 prepare 之后重新计算了 num_update_steps_per_epoch 与训练轮数——如果你在自定义调度器步数或进度条，务必在 prepare 之后再做这类计算。

三、用 accelerator.backward 替换 loss.backward

最后一步，把训练循环中惯用的 loss.backward() 替换为 [accelerator.backward]。它负责处理混合精度下的梯度缩放（grad scaler）以及 FSDP 等后端的特殊逻辑，是"唯一必须改写的反向传播调用点"：

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

四、完整的改造 diff：只需改动四处

把上述三步合在一起，对一条标准的 GLUE 微调脚本的改动如下。核心是新增 4 行、删除 3 行设备管理代码：

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

要点总结：

改动	作用
accelerator = Accelerator()	自动探测环境、初始化分布式/混合精度组件
删除 model.to(device)	设备放置交给 prepare，避免与分布式包装冲突
accelerator.prepare(...)	包装模型、切分 DataLoader、适配优化器
accelerator.backward(loss)	处理混合精度梯度缩放与后端特定逻辑
删除 batch = {k: v.to(device) ...}	批次数据移动由 prepare 后的 DataLoader 自动完成

五、启动分布式训练

5.1 脚本方式：accelerate config + accelerate launch

如果你从脚本启动训练，先用交互式命令生成并保存配置（它会询问硬件与训练设置，最终写入缓存目录下的 default_config.yaml）：

accelerate config

然后用 accelerate launch 启动训练脚本，它会按配置自动拉起所需数量的进程：

accelerate launch train.py

5.2 笔记本方式：notebook_launcher

如果你计划使用 Colaboratory 上的 TPU，可以在 Jupyter 笔记本中运行 Accelerate。做法是把所有训练逻辑封装进一个函数，然后把它传给 [notebook_launcher]：

>>> from accelerate import notebook_launcher

>>> notebook_launcher(training_function)

notebook_launcher 会在 TPU 运行时中以子进程方式重新启动该函数，使其拿到正确的分布式环境。

六、进阶细节：仓库示例中的分布式实用技巧

读完 run_glue_no_trainer.py 可以发现，改造后的脚本还用到了一批 Accelerator 提供的分布式实用 API，这里挑几个高频的说明：

• 总批量大小与日志：用 accelerator.num_processes 计算全局 batch size——total_batch_size = per_device_batch_size * accelerator.num_processes * gradient_accumulation_steps（见示例）。
• 进度条只在一个进程显示：tqdm(..., disable=not accelerator.is_local_main_process)，避免多进程刷屏。
• 多卡聚合评估指标：predictions, references = accelerator.gather((predictions, batch["labels"]))；由于多进程环境下最后一个 batch 可能产生重复，需要按 len(eval_dataloader.dataset) - samples_seen 截断（见示例）。
• 断点续训：accelerator.save_state(output_dir) 保存完整分布式状态，accelerator.load_state(checkpoint_path) 恢复；恢复后用 accelerator.skip_first_batches(train_dataloader, resume_step) 跳过已训练批次（见示例）。
• 混合精度与 Padding 对齐：示例中若启用混合精度，会把 DataCollator 的 pad_to_multiple_of 设为 8（fp8 为 16），以充分利用 NVIDIA Tensor Cores（见示例）。
• 梯度累积下的 checkpoint 时机：判断 accelerator.sync_gradients 为真时才保存，确保落在真正的优化步上（见示例）。

如果你希望更细粒度地控制后端（FSDP/DeepSpeed/DDP），仓库中对应的文档是 docs/source/en/accelerate.md，其中给出了 accelerate config 生成的 FSDP 配置示例（distributed_type、fsdp_config、mixed_precision: bf16、num_processes 等字段），以及通过 TrainingArguments 的 fsdp_config / deepspeed / DDP 参数直连各后端的写法。从 Trainer._build_accelerator_args 源码结构看，TrainingArguments 中的 mixed_precision、DDP 的 find_unused_parameters/bucket_cap_mb 等参数会被原样透传给 Accelerator，两条路径最终汇聚到同一套 Accelerate 接口上。

小结

• 改造成本极低：创建 Accelerator、prepare、backward 三处改动即可让脚本从单卡走向任意分布式环境；
• accelerate config + accelerate launch 是脚本的标准启动方式，notebook_launcher 覆盖 Colab TPU 场景；
• 断点续训、指标聚合、多进程日志等细节均可由 Accelerator 的 save_state/gather/is_local_main_process 等 API 解决，参考仓库中 16 个 no_trainer 系列示例（如 run_glue_no_trainer.py）即可覆盖分类、生成、检测、翻译等任务。
• 完整的 Accelerate 特性文档见其官方文档入口；本仓库中该指南的英文完整版为 docs/source/en/accelerate.md，本文为其阿拉伯语版 docs/source/ar/accelerate.md 所述内容的中文展开。