Accelerate load_checkpoint_and_dispatch 加载权重不匹配怎么办

如果 Accelerate load_checkpoint_and_dispatch 报出 missing keys 或 unexpected keys，先不要调整 device_map 或继续加大显存。真正要查的是：正在加载的 checkpoint 键名，是否属于当前用 config 建出来的空模型。官方大模型加载文档将这条 API 定义为“把完整或分片 checkpoint 加载到模型并分配到设备”，当前 API 还提供了 strict 参数；官方地址：https://huggingface.co/docs/accelerate/main/concept_guides/big_model_inference。

这类故障的修复顺序是“先比较键名，再核对配置和模型头，最后严格加载”。device_map="auto" 只决定模块放在哪里，不能把不属于当前模型的权重变成匹配权重。

• missing keys：空模型需要某些参数，但 checkpoint 没提供，常见于模型头或配置不一致。
• unexpected keys：checkpoint 有额外参数，常见于前缀、任务头或 adapter 边界不同。
• shape mismatch：键名相同但张量形状不同，通常是词表、隐藏层或分类头尺寸不一致。

故障现场：为什么看起来像显存问题

一次典型失败会先在 init_empty_weights() 中建立模型，然后调用 load_checkpoint_and_dispatch()。如果模型已经分配到 GPU、CPU 或磁盘，日志里常常同时出现 offload、meta device 或 hook 信息，容易让人把注意力放到设备上。但权重对应关系在分发前就已经决定：当前模型的参数名和 checkpoint 的参数名对不上，换设备不会改变这两个集合。

from accelerate import init_empty_weights, load_checkpoint_and_dispatch

# 这里必须复用保存 checkpoint 时的配置和模型类，不能只凭模型名称猜测。
with init_empty_weights():
    model = build_model_from_same_config(config)

# 先用宽松模式复现，便于得到 missing/unexpected keys；不要把它当成修复。
model = load_checkpoint_and_dispatch(
    model,
    checkpoint=checkpoint_path,
    device_map="auto",
    strict=False,
)

时间线通常是：先下载或找到分片索引，再用配置创建空模型，随后加载 checkpoint 并挂载设备钩子，最后在加载阶段暴露键名差异。官方指南也提醒，当前这条加载路径不会替读者检查 state dict 与模型是否完全正确匹配，因此 mismatch 需要在调用前主动收窄。

排查时间线：先把两套键集合摆在一起

第一步不要读取全部权重，只比较“模型侧参数名”和“分片索引里的参数名”。空模型仍然可以提供 state_dict().keys()；分片 checkpoint 通常有 *.index.json，其中的 weight_map 记录每个参数位于哪个 shard。下面的示例只读索引和元模型，适合先判断是缺少、额外还是前缀不同。

import json
from pathlib import Path

def checkpoint_keys(checkpoint_dir: str) -> set[str]:
    # 只读取分片索引，不把大模型权重整体载入内存。
    folder = Path(checkpoint_dir)
    index_files = list(folder.glob("*.index.json"))
    if not index_files:
        raise FileNotFoundError("没有找到 checkpoint 分片索引，请先确认目录格式")
    index = json.loads(index_files[0].read_text(encoding="utf-8"))
    # 新旧索引都常见 weight_map；缺失时直接报出输入问题。
    weight_map = index.get("weight_map")
    if not isinstance(weight_map, dict):
        raise ValueError("索引中没有可比较的 weight_map")
    return set(weight_map)

def compare_keys(model, checkpoint_dir: str) -> tuple[set[str], set[str]]:
    # 空模型的键名不需要真实张量，适合在 meta device 上比较。
    model_keys = set(model.state_dict().keys())
    file_keys = checkpoint_keys(checkpoint_dir)
    missing = model_keys - file_keys
    unexpected = file_keys - model_keys
    return missing, unexpected

missing, unexpected = compare_keys(model, checkpoint_path)
print("missing:", sorted(missing)[:20])
print("unexpected:", sorted(unexpected)[:20])

如果两边只差一个稳定前缀，例如一边从 model. 开始，另一边从 base_model. 开始，先确认这个前缀来自包装层，再决定是否在转换阶段统一。不要看到差异就批量删除字符串；错误的替换会把本来不同的模块伪装成同名参数。

触发条件：四类不匹配要分别处理

最容易被忽略的是“基础模型 checkpoint + 当前任务模型”组合。基础模型可能没有分类头，或 adapter 只保存增量键；如果空模型却按另一任务配置创建，就会同时看到 missing 和 unexpected。此时应该先确定要加载完整模型、基础模型，还是先加载基础模型再挂载 adapter，不能把三者的键集合直接合并。

修复动作：解释差异后再开启严格加载

把模型类、配置、checkpoint 目录和任务头统一后，再显式使用 strict=True 做一次收口。当前 API 参考将它定义为“严格要求 checkpoint state dict 的键与模型 state dict 匹配”；严格模式的价值是让已经解释过的差异不再悄悄被带进推理。

from accelerate import load_checkpoint_and_dispatch

# 差异已经解释后，再用严格模式防止错误键名被静默带过。
model = load_checkpoint_and_dispatch(
    model,
    checkpoint=checkpoint_path,
    device_map="auto",
    no_split_module_classes=["YourResidualBlock"],
    strict=True,
)

# 这里单独查看分配结果；它只用于设备布局，不负责修复权重命名。
print(model.hf_device_map)

no_split_module_classes 解决的是带残差连接的模块不要跨设备切开，device_map 和 max_memory 解决的是放置与容量，offload_folder 解决的是磁盘卸载位置。它们可以改善加载后的设备问题，但不应被当成 missing/unexpected keys 的修复开关。

防复发：把检查点契约写进加载入口

发布 checkpoint 时同时保留模型配置、任务类型、词表尺寸、是否包含任务头以及 adapter/base model 关系；加载入口先用空模型比较键集合，再执行正式加载。遇到差异时把前 20 个 missing 和 unexpected keys 保存到排障日志，避免只留下一个笼统的“权重不匹配”。这样下一次升级模型类或切换任务头时，问题会在加载边界暴露，而不是拖到第一次推理才出现。

常见问题

把 strict=False 改成 strict=True 就能解决吗？

不能。strict=True 适合在差异已经解释后收口；如果配置或任务头仍不一致，它只会更明确地报告问题，不会生成缺失权重。

device_map="auto" 会导致 missing keys 吗？

通常不会。它负责根据可用设备安排模块；missing 或 unexpected keys 应先按参数名、配置和 checkpoint 类型排查。只有键集合对齐后，才继续处理显存、offload 和模块切分。

为什么基础模型和 adapter 不能直接混着加载？

两者保存的 state dict 范围可能不同。先确认 adapter 对应的基础模型和配置，再按对应加载流程组合；不要用批量改前缀的方式掩盖来源差异。