[论文解读] Using Machine Learning at Scale in HPC Simulations with SmartSim: An Application to Ocean Climate Modeling

Sam Partee, Matthew J. Ellis+5 more|arXiv (Cornell University)|Apr 13, 2021

Oceanographic and Atmospheric Processes参考文献 31被引用 13

一句话总结

本论文提出 SmartSim，一种软件框架，可在高性能计算（HPC）模拟中实现可扩展的在线机器学习推理。该框架将 PyTorch、TensorFlow 和 ONNX 模型集成到基于 Fortran 的海洋气候模型（如 MOM6）中，实现了 12 个成员的全球海洋模拟，在 120 年的模拟过程中完成 9700 亿次推理，性能稳定且运行时开销极低。

ABSTRACT

We demonstrate the first climate-scale, numerical ocean simulations improved through distributed, online inference of Deep Neural Networks (DNN) using SmartSim. SmartSim is a library dedicated to enabling online analysis and Machine Learning (ML) for traditional HPC simulations. In this paper, we detail the SmartSim architecture and provide benchmarks including online inference with a shared ML model on heterogeneous HPC systems. We demonstrate the capability of SmartSim by using it to run a 12-member ensemble of global-scale, high-resolution ocean simulations, each spanning 19 compute nodes, all communicating with the same ML architecture at each simulation timestep. In total, 970 billion inferences are collectively served by running the ensemble for a total of 120 simulated years. Finally, we show our solution is stable over the full duration of the model integrations, and that the inclusion of machine learning has minimal impact on the simulation runtimes.

研究动机与目标

• 解决在不产生文件 I/O 瓶颈的前提下，将机器学习与用 Fortran/C/C++ 编写的传统 HPC 模拟耦合的挑战。
• 实现在长期气候模拟过程中对深度神经网络进行实时、分布式推理。
• 克服 HPC 工作负载与现代机器学习库之间在语言和数据交换方面的障碍。
• 在气候尺度分辨率下，展示 ML 增强型 HPC 模拟的可扩展性与稳定性。
• 提供一种灵活的、与库无关的接口，支持异构 HPC 系统和多种机器学习框架。

提出的方法

• 利用 SmartSim 的基础设施和客户端库，将机器学习推理与模拟逻辑解耦。
• 使用 SmartRedis API 实现从 Fortran 模拟中低延迟、远程执行机器学习模型。
• 将 JIT 转录的 PyTorch 模型集成到 MOM6 中，以替代传统的湍流动能参数化方案。
• 在 12 个集合成员中共享同一个机器学习模型，每个成员运行在 19 个节点上，并在每个时间步实现同步推理。
• 利用基于 Redis 的数据交换机制，消除文件 I/O 开销，实现实时模拟与机器学习组件之间的数据流。
• 支持异构 HPC 系统中的 CPU 和 GPU 工作负载，确保在多样化硬件和软件堆栈中的互操作性。

实验结果

研究问题

• RQ1 在不破坏模拟稳定性或性能的前提下，能否将机器学习集成到大规模、生产级的 HPC 气候模拟中？
• RQ2 单个机器学习模型在大规模环境下，能否高效地为多个 HPC 模拟实例提供分布式、实时推理？
• RQ3 在使用共享模型架构集成到全球海洋模型时，在线机器学习推理的性能开销是多少？
• RQ4 SmartSim 框架能否实现 Fortran 基础模拟与现代机器学习库（如 PyTorch 和 TensorFlow）之间的无缝、语言无关的通信？
• RQ5 通过 SmartSim 集成机器学习推理是否能在长时间气候模拟中保持数值稳定性？

主要发现

• SmartSim-EKE 集合成功完成了 12 次全球高分辨率海洋模拟，每次持续 10 年，未观察到任何数值不稳定性。
• 在 120 年的模拟期间，集合内共执行了 9700 亿次推理操作，展示了在接近百亿亿次规模下的可扩展性。
• 通过 SmartSim 集成机器学习推理引入的性能开销极小，模拟运行时间基本不受影响。
• 该框架在异构 HPC 系统上实现了数千个处理器的线性扩展，支持 CPU 和 GPU 工作负载。
• 在所有集合成员中使用共享机器学习模型，实现了无重复模型的统一、同步推理。
• 该解决方案已具备生产就绪性且完全可复现，所有源代码、数据集和训练模型均已公开发布在 GitHub 上。