AMD、MONAI Model Zoo 5バンドルをInstinct GPUで検証済み推論に対応

AMDはAMD Life Science 26.08で、MONAI Model Zooの5つのバンドルについて、AMD Instinct GPU上のROCm最適化・検証済み推論サポートをリリースしました。MI355X(gfx950)でのp50レイテンシは、spleen_deepeditの3.60 msからpancreas_ct_dintsの3,287 msまでです。

AMD InstinctはAMDのデータセンター向けAI・HPCアクセラレーター製品群で、ROCmはそのGPU計算向けのオープンなソフトウェア群です。

何が新しいか

今回のリリースはEarly Accessで、推論に焦点を置いています。最適化は各バンドルに同梱するinference_rocm.json(バンドルによりinference_rocm.yaml)というROCmオーバーレイで実現し、MONAIのネイティブな設定オーバーレイ機構によりロード時にベース設定へマージします。元のバンドル、重み、推論ロジックは変更しません。推論には通常のPyTorchを使い、オフラインのエンジンコンパイルは不要です。

オーバーレイは5バンドルすべてに同じ3つの手法を適用します。

• channels-last 3D(memory_format=torch.channels_last_3d):HIPの3D畳み込みカーネルが想定するレイアウトに合わせ、冗長な転置を除去
• BF16自動混合精度(amp_kwargs={'dtype': torch.bfloat16})
• torch.compile():Inductorバックエンドを使い、ROCm上でhipBLASLtのカーネル選択とグラフレベルの融合を行う。ネットワーク単位でtorch.compile(@network)として適用

5つのうち4つはこれらのキーだけで動作し、spleen_deepeditのみenable_gemm_transpose(True)を呼ぶ1行を追加しています。

対象バンドルとMI355Xでの性能

計測はAMD Instinct MI355X(gfx950)、ROCm 10.0で、脾臓CTボリューム(spleen_10.nii.gz)を使い、ウォームアップ5回の後に20回計測したp50値です。

バンドル	タスク	アーキテクチャ	対象	p50 (ms)	Dice
Vista3D	インタラクティブ3Dセグメンテーション	SegResNetベース+SAM風プロンプト	133構造	268	PASS
SwinUNETR (BTCV)	多臓器CTセグメンテーション	Swin Transformer U-Net	腹部13臓器	293	PASS
wholeBody_ct_segmentation	全身CTセグメンテーション	SegResNet	104構造	386	PASS
pancreas_ct_dints_segmentation	膵臓CTセグメンテーション	DiNTS	膵臓	3,287	PASS
spleen_deepedit_annotation	インタラクティブ脾臓アノテーション	DynUNet	脾臓	3.60	PASS

AMDによると、スライディングウィンドウ推論はデータ移動がコストの大半を占めるため、MI355Xの288 GB HBM3Eにより、最大級のVista3DとwholeBody_ctもワーキングセットをあふれさせずに動作します。

要件と試し方

• AMD Instinct GPU(計測はMI355X、リリースはMI300XとMI325Xもサポート)
• ROCm 10.0、Python 3.12
• MONAI 1.6.0(AMD ROCm package indexまたはROCm PyTorch Docker imageから入手)

git clone https://github.com/AMD-Ecosystem/monai-model-zoo

export PYTORCH_MIOPEN_SUGGEST_NHWC=1
export MIOPEN_FIND_MODE=1
export MIOPEN_FIND_ENFORCE=4
export TORCHINDUCTOR_MAX_AUTOTUNE=1
export TORCHINDUCTOR_MAX_AUTOTUNE_GEMM=1
export TORCHINDUCTOR_COORDINATE_DESCENT_TUNING=1
export TORCHINDUCTOR_EPILOGUE_FUSION=1
export TORCHINDUCTOR_MAX_AUTOTUNE_CONV_BACKENDS=ATEN,TRITON

cd models/BUNDLE_NAME
python -m monai.bundle run \
  --config_file "['configs/inference.json', 'configs/inference_rocm.json']" \
  --bundle_root . \
  --dataset_dir /path/to/input_dir \
  --output_dir /path/to/output_dir

注意点は次の通りです。

• 必ずバンドルディレクトリ内から実行します。MONAIはrun()開始時にバンドルルートをPythonパスに加える一方、initialize()はその前にscripts.*を解決するため、リポジトリルートから起動するとカスタムコードを持つバンドルのインポートに失敗します。
• --dataset_dirと--output_dirは絶対パスで指定します。
• 入力キーはバンドルごとに異なり、vista3dは--dataset_dirではなく--input_dictを使います。

今後

AMDは5バンドル全体での学習の検証、追加のModel Zooバンドル対応、畳み込み負荷の大きいwholeBody_ctとspleen_deepedit向けのMIOpen Composable Kernel(CK)3Dカーネルの作業を続けています。本記事はAMD Life Science 26.08医療画像スタックのシリーズの一つで、hipCIMとMONAIデータロードの記事が既に出ており、インタラクティブアノテーションの記事が今後予定されています。