Huggingface AutoTrain-Advanced 项目中的模型加载与词汇表尺寸不匹配问题解析

问题背景

在使用Huggingface的AutoTrain-Advanced项目进行大语言模型(LLM)微调时，开发者可能会遇到一个典型问题：当尝试加载经过微调的模型时，系统会报出词汇表尺寸不匹配的错误。具体表现为模型嵌入层(embed_tokens)和语言模型头层(lm_head)的权重矩阵尺寸与预期不符。

错误现象

错误信息通常会显示如下内容：

RuntimeError: Error(s) in loading state_dict for LlamaForCausalLM:
    size mismatch for model.embed_tokens.weight: copying a param with shape torch.Size([32004, 8192]) from checkpoint, the shape in current model is torch.Size([32000, 8192]).
    size mismatch for lm_head.weight: copying a param with shape torch.Size([32004, 8192]) from checkpoint, the shape in current model is torch.Size([32000, 8192])

这表明微调后的模型词汇表大小从原始的32000增加到了32004，而系统尝试加载时仍期望保持原始尺寸。

根本原因分析

经过深入调查，我们发现这一问题主要源于以下几个方面：

1. 特殊标记的自动添加：在微调过程中，系统会自动检测并添加一些特殊标记到词汇表中。在CodeLlama模型中，通常会添加四个特殊标记：▁<PRE>、▁<MID>、▁<SUF>和▁<EOT>，这导致词汇表尺寸增加了4。

2. 模型权重与词汇表不匹配：当词汇表大小改变后，模型的嵌入层和输出层需要相应调整权重矩阵的尺寸，但加载时系统仍尝试使用原始尺寸。

3. 版本兼容性问题：不同版本的transformers和peft库在处理词汇表扩展时的行为可能不一致，特别是在模型合并和加载过程中。

解决方案

针对这一问题，我们推荐以下解决方案：

1. 正确合并适配器

使用PeftModel的merge_and_unload方法正确合并适配器权重：

from peft import PeftModel

def merge_adapter(base_model_path, target_model_path, adapter_path):
    model = AutoModelForCausalLM.from_pretrained(
        base_model_path,
        torch_dtype=torch.float16,
        low_cpu_mem_usage=True,
        trust_remote_code=True,
    )
    
    tokenizer = AutoTokenizer.from_pretrained(target_model_path)
    model.resize_token_embeddings(len(tokenizer))
    
    model = PeftModel.from_pretrained(model, adapter_path)
    model = model.merge_and_unload()
    
    model.save_pretrained(target_model_path)
    tokenizer.save_pretrained(target_model_path)

2. 正确加载合并后的模型

加载合并后的模型时，需要确保使用正确的参数：

model = AutoModelForCausalLM.from_pretrained(
    model_path,
    torch_dtype=torch.float16,
    low_cpu_mem_usage=True,
    trust_remote_code=True,
    device_map="cuda"
)

3. 处理推理过程中的未知标记

对于推理过程中出现的未知标记问题，可以采取以下措施：

1. 检查训练数据中是否包含特殊标记
2. 确保tokenizer正确加载了所有特殊标记
3. 在生成文本时设置适当的参数：

output = model.generate(
    input_ids,
    max_new_tokens=200,
    pad_token_id=tokenizer.eos_token_id
)

最佳实践建议

1. 版本一致性：保持transformers、peft和autotrain-advanced版本的兼容性，推荐使用较新的稳定版本组合。

2. 词汇表检查：在微调前后对比tokenizer的配置文件和特殊标记，确认词汇表变化是否符合预期。

3. 数据预处理：在训练前彻底检查数据集，确保不包含意外的特殊字符或标记。

4. 资源管理：对于大模型如CodeLlama-34b，确保有足够的GPU内存，并使用适当的量化技术(int8)来减少内存占用。

总结

在Huggingface生态中进行大语言模型微调时，词汇表尺寸变化是一个常见但容易被忽视的问题。通过理解模型架构、正确合并适配器以及合理配置加载参数，开发者可以有效解决这类问题。本文提供的解决方案不仅适用于CodeLlama系列模型，其原理和方法也可推广到其他类似架构的大语言模型微调场景中。