Using Machine Learning at Scale in HPC Simulations with SmartSim: An Application to Ocean Climate Modeling

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, highresolution 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. 1

[1]  Noah D. Brenowitz,et al.  Prognostic Validation of a Neural Network Unified Physics Parameterization , 2018, Geophysical Research Letters.

[2]  Nikos Komodakis,et al.  Wide Residual Networks , 2016, BMVC.

[3]  Jian Sun,et al.  Deep Residual Learning for Image Recognition , 2015, 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR).

[4]  Pierre Gentine,et al.  Deep learning to represent subgrid processes in climate models , 2018, Proceedings of the National Academy of Sciences.

[5]  Nicholas Geneva,et al.  Quantifying model form uncertainty in Reynolds-averaged turbulence models with Bayesian deep neural networks , 2018, J. Comput. Phys..

[6]  M. Holzer,et al.  Recent increase in oceanic carbon uptake driven by weaker upper-ocean overturning , 2017, Nature.

[7]  P. Gent,et al.  Isopycnal mixing in ocean circulation models , 1990 .

[8]  Prabhat,et al.  CosmoFlow: Using Deep Learning to Learn the Universe at Scale , 2018, SC18: International Conference for High Performance Computing, Networking, Storage and Analysis.

[9]  P. O'Gorman,et al.  Using Machine Learning to Parameterize Moist Convection: Potential for Modeling of Climate, Climate Change, and Extreme Events , 2018, Journal of Advances in Modeling Earth Systems.

[10]  P. Gent,et al.  The NCAR Climate System Model, Version One* , 1998 .

[11]  Pierre Baldi,et al.  A Fortran-Keras Deep Learning Bridge for Scientific Computing , 2020, Sci. Program..

[12]  C. Wunsch,et al.  Ocean Circulation Kinetic Energy: Reservoirs, Sources, and Sinks , 2009 .

[13]  Jochem Marotzke,et al.  The Oceanic Eddy Heat Transport , 2002 .

[14]  Pierre Gentine,et al.  Achieving Conservation of Energy in Neural Network Emulators for Climate Modeling , 2019, ArXiv.

[15]  Zhuowen Tu,et al.  Aggregated Residual Transformations for Deep Neural Networks , 2016, 2017 IEEE Conference on Computer Vision and Pattern Recognition (CVPR).

[16]  D. Chelton,et al.  Global observations of nonlinear mesoscale eddies , 2011 .

[17]  Xiaomeng Huang,et al.  A Moist Physics Parameterization Based on Deep Learning , 2020, Journal of Advances in Modeling Earth Systems.

[18]  Nick Hand,et al.  Launching Python Applications on Peta-scale Massively Parallel Systems , 2016 .

[19]  Hrvoje Jasak,et al.  A tensorial approach to computational continuum mechanics using object-oriented techniques , 1998 .

[20]  Noah D. Brenowitz,et al.  fv3gfs-wrapper: a Python wrapper of the FV3GFS atmospheric model , 2021, Geoscientific Model Development.

[21]  Nicolas Gruber,et al.  The Oceanic Sink for Anthropogenic CO2 , 2004, Science.

[22]  Jeremy McGibbon,et al.  Machine Learning Climate Model Dynamics: Offline versus Online Performance. , 2020, 2011.03081.

[23]  R. Feely,et al.  The oceanic sink for anthropogenic CO2 from 1994 to 2007 , 2019, Science.

[24]  Andrew Stuart,et al.  Earth System Modeling 2.0: A Blueprint for Models That Learn From Observations and Targeted High‐Resolution Simulations , 2017, 1709.00037.

[25]  Robert Hallberg,et al.  Using a resolution function to regulate parameterizations of oceanic mesoscale eddy effects , 2013 .

[26]  E. Jennings,et al.  A turbulent eddy-viscosity surrogate modeling framework for Reynolds-Averaged Navier-Stokes simulations , 2020 .

[27]  Vladimir M. Krasnopolsky,et al.  Using Ensemble of Neural Networks to Learn Stochastic Convection Parameterizations for Climate and Numerical Weather Prediction Models from Data Simulated by a Cloud Resolving Model , 2013, Adv. Artif. Neural Syst..

[28]  R. Hallberg,et al.  Energy budget-based backscatter in an eddy permitting primitive equation model , 2015 .

[29]  Kevin E. Trenberth,et al.  Earth’s Energy Imbalance , 2014 .

[30]  Syukuro Manabe,et al.  Thermal Equilibrium of the Atmosphere with a Given Distribution of Relative Humidity , 1967 .

[31]  Chris Blanton,et al.  The GFDL Global Ocean and Sea Ice Model OM4.0: Model Description and Simulation Features , 2019, Journal of Advances in Modeling Earth Systems.

[32]  Veronika Eyring,et al.  Overview of the Coupled Model Intercomparison Project Phase 6 (CMIP6) experimental design and organization , 2015 .