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

TL;DR

This paper presents a framework for migrating atmospheric and oceanic AI models from PyTorch to MindSpore, optimizing for Chinese hardware, and evaluating performance improvements in speed, accuracy, and energy efficiency.

Contribution

It introduces a comprehensive migration and optimization framework tailored for Chinese chips, enabling hardware-independent climate modeling with preserved accuracy.

Findings

Migration preserves model accuracy

Significant speed improvements on Chinese chips

Enhanced energy efficiency and reduced dependency

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.