TorchTitan项目中大规模模型训练中的NaN问题分析与解决

引言

在分布式深度学习训练中，特别是在使用大规模模型如405B参数量的模型时，训练初期出现NaN（非数值）是一个常见但棘手的问题。本文将深入分析在TorchTitan项目中使用FSDP（完全分片数据并行）和TP（张量并行）在256个GPU上训练405B模型时遇到的NaN问题，并提供系统性的解决方案。

问题现象

当使用FSDP+TP组合在256个GPU上训练405B参数的模型时，训练过程在第一步就出现了NaN值。这种情况通常表明数值不稳定或硬件问题，需要系统性的排查方法。

排查方法与解决方案

1. 硬件层面检查

首先应考虑硬件问题，特别是当NaN问题在多次重启后仍然出现时。推荐使用以下环境变量进行检测：

export TORCH_NCCL_NAN_CHECK=1

这个设置会在所有集体通信操作中自动检查NaN值，并在检测到时触发断言，明确指出触发问题的节点和rank。根据经验，这通常能发现GPU硬件问题（如XID错误），简单的GPU重启往往就能解决问题。

2. 并行策略优化

对于405B参数量的模型，必须使用某种形式的2D并行策略。可以尝试以下两种方案：

• 张量并行（TP）：传统的模型并行方式
• 上下文并行（CP）：一种替代方案，有时能提供更好的稳定性

实验表明，在256个GPU上运行405B模型时，这两种策略都能稳定运行，因此如果出现NaN更可能是硬件问题而非策略问题。

3. 显存管理优化

虽然显存分配重试通常不会直接影响数值稳定性，但高显存占用（如超过95%）可能导致CUDA缓存重试，在某些大规模训练场景中可能与NaN问题相关。建议：

• 降低批次大小（如设为1）进行测试
• 监控显存使用情况，保持合理余量

4. 通信库版本问题

实际案例表明，某些情况下NaN问题源于通信库（如CCL）的版本兼容性问题。确保使用经过验证的稳定版本是关键。

性能优化建议

在解决NaN问题后，针对405B模型训练还可以考虑以下性能优化措施：

1. 激活检查点（AC）策略：可以尝试不同的AC策略（选择性AC→无AC→完全AC）来平衡内存和性能
2. 混合并行策略：结合PP（流水线并行）+FSDP+TP的配置可以进一步提升训练效率
3. 硬件选择：使用H200等具有更高HBM的GPU可以获得更好的批次大小和MFU（模型浮点运算利用率）

实际训练数据参考

在H200上的实际训练数据显示：

• 批次大小：2（使用完全AC）
• 训练速度：约14小时完成3000步
• MFU：37-40%之间（H200上的表现）

结论

大规模模型训练中的NaN问题需要系统性的排查方法。从硬件检查开始，逐步验证并行策略和显存配置，最后考虑软件栈兼容性。TorchTitan项目中的实践表明，通过科学的方法可以稳定地在256个GPU上训练405B参数的模型。对于希望进一步优化性能的用户，建议尝试不同的并行组合和硬件配置，同时密切监控训练过程中的各项指标。