Keras 3分布式训练中的NaN问题分析与解决方案

问题背景

在深度学习框架Keras 3中，开发者报告了一个关于分布式训练的重要问题：当使用tf.distribute.MirroredStrategy进行多GPU训练时，模型训练过程中会出现NaN（非数值）值。这个问题不仅出现在复杂的卷积神经网络中，甚至在最简单的恒等映射模型上也会重现。

问题重现

开发者提供了一个基于MNIST数据集的简单卷积神经网络示例。当启用分布式训练时（设置do_bug=True），模型训练过程中损失值和准确率都变成了NaN：

422/422 ━━━━━━━━━━━━━━━━━━━━ 11s 22ms/step - accuracy: nan - loss: nan - val_accuracy: 0.9697 - val_loss: 0.0839

而关闭分布式训练后（设置do_bug=False），模型训练正常：

422/422 ━━━━━━━━━━━━━━━━━━━━ 4s 6ms/step - accuracy: 0.7623 - loss: 0.7774 - val_accuracy: 0.9767 - val_loss: 0.0851

问题根源

经过深入分析，发现问题源于Keras 3中变量聚合方式的处理。在分布式训练中，Keras 3错误地将"sum"聚合方式映射到了tf.VariableAggregation.SUM，而实际上应该使用tf.VariableAggregation.NONE。

具体来说，在_map_aggregation函数中，错误的映射关系如下：

def _map_aggregation(self, aggregation):
    mapping = {
        "none": tf.VariableAggregation.NONE,
        "sum": tf.VariableAggregation.SUM,  # 这是错误的映射
        "mean": tf.VariableAggregation.MEAN,
        "only_first_replica": tf.VariableAggregation.ONLY_FIRST_REPLICA,
    }
    return mapping[aggregation]

正确的映射应该是：

def _map_aggregation(self, aggregation):
    mapping = {
        "none": tf.VariableAggregation.NONE,
        "sum": tf.VariableAggregation.NONE,  # 正确的映射
        "mean": tf.VariableAggregation.MEAN,
        "only_first_replica": tf.VariableAggregation.ONLY_FIRST_REPLICA,
    }
    return mapping[aggregation]

技术细节

在TensorFlow的分布式训练中，变量聚合方式决定了不同设备上的变量如何同步：

1. NONE：不进行聚合，每个设备保持自己的变量副本
2. SUM：对所有设备上的变量值求和
3. MEAN：对所有设备上的变量值取平均
4. ONLY_FIRST_REPLICA：只使用第一个设备上的变量值

在Keras 3中，错误的聚合方式会导致梯度计算异常，最终产生NaN值。这个问题在Keras 3.5.0版本中不存在，但在3.9.0版本中出现，说明是在某个中间提交引入的回归问题。

解决方案

目前临时的解决方案是修改_map_aggregation函数，将"sum"映射到NONE而不是SUM。Keras团队已经确认这个问题，并将在未来的版本中修复。

对于开发者来说，可以采取以下措施：

1. 暂时回退到Keras 3.5.0版本
2. 使用修改后的_map_aggregation函数
3. 等待官方发布修复版本

影响范围

这个问题影响所有使用Keras 3进行分布式训练的场景，包括但不限于：

• 多GPU训练
• 多机训练
• 任何使用tf.distribute策略的训练

值得注意的是，这个问题在tf-keras（TensorFlow内置的Keras实现）中不存在，说明是Keras 3特有的问题。

结论

分布式训练中的NaN问题揭示了Keras 3在变量聚合处理上的一个关键缺陷。通过理解问题的根源和解决方案，开发者可以避免在分布式训练中遇到类似的陷阱。Keras团队正在积极解决这个问题，预计在未来的版本中会提供官方的修复方案。