如何在Python中解决PyTorch加载大体积模型权重文件导致的系统OOM？-Python教程-PHP中文网

因为torch.load()默认将整个权重文件解压反序列化到cpu内存，峰值占用可达文件大小数倍；使用map_location='meta'可跳过实际加载，仅解析结构，大幅降低内存消耗。

如何在python中解决pytorch加载大体积模型权重文件导致的系统oom？

为什么`torch.load()`会直接爆内存？

因为默认情况下torch.load()会把整个权重文件（比如几个GB的.pt或.pth）全部解压、反序列化、加载到CPU内存，再转到GPU——哪怕你只想要其中几层参数。尤其当模型用torch.save(model.state_dict(), ...)保存时，文件是完整state dict，没有分块或延迟加载机制。

即使设了map_location='cpu'，也挡不住反序列化过程本身吃内存
使用zip格式（PyTorch 1.6+默认）时，解压缓冲区仍可能峰值占用数倍原始文件大小
Windows上还容易遇到OSError: [WinError 8]——本质是内存映射失败，不是磁盘空间问题

用`torch.load(..., map_location='meta')`跳过实际加载

PyTorch 2.0+支持map_location='meta'，它让torch.load()只解析结构、不分配真实张量内存，适合快速探查大模型结构或做子集提取。

能秒级读取10GB文件，返回一个dict，所有Tensor都是torch.Tensor(..., device='meta')
后续可按需用torch.empty(..., device='cuda') + copy_()加载特定键
注意：老版本PyTorch不支持'meta'，会报RuntimeError: device meta is not supported

ckpt = torch.load("big_model.pth", map_location="meta")
print(ckpt["model.layers.0.self_attn.q_proj.weight"].size())  # 不占内存
# 真正加载某一层：
layer0_weight = torch.empty_like(ckpt["model.layers.0.self_attn.q_proj.weight"], device="cuda")
layer0_weight.copy_(torch.load("big_model.pth", map_location="cuda")["model.layers.0.self_attn.q_proj.weight"])

用`torch._C._load_for_mobile()`或自定义`zipfile`流式读取

PyTorch的.pt文件本质是ZIP包，里面包含data.pkl和多个data/X二进制分片。绕过torch.load、直接用zipfile读指定成员，能彻底规避反序列化开销。

提示词大师-python版

图片提示词生成器？不止如此。马甲系统 —— 把脑海中的画面，翻译成AI能理解的专业表达。用得越多，它越懂你：首次需要多问几句确认方向，用久了几乎一说就懂。用得越多，它越快：缓存机制让后续对话越来越省。 RAG进化：成功案例持续入库，越跑越聪明。输入「新手指南」查看完整功能介绍

下载

适用于已知权重键名、且文件是ZIP格式（torch.save(..., _use_new_zipfile_serialization=True)）
必须手动处理pickle header和tensor layout，推荐用torch._pickle.load配合io.BytesIO
风险点：PyTorch内部格式未公开承诺稳定，_load_for_mobile是私有API，仅建议用于离线分析场景

import zipfile, io, torch
with zipfile.ZipFile("big_model.pth") as z:
    with z.open("data.pkl") as f:
        # 只读pkl头获取键列表，不加载全部
        keys = torch._pickle.load(f)
    # 再打开具体tensor文件，如"data/0"
    with z.open("data/0") as f:
        tensor_bytes = f.read()
        # 后续用torch._utils._rebuild_tensor_v2等重建，过程复杂，慎用

真正稳妥的做法：改保存方式 + 分块加载

治本之策不是硬扛大文件，而是从源头避免生成单一大权重文件。Hugging Face Transformers和DeepSpeed都默认采用分片（sharded）保存，比如pytorch_model-00001-of-00003.bin，每个文件几百MB，可按需加载。

训练时用model.save_pretrained(..., max_shard_size="2GB")代替torch.save
推理时用from_pretrained(..., device_map="auto")，它会自动调度各分片到CPU/GPU，且只驻留当前需要的层
若必须处理旧的大.pth文件，先用transformers.models.modeling_utils.shard_checkpoint()拆分，再加载

别指望靠调大swap或改torch.set_num_threads()解决OOM——那是把问题推给系统，而不是解决问题本身。真正卡住的永远是反序列化那一刻的内存峰值，不是显存或CPU时间。