PEFT项目中的LoRA模型加载与词嵌入尺寸匹配问题

概述

在使用PEFT(Parameter-Efficient Fine-Tuning)进行大语言模型微调时，特别是采用LoRA(Low-Rank Adaptation)方法时，开发者可能会遇到模型加载时的尺寸不匹配问题。本文将深入分析这一问题的成因，并提供专业解决方案。

问题现象

当使用PEFT训练一个包含可训练词嵌入层(embed_tokens)和语言模型头(lm_head)的LoRA模型后，在加载模型检查点时，系统会报告尺寸不匹配错误。具体表现为：

• 检查点中embed_tokens层的权重尺寸为[33004, 4096]
• 当前模型中的对应层尺寸为[32000, 4096]

根本原因分析

这一问题源于以下几个技术细节：

1. 词表扩展：在训练过程中，开发者通过model.resize_token_embeddings(len(tokenizer))扩展了模型的词表大小，从基础的32000增加到33004个token。

2. 模块保存设置：在PEFT配置中指定了modules_to_save = ["lm_head", "embed_tokens"]，这意味着这两个层的参数会被完整保存而非仅保存低秩适配部分。

3. 加载时尺寸不匹配：当直接加载基础模型时，其词表大小恢复为默认的32000，与检查点中保存的33004不匹配。

解决方案

正确的模型加载流程应包含以下步骤：

1. 加载基础模型：

from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-chat-hf",
    device_map="auto",
    torch_dtype=torch.bfloat16,
    trust_remote_code=True,
)

2. 加载训练时使用的tokenizer并调整模型尺寸：

from transformers import LlamaTokenizer
tokenizer = LlamaTokenizer.from_pretrained("<PATH-TO-YOUR-REPO>")
model.resize_token_embeddings(len(tokenizer))

3. 加载PEFT适配器：

from peft import PeftModel
model = PeftModel.from_pretrained(model, "<PATH-TO-YOUR-REPO>")

技术建议

1. 关于LoRA秩的选择：虽然技术上可以使用高达1024的秩(r)，但这会显著增加可训练参数数量(本例中达到总参数的21%)，几乎接近全参数微调的效果。建议从较小的秩(如8或16)开始，逐步增加以平衡效果和效率。

2. 词表扩展注意事项：当扩展词表时，新增token的嵌入会被随机初始化，需要在训练数据中给予足够多的曝光机会，以确保它们能学到有意义的表示。

3. 检查点优化：对于包含完整参数层(如embed_tokens)的PEFT模型，检查点会相对较大。可以考虑：

   • 使用量化技术减小模型大小
   • 仅在必要时保存完整参数层
   • 定期清理中间检查点

总结

PEFT框架为大型语言模型的高效微调提供了强大支持，但在处理模型结构变化(如词表扩展)时需要特别注意前后一致性。通过理解模型架构和PEFT的工作原理，开发者可以避免常见的陷阱，构建更加稳健的微调流程。对于需要大幅修改模型结构的任务，建议建立完整的训练-验证-部署流程文档，确保各环节的参数一致性。