在mlx-examples项目中优化Phi-3模型微调时的EOS标记处理

问题背景

在使用mlx-examples项目中的Phi-3模型进行二分类任务微调时，开发者发现模型在推理阶段存在一个常见问题：模型生成预测类别后不会自动停止，而是继续生成无关内容。这种情况在序列生成任务中尤为常见，特别是当模型没有被明确告知何时应该停止生成时。

技术分析

这个问题本质上与模型训练时EOS(End Of Sequence)标记的处理方式有关。在Transformer架构的语言模型中，EOS标记起着至关重要的作用：

1. 训练阶段：EOS标记帮助模型学习何时应该停止生成
2. 推理阶段：模型遇到EOS标记时会停止生成新token

当前mlx-examples项目中的ChatDataset实现没有显式添加EOS标记，这导致模型在微调过程中没有学习到何时应该停止生成，从而在推理时产生多余输出。

解决方案

通过对ChatDataset类的简单修改，可以解决这个问题。关键修改是在每个训练样本的末尾显式添加EOS标记：

class ChatDataset(Dataset):
    def __init__(self, path: Path, tokenizer: PreTrainedTokenizer):
        super().__init__(path)
        self._tokenizer = tokenizer

    def __getitem__(self, idx: int):
        messages = self._data[idx]["messages"]
        text = self._tokenizer.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )
        text = text + self._tokenizer.eos_token  # 显式添加EOS标记
        return text

技术细节

1. 为什么需要手动添加EOS标记：

   • 大多数现代分词器(如HuggingFace的tokenizer)不会自动在样本末尾添加EOS标记
   • 对于分类任务，明确的停止信号尤为重要

2. 修改后的效果：

   • 模型在训练时会学习到EOS标记的位置
   • 推理时生成目标类别后会正确停止
   • 避免了多余输出的问题

3. 适用范围：

   • 不仅适用于二分类任务
   • 对所有需要精确控制生成长度的任务都有效

最佳实践建议

1. 对于分类任务，建议始终添加EOS标记
2. 对于生成任务，可根据需求决定是否添加
3. 可以考虑将EOS标记添加作为可配置选项，以适应不同任务需求

总结

这个简单的修改显著提升了Phi-3模型在分类任务中的表现，特别是在输出控制方面。它展示了在模型微调过程中，正确处理特殊标记(如EOS)的重要性。这种技术不仅适用于mlx-examples项目，也可以推广到其他基于Transformer架构的模型微调场景中。