cuda-pytorch

Multi-architecture Docker base image for PyTorch deep learning inference with NVIDIA GPU acceleration, supporting both AMD64 desktop GPUs and ARM64 Jetson devices. Part of the Phystack CUDA image family.

Overview

cuda-pytorch provides a ready-to-use Docker image with PyTorch, Hugging Face Transformers, and supporting ML libraries pre-installed on top of the Phystack CUDA base image. It targets GPU-accelerated inference workloads on both desktop-class NVIDIA GPUs (Ampere through Blackwell) and NVIDIA Jetson edge devices (Nano, Xavier, Orin). Downstream services and projects use this as a FROM base image for their own inference containers.

Architecture

graph LR
    Base[phygrid/cuda-base] --> Image[cuda-pytorch]
    Image --> Service1[Inference services]
    Image --> Service2[ML pipelines]
    Service1 --> GPU[NVIDIA GPU]
    Service2 --> GPU
    HF[Hugging Face Hub] -->|model download| Service1

Loading

Tech Stack

Layer	Technology
Base image	phygrid/cuda-base (Python 3.11, FastAPI, Uvicorn)
ML framework	PyTorch 2.8.0 / CUDA 12.8 (AMD64), PyTorch 2.5.0 / JetPack 6.1 (ARM64)
Transformers	Hugging Face Transformers 4.36.2, Accelerate 0.25.0
Optimization	bitsandbytes 0.41.3 (AMD64 only), Optimum 1.16.1
Vision / Audio	OpenCV 4.8, torchvision, librosa 0.10, torchaudio
Scientific	scikit-learn 1.3, scipy 1.11, matplotlib 3.8
CI/CD	GitHub Actions, Docker Buildx multi-arch
Registry	Docker Hub (phygrid/cuda-pytorch)

Prerequisites

• Docker with Buildx (for multi-arch builds)
• NVIDIA GPU with driver supporting CUDA 12.8 (AMD64) or JetPack 6.1 (ARM64)
• --gpus all or --runtime nvidia capability on the Docker host

Installation

docker pull phygrid/cuda-pytorch:latest

Specific version:

docker pull phygrid/cuda-pytorch:v1.0.11

Usage

As a base image

FROM phygrid/cuda-pytorch:latest

COPY models/ /app/pytorch_models/
COPY src/ /app/src/
RUN pip install -r requirements.txt

CMD ["python", "inference_server.py"]

Interactive development

docker run -it --rm \
  --gpus all \
  -v $(pwd):/app/workspace \
  -v ~/.cache/huggingface:/app/cache/huggingface \
  -p 8000:8000 \
  phygrid/cuda-pytorch:latest bash

Hugging Face model inference

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "microsoft/DialoGPT-medium"
tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir="/app/cache/transformers")
model = AutoModelForCausalLM.from_pretrained(model_name, cache_dir="/app/cache/transformers")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)

inputs = tokenizer.encode("Hello, how are you?", return_tensors="pt").to(device)
outputs = model.generate(inputs, max_length=50, num_return_sequences=1)

Environment Variables

Variable	Default	Description
PYTORCH_CUDA_ALLOC_CONF	max_split_size_mb:128	CUDA memory allocator configuration
CUDA_LAUNCH_BLOCKING	0	Set to 1 for synchronous CUDA debugging
TRANSFORMERS_CACHE	/app/cache/transformers	Hugging Face Transformers model cache
HF_HOME	/app/cache/huggingface	Hugging Face Hub cache directory
TOKENIZERS_PARALLELISM	true	Enable parallel tokenizer processing
TORCH_CUDA_ARCH_LIST	8.0;8.6;8.9;9.0 (AMD64) / 5.3;6.2;7.2;8.7 (ARM64)	Target GPU compute capabilities

Volume Mounts

Path	Purpose
/app/pytorch_models	PyTorch model storage
/app/cache/transformers	Hugging Face Transformers cache
/app/cache/huggingface	Hugging Face Hub cache
/app/cache/torch	PyTorch hub cache
/app/data	Input/output data
/app/logs	Application logs

Project Structure

Dockerfile              # Multi-arch build with PyTorch installation
VERSION                 # Semantic version (patch auto-incremented by CI)
.github/workflows/
  docker-build-deploy.yml  # Build, push, tag, and release workflow
LICENSE                 # MIT

Testing

The image ships a built-in validation script at /app/pytorch_test.py. It verifies PyTorch, CUDA availability, and GPU tensor operations.

docker run --rm --gpus all phygrid/cuda-pytorch:latest python /app/pytorch_test.py

A Docker HEALTHCHECK runs this script every 30 seconds inside running containers.

Deployment

Built and pushed to Docker Hub by GitHub Actions on every push to main that changes Dockerfile, VERSION, or the workflow file.

The pipeline builds linux/amd64 and linux/arm64 images, pushes versioned and latest tags, creates a GitHub release, and syncs the README to Docker Hub.

Publishing

• Patch auto-increment: CI bumps the patch version and commits the updated VERSION file when the current tag already exists.
• Major/minor bump: Edit the VERSION file manually and push to main.
• Tags: Created automatically by CI in the format v1.0.11.
• Registry: Docker Hub at phygrid/cuda-pytorch.

Troubleshooting

Out of GPU memory during model loading -- Lower PYTORCH_CUDA_ALLOC_CONF to max_split_size_mb:64 for Jetson devices, or use device_map="auto" with low_cpu_mem_usage=True when loading large models.

ARM64 build fails for bitsandbytes -- This is expected. bitsandbytes is only installed on AMD64. The Dockerfile conditionally skips it on ARM64.

Related Documentation

• phygrid/cuda-base -- upstream base image
• PyTorch CUDA compatibility -- version matrix
• NVIDIA JetPack 6.1 -- ARM64 Jetson SDK
• Phystack CI/CD docs -- shared pipeline reference

License

MIT