Liger-Kernel项目中Gemma2模型收敛问题分析与解决

问题背景

在Liger-Kernel项目的测试过程中，发现Gemma2模型无法通过收敛性测试。经过初步排查，问题定位在LigerRMSNorm层的实现上。当启用RMS归一化时，模型在训练过程中参数更新出现异常，导致收敛失败。

现象分析

通过简化测试场景，设置训练步数为2步，并打印模型输出的logits及其梯度，可以观察到以下现象：

1. 第一步训练时，无论是否使用LigerRMSNorm，模型输出和梯度表现一致
2. 第二步训练时，使用LigerRMSNorm的模型输出logits与预期结果出现偏差
3. 梯度值在第二步出现异常变化，表明反向传播过程存在问题

根本原因

深入分析Gemma2模型结构后发现，该模型在解码层中使用了两次RMS归一化操作：

• 注意力机制后的归一化(post_attention_layernorm)
• 前馈网络后的归一化(post_feedforward_layernorm)

问题出在LigerRMSNorm的反向传播实现上。原始实现中对输入梯度dY进行了原地操作(in-place operations)，这在Gemma2的特殊结构下会导致梯度计算错误。

解决方案

修复方案是避免对dY进行原地操作，改为创建新的张量来存储中间计算结果。具体修改包括：

1. 在计算dX时，不再直接使用dY作为输入
2. 创建新的临时张量来存储中间计算结果
3. 确保所有运算都不会意外修改输入梯度

这种修改虽然增加了少量内存开销，但保证了梯度计算的正确性。

技术启示

1. 在实现自定义归一化层时，需要特别注意反向传播的实现细节
2. 原地操作虽然能提高效率，但在复杂模型结构中可能引发难以察觉的问题
3. 对于包含多个归一化层的模型结构，梯度计算需要更加谨慎处理
4. 单元测试中增加对中间结果的检查有助于快速定位问题

总结

通过分析Gemma2模型结构和LigerRMSNorm实现，我们定位并解决了模型收敛问题。这个案例提醒我们，在优化深度学习框架时，不仅需要考虑计算效率，还需要确保数值计算的正确性，特别是在处理复杂模型结构时。对基础算子的实现需要经过严格的测试验证，才能保证在各种模型结构下都能正常工作。