如何从PyTorch Checkpoint文件夹加载模型并转torch.load兼容格式

方法1：合并后直接保存为Torch可加载格式 #

在执行model.merge_and_unload()后，无需依赖save_pretrained()的默认保存逻辑，直接用torch.save()将模型权重保存为.pt/.pth格式，后续即可用torch.load()加载：

import torch

# 合并模型后
model = model.merge_and_unload()

# 保存为Torch权重格式
torch.save(model.state_dict(), "merged_model.pt")

# 后续加载示例
loaded_model = WhisperForConditionalGeneration.from_pretrained(
    peft_config.base_model_name_or_path,
    state_dict=torch.load("merged_model.pt"),
    device_map="auto"
)

如果需要保留save_pretrained()生成的目录结构，同时生成Torch兼容的权重文件，可调用时关闭safetensors序列化：

model.save_pretrained(<model_path>, safe_serialization=False)

此操作会生成pytorch_model.bin文件，该文件可直接用torch.load()加载。

方法2：将已保存的safetensors格式转为Torch格式 #

若已通过safetensors格式保存了模型，可先加载safetensors权重，再转存为Torch格式：

from safetensors.torch import load_file

# 加载safetensors权重文件
safetensors_weights = load_file("<model_path>/model.safetensors")

# 转存为Torch格式
torch.save(safetensors_weights, "merged_model.pt")

补充说明 #

• torch.load()加载的是模型权重字典，需结合原模型结构（如WhisperForConditionalGeneration）重建可推理模型，无法直接通过torch.load()得到完整可用模型（除非保存的是整个模型，但这种方式兼容性差，不推荐）。
• 若要直接保存包含结构和权重的完整模型，可使用以下代码，但需注意跨PyTorch版本可能出现加载问题：

torch.save(model, "full_model.pt")
# 加载时
loaded_full_model = torch.load("full_model.pt")

内容的提问来源于stack exchange，提问作者afsara_ben