Torchtune项目中知识蒸馏时处理不同词汇表大小的技术实践

背景介绍

在大型语言模型的知识蒸馏过程中，我们经常会遇到教师模型和学生模型词汇表大小不一致的情况。最近在Torchtune项目中使用Qwen2.5系列模型进行知识蒸馏时就遇到了这样的技术挑战：Qwen2.5 14B模型的词汇表大小为152064，而Qwen2.5 3B模型的词汇表大小为151936，两者存在128个token的差异。

问题分析

这种词汇表大小的差异主要源于模型设计时的不同考量。通过对比Qwen2.5 14B和3B的词汇表文件可以发现，前151664个token是完全一致的，差异仅出现在最后的填充部分。这种设计可能是出于计算效率的考虑，比如使注意力计算更加高效。

在知识蒸馏过程中，当教师模型和学生模型的词汇表大小不一致时，直接计算损失函数会遇到维度不匹配的问题，具体表现为类似"RuntimeError: The size of tensor a (152064) must match the size of tensor b (151936)"的错误。

解决方案

针对这一问题，我们提出了以下技术解决方案：

1. 词汇表对齐处理：由于差异仅出现在最后的填充token上，这些token在实际预测中很少被使用。因此，可以在计算损失函数前对教师模型的输出logits进行截断处理，使其维度与学生模型一致。

2. 代码实现：在知识蒸馏的损失计算阶段，添加以下处理逻辑：

for _, i in enumerate(teacher_logits):
    teacher_logits[i] = teacher_logits[i][:, :, :151936]

3. 多文件模型加载：在后续模型评估阶段，需要注意完整加载包含所有模型层的检查点文件。Qwen2.5 3B模型可能分多个文件保存，需要确保加载所有必要的文件，避免出现"Missing key(s) in state_dict"的错误。

技术扩展

这一案例揭示了知识蒸馏中一个更深层次的问题：当前大多数知识蒸馏方法都假设教师模型和学生模型使用相同的tokenizer。这在同系列模型间通常成立，但在跨模型系列或不同规模模型间可能不适用。

学术界已经开始关注这一问题，相关研究提出了新的蒸馏损失函数来处理不同词汇表的情况。在实际工程实践中，我们需要根据具体场景选择合适的方法：

1. 对于同系列不同规模的模型，可以采用词汇表截断的简单方法
2. 对于完全不同的模型系列，可能需要考虑更复杂的词汇映射方法
3. 未来Torchtune项目可能会增加对独立教师/学生tokenizer的支持，以提高灵活性

实践建议

基于这一案例，我们总结出以下实践建议：

1. 在进行知识蒸馏前，务必检查教师模型和学生模型的词汇表兼容性
2. 对于同系列模型的小规模差异，可以采用简单的维度对齐方法
3. 模型保存和加载时要注意检查点文件的完整性
4. 监控蒸馏过程中的损失曲线，确保知识传递的有效性
5. 考虑使用WandB等工具进行训练过程的可视化监控

通过这种方法，我们成功实现了从Qwen2.5 14B到3B的知识蒸馏，为资源受限场景下的模型部署提供了有效解决方案。