Accelerate disk_offload 怎样把模型权重映射到本地目录

如果模型已经加载到内存，Accelerate 的 disk_offload() 可以把参数转换成指定目录中的 memory-mapped 权重；前向计算时，所需参数会被取回并送到 execution_device。如果模型还在检查点阶段，则用 load_checkpoint_and_dispatch()，在 device_map 中把目标模块标为 "disk"，再通过 offload_folder 提供目录。

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

本地目录是权重的落盘位置，不是执行设备。offload_dir 或 offload_folder 要指向可持续读写的目录，execution_device 仍要指向实际承担前向计算的设备。

• 已加载的模型：使用 disk_offload(model, offload_dir, execution_device)。
• 加载检查点：使用 device_map 的 "disk"，并设置 offload_folder。
• 只支持推理：磁盘卸载不是训练内存优化方案，磁盘速度也会直接影响取权重开销。

先把本地目录和执行设备分开理解

disk_offload() 的三个关键参数很容易被混在一起：model 是要被卸载的模型对象，offload_dir 是保存 memory-mapped 权重的目录，execution_device 是执行前向计算的设备，通常是 GPU。官方文档描述的行为是：参数落在目录中，前向需要某层时再加载到执行设备，使用后再卸回去。

因此，目录不应该写成只在当前命令中存在的临时路径，也不要把它误写成 "cuda:0"。如果本地没有可用 GPU，应该先确认实际推理设备和内存预算；磁盘卸载本身不会把计算搬到磁盘上。

模型已经加载时直接调用 disk_offload

当你已经有一个完整的 PyTorch 模型对象，最小写法是把它交给 disk_offload()。下面的目录只是示例，生产环境应换成容量足够、生命周期稳定的本地盘目录。

from pathlib import Path
import torch
from accelerate import disk_offload

# 目录用于保存 memory-mapped 权重，不要指向一次性临时目录
offload_dir = Path("./model-offload")
offload_dir.mkdir(parents=True, exist_ok=True)

# model 应该是已经加载好的 PyTorch 模型
model = disk_offload(
    model,
    offload_dir=offload_dir,
    execution_device=torch.device("cuda:0"),  # 前向计算仍在 GPU 上完成
)

# 输入也要放到与执行设备匹配的位置
inputs = inputs.to("cuda:0")
with torch.no_grad():  # 推理不建立反向图，避免额外占用显存
    outputs = model(**inputs)

这个调用的重点不是返回一个新目录，而是让模型参数与目录建立卸载关系。前向时，Accelerate 的钩子会处理参数在磁盘、CPU 和执行设备之间的移动。若输入仍在 CPU，而模型的执行设备是 GPU，就会出现设备不匹配；若目录没有写权限或剩余空间不足，问题也会在卸载阶段暴露。

加载检查点时用 device_map 把模块放到 disk

如果模型还没有完整加载，可以让 load_checkpoint_and_dispatch() 根据模块名分配设备。映射值可以是 GPU 编号、"cpu" 或 "disk"；当映射里出现 "disk" 时，要同时提供 offload_folder。下面的模块名是示意，实际名称要以当前模型的模块树为准。

from pathlib import Path
import torch
from accelerate import init_empty_weights, load_checkpoint_and_dispatch

checkpoint_dir = "./checkpoint"
offload_folder = Path("./checkpoint-offload")
offload_folder.mkdir(parents=True, exist_ok=True)

# 在 meta 设备上创建空模型，避免先复制一份完整权重到 RAM
with init_empty_weights():
    model = build_model_from_config()  # 按项目配置创建模型骨架

# 只把明确的模块放到磁盘；残差块应按模型规则避免被拆开
device_map = {
    "embed": 0,                 # 输入嵌入放到第一张 GPU
    "layers.0": 0,              # 热路径层优先放在 GPU
    "layers.1": "cpu",         # CPU 作为中间层级
    "layers.2": "disk",        # 该模块参数映射到 offload_folder
    "lm_head": 0,               # 输出头回到执行设备
}

model = load_checkpoint_and_dispatch(
    model,
    checkpoint=checkpoint_dir,
    device_map=device_map,
    offload_folder=offload_folder,
    no_split_module_classes=["DecoderLayer"],  # 保持残差块的完整性
)

这里的关键是三者成套出现：模块名决定哪些权重使用磁盘，"disk" 表示落盘，offload_folder 提供落盘目录。不要照抄示例中的模块名；如果检查点键与模型结构不一致，加载器可能因为缺失或不匹配的键而失败。

目录能用不等于推理一定合适

第一次使用时，我更愿意把目录、设备和模块映射分别写出来，再根据 hf_device_map 检查结果。自动映射适合快速试验，但它会根据当时可见的 GPU 和 CPU 内存做决定；如果 CPU 内存不足，可以显式把部分模块放到 "disk"。磁盘卸载没有预取机制，慢盘或拥挤的网络盘会放大每次取权重的延迟。

还要注意它的适用范围：官方说明 disk offload 面向推理，参数会在前向所需时进入执行设备，完成后再卸载；这不是训练时保存梯度和优化器状态的方案。推理容器重启后，目录中的映射文件也应与当前模型检查点和代码版本保持一致，不能把任意目录当作可复用缓存。

常见问题

offload_dir 和 offload_folder 是同一个参数吗？

它们都表示本地权重目录，但分别属于不同入口：disk_offload() 使用 offload_dir，load_checkpoint_and_dispatch() 使用 offload_folder。不要把两个函数的参数名混写。

为什么放到 disk 后仍然需要 GPU？

disk 只负责承载权重，前向计算仍由 execution_device 或映射中的计算设备完成。只有把执行设备明确设为 CPU，才是 CPU 推理；这时性能和内存边界要另行评估。

怎样判断映射是否符合预期？

加载后查看模型的 hf_device_map，确认目标模块是否出现在 "disk"、"cpu" 或 GPU 对应位置。这个检查只说明映射结果，不代表已经测得吞吐或延迟。