Accelerate device_map 怎么给每张 GPU 和 CPU 设置 max_memory

给 Accelerate 的 max_memory 传入字典即可分别设置每张 GPU 和 CPU 的权重分配预算，例如 {0: "10GiB", 1: "12GiB", "cpu": "24GiB"}。整数键表示 GPU 设备编号，"cpu" 表示 CPU 内存预算。它参与自动生成 device_map，不能保证整个推理进程的峰值内存都限制在这些数字内。

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

本文按 Accelerate v1.15.0 文档解释配置。所有容量都是教学示例，不代表某台机器的实测结果；代码也没有在本文环境中执行。

为什么把自动可用容量改成显式预算

我更倾向于把设备预算作为部署配置单独保存。这样，模型不变而机器余量改变时，可以说明“为什么这次给 GPU 0 少分一点”，不用把决定藏在默认可用内存里。

未显式配置时，自动规划会参考设备可用容量。共享机器上，其他进程、输出处理和服务本身也需要空间。显式预算的用途，是让权重放置策略反映这些约束；预算本身不会预先占住显存，也不会让其他进程无法使用它。

设备编号、单位和 CPU 预算怎么写

值可以写带单位的字符串，也可以写以字节为单位的整数。GB 与 GiB 的换算不同，团队配置最好统一单位。不要把整数 10 当成10GB。

如果通过设备可见性配置筛选或重排了 GPU，预算键应对应当前进程看到的逻辑编号，不能直接照抄机箱标签。CPU 键也不是“让 CPU 获得额外显存”，而是给主机 RAM 中的权重留位置。

先计算映射，再按相同精度加载

下面用一个小型顺序网络展示参数怎样传递。它故意不下载模型，也不预设映射输出；模型足够小时，未必会使用第二张 GPU 或 CPU。

import torch
from torch import nn
from accelerate import init_empty_weights, infer_auto_device_map

# 示例预算，必须按当前进程可见的设备和实际余量调整。
budget = {0: "10GiB", 1: "12GiB", "cpu": "24GiB"}
weight_dtype = torch.float16

# 只创建参数结构，不在这里分配完整权重；这不是可直接推理的模型。
with init_empty_weights():
    model = nn.Sequential(
        nn.Linear(2048, 4096),
        nn.ReLU(),
        nn.Linear(4096, 2048),
    )

# dtype 是规划依据，返回值描述模块放在哪个设备。
placement = infer_auto_device_map(
    model,
    max_memory=budget,
    dtype=weight_dtype,
)
print(placement)

init_empty_weights 创建的是参数结构。此时打印映射只能了解规划，不能拿空权重执行推理。真实项目还要加载与模型结构对应的检查点。

需要一次完成加载和分配时，可以把同一个预算交给 load_checkpoint_and_dispatch：

import torch
from accelerate import load_checkpoint_and_dispatch

def load_budgeted(empty_model, checkpoint, no_split_classes):
    # 调用方提供匹配的空模型、真实本地检查点和模型实际的模块类名。
    budget = {0: "10GiB", 1: "12GiB", "cpu": "24GiB"}
    return load_checkpoint_and_dispatch(
        empty_model,
        checkpoint=checkpoint,
        device_map="auto",  # 自动规划时使用下面的每设备预算。
        max_memory=budget,
        dtype=torch.float16,  # 与规划采用同一权重精度。
        no_split_module_classes=no_split_classes,
        offload_folder="./model-offload",  # 映射含 disk 时使用该目录。
    )

这个函数是供项目集成的片段：调用方负责构建空模型、提供本地检查点，并填写模型真正需要保持完整的模块类名。含残差连接的块通常需要考虑 no_split_module_classes；不能随意写一个不存在的类名。实际加载还应满足模型自身的精度和设备支持条件。

如果选择先计算映射、再把明确的映射字典交给加载器，应在计算阶段应用预算，并保持精度等设置一致。手工指定模块位置后，不应指望额外传一个 max_memory 就自动修正不合理的手工布局。

预算没有用满与加载后仍溢出的原因

预算是上限，不是每台设备必须达到的填充目标。模块大小、不可拆分的块以及卸载时需要回搬权重的空间，都会影响最终布局。因此，把两张卡的预算相加，不能直接推断一定能装下同等大小的模型。

同样，权重可以放下，也不等于实际请求一定不会显存不足。输入长度、批量大小、缓存、临时张量和 CUDA 运行开销仍会占用空间。应给这些部分留余量，不能把显卡标称容量原样填进预算。

CPU RAM 也要容纳服务对象、输入数据及加载过程的临时开销。自动布局涉及磁盘卸载时，要提供可用的卸载目录；磁盘增加的是权重存放空间，数据搬运仍有成本。不要把增加 CPU 或磁盘预算理解成获得了同等性能的 GPU 容量。

调整配置时保留哪些约定

我会把预算、权重精度、不可拆分模块类名和卸载目录放在同一份部署说明里。调整预算时，也记录调整原因，例如给输入输出所在设备留更多空间。这比只记一个“总显存够用”更容易复用。

能给每张卡设置不同预算吗？可以，分别填写整数设备键即可。预算不同并不保证最终占用差值恰好等于预算差值。

CPU 预算是否限制整个 Python 进程？不能这样理解。它服务于权重布局，不是操作系统级进程内存配额。

这是不是多卡训练配置？本文讨论的是大模型推理的权重分配与卸载。训练时的梯度、优化器状态及并行策略需要另行设计。

实际配置的核心是：为每个可见设备给出合理预算，让规划精度和加载精度一致，并为运行时留空间。

参考：Accelerate 大模型内存布局说明；大模型加载与分配 API。