TransformerEngine中激活重计算与FP8量化配方的保存问题分析

背景介绍

在深度学习训练过程中，TransformerEngine项目采用了FP8量化技术来优化模型性能。其中，激活重计算(activation recomputation)是一种常见的内存优化技术，它通过在反向传播时重新计算前向传播的中间结果，而不是保存所有激活值，从而减少内存占用。

问题发现

在TransformerEngine的实现中，当启用激活重计算功能时(activation_recompute=True)，发现FP8GlobalStateManager的上下文状态与首次计算激活时的状态不一致。具体表现为：

1. 在反向传播调用时，FP8量化配方(recipe)没有被正确恢复
2. 对于新的量化配方，特别是per-tensor当前缩放(current scaling)的情况，激活重计算会被跳过
3. 默认的延迟缩放(delayed scaling)配方会被错误地使用

技术分析

FP8量化需要维护一系列状态信息，包括：

1. 量化配方(recipe)：定义了如何进行量化的参数和策略
2. 缩放因子(scale)和最大值(amax)：在延迟缩放策略中需要保存和恢复
3. FP8启用状态：指示当前是否应该使用FP8量化

在正常的正向-反向传播过程中，这些状态会被自动保存和恢复。然而，在激活重计算场景下，由于计算流程的特殊性，量化配方没有被正确恢复，导致计算结果不一致。

解决方案

经过讨论，确定了以下解决方案原则：

1. 不应该要求用户在调用loss.backward()时必须手动创建FP8自动转换上下文
2. 应该像保存量化器(quantizer)状态一样，在检查点中保存量化配方
3. 在从检查点加载时，需要正确恢复FP8GlobalStateManager的状态

具体实现上，需要：

1. 在保存检查点时同时保存量化配方和FP8启用状态
2. 在激活重计算阶段恢复这些状态
3. 确保恢复的状态与原始正向传播时完全一致

实现建议

一个可行的实现方案是扩展FP8GlobalStateManager的功能，使其能够：

1. 在正向传播时记录当前的量化配方和FP8启用状态
2. 将这些信息与激活值一起保存
3. 在反向传播的激活重计算阶段恢复这些状态

这样无论loss.backward()在什么上下文中被调用，都能保证使用正确的量化参数进行重计算。

总结

TransformerEngine中FP8量化与激活重计算的交互是一个需要特别注意的环节。正确保存和恢复量化配方对于保证数值一致性至关重要。通过扩展状态管理器的功能，可以确保在激活重计算时使用与原始正向传播相同的量化参数，从而保证训练过程的正确性和稳定性。