Radeon (ROCm) でAI気象モデルを実行する

Last updated at 2026-01-19Posted at 2025-12-07

はじめに

データ駆動型のAI気象モデルは、物理法則に基づく物理気象モデルと比べて、いくつかの指標において同程度やそれ以上の精度を実現できることがわかり、ここ数年で注目を集めている。AI気象モデルによる予報を現業利用するセンターもあり、その予報結果は物理モデルと同様に無料でダウンロードできるものがある。

しかし、特定のモデルと初期値から計算した予報が欲しい場合、自分で計算する必要が出てくる。この際、CPUでの計算は時間がかかるため、実質的にGPUを使用することが必須である。本記事ではAMD社製のGPUであるRadeonを使用したAI気象モデルの実行について解説する。

AI気象モデルの実行に使用できるGPU

製品比較

AI気象モデルを実行するために使用できるGPUには、いくつか条件がある。

最も大きな制約条件はGPUのメモリ(VRAM)である。AI気象モデルの推論には、VRAMは32GBあるとよく、多くのモデルは最低でも24GBを必要とする。そのためVRAMが8GBや16GBである多くのコンシューマー向けGPUは使用することができない。

また、AI気象モデルに用いられる主要な深層学習フレームワーク（PyTorch, JAX, ONNX Runtime）のバックエンドとしての対応状況を考えると、実質的に動かせるのはNVIDIAのCUDAまたはAMDのROCmになる。私が試した限りでは、Apple (Mac) のCoreMLは対応していない命令や配列長の制限で動かないことが多い。IntelのOpenVINOは未調査だが、VRAMが多くAMDやNVIDIAと比べて性能やコスパで優れている製品がないのが実情だ。

これらを踏まえて、一般のご家庭でも導入できる、現在新品を購入可能なVRAM 32GB以上の製品をまとめると以下のようになる¹。

GPU	VRAM (GB)	FP32性能 (TFLOPS)	メモリ帯域幅 (GB/s)	消費電力(W)	価格 (万円)
Radeon PRO AI R9700	32	47.84	644.6	300	25
Geforce RTX 5090	32	104.8	1790	575	50
Ryzen AI Max+ 395	最大96	14.85	256	230	30〜40
NVIDIA GB10	共有128	29.71	273.2	240	60〜90

上記のうち、AMD の Radeon PRO AI R9700 と NVIDIA の Geforce RTX 5090 はグラフィックボードと呼ばれるもので、PCIeスロットに刺して使用する。そのため別途グラフィックボードが刺さる自作PC本体を用意する必要があり、使用するパーツのスペックにもよるが上記の価格に追加で20万円程度かかると考えたほうがよい。自作に不安がある場合にはBTOで購入することも可能である。この場合、メーカーによる動作の保証がある代わりに、自作と比べて20〜30万円程度は高くなる。

一方 AMD の Ryzen AI Max+ 395 や NVIDIA の GB10 Grace Blackwell Superchip はSystem-on-a-chip (SoC) と呼ばれるもので、一つのチップにCPUとGPUが搭載されている。元々パラメータ数の多いローカルLLMを実行することを目的とした製品であるため、VRAMは多くAI気象モデルの実行に十分である。ただしグラフィックボードと比べて演算性能が劣る点に注意が必要である。

以下、私が所有している4製品について、それぞれの特徴と利点・欠点をまとめる。

Ryzen AI Max+ 395

Ryzen AI Max+ 395 は16コア32スレッドのCPUと、Radeon 8060Sと呼ばれるGPUを内蔵したSoCである。コードネームはStrix Halo。メモリは最大128GBで、このうち96GBまでGPU用に割り当てることができる。日本で入手可能な製品は、ミニPC（ミニワークステーション）の形で GMKtec EVO-X2 (30〜40万円) と Minisforum MS-S1 MAX (40万円) がある。私はEVO-X2の128GBを所有している。主な利点と欠点は次の通りである。

利点
- VRAMが最大96GBのGPUを安く入手することができる
- 電源アダプタは230W、推論時の消費電力は150W程度と低い
- 大きさがコンパクト
欠点
- ROCmの対応が不十分な点があり、設定に一手間必要とする（この記事で解説）
- 演算性能は高くないので、演算機律速なタスクには向かない
- ファンの回転時はそれなりにうるさい。エアコンの音に近く、筆者の主観では添い寝は可能

GB10

NVIDIAからはDGX Sparkという名称で発売されているほか、各社から同じGB10 Grace Blackwellスーパーチップを用いたOEM製品が発売されている。メモリはCPUとGPUの共有で128GBである。私はASUSのAscent GX10を所有している。主な利点と欠点は次の通りである。

利点
- 電源アダプタは240W、推論時の消費電力は170Wと低い
- 大きさがコンパクト
- ファンの音はかなり静か
欠点
- 価格が高い
- CPUのアーキテクチャがx86-64(amd64)ではなくARM(aarch64)であるため、onnxruntimeや一部のパッケージをインストールする際に自前でビルドする必要がある

Radeon AI PRO R9700

Radeon AI PRO R9700 は最新世代のAMDのプロ向けのGPUである。中身は実質的にコンシューマー向けの最上位機種のRadeon 9070 XT のVRAMを16GBから32GBにしたものである。シロッコファンを搭載したブロワータイプ（外排気）のGPUであるのが特徴。私はPowerColor製のものを所有している。

利点
- VRAM 32GBのグラフィックボードを安く入手することができる
- 2スロット厚のため、他のカードとも物理的に干渉しにくく多くのPCやラックサーバーに搭載できる
欠点
- シロッコファンの高音がうるさいので、添い寝には不向き

Geforce RTX 5090

Geforce RTX 5090はNVIDIAのコンシューマー向けの最上位機種である。その演算性能自体はワークステーション向けの最上位機種であるRTX PRO 6000 Blackwell に匹敵する割に、150万円する6000 BWと比べて安く入手することができる。私は3スロット厚空冷のINNO3D GeForce RTX 5090 X3 OCを所有している。

利点
- 演算性能が高く、GPU 1枚に乗るモデルであれば学習・推論などで困ることはない
欠点
- 基本的に空冷であれば3スロット厚を超えるものが多く、物理的に設置できるPCが限られ、ラックサーバーには搭載できない
- デフォルトのTDPは575Wと高いため、消費電力が高く発熱も大きいため、電源容量や冷却に気を使う必要がある
- ファンの回転数が80%近くまで上がるとそれなりにうるさい。エアコンの音に近く、筆者の主観では添い寝は可能

VRAM消費量・実行時間

240時間(40ステップ)の単一メンバの予報のうち、VRAM使用量と推論にかかった時間をまとめる（WIP）。推論にはデータのダウンロードや出力は含んでいない。コンパイルが走る場合は、結果がキャッシュされた2回目以降の時間である。空欄は未計測である。

Ryzen AI Max+ 395およびR9700で使用するROCmのバージョンは7.1.1である。ただしONNX runtimeを使用する場合は6.4.4である。GB10ではCUDA13.0を使用し、RTX 5090ではCUDA12.8を使用している。

いずれのモデルもEarth2Studioラッパーを使用して実行した。

モデル	FW	VRAM	Ryzen 395	R9700	GB10	RTX 5090
SFNO	Torch	23GB	3分05秒	1分01秒	47秒	9秒
FCN3	Torch	95GB	5時間	実行不能	3分38秒	実行不能
Aurora	Torch	39(32)GB	56分	25分	8分30秒	1分45秒
AIFS	Torch	23-51GB	2分08秒	1分07秒	1分32秒	24秒
Pangu	ONNX	27GB	2分20秒	2分57秒	3分08秒	41秒
FuXi	ONNX	9GB	2分27秒	21分18秒	1分59秒	28秒
FengWu	ONNX	24GB	3分36秒	7分28秒	2分20秒	21秒
GraphC	JAX	75%	1分45秒	1分08秒	1分25秒	49秒

まとめ

上記を踏まえると、AI気象モデルの実行のみを行う場合、AMD製品のRyzen AI Max+ 395やRadeon PRO AI R9700のコストパフォーマンスが良い。ただし、AI気象モデルの多くはNVIDIA GPUを用いたCUDAを前提としているため、AMD RadeonのROCm環境での実行についての直接的な情報は少ない²。

AI気象モデルの実行環境について

様々な用語が出てくるが、ざっくり次のような構造になっている。

ラッパーライブラリ
- Earth2Studio (NVIDIA)
- ai-models (ECMWF、開発終了)
深層学習フレームワーク
- PyTorch
  - Anemoi (ECMWF、AIFS用)
- JAX (Google系モデル)
- ONNX runtime (中国系のモデル、推論専用)
GPU計算プラットフォーム
- CUDA (NVIDIA)
- ROCm (AMD)

複数のAI気象モデルの重みデータや入力データの形式を揃えるのが面倒であるため、ラッパーライブラリを利用すると便利である。従来はECMWFのai-modelsが使われていたが、現在はArchiveとなり開発が終了している。近年活発に開発されているのがNVIDIAのEarth2Studioであり、今後はこちらを使うと良いだろう。

AI気象モデルによって、使われている深層学習フレームワークが異なる。Earth2Studioでインストールした場合、CUDA環境を前提としたパッケージがインストールされる。これをRadeonで実行したい場合はROCmに対応したPyTorch、JAX、ONNX runtimeにそれぞれ差し替える必要がある。

Ryzen AI Max+ 395やR9700で利用可能なROCmのバージョンには、6系の6.4.4または7系の7.1.1がある。ここでは、基本的に不具合の修正が行われた7.1.1を利用する。ただし、7.1.1ではONNX runtimeでAI気象モデルの実行に必要なROCMExecutionProviderが廃止された³ため、ONNXモデルは6.4.4を利用する。以下の記事のように両者が共存した環境を前提とする。

Earth2Studio

共通事項

uvを使用したEarth2Studioの環境構築について説明する。基本的には、依存関係をpyproject.tomlに記載した後、uv syncを実行することで環境構築を行う（ように書き直し中）。pyproject.tomlではCUDA向けのパッケージをROCm向けのものに置き換える処理を行う。

また、Earth2Studioのコード自体にいくつか修正を加える必要がある。このために、まずは適当なリポジトリにcloneしたものを用意する。

git clone https://github.com/NVIDIA/earth2studio.git
cd earth2studio

ここでは複数のモデルに共通する例として、ONNX runtime向けの変更を行おう。1つ目の変更としてort.preload_dlls()が存在せず警告またはエラーになるため、earth2studio/models/utils.py の20行目付近に次のようにガードを入れる。

try:
    import onnxruntime as ort
    
    if hasattr(ort, "preload_dlls"):
        ort.preload_dlls()
    from onnxruntime import InferenceSession

さらにROCMExecutionProviderが使用されるように、63行目付近を次のように書き換える。

    if device.type == "cuda":
        if device.index is None:
            device_index = torch.cuda.current_device()
        else:
            device_index = device.index
        available = ort.get_available_providers()
        if "CUDAExecutionProvider" in available:
            providers = [
                (
                    "CUDAExecutionProvider",
                    {
                        "device_id": device_index,
                    },
                ),
                "CPUExecutionProvider",
            ]
        elif "ROCMExecutionProvider" in available:
            providers = [
                (
                    "ROCMExecutionProvider",
                    {
                        "device_id": device_index,
                    },
                ),
                "CPUExecutionProvider",
            ]
        else:
            providers = [
                "CPUExecutionProvider",
            ]

そうしたらモデルの実行プログラムのあるディレクトリで、以下を追記する（/path/to/earth2studioは上記リポジトリを参照するように書き換える）。editable = trueによりリポジトリに加えた修正がすぐに反映されるようになる。

[tool.uv.sources]
earth2studio = { path = "/path/to/earth2studio", editable = true }

pyproject.tomlにこの他の設定を書き加えた上でuv syncを実行することで、仮想環境にパッケージのインストールする。

なお、Earth2Studioのライブラリの使い方自体については別の記事にまとめたのでそちらを参照のこと。

FourCastNet v2 / SFNO (PyTorch)

pyproject.toml

[project]
name = "sfno-demo"
version = "0.1.0"
description = "Add your description here"
readme = "README.md"
requires-python = ">=3.12,<3.13"
dependencies = [
    "earth2studio[data,sfno]",
    "torch",
    "torchvision",
    "pytorch-triton-rocm",
]

[[tool.uv.index]]
name = "pytorch-rocm"
url = "https://download.pytorch.org/whl/nightly/rocm7.1"
explicit = true

[tool.uv.sources]
earth2studio = { path = "/path/to/earth2studio", editable = true }
torch = { index = "pytorch-rocm" }
torchvision = { index = "pytorch-rocm" }
pytorch-triton-rocm = { index = "pytorch-rocm" }

[tool.uv]
override-dependencies = ["numba==0.63.1", "llvmlite==0.46.0"]

FourCastNet v3 (PyTorch)

FourCastNet v2 の pyproject.toml の sfno を fcn3 に書き換えれば良い。

内部で用いられるtorch_harmonics（DISCO）の球面畳み込みが巨大な中間行列を要求するため、他のモデルと比べてVRAMを多く使用する。そのため、Ryzen AI Max+ 395でVRAM 96GBの割り当てが必要。結論から言うと、このようにVRAMを割り当てたとしても実用的な速度で動作しない。

サンプルコードを実行したところ、96GBのVRAMのうち96%を使用し、推論時間は1ステップあたり7分42秒（ROCm 6.4.4の場合は15分）であった。これは1メンバ40ステップ予報の場合5時間かかる計算になる。NVIDIAのGH200で実行したところ1ステップあたり1秒、全体で1分程度であり、これと比較してとても遅い。

原因としては、torch-harmonics disco CUDA extension is not available. FCN3 run on GPU/CUDA will be slower.と警告が出ており、torch-harmonicsのROCm向け最適化が存在しないことだと考えている。

Aurora (PyTorch)

FourCastNet v2 の pyproject.toml の sfno を aurora に書き換えれば良い。

デフォルトだとVRAMは39GB程度使用するため、R9700やRTX 5090のような32GBのGPUでは動かない。しかし、ロールアウト（時間ループ）で毎回VRAMの解放処理を入れることで、32GB以内にギリギリ抑えて実行することができる。Earth2Studioのearth2studio/models/px/aurora.py以下のように編集する。

まず、gcをインポートする。

import gc

次に、_forwardメソッドのforループの末尾に、メモリ解放用のコードを追加する。

        for t in range(coords["time"].shape[0]):
            batch_coords = coords.copy()
            batch_coords["time"] = batch_coords["time"][t : t + 1]
            # x shape: [b,1,2,69,720,1440], batch, lead_time, variables, lat, lon
            input_batch = self._prepare_input(x[:, t : t + 1], batch_coords)
            # Convert tensor to Batch, atmos_vars is the first 65 variables, surf_vars is the last 4 variables of x
            output_batch = self.model(input_batch)
            # Convert Batch to tensor
            out[:, t : t + 1] = self._prepare_output(output_batch, coords)
            
            # memory release
            del input_batch, output_batch
            gc.collect()
            torch.cuda.empty_cache()

_default_generatorメソッドのwhileループの末尾にも、同様にして追加する

        while True:
            ### (略) ###

            yield x[:, :, 1:], coords_out
            
            # memory release
            del init_x
            gc.collect()
            torch.cuda.empty_cache()

AIFS-Single v1.0 (PyTorch, Flash Attention, Anemoi)

AIFS-Singleの最新版はv1.1だが、Earth2Studioの現在の最新版であるv0.11.0ではAIFS-Single v1.0にしか対応していない。

uv init --python=3.12
uv sync
uv pip install torch torchvision --index-url https://download.pytorch.org/whl/rocm6.4
uv add "earth2studio @ git+https://github.com/NVIDIA/earth2studio.git@0.11.0"
uv add earth2studio --extra aifs --no-build-isolation-package flash-attn

ここでFlash Attention はROCm用のものをインストールする必要がある。そのままでも動作するが、実行がとても遅くなる。

まず別のディレクトリにリポジトリをcloneする。

git clone --recursive https://github.com/ROCm/flash-attention.git
cd flash-attention
git checkout main_perf

次に元のディレクトリに戻って、Flash Attentionをインストールする。/path/to/flash-attentionの部分はローカルリポジトリのパスに書き換える。

export FLASH_ATTENTION_TRITON_AMD_ENABLE="TRUE"
uv add --editable --no-build-isolation /path/to/flash-attention

FLASH_ATTENTION_TRITON_AMD_ENABLEの環境変数の設定は実行時にも必要になる。忘れないように~/.bashrcに同じ設定を追加しておくとよい。

R9700で推論する場合、使用するVRAMを32GB以内に収めるため、以下のように環境変数を指定する必要がある（Anemoiの節で詳しく述べる）。

os.environ['ANEMOI_INFERENCE_NUM_CHUNKS']='16'

AIFS-ENS v1.0

AIFS-Singleの場合と基本的に同じで、--extra aifsensにすればよい。

uv add earth2studio --extra aifsens --no-build-isolation-package flash-attn

1メンバあたりの推論時間はAIFS-Singleとほぼ同じである。

Pangu-Weather, FuXi, FengWu (ONNX runtime)

はじめに次のようなpyproject.tomlを作成する。ここでデバイスの判定に必要なためPyTorchも合わせてインストールしている。FengWuの場合はpanguをfengwuに置き換えれば良い。

[project]
name = "pangu"
version = "0.1.0"
description = "Add your description here"
readme = "README.md"
requires-python = ">=3.12,<3.13"
dependencies = [
    "earth2studio[data,pangu]",
    "onnxruntime-rocm",
    "torch",
    "torchvision",
    "pytorch-triton-rocm",
]

[tool.uv]
exclude-dependencies = ["onnxruntime-gpu"]

[tool.uv.sources]
earth2studio = { git = "https://github.com/NVIDIA/earth2studio.git", rev = "0.11.0" }
onnxruntime-rocm = { url = "https://repo.radeon.com/rocm/manylinux/rocm-rel-6.4.4/onnxruntime_rocm-1.21.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl" }
torch = { index = "pytorch" }
torchvision = { index = "pytorch" }
pytorch-triton-rocm = { index = "pytorch" }

[[tool.uv.index]]
name = "pytorch"
url = "https://download.pytorch.org/whl/rocm6.4"
explicit = true

次にインストールを行う。

uv sync

FuXiやFengWuの実行で、MIOpen(HIP): Warning [IsEnoughWorkspace] [GetSolutionsFallback AI] Solver <GemmBwdRest>といった警告が出ることがある。これは以下の記事のように設定を変更してonnxruntimeをビルドし直すことで解消できる。ただし、実行時間はそれほど変わらないので、手間をかけて対応するかは各自の判断となる。

GraphCast (JAX)

以下のpyproject.tomlを作成する。

[project]
name = "graphcast-demo"
version = "0.1.0"
description = "Add your description here"
readme = "README.md"
requires-python = ">=3.12,<3.13"
dependencies = [
    "earth2studio[data,graphcast]",
    "jax==0.7.1",
    "jaxlib==0.7.1",
    "jax-rocm7-pjrt==0.7.1",
    "jax-rocm7-plugin==0.7.1",
    "torch",
    "torchvision",
    "triton==3.5.1+rocm7.1.1.gita272dfa8",
]

[tool.uv.sources]
earth2studio = { git = "https://github.com/NVIDIA/earth2studio.git", rev = "0.11.0" }
jaxlib = { url = "https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/jaxlib-0.7.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl" }
jax-rocm7-pjrt = { url = "https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/jax_rocm7_pjrt-0.7.1-py3-none-manylinux_2_28_x86_64.whl" }
jax-rocm7-plugin = { url = "https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/jax_rocm7_plugin-0.7.1-cp312-cp312-manylinux_2_28_x86_64.whl" }
torch = { url = "https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/torch-2.9.1%2Brocm7.1.1.lw.git351ff442-cp312-cp312-linux_x86_64.whl" }
torchvision = { url = "https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/torchvision-0.24.0%2Brocm7.1.1.gitb919bd0c-cp312-cp312-linux_x86_64.whl" }
triton = { url = "https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/triton-3.5.1%2Brocm7.1.1.gita272dfa8-cp312-cp312-linux_x86_64.whl" }

[tool.uv]
# earth2studioのjax[cuda12]依存を無視してROCm版を使う
override-dependencies = ["jax==0.7.1", "jaxlib==0.7.1"]

インストールする。

uv sync

実行時の注意点として、JAXのコンパイルには時間がかかるので、永続キャッシュを指定しておくと良い。初回はコンパイルに10分以上かかるが、2回目以降は10秒程度で済むようになる。

export JAX_COMPILATION_CACHE_DIR="/path/to/jax_cache"

VRAM使用量はいずれも全体の75%である（JAXの仕様）。

ai-models

FourCastNet v2

Pythonのバージョンは3.10-3.12系であれば動くはず。

ai-modelsをインストールする

pip install ai-models

ai-models-fourcastnetv2 をインストールする。最新のPyTorchを使用する場合、ai_models_fourcastnetv2/model.pyの重みデータの読み込み部分でweights_only=Falseを指定しないとエラーになる。

- checkpoint = torch.load(checkpoint_file, map_location=self.device)
+ checkpoint = torch.load(checkpoint_file, map_location=self.device, weights_only=False)

自分でソースコードを変更してインストールしても良いし、以下のリポジトリに修正済みのものをインストールすることもできる。

pip install git+https://github.com/wm-ytakano/ai-models-fourcastnetv2.git

ai-models-fourcastnetv2の依存関係としてCUDA版のPyTorchがインストールされる。ROCm対応版のPyTorchで上書きする。

pip install --force-reinstall torch torchvision --index-url https://download.pytorch.org/whl/rocm6.4

Numpyのバージョンが2系だとエラーになるので、1.26.4に戻しておく

pip install --force-reinstall "numpy==1.26.4"

初回は実行するディレクトリでモデルのウェイトをダウンロードしておく。

cd assets
ai-models --download-assets fourcastnetv2-small
cd ..

実行する。--dateや--timeは予報データが存在する日時に置き換えること。

ai-models --input ecmwf-open-data --date 20251208 --time 0000 --assets assets --lead
-time 240 fourcastnetv2-small

Ryzen AI Max+ 395での240時間(40step)の実行時間は以下の通りであった。

2025-12-09 17:00:15,777 INFO Elapsed: 46 seconds.
2025-12-09 17:00:15,777 INFO Average: 1 second per step.

Pangu-Weather

Pythonのバージョンは3.10-3.12系であれば動くはず。

ai-models はCUDAを使用することを前提としているため、ROCmにも対応するようにソースコードを修正してからpipでインストールしよう。

適当な作業ディレクトリで、リポジトリをクローンする。

git clone https://github.com/ecmwf-lab/ai-models.git
cd ai-models

以下のように"ROCMExecutionProvider"を追加する。

--- a/src/ai_models/model.py
+++ b/src/ai_models/model.py
@@ -220,7 +220,7 @@ class Model:

         available_providers = ort.get_available_providers()
         providers = []
-        for n in ["CUDAExecutionProvider", "CPUExecutionProvider"]:
+        for n in ["ROCMExecutionProvider", "CUDAExecutionProvider", "CPUExecutionProvider"]:
             if n in available_providers:
                 providers.append(n)

pip install .

Pangu-Weatherの実行にはai-models-panguweatherのインストールが追加で必要になる。こちらはソースコードの修正の必要がないので、通常通りインストールする。

pip install ai-models-panguweather

ai-models-panguweatherの依存関係としてonnxruntime-gpuがインストールされる。これはCUDA用のためアンインストールする。その後、ROCm用のonnxruntime-rocmをインストールする。

pip uninstall onnxruntime-gpu
pip install onnxruntime-rocm -f https://repo.radeon.com/rocm/manylinux/rocm-rel-6.4/

初回は実行するディレクトリでモデルのウェイトをダウンロードしておく。

cd assets
ai-models --download-assets panguweather
cd ..

Pangu-Weatherを実行する。--dateや--timeは予報データが存在する日時に置き換えること。

ai-models --input ecmwf-open-data --date 20251201 --time 0000 --assets assets --lead-time 240  panguweather

もしここで以下のようなエラーが出る場合、環境変数の設定export HSA_OVERRIDE_GFX_VERSION=11.0.0を忘れていないか確認する (Ryzen AI Max+ 395)。

longer_pathname_so_that_rpms_can_support_packaging_the_debug_info_for_all_os_profiles/src/clr/hipamd/src/hip_global.cpp:158 : 3448962776 us:  Module not initialized

GraphCast

Pythonの仮想環境を作成する。Pythonのバージョンは3.10-3.12系であれば動くはず。

まずai-modelsをそのままインストールする

pip install ai-models

次に、ai-models-graphcastをインストールする。公式のものは依存関係としてdm-haiku@0.0.10をインストールするが、これはJAX 0.4.24 以降で削除されたjax.linear_utilを呼び出そうとしてエラーになる。pyproject.tomlでdm-haiku@0.0.14に変更したPull Requestが投げられているので、こちらのブランチを使用する。

pip install  "git+https://github.com/kashif/ai-models-graphcast.git@patch-1"

GraphCastとその依存性を公式からインストールする。ここでNumpyのバージョンが2系に上がってしまうが、ai-modelsで使用する場合エラーになるので、元のnumpy@1.26.4に戻しておく。

pip install git+https://github.com/deepmind/graphcast.git
pip install numpy==1.26.4

ROCmに対応したJAXをインストールする。

pip install jax[rocm]

初回は実行するディレクトリでモデルのウェイトをダウンロードしておく。

mkdir assets
cd assets
ai-models --download-assets graphcast
cd ..

実行前にいくつか設定する必要がある。JAXが実行前にJITコンパイルを行う時にRAMを50GB程度消費するため、RAMが64GB以上用意するのが望ましい。Ryzen AI Max+ 395であればRAM 64GB, VRAM 64GBにするとバランスが良い。不足する場合には以下のようにswapを用意する必要がある。

sudo fallocate -l 100G /home/swapfile
sudo chmod 600 /home/swapfile
sudo mkswap /home/swapfile
sudo swapon /home/swapfile
swapon --show

またXLAの最適化をOFFにするように環境変数を設定する。これをやらないとXLAコンパイルに長時間かかってしまう。

export XLA_FLAGS="--xla_gpu_autotune_level=0"

またR9700のようにVRAMが32GBの場合、VRAMが不足して実行できなくなる。これを避けるために、メモリ確保をJAXではなくROCm側に委ねることでVRAM使用量を削減し実行可能になるが、計算速度は遅くなる。

export XLA_PYTHON_CLIENT_ALLOCATOR=platform

実行する

ai-models --input ecmwf-open-data --date 20251207 --time 0000 --assets assets --lead-time 240  graphcast

Ryzen AI Max+ 395での実行時間の出力を抜き出したのがこちら。初期値データはキャッシュされた状態で計測している。実際の計算時間(full rollout prediction)よりも、前処理や後処理に多くの時間がかかっている。

Loading params/GraphCast_operational - ERA5-HRES 1979-2021 - resolution 0.25 - pressure levels 13 - mesh 2to6 - precipitation output only.npz: 0.2 second.
Building model: 0.2 second.
Creating forcing variables: 7 seconds.
Converting GRIB to xarray: 1 second.
Reindexing: 1 second.
Creating training data: 10 seconds.
Extracting input targets: 2 seconds.
Creating input data (total): 12 seconds.
Doing full rollout prediction in JAX: 55 seconds.
Converting output xarray to GRIB and saving
Saving output data: 1 minute 41 seconds.
Total time: 2 minutes 52 seconds.

R9700での実行結果。export XLA_PYTHON_CLIENT_ALLOCATOR=platformにより、Ryzen AI Max+ 395よりも遅くなった。

precipitation output only.npz: 0.3 second.
Building model: 0.3 second.
Creating forcing variables: 5 seconds.
Converting GRIB to xarray: 2 seconds.
Reindexing: 3 seconds.
Creating training data: 11 seconds.
Extracting input targets: 8 seconds.
Creating input data (total): 19 seconds.
Doing full rollout prediction in JAX: 6 minutes 50 seconds.
Converting output xarray to GRIB and saving
Saving output data: 1 minute 25 seconds.
Total time: 8 minutes 39 seconds.

ラッパーを使用しない推論

AIFS-ENS (Anemoi)

AIFS ENS v1.0 の実行方法について説明する。以下のHugging FaceのリポジトリにあるJupyter Notebookをベースに環境構築と実行を行う。

Python仮想環境

AIFSが使用するAnemoiというAI気象モデルフレームワークは、学習時と推論時のバージョンチェックが合わないとエラーになるようになる。AIFS-ENS v1.0の場合、Python 3.11.6にする必要がある。

PyTorchのインストール

Stable (2.9.1) > Linux > Pip > Python > ROCm 6.4 で出てくる以下のコマンドを実行。PyTorchのバージョンはあまり気にせずインストールしたが、実行できた。

pip install torch torchvision --index-url https://download.pytorch.org/whl/rocm6.4

Anemoiのインストール

pip install anemoi-inference[huggingface]==0.6.0 anemoi-models==0.6.0 anemoi-graphs==0.6.0 anemoi-datasets==0.5.23
pip install earthkit-regrid==0.4.0 'ecmwf-opendata>=0.3.19'

Flash Attention のインストール

Flash Attention はROCm用のものをインストールする必要がある。

git clone --recursive https://github.com/ROCm/flash-attention.git
cd flash-attention
git checkout main_perf

pip install packaging ninja triton==3.2.0

export FLASH_ATTENTION_TRITON_AMD_ENABLE="TRUE"
pip install --no-build-isolation -v .

FLASH_ATTENTION_TRITON_AMD_ENABLEの環境変数の設定は実行時にも必要になる。忘れないように~/.bashrcに同じ設定を追加しておくとよい。

export FLASH_ATTENTION_TRITON_AMD_ENABLE="TRUE"

実行

公式のJupyter Notebookを実行すればよい。VSCode等であればipykernelのインストール、Jupyter Labを使用する場合はさらにjupyterlabのインストールやjupyterのkernelの設定を行う。詳細は割愛。

R9700のようにVRAMが32GBの場合は、Notebook内に書かれている通り環境変数を指定してVRAM使用量を削減する必要がある。

# os.environ['PYTORCH_CUDA_ALLOC_CONF']='expandable_segments:True' 
os.environ['ANEMOI_INFERENCE_NUM_CHUNKS']='16'

PYTORCH_CUDA_ALLOC_CONFの方はUserWarning: expandable_segments not supported on this platformと出力されるためROCm環境で無効のようだ。

ANEMOI_INFERENCE_NUM_CHUNKSのチャンク分割数が多くなるほど、VRAM使用量は削減されるが、実行時間は長くなる。Ryzen AI Max+ 395で lead_time=72 (12steps) の推論を行った場合の比較結果は以下の通り。VRAMが32GBなら3、24GBなら16を指定すればよさそう。

`ANEMOI_INFERENCE_NUM_CHUNKS`	Wall time (s)	VRAM (MiB)
1 (無指定)	53	51800
2	54.4	37902
3	55	31760
4	55.5	29466
8	57.8	26690
16	62	24142
32	70	23242

GenCast

GenCastは現在はEarth2Studioには実装されておらず、ai-modelまたはオリジナルのGraphCastリポジトリを元に実行する必要がある。ここではGraphCastリポジトリの例をベースに実行する方法を説明する。

インストール

次のpyproject.tomlを作成し環境構築(uv sync)する。ROCm 7.1.1 向けに、JAX 0.7.1を使用している。

[project]
name = "gencast-demo"
version = "0.1.0"
description = "Add your description here"
readme = "README.md"
requires-python = ">=3.12,<3.13"
dependencies = [
    "jax==0.7.1",
    "jaxlib==0.7.1",
    "jax-rocm7-pjrt==0.7.1",
    "jax-rocm7-plugin==0.7.1",
    "triton==3.5.1+rocm7.1.1.gita272dfa8",
    "graphcast",
    "netcdf4>=1.7.3",
]

[tool.uv.sources]
jaxlib = { url = "https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/jaxlib-0.7.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl" }
jax-rocm7-pjrt = { url = "https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/jax_rocm7_pjrt-0.7.1-py3-none-manylinux_2_28_x86_64.whl" }
jax-rocm7-plugin = { url = "https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/jax_rocm7_plugin-0.7.1-cp312-cp312-manylinux_2_28_x86_64.whl" }
triton = { url = "https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/triton-3.5.1%2Brocm7.1.1.gita272dfa8-cp312-cp312-linux_x86_64.whl" }
graphcast = { path = "../../projects/graphcast", editable = true }

ここでGraphCastはローカルにgit clone git@github.com:google-deepmind/graphcast.gitした上で、後述のように修正したものを参照する。pyproject.yamlの相対パスは各自の環境に合わせて適宜書き換える。

GenCastが使用するsplash_attentionはTPUにしかないため、GPUで動かす場合は修正が必要になる。

はじめに、graphcast/denoiser.pyのSparseTransformerConfigを変更可能になるようにデコレーターのfrozenをFalseに変更する。

@chex.dataclass(frozen=False, eq=True)
class SparseTransformerConfig:

次に、graphcast/sparse_transformer.pyで、以下のimport文をコメントアウトする。

from jax.experimental.pallas.ops.tpu import splash_attention

最後に、型ヒントのsplash_attention.splash_attention_mask.Maskを削除する。WeatherMeshMaskはsplash_attention.splash_attention_mask.Maskの継承をやめる。

実行

実行方法については詳しく解説しないが、graphcastリポジトリのgencast_demo_cloud_vm.ipynb をベースに行う。

ただしGPUでの実行に対応するため、sparse_transformer_configのattention_typeとmask_typeを以下のように変更する。

denoiser_architecture_config = ckpt.denoiser_architecture_config
denoiser_architecture_config.sparse_transformer_config.attention_type = "triblockdiag_mha"
denoiser_architecture_config.sparse_transformer_config.mask_type = "full"

実行は初回にJAXのコンパイルで10分程度かかるが、キャッシュにより2回目以降は早くなる。
1°モデル(GenCast 1p0deg <2019.npz)の場合、1ステップ(12時間予報)あたり、Ryzen AI Max+ 395で49秒、R9700で36秒であった。

0.25°モデル(GenCast 0p25deg Operational <2022.npz)の場合、Ryzen AI Max+ 395ではVRAMを80GB近く使用し、1ステップ(12時間予報)あたり10分05秒かかる。R9700ではVRAMを30GB使用し、1ステップあたり9分12秒かかる。参考としてRTX 5090では最初の1ステップに毎回コンパイルが走るため3分かかり、それ以降は1ステップ1分である。

FuXi-ENS

重みとコードはZenodoからダウンロードできる。

以下のようなpyproject.tomlを作成すれば実行できる。

pyproject.toml

[project]
name = "fuxi-ens"
version = "0.1.0"
description = "Add your description here"
readme = "README.md"
requires-python = ">=3.12"
dependencies = [
    "bottleneck>=1.6.0",
    "dask>=2025.12.0",
    "netcdf4>=1.7.4",
    "onnxruntime-rocm",
    "xarray>=2025.12.0",
]

ただしこのままだと MIOpen(HIP): Warning [IsEnoughWorkspace] [GetSolutionsFallback AI] Solver <GemmBwdRest>, workspace requiredという警告が出て、演算をフォールバックするため遅い。Ryzen AI Max+ 395では1ステップ目が860秒、2ステップ目以降が430秒かかる。R9700の場合も1ステップ目が660秒、2ステップ目以降337秒かかる。なお参考としてRTX 4090では2.0秒/ステップ (VRAM 22GB)、RTX 5090では1.6秒/ステップ (VRAM 29GB)である。

これに対処するには、以下の記事のように設定を変更してonnxruntimeをビルドし直す必要がある。

ビルドしたonnxruntime-rocmを参照するようにpyproject.tomlに追記する。

pyproject.toml

[tool.uv.sources]
onnxruntime-rocm = { path = "/path/to/onnxruntime/build/Linux/Release/dist/onnxruntime_rocm-1.22.2-cp312-cp312-linux_x86_64.whl" }

これによりR9700では初回の1ステップ目が243秒かかるが、それ以降は1ステップで24秒で実行できるようになる。

Ryzen AI Max+ 395については、HW Exception by GPU node-1 (Agent handle: 0xd851680) reason :GPU Hangというエラーが出ている。これは別の問題のようだが、詳しくは調査していない。

入手可能と書いたが、RTX 5090とGB10の国内在庫はかなり少なくなってきている(2026/1/9) ↩
データセンター向けのAMD Instinct についてはROCm公式ブログに記事がある https://rocm.blogs.amd.com/artificial-intelligence/ai-weather-forecasting/README.html ↩
RadeonでONNX runtimeを実行する場合、ROCMExecutionProvider (ROCm EP)とMIGraphXExecutionProvider(MIGraphX EP)の2つの選択肢がある。ただし、AI気象モデルの場合モデルが巨大で2GBを超えるため、、MIGraphX EPではエラーが出る（[libprotobuf ERROR /onnxruntime/build/Linux/Release/_deps/protobuf-src/src/google/protobuf/message_lite.cc:449] onnx.ModelProto exceeded maximum protobuf size of 2GB: 2268174223）。これを無理やり回避することもできるが、モデルのコンパイルに非常に時間がかかるなどの問題があり、実質的にMIGraphX EPでは動作しないと考えたほうがよい。 ↩

You get articles that match your needs
You can efficiently read back useful information
You can use dark theme

What you can do with signing up