Contentores ARM64 e DGX Spark

O Dockerfile compila para Linux AMD64 e ARM64. A CPU continua a ser a predefinição em ambos; uma CPU ARM64 não implica uma GPU NVIDIA.

CPU ARM64

Compila no anfitrião de destino. O Docker seleciona a sua arquitetura nativa:

docker compose run --build --rm laya

As compilações cruzadas com docker buildx build --platform linux/arm64 --load -t laya:arm64 . precisam de um builder ARM64 ou de emulação configurada. Uma compilação emulada não mostra o desempenho de inferência nativo.

DGX Spark

Usa o anfitrião Linux do Spark com o seu controlador NVIDIA suportado e o NVIDIA Container Toolkit. O override seleciona ARM64, as wheels CUDA 13.0 e a GPU 0. Usa-o em vez do compose.cuda.yaml, e não com ele:

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

Para servir a API HTTP no Spark, adiciona compose.http.yaml. As definições de serviço HTTP aplicam-se sem alterações:

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

Define LAYA_GPU_ID para selecionar outro dispositivo.

TORCH_VERSION fixa o PyTorch em todas as compilações, tanto em CPU como em CUDA. O Compose lê LAYA_TORCH_VERSION; as compilações diretas usam --build-arg TORCH_VERSION=2.14.0 --build-arg TORCH_INDEX=cu130. Alterar qualquer um deles exige uma recompilação, porque uma variável de ambiente de runtime não pode substituir a wheel instalada.

As compilações CUDA 13.0 do PyTorch para ARM64 dependem do cuSPARSELt 0.8.0 (PyTorch 2.11) ou 0.8.1 (PyTorch 2.14). As wheels AArch64 da NVIDIA para essas duas versões declaram manylinux2014_sbsa no seu ficheiro WHEEL, que o pip check rejeita; a 0.9.0 corrige isso. A compilação verifica que a biblioteca é ELF64 AArch64 e carrega, e depois corrige essa etiqueta e o hash do seu RECORD. Qualquer outra versão com o mesmo defeito falha a compilação em vez de receber a reparação, e o pip check continua a correr. Isto segue a descoberta de compatibilidade do @TheIrritainer no FastLaya.

Reportar resultados do Spark

A CI não tem GPU, por isso a inferência no Spark precisa de um relatório a partir de hardware real. Executa isto no Spark e inclui o seu resultado com o nvidia-smi, as versões do SO e do controlador, e a revisão da imagem:

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

Repete com cada checkpoint que pretendas executar. Um teste numa CPU ARM64 nativa não estabelece suporte a kernels Blackwell nem a inferência em GPU. A stack CUDA específica da plataforma Jetson não é abrangida pelo override do Spark.

Apple Silicon

A aceleração por GPU da Apple precisa de PyTorch nativo do macOS com MPS. O Docker Desktop corre um contentor Linux, que não tem backend MPS, por isso o contentor usa a CPU. O Laya já tem um caminho de dispositivo MPS, e o PR #51 e o PR #109 tratam da compatibilidade e do desempenho do MPS fora do Docker.