ONNX Runtime GPU: Почему и как нужно собирать самостоятельно

ONNX Runtime GPU: Почему стоит собрать собственный билд?

Прежде чем перейти к практической части, уточним несколько ключевых понятий.

Что такое onnxruntime_gpu?

ONNX Runtime (ORT) – оптимизированный движок для инференса ИИ. Модели, созданные в PyTorch или TensorFlow, конвертируются в формат .onnx, после чего ORT обеспечивает их максимально быструю работу на разных аппаратных платформах. Версия onnxruntime_gpu использует CUDA и cuDNN от NVIDIA для ускорения вычислений на GPU.

Почему эта версия популярна в генеративных моделях (изображения/видео)?

Модели вроде Stable Diffusion или Whisper обладают огромным числом параметров и требуют огромных вычислительных ресурсов. ORT предлагает следующие оптимизации:

• Graph Optimization: удаляет лишние операции и объединяет их (Operator Fusion).
• Memory Management: эффективно управляет памятью GPU, позволяя запускать крупные модели без задержек.
• Hardware Acceleration: в сочетании с TensorRT обеспечивает в несколько раз большую скорость по сравнению с чистым PyTorch.

Почему пользователям aarch64 (ARM) необходимо собирать из исходников?

Обычно достаточно выполнить pip install onnxruntime-gpu, но для aarch64 ситуация иная.

1. Отсутствие готовых бинарников: официальные колёса в PyPI в основном предназначены для x86_64.
2. Поддержка новейших архитектур: такие версии CUDA, как 13.0, и Compute Capability 12.1 (Blackwell) часто не покрываются официальными релизами.
3. Оптимальная настройка: если требуется учесть специфические особенности CPU/GPU конкретного узла (например, DGX Spark), единственный путь – собрать библиотеку самостоятельно.

Сборка ONNX Runtime GPU на DGX‑Spark (aarch64)

Ниже представлен пошаговый процесс, основанный на моём опыте.

1. Проверка окружения (Environment Check)

Сначала убедимся в характеристиках системы. Я использовал DGX‑Spark с GPU Grace‑Blackwell.

# Информация о GPU
nvidia-smi
# Версия CUDA
nvcc --version

2. Установка недостающего компонента: cuDNN

Поскольку cuDNN отсутствовал, я установил его через pip в виде Python‑пакета и указал путь вручную.

# cuDNN для CUDA 13
pip install nvidia-cudnn-cu13

# Проверка пути установки
python3 -c "import site, os; print(os.path.join(site.getsitepackages()[0], 'nvidia/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"

Совет: добавление CMAKE_CUDA_FLAGS помогает избежать ошибок с включением заголовков CCCL.

4. Проверка артефактов и установка

После завершения сборки в build/Linux/Release/dist/ появится .whl‑файл, предназначенный только для aarch64.

# Показать полученный 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!