解决HuggingFace PEFT库中LoRA微调时的Embedding层尺寸不匹配问题

背景介绍

在使用HuggingFace的PEFT(Parameter-Efficient Fine-Tuning)库进行LoRA(Low-Rank Adaptation)微调时，开发者可能会遇到一个常见问题：当原始预训练模型的Embedding层词汇表大小与当前tokenizer的实际词汇量不一致时，会导致模型加载失败。这种情况通常表现为"size mismatch"错误，特别是在处理自定义数据集或特定领域任务时。

问题分析

以Mamba模型为例，原始预训练模型的Embedding层可能配置为Embedding(50280, 2560)，而实际tokenizer的词汇量只有50277。这种差异通常源于：

1. 预训练阶段使用了更大的词汇表
2. 模型发布后tokenizer经过了调整或精简
3. 用户自定义tokenizer导致词汇量变化

当使用PEFT的LoRA配置对Embedding层进行适配时，这种尺寸不匹配会导致无法加载预训练权重，抛出类似以下的错误：

RuntimeError: size mismatch for base_model.model.backbone.embeddings.lora_embedding_A.default: copying a param with shape torch.Size([64, 50280]) from checkpoint, the shape in current model is torch.Size([64, 50277])

解决方案

方法一：单独加载适配器

最可靠的解决方案是分别加载基础模型和适配器权重：

1. 首先加载基础模型：

model = AutoModelForCausalLM.from_pretrained(model_dir, torch_dtype=torch.bfloat16)

2. 然后单独加载适配器：

model.load_adapter(adapter_path, "default")

这种方法避免了直接使用AutoPeftModelForCausalLM.from_pretrained()时遇到的尺寸检查问题。

方法二：调整LoRA配置

如果不需要对Embedding层进行适配，可以在LoRA配置中排除Embedding层：

lora_config = LoraConfig(
    r=64,
    target_modules=["x_proj", "in_proj", "out_proj"],  # 移除了embeddings
    task_type="CAUSAL_LM",
    bias="none",
    use_rslora=True,
)

方法三：词汇表对齐

对于高级用户，可以考虑：

1. 扩展当前tokenizer的词汇表以匹配原始模型
2. 或者修改模型Embedding层尺寸以匹配当前tokenizer

但这种方法需要谨慎处理，可能会影响模型性能。

最佳实践建议

1. 在使用PEFT进行微调前，先检查模型和tokenizer的词汇表尺寸是否匹配
2. 对于Embedding层的适配要特别小心，通常其他层的适配已经能带来足够好的效果
3. 考虑使用model.resize_token_embeddings(len(tokenizer))来显式调整Embedding层尺寸
4. 保存完整的微调配置，包括LoRA参数和目标模块选择，便于复现和部署

总结

处理PEFT中Embedding层尺寸不匹配问题时，单独加载适配器是最稳健的解决方案。理解模型架构和tokenizer之间的关系对于成功进行参数高效微调至关重要。通过合理配置LoRA目标模块和采用分步加载策略，可以有效规避这类尺寸不匹配问题，实现模型的顺利微调和部署。