Conteneurs ARM64 et DGX Spark

Le Dockerfile compile pour Linux AMD64 et ARM64. Le CPU reste la valeur par défaut sur les deux ; un CPU ARM64 n’implique pas un GPU NVIDIA.

CPU ARM64

Compile sur l’hôte cible. Docker sélectionne son architecture native :

docker compose run --build --rm laya

Les compilations croisées avec docker buildx build --platform linux/arm64 --load -t laya:arm64 . ont besoin d’un builder ARM64 ou d’une émulation configurée. Une compilation émulée ne montre pas les performances d’inférence natives.

DGX Spark

Utilise l’hôte Linux du Spark avec son pilote NVIDIA pris en charge et le NVIDIA Container Toolkit. L’override sélectionne ARM64, les roues CUDA 13.0 et GPU 0. Utilise-le à la place de compose.cuda.yaml, pas avec :

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

Pour servir l’API HTTP sur le Spark, ajoute compose.http.yaml. Les réglages de service HTTP s’appliquent sans changement :

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

Définis LAYA_GPU_ID pour sélectionner un autre appareil.

TORCH_VERSION épingle PyTorch pour chaque build, CPU comme CUDA. Compose lit LAYA_TORCH_VERSION ; les builds directs prennent --build-arg TORCH_VERSION=2.14.0 --build-arg TORCH_INDEX=cu130. Changer l’un ou l’autre exige une recompilation, car une variable d’environnement de runtime ne peut pas remplacer la roue installée.

Les builds CUDA 13.0 de PyTorch pour ARM64 dépendent de cuSPARSELt 0.8.0 (PyTorch 2.11) ou 0.8.1 (PyTorch 2.14). Les roues AArch64 de NVIDIA pour ces deux versions déclarent manylinux2014_sbsa dans leur fichier WHEEL, ce que pip check rejette ; 0.9.0 corrige cela. Le build vérifie que la bibliothèque est ELF64 AArch64 et se charge, puis corrige ce tag et son hash RECORD. Toute autre version avec le même défaut fait échouer le build au lieu de recevoir la réparation, et pip check tourne quand même. Cela suit la découverte de compatibilité de @TheIrritainer dans FastLaya.

Rapporter des résultats Spark

La CI n’a pas de GPU, donc l’inférence Spark a besoin d’un rapport depuis du matériel réel. Exécute ceci sur le Spark et joins sa sortie avec nvidia-smi, les versions de l’OS et du pilote, et la révision de l’image :

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())
'

Répète avec chaque checkpoint que tu comptes exécuter. Un test CPU ARM64 natif n’établit pas le support des kernels Blackwell ni de l’inférence GPU. La pile CUDA spécifique à la plateforme Jetson n’est pas couverte par l’override Spark.

Apple Silicon

L’accélération GPU Apple nécessite PyTorch macOS natif avec MPS. Docker Desktop exécute un conteneur Linux, qui n’a pas de backend MPS, donc le conteneur utilise le CPU. Laya a déjà une voie d’appareil MPS, et PR #51 et PR #109 traitent la compatibilité et les performances MPS hors Docker.