Hugging Face Transformers 怎样只加载模型部分权重

如果显存只够放下模型的一部分，先不要急着把 device_map="auto" 当成“只加载部分权重”。两者解决的问题不同：真正的部分加载是筛掉不需要的键，让对应模块保持未初始化；显存治理则是保留完整权重，再把模块分配到 GPU、CPU 或磁盘。

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

本文以 Transformers 的 state_dict 入口为主线，说明如何按模块读取 safetensors、判断哪些参数没有预训练值，以及什么时候应该改用设备分配。

先分清：丢弃权重和分配显存不是一回事

from_pretrained() 默认会把检查点权重装配到模型结构中。若只需要骨干网络中的少数层，可以准备一个只包含目标键的 state_dict，再配合配置对象创建模型。没有出现在字典里的参数不会 magically 变成预训练参数，通常会留下“未从预训练模型初始化”的提示。

如果目标只是让完整模型适配有限显存，应该保留全部权重，使用 device_map、max_memory 和必要的 offload_folder。这条路线改变的是设备位置，不改变模型参数集合。

按键名筛选 state_dict，再交给 from_pretrained

筛选的核心不是凭感觉写“第几层”，而是先根据检查点实际命名规则定义前缀。下面示例用 safe_open 按需读取 safetensors 中的键，只保留嵌入层、第一层 Transformer block 和输出头。不同模型的前缀可能不同，使用前应以该模型的权重键名为准。

from safetensors import safe_open
from transformers import AutoConfig, AutoModelForCausalLM

model_id = "your-org/your-model"
config = AutoConfig.from_pretrained(model_id)

# 只选择本次任务需要的模块；前缀必须和检查点中的键名一致。
wanted_prefixes = (
    "model.embed_tokens.",
    "model.layers.0.",
    "lm_head.",
)
partial_state_dict = {}

# safe_open 按键读取张量，避免先把所有张量都复制到 Python 字典。
with safe_open("/path/to/model.safetensors", framework="pt", device="cpu") as checkpoint:
    for key in checkpoint.keys():
        if key.startswith(wanted_prefixes):
            partial_state_dict[key] = checkpoint.get_tensor(key)

# None 表示从提供的 config 和 state_dict 组装模型，而不是再指定一个完整检查点。
model, loading_info = AutoModelForCausalLM.from_pretrained(
    None,
    config=config,
    state_dict=partial_state_dict,
    output_loading_info=True,
)

# 缺失键对应仍未加载的参数；不要把它们当成已经训练好的权重。
print("missing:", len(loading_info["missing_keys"]))
print("unexpected:", len(loading_info["unexpected_keys"]))

这里的关键关系是：config 决定完整模型骨架，partial_state_dict 只提供选中的张量，from_pretrained 负责把二者拼起来。未被选中的模块仍然存在，但没有对应的预训练值。

未加载参数怎么处理，才不会得到“看起来能跑”的假模型

部分加载适合做结构迁移、权重分析、按模块初始化或自定义微调准备，不适合直接把缺失层当成可推理成品。先看 missing_keys，再决定是补齐权重、冻结已加载模块，还是为新模块安排训练。

# 只冻结已经加载的骨干模块；新模块需要明确训练策略。
for name, parameter in model.named_parameters():
    if any(name.startswith(prefix) for prefix in wanted_prefixes):
        parameter.requires_grad = False

# 训练或推理前，显式确认缺失键范围，避免把随机层悄悄带入结果。
missing_keys = loading_info["missing_keys"]
if missing_keys:
    print("仍需初始化或训练的参数数量:", len(missing_keys))

ignore_mismatched_sizes=True 只适合处理形状不一致，例如分类标签数量变化；它不是“安全忽略所有缺失权重”的开关。若键名不匹配，应优先修正前缀、模型配置或权重来源。

显存不足时用 device_map 分配，而不是删除权重

如果你的真实需求是“完整模型仍要推理，只是单张 GPU 放不下”，可以让 Accelerate 根据容量分配模块。device_map="auto" 会计算设备布局，max_memory 可以限制设备上限；当映射使用 "disk" 时，需要指定 offload_folder 保存卸载权重。

from transformers import AutoModelForCausalLM

# 这条路径保留完整权重，只改变模块的设备位置。
model = AutoModelForCausalLM.from_pretrained(
    "your-org/your-model",
    device_map="auto",
    max_memory={0: "10GB", "cpu": "30GB"},
    offload_folder="./offload",
    dtype="auto",
)

# 读取实际分配，确认模型没有被误解读成“只加载了部分权重”。
print(model.hf_device_map)

这条方案的判断标准很简单：想减少参数集合，就筛选 state_dict；想保留完整能力但降低单设备压力，就用 device_map。两者也可以组合，但应先明确哪些模块真的可以没有预训练值。

按场景选择加载方式

常见问题

只传入几个权重键，模型会自动补齐其它层吗？

不会自动补成预训练值。模型骨架可以由配置创建，但没有对应键的参数仍需要初始化或训练，必须结合 missing_keys 判断。

device_map=auto 是不是等于只加载需要的层？

不是。它主要决定模块放在哪个设备，目标仍是让完整模型参与推理；要减少权重集合，应使用筛选后的 state_dict。

为什么按层前缀筛选后仍然报键名错误？

常见原因是模型架构的命名空间不同，或者权重来自另一种模型头。先对照检查点键名和 config 对应的模型类，再调整前缀，不要用关闭提示的方式掩盖映射错误。

总结

Transformers 的部分权重加载，本质是“配置创建模型骨架 + 筛选后的 state_dict 提供少量参数”。它能减少实际装入的参数，但也会留下必须处理的未初始化模块；显存不够时则应使用 device_map、容量上限和 offload，把完整权重拆到不同设备。先确定目标，再选择入口，才不会得到一个能实例化、却没有完整预训练能力的模型。