GraphCast模型在AMD CPU上的运行问题分析与解决方案
背景介绍
GraphCast是Google DeepMind开发的一款基于图神经网络的天气预报模型。该模型通常运行在GPU或TPU上以获得最佳性能,但最近有开发者尝试在AMD CPU上运行其演示案例gencast_mini_demo.ipynb时遇到了NaN(非数值)输出的问题。
问题现象
当尝试在CPU上运行GraphCast的预测流程时,模型输出全部变为NaN值。经过调试发现,这一问题出现在模型的自回归预测过程中,特别是在处理噪声水平计算时。
技术分析
根本原因
深入调试后发现,问题出在apply_stochastic_churn
函数中的噪声水平计算部分。具体来说,当计算new_noise_level^2 - noise_level^2
时,由于浮点数精度问题,结果可能变为一个极小的负值(如-6.66e-11)。随后对这个负值取平方根操作导致了NaN的产生。
数值稳定性问题
在机器学习模型中,数值稳定性是一个常见挑战。特别是在涉及平方根、对数等数学运算时,输入值的微小变化可能导致输出结果完全无效。在本案例中,虽然理论上new_noise_level^2 - noise_level^2
应该总是非负的,但浮点数运算的精度限制导致了负值的出现。
JIT编译的影响
有趣的是,当禁用JIT编译(通过设置jax_disable_jit=True
)时,模型能够正常运行。这表明问题可能与XLA编译器对浮点运算的优化处理方式有关。JIT编译通常会重组计算图以提高性能,但有时会引入数值稳定性问题。
解决方案
临时修复方案
最直接的解决方案是在计算平方根前确保被开方数非负:
diff = new_noise_level**2 - noise_level**2
extra_noise_stddev = (jnp.sqrt(jnp.where(diff < 0., 0., diff))
* noise_level_inflation_factor)
这种方法简单有效,通过显式处理负值情况,避免了NaN的产生。
更稳健的数值处理
更稳健的做法是使用jnp.maximum
函数:
extra_noise = jnp.maximum(0, new_noise_level**2 - noise_level**2)
extra_noise_stddev = jnp.sqrt(extra_noise) * noise_level_inflation_factor
这种方法在数学上更清晰,表达了"我们只对非负部分感兴趣"的意图。
性能考量
虽然这些修改解决了NaN问题,但在CPU上运行GraphCast仍然面临性能挑战:
- 计算速度显著慢于GPU/TPU
- 内存占用较高(约17GB)
- 并行效率受限
对于实际应用,建议在完成调试后重新启用JIT编译以获得更好的性能,或者考虑使用云端的GPU/TPU资源进行大规模预测。
结论
本次调试揭示了GraphCast模型在CPU环境下运行时的一个数值稳定性问题,并提供了有效的解决方案。这一经验也提醒我们:
- 在跨平台部署机器学习模型时,数值稳定性需要特别关注
- 浮点数运算的精度问题可能在看似无害的数学表达式中突然出现
- 调试工具如
jax.debug.print
和jax_debug_nans
对于定位数值问题非常有用
通过这些解决方案,开发者现在可以在AMD CPU上成功运行GraphCast模型,为没有GPU/TPU硬件资源的用户提供了更多可能性。
HunyuanImage-3.0
HunyuanImage-3.0 统一多模态理解与生成,基于自回归框架,实现文本生成图像,性能媲美或超越领先闭源模型00ops-transformer
本项目是CANN提供的transformer类大模型算子库,实现网络在NPU上加速计算。C++045Hunyuan3D-Part
腾讯混元3D-Part00GitCode-文心大模型-智源研究院AI应用开发大赛
GitCode&文心大模型&智源研究院强强联合,发起的AI应用开发大赛;总奖池8W,单人最高可得价值3W奖励。快来参加吧~0289Hunyuan3D-Omni
腾讯混元3D-Omni:3D版ControlNet突破多模态控制,实现高精度3D资产生成00GOT-OCR-2.0-hf
阶跃星辰StepFun推出的GOT-OCR-2.0-hf是一款强大的多语言OCR开源模型,支持从普通文档到复杂场景的文字识别。它能精准处理表格、图表、数学公式、几何图形甚至乐谱等特殊内容,输出结果可通过第三方工具渲染成多种格式。模型支持1024×1024高分辨率输入,具备多页批量处理、动态分块识别和交互式区域选择等创新功能,用户可通过坐标或颜色指定识别区域。基于Apache 2.0协议开源,提供Hugging Face演示和完整代码,适用于学术研究到工业应用的广泛场景,为OCR领域带来突破性解决方案。00- HHowToCook程序员在家做饭方法指南。Programmer's guide about how to cook at home (Chinese only).Dockerfile09
- PpathwayPathway is an open framework for high-throughput and low-latency real-time data processing.Python00
热门内容推荐
最新内容推荐
项目优选









