ONNX Runtime GPU:なぜ自前でビルドすべきか?

本格的なビルドに入る前に、いくつかの概念を整理しておきます。

onnxruntime_gpuとは?



ONNX Runtime(ORT) は最適化された AI 推論エンジンです。PyTorch や TensorFlow で作成したモデルを .onnx 形式に変換し、さまざまなハードウェア上で最速で実行できるようにします。その中でも onnxruntime_gpu は NVIDIA GPU の CUDAcuDNN を活用し、演算速度を最大化したバージョンです。

image of onnxruntime

なぜ AI 生成モデル(画像・動画)で頻繁に使われるのか?

Stable Diffusion や Whisper といったモデルはパラメータが巨大で、計算量が膨大です。ORT は以下の最適化を提供します:

  • Graph Optimization: 不要な演算を削減し、複数の演算を一つに統合(Operator Fusion)します。
  • Memory Management: GPU メモリ割り当てを効率的に管理し、大規模モデルでもスムーズに動作させます。
  • Hardware Acceleration: TensorRT などのアクセラレータと組み合わせ、純粋な PyTorch に比べて数倍の速度向上を実現します。

なぜ aarch64(ARM)ユーザーはソースビルドが必要なのか?



通常は pip install onnxruntime-gpu でインストールできますが、aarch64 では状況が異なります。

  1. バイナリが提供されていない: PyPI にある公式 Wheel は主に x86_64 用です。
  2. 最新アーキテクチャへの対応: 最近の CUDA 13.0Compute Capability 12.1(Blackwell) などは、正式リリース版がまだ対応していないことがあります。
  3. 最適化の追求: 特定のマシン(DGX Spark など)の CPU/GPU 特性に合わせた最適化オプションを適用したい場合、ソースからビルドするのが唯一の方法です。

DGX‑Spark(aarch64)で ONNX Runtime GPU をビルドする手順

以下に、私の実体験をもとにしたビルド手順をまとめます。

1. 環境確認(Environment Check)

まずシステムのスペックを確認します。私は DGX‑Spark(Grace Blackwell GPU)を使用しました。

項目 情報
OS/Arch Linux aarch64
GPU NVIDIA GB10(Blackwell)
CUDA 13.0(V13.0.88)
Python 3.12.3
Compute Cap 12.1
# GPU 情報確認
nvidia-smi
# CUDA バージョン確認
nvcc --version

2. cuDNN のインストール(不足コンポーネントの追加)

システムに cuDNN が無かったため、まず pip で Python パッケージとしてインストールし、パスを設定しました。

# CUDA 13 用 cuDNN インストール
pip install nvidia-cudnn-cu13

# インストール先パス確認(Python スクリプト利用)
python3 -c "import site, os; print(os.path.join(site.getsitepackages()[0], 'nvidia/cudnn'))"

その後、環境変数を設定してビルドスクリプトが cuDNN を見つけられるようにします。

export CUDA_HOME=/usr/local/cuda
export CUDNN_HOME=/home/jesse/onnxruntime/venv/lib/python3.12/site-packages/nvidia/cudnn
export PATH=$CUDA_HOME/bin:$PATH
export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$CUDNN_HOME/lib:$LD_LIBRARY_PATH
export CUDACXX=$CUDA_HOME/bin/nvcc

3. ONNX Runtime のソースビルド

いよいよ核心部分です。build.sh スクリプトに Blackwell アーキテクチャ(121)と新しい CUDA パスを指定します。

./build.sh \
  --config Release \
  --update --build \
  --parallel \
  --build_wheel \
  --use_cuda \
  --cuda_home $CUDA_HOME \
  --cudnn_home $CUDNN_HOME \
  --skip_tests \
  --cmake_generator Ninja \
  --cmake_extra_defines CMAKE_CUDA_ARCHITECTURES=121 \
  --cmake_extra_defines CMAKE_CUDA_FLAGS="-I/usr/local/cuda/include/cccl" \
  --cmake_extra_defines CMAKE_CXX_FLAGS="-I/usr/local/cuda/include/cccl"

Tip: CCCL 関連のインクルードパスエラーを防ぐために CMAKE_CUDA_FLAGS を追加すると効果的です。

4. ビルド成果物の確認とインストール

ビルドが完了すると build/Linux/Release/dist/ 配下に aarch64 用の .whl ファイルが生成されます。

# 生成された Wheel ファイル確認
ls -lh build/Linux/Release/dist/onnxruntime_gpu-*.whl

# インストール
pip install build/Linux/Release/dist/onnxruntime_gpu-1.25.0-cp312-cp312-linux_aarch64.whl

5. 最終検証

ライブラリが正しく GPU を認識するか確認します。

import onnxruntime as ort
print("ORT version:", ort.__version__)
print("Available providers:", ort.get_available_providers())

結果: ['CUDAExecutionProvider', 'CPUExecutionProvider'] と表示されれば成功です! 🎉


まとめ

ビルドした Wheel ファイルは後で Docker イメージを作成する際に COPY 命令で組み込めるので再利用が可能です。aarch64 と最新 CUDA 環境は構築に手間がかかりますが、一度正しく構築すれば生成 AI モデルの運用で圧倒的な性能を引き出せます。

この手順が DGX‑Spark ユーザーの皆さんにとって一筋の光となりますように!