Torchtitan项目中FP8行量化训练出现NaN损失问题的分析与解决

问题背景

在Torchtitan项目（一个基于PyTorch的大模型训练框架）中，用户报告在使用Llama3 8B模型进行FP8行量化训练时出现了NaN（非数字）损失值的问题。该问题特别出现在使用rowwise量化方案、模型编译（torch.compile）和RMSNorm层组合的情况下。

问题现象

当用户尝试使用以下配置运行训练时：

• 模型：Llama3 8B
• 量化：FP8行量化(rowwise)
• 训练配置：启用torch.compile
• 归一化层：RMSNorm

训练过程中会出现NaN损失值，导致训练失败。值得注意的是，同样的配置在几周前可以正常工作。

技术分析

FP8量化简介

FP8（8位浮点数）是一种新兴的深度学习量化格式，它能在保持模型精度的同时显著减少内存占用和计算开销。行量化(rowwise)是一种特殊的量化策略，它为矩阵的每一行使用独立的缩放因子，相比全局量化能更好地保留精度。

问题根源追踪

经过深入排查，发现问题源于PyTorch核心代码中的一个特定提交。这个提交修改了RMSNorm的实现，与torch.compile和FP8行量化产生了不良交互。具体表现为：

1. 在FSDP（完全分片数据并行）的foreach_all_gather_copy_out函数中出现了NaN值
2. 问题仅出现在RMSNorm层，替换为LayerNorm后问题消失
3. 问题具有非确定性，有时表现为NaN，有时表现为CUDA内核错误

问题复现范围

通过二分法排查PyTorch提交历史，确定了问题引入的具体版本范围：

• 正常工作的最后版本：PyTorch 2.6.0（2025年1月29日发布）
• 问题首次出现的提交：2025年3月8日的某个特定修改

临时解决方案

在等待PyTorch核心团队修复根本问题的同时，Torchtitan项目提供了两种临时解决方案：

1. 环境变量法：设置TORCHINDUCTOR_EMULATE_PRECISION_CASTS=1，强制模拟精度转换
2. 代码修改法：将RMSNorm替换为LayerNorm（已在Torchtitan #1108中实现）

技术启示

1. 量化训练的敏感性：低精度训练对数值稳定性要求极高，任何微小的数值处理不当都可能导致NaN
2. 编译优化的复杂性：torch.compile的优化可能暴露底层数值计算问题
3. 版本兼容性：深度学习框架的快速迭代可能引入意外的回归问题

长期解决方案

PyTorch核心团队正在处理这个问题的根本修复，涉及以下几个方面：

1. 改进FP8在FSDP中的处理逻辑
2. 增强RMSNorm的数值稳定性
3. 完善torch.compile对量化操作的支持

最佳实践建议

对于使用Torchtitan进行FP8量化训练的用户：

1. 密切关注PyTorch版本更新
2. 在启用新特性时进行充分的稳定性测试
3. 保持对训练过程的监控，特别是损失值和梯度变化
4. 考虑使用TORCHINDUCTOR_EMULATE_PRECISION_CASTS=1作为预防措施

这个问题展示了深度学习系统底层优化与高层抽象之间复杂的交互关系，也提醒我们在追求性能优化的同时需要保持对数值稳定性的高度关注。