亚马逊新论文指出，KV缓存策略不只是推理阶段的优化手段，它还会反过来影响模型该怎么训练。如果模型在推理时会丢掉一部分上下文，那么训练时也应该模拟这种遗忘。

匹配缓存策略，防止长上下文失效

打开网易新闻 查看精彩图片

把微调方式和KV缓存策略对齐，就能避免长上下文能力在推理时崩掉。在128k token的测试里，全注意力微调的模型一旦换成稀疏推理，经常输出又长又没意义的回答。

而匹配缓存策略微调的模型，不仅能正常作答，还知道什么时候该停下来。这个方法支持任意缓存策略，并且可以在40 GB A100上计算4B模型的梯度。