Transformers项目中RecurrentGemma模型长文本推理崩溃问题分析

问题背景

在Transformers项目的最新版本(v4.50.3)中，研究人员发现RecurrentGemma-9b-it模型在处理长文本输入时会出现崩溃现象。具体表现为：当输入文本长度超过模型滑动窗口注意力机制设定的2048个token限制时，程序会抛出张量维度不匹配的运行时错误。

技术细节分析

RecurrentGemma模型采用了滑动窗口注意力机制，这种设计原本是为了处理长序列而优化的。然而在实际使用中发现，当输入序列超过预设的2048个token窗口大小时，模型在生成阶段会触发以下关键错误：

RuntimeError: The size of tensor a (2048) must match the size of tensor b (2402) at non-singleton dimension 3

深入分析表明，这个问题源于模型内部对注意力掩码(attention_mask)的处理逻辑。在解码阶段，模型假设接收到的attention_mask已经按照滑动窗口大小进行了裁剪，但实际上当前的通用输入准备函数并未执行这一裁剪操作。

问题演变历史

通过代码审查发现，这个问题经历了两个关键阶段：

1. 在v4.46.0版本中，移除了模型特定的prepare_inputs_for_generation方法，同时错误地中断了past_key_values参数的传递，这曾导致模型生成功能完全失效。

2. 在v4.48.3版本中，恢复了kwargs参数的传递，虽然修复了生成功能，但长文本处理的问题仍然存在。

解决方案

针对这个问题，技术团队提出了明确的修复方案：在_update_causal_mask()方法内部添加对attention_mask的裁剪逻辑，确保其长度不超过因果掩码的尺寸。这一修改与原始prepare_inputs_for_generation方法的处理方式保持一致，能够有效解决长文本输入时的张量维度不匹配问题。

对开发者的建议

对于需要使用RecurrentGemma模型处理长文本的开发者，建议：

1. 暂时将输入文本控制在2048个token以内
2. 关注Transformers项目的后续更新，及时获取修复版本
3. 如需立即使用，可以考虑回退到v4.42.4等早期版本

这个问题凸显了在模型优化过程中保持向后兼容性的重要性，特别是在处理具有特殊架构要求的模型时。技术团队正在通过回归测试确保类似问题不会再次出现。