Accelerate device_map 怎么把大模型分配到多设备

当一个模型的权重超过单张 GPU 显存时，Accelerate 的 device_map 可以把不同模块放到多张 GPU，并在需要时把一部分权重放到 CPU 或磁盘。最容易误解的一点是：这首先是“模型如何落位”的问题，不是把一次推理自动变成多卡同时计算。

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

先把“多设备”理解成权重分片，而不是多卡并行

device_map 描述的是模型模块到设备的映射，例如某些 Transformer block 在 cuda:0，后续 block 在 cuda:1，也可以有一部分在 cpu。Accelerate 会在模块之间安装对齐设备的处理逻辑，让中间张量随着模型层移动。

这种方式适合“模型放不进单卡，但希望先完成推理”的场景。官方 Big Model Inference 文档明确把它定位为推理加载能力；一次前向计算通常仍按层向前推进，跨 GPU 的传输和 CPU/磁盘卸载都会带来额外开销。如果目标是训练，应该另外评估数据并行、张量并行或流水线并行方案。

用 device_map="auto" 让权重按 GPU、CPU 和磁盘落位

最小方案是把 device_map="auto" 传给 Transformers 的 from_pretrained。如果使用 Accelerate 的底层加载接口，也可以在空模型骨架上调用 load_checkpoint_and_dispatch。前者适合常见 Transformers 模型，后者适合自己控制 checkpoint 和空权重初始化的项目。

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

model_id = "your-org/your-causal-model"

# auto 会根据可见显存决定层的落位，必要时才向 CPU 或磁盘卸载。
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="auto",
    torch_dtype=torch.float16,
    low_cpu_mem_usage=True,
)
tokenizer = AutoTokenizer.from_pretrained(model_id)

# 先查看实际分配，避免把“加载成功”误认为“全部在 GPU”。
print(model.hf_device_map)

text = "请用一句话解释 device_map 的作用。"
# 输入应送到模型首个可执行设备，而不是盲目写死 cuda:0。
input_device = next(model.parameters()).device
inputs = tokenizer(text, return_tensors="pt").to(input_device)
with torch.inference_mode():
    # 推理模式减少 autograd 开销，但不改变模型的多设备落位。
    output = model.generate(**inputs, max_new_tokens=64)
print(tokenizer.decode(output[0], skip_special_tokens=True))

这里的 torch_dtype 只表示加载时采用更低精度的意图，实际可用精度仍要结合模型、硬件和 Transformers 版本确认。若 hf_device_map 中出现 disk，还要给支持磁盘卸载的加载方式提供 offload_folder，并为临时索引和磁盘读写留出空间。

用 max_memory 与 no_split_module_classes 收紧分配边界

自动分配是起点，不一定是生产配置。GPU 还要承载输入、KV cache 和生成过程，直接把显存填满会让模型刚加载完成就因生成阶段的额外占用而失败。可以用 max_memory 显式给每张 GPU、CPU 设置预算；设备编号要以当前进程可见设备为准。

max_memory = {
    0: "20GiB",   # 给第一张 GPU 留出输入与生成缓存空间。
    1: "22GiB",   # 第二张 GPU 的预算按实际显存和并发需求调整。
    "cpu": "48GiB",  # CPU 预算不是越大越好，还要保留系统内存。
}

# 这个名字必须替换成模型真实的层类名，不能凭空填写。
no_split = ["YourDecoderLayer"]

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="auto",
    max_memory=max_memory,
    no_split_module_classes=no_split,
    torch_dtype=torch.float16,
    low_cpu_mem_usage=True,
)

# 分组后的映射能帮助定位是哪一段模块落到了 CPU 或 disk。
for module_name, device in model.hf_device_map.items():
    print(module_name, "->", device)

no_split_module_classes 适用于带残差连接等不适合跨设备拆开的层。它不是“把每层都固定到一张卡”的开关；如果某个完整层本身大于可用设备预算，继续调低 max_memory 并不能解决问题，必须先换精度、量化、减少模型规模，或让最大层拥有足够的设备空间。

用 hf_device_map 和输入设备确认加载结果

不要只看进程没有报错。先遍历 model.hf_device_map，确认关键模块是否落在预期设备；再检查 next(model.parameters()).device，把 tokenizer 产生的张量移动到这个设备。若模型包含自定义输入嵌入、视觉编码器或额外投影层，还要把它们列入自己的检查清单。

验证时重点看三个边界：第一，映射是否覆盖模型模块；第二，最大的不可拆分层是否有完整设备预算；第三，输入和生成参数是否与模型的首个执行设备一致。只要出现大量 cpu 或 disk，就应把性能下降解释为数据搬运成本，而不是简单归因于“多卡没有生效”。

常见问题

device_map="auto" 会让两张 GPU 同时计算吗？

不保证。它主要负责模块落位与推理调度，官方说明的典型路径是层之间传递结果，单次前向并不等同于流水线并行。

为什么加了 auto 仍然显存不足？

常见原因是最大层本身放不下，或者生成时的缓存没有预留空间。先看每张卡的 max_memory 和不可拆分层，再考虑量化或更小模型。

什么时候需要自己写 device_map 字典？

当某些模块必须固定到特定设备、需要避开一张正在承担其他任务的 GPU，或要控制 CPU/磁盘卸载范围时，再从 auto 生成的映射出发做人工调整。