Accelerate mixed precision 训练中统一设备与精度的配置

我在把一段单卡 PyTorch 训练脚本交给 Accelerate 管理时，最容易出问题的不是启动命令，而是“谁负责设备、谁负责精度”没有统一：模型手动放到了 CUDA，DataLoader 仍按原方式出 batch，autocast 又在另一个地方单独配置。结果就是代码能启动，但换设备或换启动方式后，配置很难解释。

更稳妥的做法是只保留一个 Accelerator 实例作为配置入口。让 prepare 接管模型、优化器、DataLoader 和调度器的设备准备；让 mixed_precision 明确选择 no、fp16、bf16 或 fp8；训练计算放进 accelerator.autocast()，反向传播交给 accelerator.backward()。

官方地址：https://huggingface.co/docs/accelerate/

如果目标是让设备和精度在单卡、多卡或不同启动方式之间保持一致，优先把配置收口到 Accelerator(mixed_precision=...)，并在创建训练对象后统一调用一次 prepare。不要在后续循环里继续散落固定的 .to("cuda")。

先确定唯一的设备与精度入口

Accelerator 的 mixed_precision 参数是最直观的入口。当前官方参考文档列出的选项包括 no、fp16、bf16 和 fp8；如果代码里没有显式传入，它还会读取 ACCELERATE_MIXED_PRECISION，并与本机 Accelerate 配置或启动参数协同工作。

设备也不应由业务代码猜测。创建实例后用 accelerator.device 作为需要手动创建张量时的设备来源，模型、优化器和 DataLoader 则交给 prepare。这样，训练脚本中的“设备选择”和“精度选择”是同一个运行时对象的属性。

用 prepare、autocast 和 backward 组成训练骨架

下面这段骨架只保留和设备、精度有关的部分。示例使用 bf16，实际项目可以按硬件支持改成 fp16 或 no。代码中的中文注释说明配置边界；它是文章示例，不代表本文在本机执行过。

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


def build_loader():
    # 用小型张量数据集说明 DataLoader 也交给 Accelerator 准备
    features = torch.randn(64, 8)
    labels = torch.randint(0, 2, (64,))
    return DataLoader(TensorDataset(features, labels), batch_size=8, shuffle=True)


def main():
    # 精度在这里集中配置；换成 "no" 就是全精度路径
    accelerator = Accelerator(mixed_precision="bf16")
    model = nn.Sequential(nn.Linear(8, 16), nn.ReLU(), nn.Linear(16, 2))
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
    scheduler = torch.optim.lr_scheduler.LinearLR(optimizer, total_iters=10)
    train_loader = build_loader()

    # prepare 会按当前运行方式准备模型、优化器、数据和调度器
    model, optimizer, train_loader, scheduler = accelerator.prepare(
        model, optimizer, train_loader, scheduler
    )
    accelerator.print(
        f"device={accelerator.device}, mixed_precision={accelerator.mixed_precision}"
    )

    for features, labels in train_loader:
        optimizer.zero_grad()
        # autocast 只包住需要混合精度的前向和损失计算
        with accelerator.autocast():
            logits = model(features)
            loss = nn.functional.cross_entropy(logits, labels)
        # backward 负责按当前精度和分布式配置执行反向传播
        accelerator.backward(loss)
        optimizer.step()
        scheduler.step()


if __name__ == "__main__":
    main()

这里有三个容易被忽略的点。第一，prepare 的返回值要覆盖原对象，后面的循环使用的是准备后的 DataLoader 和模型。第二，autocast 只包住前向与损失计算，优化器更新仍然放在上下文外。第三，反向传播使用 accelerator.backward，不要同时再调用一次 loss.backward()。

用三种方式切换 mixed precision

为了避免代码、环境和启动命令互相打架，我通常只选一种作为当前项目的主入口，其他方式只用于覆盖或临时实验。

1. 代码显式指定：Accelerator(mixed_precision="bf16")，适合把训练脚本的行为写得最清楚。
2. 环境变量指定：在启动前设置 ACCELERATE_MIXED_PRECISION=fp16，适合容器或任务平台注入配置。
3. 命令行指定：先用 accelerate config 保存本机默认配置，运行时再用 accelerate launch --mixed_precision bf16 train.py 覆盖。

# 先生成或更新本机 Accelerate 配置
accelerate config

# 运行时显式选择 BF16，避免依赖旧的默认配置
accelerate launch --mixed_precision bf16 train.py

# 容器任务也可以从环境变量注入 FP16
ACCELERATE_MIXED_PRECISION=fp16 accelerate launch train.py

如果代码已经明确传入 mixed_precision，就不要再把另一个值当成“隐式真相”。工程上可以保留启动参数，但应在日志里打印 accelerator.mixed_precision，让最终生效的模式可见。

设备不一致时按对象边界排查

最常见的症状是 Expected all tensors to be on the same device。先检查是否在 prepare 之后又把模型或输入强行放到固定设备，再检查是否在循环里手动创建了没有设备参数的新张量。

# 手动创建和输入同形状的张量时，沿用 Accelerator 的设备
mask = torch.ones(features.shape[0], device=accelerator.device)

# 不要把准备后的对象重新固定到某一张 GPU
# model = model.to("cuda:0")  # 会破坏多进程或 CPU 回退的统一入口

如果 batch 来自已经准备好的 DataLoader，通常不需要再写一层 features.to("cuda")。只有脱离 DataLoader、在业务逻辑里临时创建的张量，才需要使用 accelerator.device，并且要确认它和模型当前所在设备一致。

精度没有生效时检查这四件事

官方参考还提供了 optimizer_step_was_skipped 属性，可用于识别混合精度下因梯度溢出而跳过的优化器更新。这个信号比只看 loss 是否下降更适合定位“训练在跑但参数没有更新”的情况。

评估、保存和回退的边界

训练阶段启用混合精度，不等于所有计算都必须降精度。评估中对数值稳定性敏感的部分可以在 accelerator.autocast() 外执行，或给 autocast 传入专门的处理器覆盖设置。保存模型时，按项目需要使用 accelerator.unwrap_model(model) 取得未包装模型，并在多进程场景下让主进程负责写文件。

如果项目只是单卡推理、没有混合精度和分布式需求，就没有必要为了“统一”而强行把每一个对象都交给 Accelerate。相反，只要训练脚本需要切换 CPU/GPU、多卡或精度模式，集中到 Accelerator 的收益才明显。

总结：把配置收口，排查就有边界

这次配置的重点不是记住某个启动参数，而是建立单一职责：Accelerator 决定设备和精度，prepare 准备训练对象，autocast 包住前向计算，backward 处理反向传播。出现问题时，先看最终生效的 accelerator.device 和 accelerator.mixed_precision，再查手动设备迁移与硬件支持，通常比逐行猜测更快。

常见追问

• Accelerate 的 mixed precision 应该写在代码还是命令行？项目默认行为建议固定在代码或配置文件，临时实验再用 launch 参数覆盖。
• 用了 prepare 还需要手动把 batch 放到 GPU 吗？由 prepare 处理的 DataLoader 通常不需要；临时新建的张量使用 accelerator.device。
• fp16 和 bf16 可以只改参数名切换吗？可以切换配置值，但仍要确认当前硬件、PyTorch 和模型数值稳定性是否支持。
• 为什么训练代码里要用 accelerator.backward？它会按当前的混合精度与分布式配置执行正确的反向传播路径。