Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Deploying Atmospheric and Oceanic AI Models on Chinese Hardware and Framework: Migration Strategies, Performance Optimization and Analysis

Yuze Sun, Wentao Luo, Yanfei Xiang, Jiancheng Pan, Jiahao Li, Quan Zhang, Xiaomeng Huang

TL;DR

This work tackles GPU-centric limitations in climate modeling by migrating large atmospheric and oceanic AI models from PyTorch to MindSpore and optimizing for Chinese hardware such as Ascend and DCU. It develops a migration-and-optimization framework, leveraging static graphs, operator mapping, mixed precision, and distributed training to achieve good functional parity with GPUs while improving energy efficiency. Empirical results on FourCastNet, GraphCast, and AI-GOMS across ERA5 and HYCOM show accuracy within $<5\%$ of the original and GPU-like throughput with favorable energy profiles, especially when using MindSpore on Ascend. The study offers practical guidance for deploying large-scale climate models on domestic hardware, supporting greater technological independence and ecosystem development.

Abstract

With the growing role of artificial intelligence in climate and weather research, efficient model training and inference are in high demand. Current models like FourCastNet and AI-GOMS depend heavily on GPUs, limiting hardware independence, especially for Chinese domestic hardware and frameworks. To address this issue, we present a framework for migrating large-scale atmospheric and oceanic models from PyTorch to MindSpore and optimizing for Chinese chips, and evaluating their performance against GPUs. The framework focuses on software-hardware adaptation, memory optimization, and parallelism. Furthermore, the model's performance is evaluated across multiple metrics, including training speed, inference speed, model accuracy, and energy efficiency, with comparisons against GPU-based implementations. Experimental results demonstrate that the migration and optimization process preserves the models' original accuracy while significantly reducing system dependencies and improving operational efficiency by leveraging Chinese chips as a viable alternative for scientific computing. This work provides valuable insights and practical guidance for leveraging Chinese domestic chips and frameworks in atmospheric and oceanic AI model development, offering a pathway toward greater technological independence.

Deploying Atmospheric and Oceanic AI Models on Chinese Hardware and Framework: Migration Strategies, Performance Optimization and Analysis

TL;DR

This work tackles GPU-centric limitations in climate modeling by migrating large atmospheric and oceanic AI models from PyTorch to MindSpore and optimizing for Chinese hardware such as Ascend and DCU. It develops a migration-and-optimization framework, leveraging static graphs, operator mapping, mixed precision, and distributed training to achieve good functional parity with GPUs while improving energy efficiency. Empirical results on FourCastNet, GraphCast, and AI-GOMS across ERA5 and HYCOM show accuracy within of the original and GPU-like throughput with favorable energy profiles, especially when using MindSpore on Ascend. The study offers practical guidance for deploying large-scale climate models on domestic hardware, supporting greater technological independence and ecosystem development.

Abstract

With the growing role of artificial intelligence in climate and weather research, efficient model training and inference are in high demand. Current models like FourCastNet and AI-GOMS depend heavily on GPUs, limiting hardware independence, especially for Chinese domestic hardware and frameworks. To address this issue, we present a framework for migrating large-scale atmospheric and oceanic models from PyTorch to MindSpore and optimizing for Chinese chips, and evaluating their performance against GPUs. The framework focuses on software-hardware adaptation, memory optimization, and parallelism. Furthermore, the model's performance is evaluated across multiple metrics, including training speed, inference speed, model accuracy, and energy efficiency, with comparisons against GPU-based implementations. Experimental results demonstrate that the migration and optimization process preserves the models' original accuracy while significantly reducing system dependencies and improving operational efficiency by leveraging Chinese chips as a viable alternative for scientific computing. This work provides valuable insights and practical guidance for leveraging Chinese domestic chips and frameworks in atmospheric and oceanic AI model development, offering a pathway toward greater technological independence.
Paper Structure (20 sections, 3 equations, 4 figures, 5 tables)

This paper contains 20 sections, 3 equations, 4 figures, 5 tables.

Table of Contents

  1. Introduction
  2. Method
  3. Model Migration Framework Design.
  4. Optimization Strategies for Adapting to Chips.
  5. Evaluation Benchmarks.
  6. Experimental Setup
  7. Hardware Configuration.
  8. Datasets and Models.
  9. Training and Inference Configuration.
  10. Experimental Results
  11. Training Efficiency Analysis.
  12. Training Accuracy Analysis.
  13. Inference Efficiency Analysis.
  14. Inference Accuracy Analysis.
  15. Energy Consumption Analysis.
  16. ...and 5 more sections

Figures (4)

  • Figure 1: The migration and optimization framework for atmospheric and oceanic models from PyTorch to MindSpore based on different hardware.
  • Figure 2: Comparison of Inference Accuracy between MindSpore and PyTorch Frameworks
  • Figure S1: Training Performance of AI-GOMS under different chip and frameworks
  • Figure S2: The migration process from a model based on the Pytorch framework to the Mindspore framework