2026实战：用 accelerate 库加速 PyTorch 训练，核心语法与避坑指南

本文是一篇 Python 库用法教程，结合 2026年09月的热门 / 新发布 / 高频实用包，讲清安装、核心语法和能直接抄走的写法。

1. 安装与导入：一条命令搞定多设备环境

安装 accelerate 只需要一条命令，推荐使用 pip 或更快的 uv：

bash 复制代码

pip install accelerate
# 或使用 uv（速度更快，依赖解析更严谨）
uv pip install accelerate

若需要与特定 PyTorch 版本对齐，可加 --extra-index-url 指定 CUDA 轮子源。安装完成后，用 accelerate config 生成配置文件，这是整个库的"设备调度中心"：

bash 复制代码

accelerate config

交互式问答会依次询问：计算环境（单机单卡 / 单机多卡 / 多机多卡）、是否使用 CPU、混合精度模式（no / fp16 / bf16 / fp8）、是否启用 DeepSpeed 或 FSDP。回答后会在 ~/.cache/huggingface/accelerate/default_config.yaml 写入配置。若想跳过问答直接生成，可用 accelerate config default 生成一份默认单卡配置；用 accelerate config --config_file my.yaml 指定路径。

检查配置是否生效：

bash 复制代码

accelerate env

它会打印当前检测到的设备、PyTorch 版本、CUDA 可用性以及默认配置内容，是排查"为什么没走 GPU"的第一入口。

接下来在 Python 中导入并初始化。核心类只有一个：Accelerator。

python 复制代码

from accelerate import Accelerator
import torch

accelerator = Accelerator(
    mixed_precision="fp16",   # 覆盖配置文件，可选 "no"/"fp16"/"bf16"/"fp8"
    gradient_accumulation_steps=4,
    log_with="tensorboard",   # 可选 "wandb"、"all"
)

print("device:", accelerator.device)
print("process index:", accelerator.process_index)
print("num processes:", accelerator.num_processes)
print("mixed precision:", accelerator.mixed_precision)
print("distributed type:", accelerator.distributed_type)

# 准备模型、优化器、数据加载器
model = torch.nn.Linear(10, 2)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
model, optimizer = accelerator.prepare(model, optimizer)

Accelerator 构造参数中，mixed_precision 优先级高于配置文件，临时调试时常用；gradient_accumulation_steps 让 accelerator.backward(loss) 自动累积梯度，无需手写缩放；log_with 指定实验追踪后端。accelerator.device 返回当前进程应使用的设备（如 cuda:0 或 cpu），process_index 与 num_processes 用于区分多进程中的当前 rank。

关键易错点：accelerator.prepare() 会就地包装模型、优化器和 DataLoader，返回的是包装后的对象，必须用返回值覆盖原变量，否则多卡下模型不会同步梯度。另外，混合精度下不要手动调用 loss.half()，Accelerator 已在 backward 中处理；若在 CPU 上设置 fp16，运行时会回退到 no 并给出警告。最后，所有分布式初始化代码都应放在 if __name__ == "__main__": 之下，避免多进程重复执行。

2. 核心对象 Accelerator：初始化与关键参数

安装与导入是第一步，accelerate 的包名与导入名一致：

bash 复制代码

pip install accelerate

python 复制代码

from accelerate import Accelerator

Accelerator 是整个库的中枢，它负责探测当前运行环境（单卡、多卡、TPU、混合精度），并把模型、优化器、数据加载器统一搬运到正确设备上。最简初始化只需一行：

python 复制代码

accelerator = Accelerator()

但真实训练往往要显式控制行为，构造参数才是重点。

关键构造参数

device_placement=True 是默认值。开启后，accelerator.prepare() 会自动把模型和 batch 放到对应设备；若你已有自定义的设备搬运逻辑，可设为 False 手动接管，否则会出现设备重复搬运或张量错位。

mixed_precision 控制混合精度，取值为 "no"、"fp16"、"bf16"。它决定 autocast 的 dtype 以及 GradScaler 是否启用。注意："fp16" 在部分 CPU 或老卡上不可用，"bf16" 需要硬件支持，写错会在首次反向传播时报 dtype 错误。

gradient_accumulation_steps 指定梯度累积步数。设置后，accelerator.accumulate(model) 配合 with 语句使用，只有累积到指定步数才真正 optimizer.step()，等效放大 batch。

split_batches 与 gradient_accumulation_steps 配合使用：当 split_batches=True 时，传入的 batch 会被切分到各进程，而不是每个进程各取一份完整数据。多卡场景下容易混淆，需按数据规模决定。

主进程日志

分布式训练中每个进程都会执行代码，直接 print 会输出 N 份。accelerator.print 只在主进程打印：

python 复制代码

from accelerate import Accelerator

accelerator = Accelerator(
    device_placement=True,
    mixed_precision="bf16",
    gradient_accumulation_steps=4,
    split_batches=False,
)

accelerator.print(f"进程数: {accelerator.num_processes}")
accelerator.print(f"当前设备: {accelerator.device}")
accelerator.print(f"是否主进程: {accelerator.is_main_process}")

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

for step, batch in enumerate(dataloader):
    with accelerator.accumulate(model):
        outputs = model(**batch)
        loss = outputs.loss
        accelerator.backward(loss)
        optimizer.step()
        optimizer.zero_grad()

accelerator.backward(loss) 替代 loss.backward()，它内部处理了混合精度的梯度缩放。accelerator.accumulate(model) 返回上下文管理器，控制累积节奏。accelerator.print 的签名与内置 print 一致，支持 sep、end 等参数。

常见坑

一是 mixed_precision 传布尔值而非字符串，会静默失效；二是忘记用 accelerator.backward，导致 fp16 下梯度未反缩放；三是 prepare 必须在创建优化器之后调用，顺序颠倒会让参数设备不一致。初始化阶段把这些参数定清楚，后续训练循环才能稳定运行。

3. 常用 API：准备模型、优化器、数据加载器

Accelerator 实例创建之后，真正的"包装"动作几乎都集中在 prepare 一个方法里。它的签名是 prepare(*args, device_placement=True, evaluation_mode=False)，传入的对象可以是 torch.nn.Module、torch.optim.Optimizer、torch.utils.data.DataLoader 或学习率调度器，返回的是被包装后的同类型对象，因此必须接收返回值，不能原地修改后就丢弃。

python 复制代码

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

accelerator = Accelerator()
device = accelerator.device

model = torch.nn.Linear(10, 1)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1)
loader = DataLoader(TensorDataset(torch.randn(64, 10), torch.randn(64, 1)),
                    batch_size=8, shuffle=True)

model, optimizer, loader, scheduler = accelerator.prepare(
    model, optimizer, loader, scheduler
)

for x, y in loader:
    optimizer.zero_grad()
    loss = torch.nn.functional.mse_loss(model(x), y)
    accelerator.backward(loss)   # 不能用 loss.backward()
    optimizer.step()
scheduler.step()

关键点：prepare 返回的顺序与传入顺序一致。model 会被移动到 accelerator.device 并自动套上 DistributedDataParallel；optimizer 会被替换为支持混合精度的封装；loader 会换成分布式采样器，保证各进程分到不同数据；scheduler 则被适配到正确的步进节奏。若只传一部分对象，未包装的模型不会自动上设备，容易报设备不一致错误。

反向传播必须换成 accelerator.backward(loss)。它内部会按混合精度状态对 loss 做缩放，并在多卡下处理梯度同步。直接写 loss.backward() 在 fp16 下会产生梯度溢出，且绕过了 accelerate 的梯度规约逻辑，是新手最常见的坑。

跨设备收集张量用 accelerator.gather(tensor) 或 accelerator.gather_for_metrics(tensor)。前者把所有进程上的张量按第一维拼接后返回，后者在分布式评估时会自动去重，避免 DistributedSampler 补齐导致的重复样本。

python 复制代码

preds = model(x)
all_preds = accelerator.gather_for_metrics(preds)
if accelerator.is_main_process:
    print(all_preds.shape)

注意 gather 要求各进程张量形状一致，否则需先 pad_across_processes；另外只有主进程才应执行打印、写文件等操作，用 accelerator.is_main_process 判断。

4. 完整小例子：单卡与多卡通用的训练脚本

下面这个脚本把设备放置、混合精度、梯度累积全部交给 Accelerator，单卡直接 python train.py，多卡用 accelerate launch train.py，代码完全不用改。

python 复制代码

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

def main():
    accelerator = Accelerator(
        mixed_precision="fp16",          # 单卡无 GPU 时自动降级为 no
        gradient_accumulation_steps=4,   # 累积 4 步再更新
        log_with=None,
    )

    # 1. 构造一个极简数据集：128 个样本，10 维特征，3 分类
    x = torch.randn(128, 10)
    y = torch.randint(0, 3, (128,))
    loader = DataLoader(TensorDataset(x, y), batch_size=16, shuffle=True)

    # 2. 模型与优化器
    model = torch.nn.Sequential(
        torch.nn.Linear(10, 32), torch.nn.ReLU(), torch.nn.Linear(32, 3)
    )
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
    criterion = torch.nn.CrossEntropyLoss()

    # 3. 关键一步：prepare 负责设备放置、DDP 包装、AMP、DataLoader 分片
    model, optimizer, loader = accelerator.prepare(model, optimizer, loader)

    model.train()
    for epoch in range(3):
        for step, (bx, by) in enumerate(loader):
            # 4. 梯度累积上下文：内部自动 no_sync 与缩放
            with accelerator.accumulate(model):
                logits = model(bx)
                loss = criterion(logits, by)
                accelerator.backward(loss)     # 替代 loss.backward()
                optimizer.step()
                optimizer.zero_grad()

            # 5. 只在累积边界和主进程打印
            if accelerator.sync_gradients and accelerator.is_main_process:
                if step % 10 == 0:
                    accelerator.print(
                        f"epoch={epoch} step={step} loss={loss.item():.4f}"
                    )

    # 6. 保存：unwrap_model 去掉 DDP 包装再存
    accelerator.wait_for_everyone()
    if accelerator.is_main_process:
        unwrapped = accelerator.unwrap_model(model)
        torch.save(unwrapped.state_dict(), "ckpt.pt")

if __name__ == "__main__":
    main()

关键语法逐条说明：

Accelerator(mixed_precision="fp16", gradient_accumulation_steps=4)：两个最常用参数。mixed_precision 可选 "no"、"fp16"、"bf16"；gradient_accumulation_steps 决定多少个小 batch 累积成一次等效大 batch。
accelerator.prepare(...)：返回值顺序必须与传入顺序一致，常见错误是写成 loader, model, optimizer = ...，会静默错位。它会把模型搬到正确设备、用 DDP 包装、把优化器换成 AMP 版本、并给 DataLoader 注入分布式采样器。
with accelerator.accumulate(model):：只有把前向、backward、step、zero_grad 全部放进这个块内，梯度累积才正确；漏掉 step 会导致累积失效。
accelerator.backward(loss)：替代 loss.backward()，负责混合精度下的 loss 缩放。直接调用原生 backward 在 fp16 下容易梯度溢出。
accelerator.sync_gradients：布尔值，仅在真正执行优化器更新的那一步为 True，用它控制日志频率可避免多进程重复打印。
accelerator.is_main_process 与 accelerator.print：多卡时只有主进程输出，防止日志刷屏。
accelerator.unwrap_model(model)：保存权重前必须调用，否则 state_dict 的 key 会带上 module. 前缀，加载时报错。

易错点：不要在 prepare 之前把模型 .to(device)，prepare 会自己处理；optimizer.zero_grad() 要放在 step() 之后且仍在 accumulate 块内；多卡运行务必用 accelerate launch，直接 torchrun 会绕过配置。

5. 进阶写法：混合精度、梯度累积与分布式推理

混合精度：fp16 还是 bf16

Accelerator(mixed_precision=...) 接受 "no"、"fp16"、"bf16" 三种值。bf16 动态范围与 fp32 相同，几乎不会溢出，在 Ampere 及更新的 GPU 上优先选它；fp16 精度更高但需要配合 GradScaler，accelerate 会自动帮你挂上，无需手动写 torch.cuda.amp.GradScaler。判断硬件是否支持：

python 复制代码

import torch
from accelerate import Accelerator

use_bf16 = torch.cuda.is_bf16_supported()
accelerator = Accelerator(mixed_precision="bf16" if use_bf16 else "fp16")
print(accelerator.mixed_precision, accelerator.device)

注意：CPU 上 bf16 支持有限，mixed_precision 若设为 fp16 在纯 CPU 环境会退化为 no。若你用 accelerate config 生成配置文件，命令行参数会覆盖默认值，但显式传入构造函数的参数优先级最高。

梯度累积与 accumulate 上下文

gradient_accumulation_steps=N 让 accelerator 每 N 步才真正执行 optimizer.step() 并清零梯度。关键在于：不要在循环里手动判断步数 ，用 accelerator.accumulate(model) 上下文管理器，它会自动处理 no_sync、梯度缩放和 step 时机：

python 复制代码

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

accelerator = Accelerator(mixed_precision="bf16", gradient_accumulation_steps=4)
model = torch.nn.Linear(16, 4)
opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
loader = DataLoader(TensorDataset(torch.randn(64, 16), torch.randn(64, 4)), batch_size=8)
model, opt, loader = accelerator.prepare(model, opt, loader)

for epoch in range(2):
    for x, y in loader:
        with accelerator.accumulate(model):
            loss = torch.nn.functional.mse_loss(model(x), y)
            accelerator.backward(loss)
            opt.step()
            opt.zero_grad()

accelerator.accumulate 内部会在累积期间调用 model.no_sync()（分布式下跳过梯度 all-reduce），并在最后一步同步。常见错误是自己在外面写 if step % N == 0: opt.step()，这会导致 DDP 下梯度通信次数错乱。另外 accelerator.backward(loss) 已替代 loss.backward()，fp16 时它会自动做 loss scaling。

unwrap_model 保存与分布式推理

accelerator.prepare 会把模型包成 DistributedDataParallel 或 DataParallel，直接 torch.save(model.state_dict()) 会带上 module. 前缀。正确做法是用 accelerator.unwrap_model(model) 取回原始模型，再配合 accelerator.save 或 accelerator.wait_for_everyone()：

python 复制代码

accelerator.wait_for_everyone()
unwrapped = accelerator.unwrap_model(model)
if accelerator.is_main_process:
    accelerator.save(unwrapped.state_dict(), "model.pt")

推理阶段同样用 prepare 包装模型和 dataloader，accelerator 会自动切分数据到各进程。若只需主进程输出，用 accelerator.is_main_process 判断；若想合并所有进程结果，用 accelerator.gather_for_metrics(preds)，它会自动处理各进程张量形状不一致的边界情况。注意 unwrap_model 只去掉包装层，不改变权重，所以保存的 state_dict 可以直接被原始模型 load_state_dict 加载。

6. 注意事项：常见坑与避坑指南

反向传播必须走 accelerator.backward

最常见也最隐蔽的错误，是仍然调用 loss.backward()。在单卡下它未必立刻报错，但一旦启用混合精度或 DeepSpeed，梯度缩放器不会被触发，训练直接失效。

python 复制代码

from accelerate import Accelerator
import torch

accelerator = Accelerator(mixed_precision="fp16")
model = torch.nn.Linear(10, 2)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
model, optimizer, loader = accelerator.prepare(model, optimizer, loader)

for x, y in loader:
    out = model(x)
    loss = torch.nn.functional.cross_entropy(out, y)
    accelerator.backward(loss)   # 不是 loss.backward()
    optimizer.step()
    optimizer.zero_grad()

accelerator.backward(loss) 内部会根据 mixed_precision 调用对应的 GradScaler.scale，并处理梯度裁剪，返回值为 None。若你自定义了梯度裁剪，应使用 accelerator.clip_grad_norm_(model.parameters(), max_norm)，而不是 torch.nn.utils.clip_grad_norm_。

prepare 之前不要手动 to(device)

accelerator.prepare() 会自动把模型和数据放到正确的设备上。如果你提前写了 model.to(accelerator.device)，多进程下反而可能把模型放到错误设备或重复搬运。

python 复制代码

# 错误：prepare 之前手动搬设备
model = model.to(accelerator.device)
model, optimizer, loader = accelerator.prepare(model, optimizer, loader)

# 正确：交给 prepare
model, optimizer, loader = accelerator.prepare(model, optimizer, loader)

如果确实需要查询设备，用 accelerator.device 只读属性即可，不要用它做搬运。

多进程下日志重复打印

accelerator.print() 只在主进程输出，而 print() 会在每个进程各打印一次，日志被放大 N 倍。同理，tqdm 需用 accelerator.is_main_process 包裹，或直接使用 accelerator.gather_for_metrics 汇总指标。

python 复制代码

if accelerator.is_main_process:
    print(f"epoch {epoch} loss {loss.item():.4f}")
accelerator.print("only on main process")

保存模型要 unwrap

prepare() 会把模型包进 DistributedDataParallel 或 DeepSpeedEngine，直接 torch.save(model.state_dict(), ...) 会保存带 module. 前缀的键，加载时对不上。正确做法是 accelerator.unwrap_model(model) 后再取 state_dict，并用 accelerator.save() 落盘。

python 复制代码

unwrapped = accelerator.unwrap_model(model)
accelerator.save(unwrapped.state_dict(), "ckpt.pt")

随机种子必须统一

多进程下每个进程的随机状态独立，数据增强、dropout 会不一致。应在 prepare() 之后调用 accelerator.wait_for_everyone() 再设置种子，或直接使用 Accelerator(..., seed=42) 让库内部统一处理。

python 复制代码

accelerator = Accelerator(seed=42)
# 或者
accelerator.wait_for_everyone()
torch.manual_seed(42)

最后一点：accelerator.gather() 返回的张量在所有进程上形状一致，但 gather_for_metrics() 会去掉 padding 产生的重复样本，计算准确率时优先选后者，否则指标会被拉低。

7. 适用场景：什么时候该用 accelerate

四种典型适用场景

第一种是单卡代码迁移到多卡。你写好的训练循环里只有 model、optimizer、dataloader 和 loss.backward()，不想重构成 DistributedDataParallel 那一套 init_process_group、DistributedSampler、local_rank 传参。accelerate 只需三处改动就能跑多卡：

python 复制代码

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

accelerator = Accelerator(mixed_precision="fp16")
model = torch.nn.Linear(10, 1)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
loader = DataLoader(TensorDataset(torch.randn(64, 10), torch.randn(64, 1)), batch_size=8)

model, optimizer, loader = accelerator.prepare(model, optimizer, loader)

for x, y in loader:
    optimizer.zero_grad()
    loss = torch.nn.functional.mse_loss(model(x), y)
    accelerator.backward(loss)   # 替代 loss.backward()
    optimizer.step()

accelerator.save_state("ckpt")   # 只保存主进程权重

Accelerator 构造时传入 mixed_precision="fp16" 或 "bf16"，prepare() 返回的对象会自动包装模型、优化器和 dataloader，backward() 内部处理梯度缩放与跨卡同步。注意 accelerator.backward(loss) 不能换成 loss.backward()，否则混合精度下梯度会出错。

第二种是混合精度训练。原生 torch.cuda.amp 需要手动维护 GradScaler、autocast 上下文和 unscale 逻辑，accelerate 把这三件事收进 Accelerator 和 backward() 里，切换精度只改一个字符串参数，CPU 上还能自动降级为 fp32 继续跑。

第三种是同一套代码跨硬件复用。本地 1 张卡、服务器 8 张卡、TPU 或纯 CPU 调试，代码不用动，改的是启动方式：

bash 复制代码

accelerate config      # 交互式生成 ~/.cache/huggingface/accelerate/default_config.yaml
accelerate launch train.py

accelerate config 会询问设备类型、卡数、是否混合精度，写入 YAML；accelerate launch 读取该文件并设置好环境变量。Accelerator() 无参构造时就从这份配置读取，所以代码里不必写死 device。

与原生 DDP、Lightning 的取舍

原生 DDP 的写法是 torch.distributed.init_process_group、DistributedSampler、torchrun --nproc_per_node，控制粒度最细，但每个脚本都要重复这些样板代码，且混合精度要自己接 GradScaler。accelerate 的定位是薄封装：不接管训练循环，只替换设备放置、梯度同步和精度这三个横切关注点，所以已有训练脚本改造成本最低。

PyTorch Lightning 走的是另一条路，它要求把逻辑收进 LightningModule 的 training_step、configure_optimizers，再由 Trainer 驱动。好处是 checkpoint、日志、回调、多卡策略都成体系，代价是必须按它的框架重写代码。判断标准很直接：如果你已经有成熟的自定义训练循环、只想加多卡和混合精度，用 accelerate；如果是从零开始、愿意接受框架约束，Lightning 的配套更全。accelerate 与 transformers 的 Trainer 深度集成，这也是它在微调场景里被大量使用的原因。

8. 总结与下一步学习建议

bash 复制代码

pip install accelerate
accelerate config          # 交互式生成 ~/.cache/huggingface/accelerate/default_config.yaml
accelerate launch train.py # 按配置启动多卡/混合精度

回顾一下，accelerate 的核心就三层：Accelerator 负责设备放置与混合精度，prepare() 把模型、优化器、数据加载器、调度器一次性接管，gather_for_metrics / reduce / main_process_first 处理跨进程聚合与主进程逻辑。训练循环里只需 with accelerator.accumulate(model): 包住 backward、step、zero_grad，梯度累积、AMP、多卡同步都会自动处理；保存用 accelerator.save_state() / load_state() 而不是裸 torch.save，否则随机种子、优化器状态、数据进度会错位。下面这段是最小可运行骨架，可直接对照自己的脚本改写：

python 复制代码

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

accelerator = Accelerator(mixed_precision="fp16", gradient_accumulation_steps=4)
model = torch.nn.Linear(16, 1)
opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
loader = DataLoader(TensorDataset(torch.randn(256, 16), torch.randn(256, 1)), batch_size=32)

model, opt, loader = accelerator.prepare(model, opt, loader)

for epoch in range(3):
    for x, y in loader:
        with accelerator.accumulate(model):
            loss = torch.nn.functional.mse_loss(model(x), y)
            accelerator.backward(loss)
            opt.step()
            opt.zero_grad()
    accelerator.save_state(f"ckpt/epoch_{epoch}")

accelerator.prepare 返回的对象顺序必须与传入一致；accumulate 在梯度累积步数未到时会自动跳过 step，所以不要在循环外手动判断。save_state 的路径在分布式下每个进程都会写，通常配合 accelerator.wait_for_everyone() 使用。

下一步建议按三条线深入。第一，读官方文档与示例：accelerate 仓库的 examples/ 目录包含 complete_cv_example.py、complete_nlp_example.py，以及 accelerate launch --multi_gpu、--mixed_precision 等 CLI 参数说明，docs/source/usage_guides/ 下的 fsdp.md、deepspeed.md 是最权威的集成入口。第二，结合 transformers.Trainer：当 TrainingArguments 检测到 accelerate 配置后会自动启用分布式与 AMP，你只需 accelerate launch train.py，无需改一行训练代码；也可以把上面的骨架封装成自己的 Trainer 类，把 prepare、backward、save_state 抽成方法。第三，探索 FSDP 与 DeepSpeed：accelerate config 中选择 FSDP 或 DEEPSPEED，再用 accelerator.state.deepspeed_plugin 读取 zero_stage、offload_optimizer_device 等参数，即可把大模型分片到多卡。常见坑是 DeepSpeed 与 torch.compile 同时开启时的兼容问题，以及 FSDP 下 save_state 需要 FullyShardedDataParallel 的 state_dict 类型，务必先用小模型跑通再上生产。

9. 从 accelerate 配置到生产服务：与 FastAPI、websockets 的组合用法

训练脚本跑通只是第一步，真正落地时往往需要把 accelerate 管理的模型暴露成服务。这里最容易踩的坑是：accelerate launch 启动的是多进程，而 FastAPI 默认是单进程 uvicorn，两者直接混写会出现在每个 rank 上各起一个 Web 服务的尴尬局面。正确做法是把模型推理逻辑封装成独立函数，让 Accelerator 只在主进程生效，再用 FastAPI 做一层薄封装。

python 复制代码

import torch
from fastapi import FastAPI
from pydantic import BaseModel
from accelerate import Accelerator

app = FastAPI()
accelerator = Accelerator()
model = torch.nn.Linear(8, 2)
model = accelerator.prepare(model)
model.eval()

class Payload(BaseModel):
    features: list[float]

@app.post("/predict")
@torch.no_grad()
def predict(payload: Payload):
    x = torch.tensor([payload.features], device=accelerator.device)
    logits = model(x)
    # 多卡时只在主进程返回，避免重复响应
    if not accelerator.is_main_process:
        return {"skipped": True}
    return {"logits": logits.cpu().tolist()}

如果需要流式输出 token，可以用 websockets 替代 HTTP 轮询。accelerate 的 gather_for_metrics 在这里很有用，它能把各 rank 的生成结果汇总到主进程再推送：

python 复制代码

import asyncio, json, websockets
from accelerate import Accelerator
from accelerate.utils import gather_object

accelerator = Accelerator()

async def handler(ws):
    async for msg in ws:
        prompt = json.loads(msg)["prompt"]
        local = generate(prompt)          # 每个 rank 各自生成
        merged = gather_object([local])   # 汇总到主进程
        if accelerator.is_main_process:
            await ws.send(json.dumps({"tokens": merged}))

async def main():
    async with websockets.serve(handler, "0.0.0.0", 8765):
        await asyncio.Future()

if __name__ == "__main__":
    asyncio.run(main())

注意 gather_object 要求所有进程都调用，否则会死锁；这一点和 gather_for_metrics 的语义一致，但后者会自动去重 padding，更适合评测场景。

10. 配置与排错：用 tomlkit 管理 accelerate 配置，用 jsonschema 校验参数

accelerate config 生成的是 YAML，但很多团队更习惯 TOML，因为它在手写和 diff 时更清晰。tomlkit 的优势是保留注释和格式，适合做「程序化修改但不破坏可读性」的场景。下面这段代码读取 accelerate 的配置文件，把 mixed_precision 改成 bf16 并写回，注释不会丢：

python 复制代码

import tomlkit
from pathlib import Path

cfg_path = Path.home() / ".cache/huggingface/accelerate/default_config.yaml"
# 这里演示 TOML 侧，真实项目可维护一份 train.toml
doc = tomlkit.parse(Path("train.toml").read_text())
doc["mixed_precision"] = "bf16"
doc["gradient_accumulation_steps"] = 4
Path("train.toml").write_text(tomlkit.dumps(doc))

参数校验则推荐用 jsonschema 配合 jsonschema-specifications，后者提供了 2020-12 等最新规范，避免旧版 draft 校验不通过。把训练超参当成一份 JSON Schema，启动前先校验，能拦掉大量「跑了两小时才发现 lr 传成字符串」的事故：

python 复制代码

from jsonschema import validate
from jsonschema_specifications import REGISTRY
from referencing import Registry

schema = {
    "$schema": "https://json-schema.org/draft/2020-12/schema",
    "type": "object",
    "properties": {
        "learning_rate": {"type": "number", "exclusiveMinimum": 0},
        "mixed_precision": {"enum": ["no", "fp16", "bf16"]},
        "gradient_accumulation_steps": {"type": "integer", "minimum": 1},
    },
    "required": ["learning_rate", "mixed_precision"],
}

registry = Registry().with_contents(REGISTRY)
validate(instance={"learning_rate": 2e-5, "mixed_precision": "bf16"},
         schema=schema, registry=registry)

最后提一个高频报错：ValueError: You can't use gradient_accumulation_steps with ...。多数情况是 accelerate config 里开了 accumulation，代码里又手动调了 optimizer.step()。排错顺序建议是------先打印 accelerator.gradient_accumulation_steps，再确认 accelerator.accumulate(model) 上下文是否包住了整个 forward/backward，最后检查是否在 accelerator.prepare 之前就创建了 scheduler。这三步能覆盖九成以上的 accumulation 相关问题。