[論文レビュー] Proof-of-concept: Using ChatGPT to Translate and Modernize an Earth System Model from Fortran to Python/JAX

Anthony Zhou, Linnia Hawkins+1 more|arXiv (Cornell University)|Feb 13, 2024

Computational Physics and Python Applications被引用数 5

ひとこと要約

本論文は、GPT-4を用いてFortran Earth System ModelのコンポーネントをPython/JAXへ翻訳する semi-automated workflow を示し、GPU 加速の速度アップを達成し、パラメータ推定の自動微分を可能にする。CESMのCLMにおける葉レベルの光合成に焦点を当てている。

ABSTRACT

Earth system models (ESMs) are vital for understanding past, present, and future climate, but they suffer from legacy technical infrastructure. ESMs are primarily implemented in Fortran, a language that poses a high barrier of entry for early career scientists and lacks a GPU runtime, which has become essential for continued advancement as GPU power increases and CPU scaling slows. Fortran also lacks differentiability - the capacity to differentiate through numerical code - which enables hybrid models that integrate machine learning methods. Converting an ESM from Fortran to Python/JAX could resolve these issues. This work presents a semi-automated method for translating individual model components from Fortran to Python/JAX using a large language model (GPT-4). By translating the photosynthesis model from the Community Earth System Model (CESM), we demonstrate that the Python/JAX version results in up to 100x faster runtimes using GPU parallelization, and enables parameter estimation via automatic differentiation. The Python code is also easy to read and run and could be used by instructors in the classroom. This work illustrates a path towards the ultimate goal of making climate models fast, inclusive, and differentiable.

研究の動機と目的

• Earth System Models (ESMs) を GPU に優しく、かつ微分可能にするための近代化の必要性を喚起する。
• 静的解析と GPT-4 を用いて Fortran コード単位を Python/JAXへ翻訳する divide-and-conquer ワークフローを提案する。
• CESM CLM の葉レベルの光合成モジュールにおける実行時間改善とパラメータ推定機能を実証する。
• より包括的で高速かつ微分可能な気候モデリングへ向けたオープンソースツールと概念実証的な道筋を提供する。

提案手法

• 静的解析とトポロジカルソートを用いて Fortran コードを依存関係順のユニットに分割する。
• 各ユニットを Python/JAXへ翻訳し、反復を導くユニットテストを伴う GPT-4 を用いる。
• 複数の Python バリアント（NumPy、Numba、SciPy、JAX）を実行し、GPU と CPU での実行時間を比較する。
• 自動微分と勾配ベースの最適化を用いた Vcmax の微分可能なパラメータ推定を披露する。
• 翻訳の品質と信頼性を高めるためにユニットテストとともに反復する。

Figure 2: Comparing runtime of leaf-level photosynthesis in several Python translations with the original Fortran version. Runtime was measured on an Amazon EC2 G5.4xlarge instance with one NVIDIA A10G GPU.

実験結果

リサーチクエスチョン

• RQ1 大規模言語モデル主導の divide-and-conquer アプローチは、Fortran の気候モデルのコンポーネントを信頼性高く Python/JAXへ翻訳できるか。
• RQ2 GPU 加速を介した Python/JAX の翻訳は、選択したモデルコンポーネントに対して native Fortran より顕著な実行時間の改善を提供するか。
• RQ3 翻訳後の Python/JAX コードにおいて自動微分は効率的なパラメータ推定を可能にするか。
• RQ4 この翻訳をより大規模な Fortran コードベースへスケールする際の実践的な課題と制約は何か。

主な発見

• JAX-GPU翻訳は、テストした Python バリアントの中で最速の実行時間を達成し、Fortranは約100倍遅かった。
• 同じモジュールに対して、GPU並列化された Python/JAX はCPUベースのFortranを大幅に上回る。
• jit-compiled Python アプローチ（JAX, Numba）は Fortran にかなり近い性能を示し、コンパイルが速度向上に大きく寄与することを示唆している。
• 自動微分は Vcmax の勾配ベースのパラメータ推定を可能にし、葉モデルにおける均一サンプリングと比較して反復回数を削減する。
• 翻訳ワークフローは、微分可能でGPU対応の気候モデルへの道を効果的に示し、クラスルームで使える Python コードをサポートしている。

Figure 3: Measured (points) and modeled (lines) relationship between the internal partial pressure of CO2 (Pa) and the rate of assimilation (umol/m2/s). The modeled values use the Vcmax parameter value selected using either uniform sampling (orange) or gradient descent (green).