Accelerate 怎样准备多进程数据加载器

我在把单卡训练脚本改成多进程时，最容易先改错的不是模型，而是 DataLoader：每个进程到底拿一份完整数据、一个批次的哪一段，最后不足一个整批时是否会补样本，都不能靠猜。Accelerate 的推荐入口是把 DataLoader 和模型、优化器一起交给 Accelerator.prepare()，让它返回适合当前分布式环境的对象。

本文只聚焦“怎样准备多进程 DataLoader”：不展开梯度累积，也不把 num_workers 的 CPU 预取并行误当成分布式进程。官方文档入口是 https://huggingface.co/docs/accelerate/，可直接复制打开。

prepare 到底替换了什么

accelerator.prepare(loader) 不会把数据集内容神秘地复制到每个进程，而是返回一个经过分布式适配的 DataLoader。对于常见的 map-style 数据集，核心关系可以理解为：原始采样批次被按进程切分或分配，当前进程的 DataLoader 只产生自己应该消费的部分；如果开启设备放置，批次中的张量还会被放到对应设备。

当模型、优化器和调度器也一起传入时，返回顺序与传入顺序一致。这个细节很小，却能避免把旧的 loader 继续用于循环：

from accelerate import Accelerator

# 让 Accelerator 根据当前启动环境决定设备和进程信息。
accelerator = Accelerator()

# 返回值与传入顺序一致，后续循环必须使用准备后的 training_loader。
model, optimizer, training_loader = accelerator.prepare(
    model, optimizer, training_loader
)

for batch in training_loader:
    inputs, labels = batch
    # 准备后的批次由 Accelerate 负责放到合适设备，不要再次盲目搬运。
    outputs = model(inputs)
    loss = loss_function(outputs, labels)
    # 使用 Accelerate 的反向入口以适配当前分布式后端。
    accelerator.backward(loss)
    optimizer.step()
    optimizer.zero_grad()

如果只准备 DataLoader，也可以调用 accelerator.prepare(training_loader)；官方参考仍推荐优先使用 prepare() 统一准备多个 PyTorch 对象。专门的方法 prepare_data_loader() 适合需要单独处理 loader 的场景，但它不是另一套分片规则。

先决定批次是“全局大小”还是“每进程大小”

我觉得最值得先写在配置旁边的一句话是：split_batches 决定脚本里的 batch_size 如何解释。它为 False 时，各进程通常拿不同的完整批次，所有进程合起来的有效批次会随着进程数增长；它为 True 时，一个原始批次会切成多份，每个进程观察到的批次大小保持不变，但原始批次大小必须能被进程数整除。

用 DataLoaderConfiguration 把意图写清楚

当前稳定文档把 DataLoader 相关选项集中在 DataLoaderConfiguration 中。这样写的好处是，批次语义不再散落在训练循环里；读代码的人可以先看配置，再判断每个进程会拿到什么。

import torch
from torch.utils.data import DataLoader, TensorDataset
from accelerate import Accelerator, DataLoaderConfiguration

# 这里用可整除的批次，方便 split_batches=True 在多进程下切分。
dataset = TensorDataset(torch.randn(32, 8), torch.randint(0, 2, (32,)))
loader = DataLoader(dataset, batch_size=8, shuffle=True, pin_memory=True)

# split_batches=True 表示一个原始批次按进程切片；
# even_batches=True 允许尾部不足整除时补齐，避免不同进程提前结束。
loader_config = DataLoaderConfiguration(
    split_batches=True,
    dispatch_batches=False,
    even_batches=True,
)
accelerator = Accelerator(dataloader_config=loader_config)

# 只替换 loader 引用，训练循环继续消费准备后的对象。
loader = accelerator.prepare(loader)
for inputs, labels in loader:
    # 这里的 inputs 和 labels 已经是当前进程负责的那一份批次。
    pass

这个例子没有把 dispatch_batches 打开，因为它展示的是普通可索引数据集的批次切分。官方参考说明：当 dispatch_batches=True 时，DataLoader 只在主进程迭代，再把批次切分并广播给各进程；对底层是 IterableDataset 的 loader，它默认更倾向于使用这种派发方式。要不要打开，取决于数据读取是否只能由一个入口完成。

尾部样本为什么会影响“每个进程都能结束”

假设数据集长度不能被全局批次整除，多进程 DataLoader 需要决定最后一组数据怎样让所有进程保持一致。even_batches=True 的含义不是“删除多余数据”，而是必要时从数据集开头重复样本，让批次可以平均分给各 worker。它适合训练循环需要所有进程保持相同迭代节奏的情况，但重复样本会改变一次 epoch 的样本组成。

如果业务更在意不重复尾部样本，可以考虑 even_batches=False 或配合 drop_last 重新定义尾批语义，但必须同时确认模型、指标聚合和不同进程的结束条件。对于 Distributed Data Parallel，Accelerate 还提供了处理不等长输入的上下文；那是另一层控制，不应把它和 DataLoader 的普通分片混成一个开关。

启动和最小验证只看三个信号

代码准备好以后，用 Accelerate 的启动方式让它创建多进程：

# 让 Accelerate 按本机配置启动训练脚本。
accelerate launch train.py

# 需要显式指定进程数时，再把配置写在启动参数中；
# 具体参数以当前安装版本的 accelerate launch --help 为准。
accelerate launch --num_processes 2 train.py

最小验证不需要把日志打成一大片，先确认三个信号即可：

1. 对象信号：训练循环使用的是 prepare() 返回的 loader，而不是原始引用。
2. 规模信号：根据 split_batches 判断每进程 batch 和全局 batch，确认它与显存、学习率设定一致。
3. 尾部信号：数据集长度不能整除时，明确记录是补齐、丢弃，还是允许不等长结束。

如果看到每个进程都重复读取完整数据，先检查是否绕过了准备后的 loader；如果只在 IterableDataset 上出现阻塞，再检查 dispatch_batches 和数据源是否允许主进程单点读取。不要先改 num_workers，它解决的是单进程内的数据加载并行，不替代 Accelerate 的跨进程分片。

我会保留的配置判断

• 普通 DataLoader：从 accelerator.prepare(loader) 开始，除非确实需要单独准备。
• 批次语义：先决定 batch_size 是每进程还是全局，再选择 split_batches。
• 流式输入：确认是否需要主进程读取后派发，不要把默认值当成业务保证。
• 不整除数据：even_batches=True 可能重复开头样本；关闭它前先设计尾批和指标逻辑。
• 设备搬运：准备后的批次通常已由 Accelerate 处理设备放置，避免在循环里重复写一套设备分支。

结论

Accelerate 准备多进程 DataLoader 的最小可靠做法，是创建 Accelerator，把 DataLoader（通常连同模型和优化器）交给 prepare()，再围绕 split_batches、dispatch_batches 和 even_batches 明确批次、读取入口与尾部样本策略。真正需要迁移的不是 for batch 这一行，而是“一个批次由谁读取、怎样切给进程、尾部如何收束”这三件事。

相关问题

prepare() 后还要给输入调用 .to(device) 吗？

通常不需要对 DataLoader 产出的张量再写一套固定设备搬运；Accelerate 的 DataLoader 可以负责设备放置。若关闭了设备放置或使用自定义批次结构，应按当前配置确认，而不是机械删除所有设备代码。

split_batches=True 时 batch_size 有什么限制？

原始 DataLoader 的 batch_size 需要能被进程数整除，否则无法把一个批次平均切给各进程。若希望每进程拿完整批次，应重新考虑 split_batches=False 对全局批次的影响。

dispatch_batches=True 适合所有数据集吗？

不适合一概而论。它让主进程负责迭代，再分发批次；流式或只能单点读取的数据源可能更适合，但普通可索引数据集通常可以直接使用分片后的采样策略。