ARM64 と DGX Spark のコンテナ

ARM64 と DGX Spark のコンテナ

Dockerfile は Linux AMD64 と ARM64 向けにビルドします。どちらも既定は CPU のままで、ARM64 の CPU が NVIDIA GPU を意味するわけではありません。

ARM64 CPU

対象のホスト上でビルドします。Docker がネイティブのアーキテクチャを選びます：

docker compose run --build --rm laya

docker buildx build --platform linux/arm64 --load -t laya:arm64 . によるクロスビルドには、 ARM64 のビルダーか構成済みのエミュレーションが必要です。エミュレーションでのビルドは、ネイティブ の推論性能を示しません。

DGX Spark

Spark の Linux ホストを、対応する NVIDIA ドライバと NVIDIA Container Toolkit とともに使います。この override は ARM64、CUDA 13.0 の wheel、GPU 0 を選びます。 compose.cuda.yaml の代わりに使い、両方は使いません：

docker compose -f compose.yaml -f compose.spark.yaml run --build --rm laya

Spark で HTTP API を提供するには、compose.http.yaml を追加します。 HTTP サービス の設定がそのまま当てはまります：

docker compose -f compose.yaml -f compose.http.yaml -f compose.spark.yaml up --build laya-serve

LAYA_GPU_ID を設定すると別のデバイスを選べます。

TORCH_VERSION は CPU でも CUDA でも、すべてのビルドで PyTorch を固定します。Compose は LAYA_TORCH_VERSION を読み、直接ビルドは --build-arg TORCH_VERSION=2.14.0 --build-arg TORCH_INDEX=cu130 を取ります。どちらを変えるにも 再ビルドが必要です。実行時の環境変数では、インストール済みの wheel を置き換えられないからです。

PyTorch の ARM64 向け CUDA 13.0 ビルドは、cuSPARSELt 0.8.0（PyTorch 2.11）または 0.8.1 （PyTorch 2.14）に依存します。NVIDIA のこの 2 バージョンの AArch64 wheel は、WHEEL ファイルの 中で manylinux2014_sbsa を宣言しており、pip check がそれを拒否します。0.9.0 がそれを修正して います。ビルドは、そのライブラリが ELF64 AArch64 で読み込み可能かを確認し、その後そのタグと RECORD ハッシュを修正します。同じ不具合を持つ他のバージョンは、修理されずにビルドが失敗し、 pip check は引き続き実行されます。これは FastLaya での @TheIrritainer の互換性の指摘に従ったものです。

Spark の結果を報告する

CI には GPU がないので、Spark での推論には実機からの報告が必要です。Spark 上で次を実行し、その 出力を nvidia-smi、OS とドライバのバージョン、イメージのリビジョンとともに提出してください：

docker compose -f compose.yaml -f compose.spark.yaml run --build --rm laya python -c '
import json, platform, torch
from pathlib import Path
from laya import load
assert platform.machine() == "aarch64"
assert torch.cuda.is_available()
print(torch.__version__, torch.version.cuda, torch.cuda.get_device_name(0))
print(torch.cuda.get_device_capability(0), torch.cuda.get_arch_list())
agent = load("convaiinnovations/laya", device="cuda")
request = json.loads(Path("/opt/laya/examples/request.json").read_text())
result = agent.predict(request["state"], request["questions"])
assert next(agent.model.parameters()).device.type == "cuda", "fell back to CPU"
assert set(result["answers"]) == set(request["questions"])
print("CUDA inference passed", torch.cuda.max_memory_allocated())
'

実行する予定のチェックポイントごとに繰り返してください。ネイティブ ARM64 の CPU テストは、 Blackwell カーネルや GPU 推論への対応を裏付けるものではありません。Jetson のプラットフォーム 固有の CUDA スタックは、Spark の override の対象外です。

Apple Silicon

Apple GPU の高速化には、MPS を備えたネイティブ macOS の PyTorch が必要です。Docker Desktop は MPS バックエンドを持たない Linux コンテナを動かすので、コンテナは CPU を使います。Laya には すでに MPS のデバイスパスがあり、PR #51 と PR #109 が Docker の外での MPS の互換性と 性能を扱っています。