Liger-Kernel项目中的交叉熵损失函数收敛测试问题分析

问题背景

在Liger-Kernel项目中，开发团队设计了一个名为test_mini_models_with_logits的收敛性测试，用于验证经过monkey patch修改后的模型与原始模型在行为和输出上的一致性。这个测试主要比较两种模型（使用和不使用Liger补丁）的损失值和最后一步的logits输出。

问题现象

在最新版本的transformers库中，这个收敛测试开始失败。具体表现为测试比较的两个张量差异显著——实际上比较的是原始模型的logits和补丁模型的logits梯度，而非预期的两个logits张量。

技术原理

问题的根源在于LigerCrossEntropy的内存优化设计：

1. 内存优化机制：LigerCrossEntropy为了节省内存，直接在logits张量上执行原地操作(in-place operation)，将logits梯度存储在logits张量自身中。

2. 原始测试设计：在CausalLMLoss上下文中，测试通过创建shift_logits和shift_labels临时张量来保存移位后的数据，然后将这些临时张量传递给LigerCrossEntropy。这样，原地操作只会影响临时张量，原始logits仍可被测试捕获。

3. transformers优化：transformers的最新优化(PR #35646)移除了额外的内存分配，但带来了副作用——不再分配新张量来跟踪原始logits，导致测试实际上比较的是logits和其梯度。

解决方案

开发团队提出了几种解决方案思路：

1. 比较logits梯度：既然无法直接获取logits，可以计算原始CrossEntropy的logits梯度，然后与LigerCrossEntropy的logits梯度进行比较。

2. 优化BF16场景：发现BF16场景下由于logits.float()调用，LigerCrossEntropy将梯度保存在FP32张量而非原始BF16张量中，这意味着可以进一步减少内存分配。

3. 封装ForCausalLMLoss：建议编写一个包装函数来修补整个ForCausalLMLoss，而不仅仅是nn.functional.CrossEntropy，这样可以更全面地优化内存使用。

技术影响

这一问题的解决过程揭示了几个重要技术点：

1. 内存优化与测试的平衡：性能优化可能会影响测试的可观测性，需要在两者之间找到平衡。

2. 原地操作的风险：使用原地操作虽然节省内存，但会破坏原始数据，可能影响调试和验证。

3. 类型转换的影响：BF16与FP32之间的转换不仅影响计算精度，还会影响内存分配策略。

最终方案

项目团队最终决定：

1. 修改收敛测试，改为比较logits梯度而非logits本身
2. 优化LigerForCausalLMLoss实现，减少内存分配
3. 保持测试的严格性，通过比较梯度、损失和所有模型参数来确保补丁的正确性

这一解决方案既保持了LigerCrossEntropy原始设计理念，又确保了测试的有效性，同时进一步优化了内存使用效率。