XTuner项目中LLaVA模型添加特殊令牌的技术实践

背景介绍

在XTuner项目中使用LLaVA模型进行监督微调(SFT)时，开发者经常需要为特定任务添加特殊令牌(special tokens)。这些特殊令牌可以帮助模型更好地理解输入数据的结构和语义，提升模型在特定任务上的表现。然而，在已有模型权重的基础上添加新令牌会带来技术挑战，特别是维度不匹配问题。

问题分析

当我们在已训练好的LLaVA模型上添加特殊令牌时，主要会遇到两个关键问题：

1. 词汇表扩展：需要将新令牌添加到tokenizer的词汇表中
2. 模型权重调整：需要相应调整语言模型部分的词嵌入层大小

具体表现为尝试加载已有权重时会出现维度不匹配错误，例如输出层维度从[92544, 512]变为[92547, 512]的情况。

解决方案

1. 添加特殊令牌到tokenizer

首先需要在数据处理阶段将特殊令牌添加到tokenizer中。核心代码如下：

for special_token in special_tokens:
    if special_token not in tokenizer.get_vocab():
        tokenizer.add_tokens([special_token], special_tokens=True)

这段代码会检查每个特殊令牌是否已存在于词汇表中，如果不存在则添加，并标记为特殊令牌。

2. 调整模型词嵌入层大小

添加令牌后，必须相应调整语言模型部分的词嵌入层大小。在LLaVA模型初始化时需要进行以下操作：

self.llm = self._build_from_cfg_or_module(llm)
self.llm.resize_token_embeddings(new_vocab_size)  # 调整为新词汇表大小
self.visual_encoder = self._build_from_cfg_or_module(visual_encoder)

其中new_vocab_size是添加特殊令牌后的新词汇表大小。

3. 处理预训练权重加载

当从已有检查点加载权重时，需要特别注意维度转换问题。可以参考XTuner项目中的权重转换逻辑，在加载预训练权重时进行适当的维度调整，确保新旧词汇表能够正确对应。

技术细节

1. 词汇表扩展机制：添加新令牌时，tokenizer会扩展其内部词汇表，而模型需要相应扩展其词嵌入矩阵。新添加的令牌对应的词向量通常会被初始化为随机值。

2. 维度一致性检查：在模型加载阶段，系统会严格检查各层权重张量的维度是否匹配。任何不匹配都会导致运行时错误。

3. 微调策略：添加新令牌后，建议对模型进行全面微调，特别是新添加的令牌对应的词向量部分，以确保它们能够学习到有意义的表示。

最佳实践

1. 在添加特殊令牌前，仔细规划所需的令牌集合，避免频繁修改词汇表
2. 记录词汇表变更情况，确保训练和推理阶段使用相同的tokenizer配置
3. 对于生产环境，建议将添加特殊令牌后的模型重新导出为完整检查点
4. 监控模型在新令牌上的表现，必要时进行针对性训练

通过以上方法，开发者可以灵活地为XTuner项目中的LLaVA模型添加所需特殊令牌，同时确保模型能够正确加载和训练。