Llama-recipes项目中新增词汇表未正确保存的问题分析

问题背景

在Llama-recipes项目进行模型微调时，开发者发现当向tokenizer添加新词汇并调整模型嵌入层大小后，保存的模型未能正确保留这些修改。具体表现为：尽管在训练前调用了tokenizer.add_tokens()和model.resize_token_embeddings()，但最终保存的模型嵌入层维度仍保持原始大小。

技术细节分析

问题重现流程

1. 在微调脚本中，开发者按照标准流程：

   • 使用tokenizer.add_tokens()添加新词汇
   • 调用model.resize_token_embeddings(len(tokenizer))调整嵌入层大小

2. 训练完成后，通过FSDP到HuggingFace格式的转换工具保存模型

3. 检查保存的模型时发现：

   • model.get_input_embeddings().weight.shape[0]仍为原始词汇表大小
   • 新添加的词汇对应的嵌入层参数未被保存

根本原因

问题出在权重转换环节。Llama-recipes项目中使用的convert_hf_weights_to_llama.py脚本在转换权重时，没有考虑词汇表扩展的情况。该脚本直接使用原始模型的配置参数，导致新添加的词汇对应的嵌入层参数在转换过程中丢失。

解决方案

临时解决方案

开发者可以手动修改转换脚本，在权重转换前确保：

1. 正确加载扩展后的tokenizer
2. 根据实际词汇表大小调整模型配置
3. 确保嵌入层权重矩阵的维度与新词汇表大小匹配

长期改进建议

项目维护者应考虑以下改进：

1. 在权重转换脚本中加入词汇表大小验证
2. 自动检测并处理词汇表扩展情况
3. 提供明确的文档说明如何处理自定义词汇表的情况

最佳实践建议

对于需要在Llama模型中添加自定义词汇的开发者，建议：

1. 在训练前后都验证词汇表大小是否一致
2. 检查转换后的模型是否保留了所有自定义词汇
3. 考虑在转换后手动调整模型配置中的词汇表相关参数

这个问题凸显了在大型语言模型微调过程中，保持各组件(特别是tokenizer和模型架构)之间一致性的重要性。开发者需要特别注意模型配置与实际参数之间的匹配关系，特别是在进行词汇表扩展这类操作时。