Accelerate 大模型推理怎么自动分配到多块设备

模型参数总量超过单块显卡显存时，Accelerate 的常用解法是让 device_map="auto" 计算一份设备映射。它会先尽量使用可用 GPU，再把剩余层放到 CPU，必要时才回退到磁盘；加载过程中配合空权重骨架，可以避免先在内存中完整实例化一份模型。

官方文档：https://huggingface.co/docs/accelerate/usage_guides/big_modeling

自动分配到底做了什么

在普通加载方式里，模型对象和权重往往会先占用一整块主机内存，然后再迁移到 GPU。Accelerate 的大模型推理路径先用 init_empty_weights() 创建只有结构、没有实际参数的骨架，再由 load_checkpoint_and_dispatch() 一边读取 checkpoint，一边把层放到计算设备。

默认策略的顺序可以概括为“高速设备优先，慢速设备兜底”：先填 GPU，空间不足时使用 CPU，仍放不下才使用磁盘。它按模块和层组织映射，不是把一个张量随机切片。某些带残差连接的模块不适合跨设备拆分，就要通过 no_split_module_classes 保持完整。

最小写法：让 Transformers 自动计算 device map

如果模型来自 Transformers，通常不需要手动调用底层 dispatch 函数，直接在 from_pretrained 中传入 device_map="auto" 即可。低精度加载可以减少权重占用，但它只改变存储类型，不会自动解决算子、输入和显存碎片的所有问题。

from transformers import AutoModelForCausalLM, AutoTokenizer

# 用较低精度加载权重，具体 dtype 要和硬件及模型支持情况匹配
model = AutoModelForCausalLM.from_pretrained(
    "your-org/your-model",
    device_map="auto",
    torch_dtype="auto",
)
tokenizer = AutoTokenizer.from_pretrained("your-org/your-model")

# 查看 Accelerate 为模块计算出的设备映射，先确认是否出现 cpu 或 disk
device_map = getattr(model, "hf_device_map", {})
print(device_map)

# 输入通常送到首个执行设备；不要盲目写死 cuda:0
inputs = tokenizer("请给出一句简短的问候。", return_tensors="pt")
first_device = next(model.parameters()).device
inputs = {name: value.to(first_device) for name, value in inputs.items()}
outputs = model.generate(**inputs, max_new_tokens=32)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

这里的映射应先当作“能否装下”的结果来读：如果出现 cpu 或 disk，说明权重已经超出当前 GPU 预算，推理时会发生跨介质搬运。它并不表示已经获得最佳延迟，也不代表多块 GPU 会同时满载。

用预算和不可拆分规则控制结果

多卡显存大小不一致、驱动预留不同或某个大层无法拆分时，完全交给默认探测容易得到不稳定的落点。可以用 max_memory 给设备设置上限，并把模型中的关键层类加入不可拆分列表。

import torch
from transformers import AutoModelForCausalLM

# 显式给两块 GPU 和 CPU 留出预算，避免把全部显存都交给权重
max_memory = {
    0: "20GiB",
    1: "20GiB",
    "cpu": "64GiB",
}

# 残差连接模块需要整体留在一个设备，类名必须按实际模型结构填写
model = AutoModelForCausalLM.from_pretrained(
    "your-org/your-model",
    device_map="auto",
    max_memory=max_memory,
    no_split_module_classes=["GPTJBlock"],
    torch_dtype=torch.float16,
)

# 只有映射包含 disk 时才准备磁盘 offload 目录，并确保目录可写
print(getattr(model, "hf_device_map", {}))

如果映射中确实包含 disk，底层加载接口还需要提供 offload_folder，让权重以磁盘方式暂存。磁盘回退能扩大可加载模型的范围，却会明显增加访问延迟；生产环境应先确认本地 SSD、缓存目录和并发访问是否能承受。

能运行不等于吞吐最优

Accelerate 的 Big Model Inference 在多 GPU 场景下更接近模型并行：某一时刻可能只有一个 GPU 执行当前层，下一层再把结果传给另一个设备。这样做的价值是突破单卡容量限制，而不是自动把多卡变成并行流水线。若目标是高吞吐，应继续比较 pipeline parallelism、张量并行或专用推理服务。

排查时先打印设备映射，再观察显存占用和单次生成耗时。不要只看到“没有 CUDA out of memory”就判断配置成功：映射是否符合预期、输入是否落在首个执行设备、CPU 或磁盘是否成为瓶颈，才决定这份自动分配是否值得保留。

常见问题

device_map="auto" 会让所有 GPU 同时工作吗？

不会保证这一点。它主要把模型层分布到可用设备，Big Model Inference 可能按层串行执行。想提高吞吐，需要单独设计流水线或其他并行策略。

为什么明明有两块 GPU，结果仍然出现 CPU？

自动映射按可用预算和不可拆分层计算，模型总权重、临时激活或保留空间超过两块 GPU 的可用容量时，就会把部分层放到 CPU。先看 hf_device_map 和 max_memory，再决定是否换 dtype 或缩小输入。

什么时候应该手写 device map？

当 GPU 显存不对称、关键层必须相邻、或默认映射把大量计算推到 CPU 时，才值得从自动映射改为显式映射。先用自动结果定位瓶颈，再做小范围调整，通常比一开始手写整张映射更稳。