PyTorch torch.load 开启 mmap 后为什么读取参数时才占用内存

我第一次处理很大的 PyTorch 检查点时，最容易被误导的是“加载完成了，但内存没有立刻涨满”。后来把 mmap 拆成文件映射、虚拟地址和实际驻留页三个概念，现象就清楚了：mmap=True 先让序列化存储可以映射到进程地址空间，操作系统通常在对应页被访问时才把它带入内存，所以第一次读取某些参数时，RSS 或页面驻留才会明显变化。

官方文档：https://docs.pytorch.org/docs/stable/generated/torch.load.html

我先把 mmap 的三个内存概念分开

文件映射像是给文件内容建立一段可寻址的窗口。这个窗口属于进程的虚拟地址空间，和“当前有多少页真的驻留在 RAM”不是同一个量。参数 storage 通过映射找到对应文件区域后，读到某一页时才更可能触发页载入；因此“加载函数返回”和“所有参数页都被访问”之间可以存在明显间隔。

这也解释了为什么只打印一个小参数时，内存观察结果和遍历整个 state_dict 不一样。前者只会碰到少量页，后者会把更多文件页带入访问路径。图示只表达静态关系，不是某次机器运行的内存曲线。

我用最小示例确认 torch.load 的职责

下面的写法把三个选择放在一起：mmap=True 负责文件映射，map_location="cpu" 明确主机侧落点，weights_only=True 表示按权重加载的意图。实际项目要根据检查点内容和当前 PyTorch 版本确认支持范围；这里不把它们混成一个“省内存开关”。

import torch

# mmap 让检查点存储采用文件映射路径，返回时不等于所有页已经被访问。
state_dict = torch.load(
    "model.pt",
    mmap=True,
    map_location="cpu",
    weights_only=True,
)

# 这里只读取一个参数，触碰的页通常少于遍历整个参数集合。
first_name, first_param = next(iter(state_dict.items()))
print(first_name, tuple(first_param.shape), first_param.device)

我会把“函数返回时间”和“第一次真正读取参数的时间”分开记录。尤其不要只用 torch.load 前后的一个内存采样点下结论，因为 Python 对象、文件映射和操作系统缓存的统计口径并不完全相同。

第一次访问参数时，为什么看起来才占用内存

这里的“占用”通常是观察指标发生了变化，而不是 mmap 在那一刻才创建。映射本身已经建立；只是相应页此前可能还没有成为进程的常驻工作集。访问某个大权重、把它参与运算，或者遍历更多参数，都会扩大被触碰的页范围。

import torch

# 先加载映射后的权重，但不把遍历参数误当成加载动作的一部分。
state_dict = torch.load("model.pt", mmap=True, map_location="cpu", weights_only=True)

# 这个读取只观察一个张量的元信息，不主动复制整个 state_dict。
name, tensor = next(iter(state_dict.items()))
print("selected:", name, "numel:", tensor.numel())

# 只有在明确需要数值参与计算时，才让对应数据进入计算路径。
checksum = tensor.float().mean().item()
print("mean:", checksum)

这个例子不是内存基准测试，也不能据此声称固定的节省比例。它只帮助我定位现象：映射建立、对象访问、数值计算和设备迁移是四个不同阶段。若要做工程测量，应统一 RSS、页错误、CPU 内存和 GPU 显存的采样口径，并记录访问了哪些参数。

处理 map_location 与设备迁移时，我会保留两条边界

map_location="cpu" 的意义是把反序列化出来的 storage 放在 CPU 侧；它不等于模型已经在 GPU 上可用。将参数交给模型前，通常还要明确执行 load_state_dict、to(device) 或项目自己的设备迁移逻辑。CPU 映射存储和 GPU 显存中的参数副本，应该当作两条不同的资源生命周期来管理。

import torch

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 先把检查点固定在 CPU 侧，避免加载阶段隐式改变设备边界。
state_dict = torch.load("model.pt", mmap=True, map_location="cpu", weights_only=True)

model = build_model()  # 这里代表项目已有的模型构造函数。
# 显式迁移模型和参数，避免把 CPU 映射存储当成 GPU 参数。
model.load_state_dict(state_dict)
model.to(device)

# 生命周期结束后，若不再需要 CPU 侧对象，可以按项目逻辑释放引用。
del state_dict

如果 map_location 直接把存储导向其他设备，实际行为还会受到设备可用性、检查点内容和 PyTorch 支持范围影响。对大模型而言，我更倾向于先在 CPU 侧建立清晰的加载边界，再在确定的时机迁移，便于解释峰值来自哪里。

支持范围和降级方案：不是所有检查点都能随意 mmap

官方文档对 mmap 有明确限制：它服务于采用默认 zipfile 序列化方式保存的文件；旧序列化格式或不满足条件的文件不能简单套用同一加载路径。遇到不兼容时，降级方案是去掉 mmap=True，先确认普通 torch.load 能正确还原，再决定是否重新保存检查点。

另一个常见误区是把 mmap=True 当作“永远低内存”。如果后续代码遍历并计算全部参数，或者很快把权重复制到 GPU，访问页和设备副本仍会产生相应资源压力。它改变的是读取和映射时机，不会取消模型参数本身的容量需求。

我的加载取舍清单

相关问题

mmap=True 会不会把文件全部读进内存？

它首先建立文件映射；实际页是否驻留，取决于访问模式和操作系统。遍历或计算全部参数时，仍可能触碰大量页。

为什么 torch.load 返回后内存还没有明显变化？

返回说明反序列化对象已经建立，不代表所有映射页都被读取。观察具体参数或进行计算后，统计值才可能变化。

map_location="cpu" 后还需要 model.to(device) 吗？

如果模型要在 GPU 或其他设备执行，通常仍需明确迁移模型或参数。CPU 侧映射存储与设备侧副本不是同一个资源。