5.1. GPU Setup

This guide covers installing and configuring JAX for GPU acceleration.

Important

GPU acceleration is supported on Linux only. macOS and Windows automatically use the CPU backend — NLSQ enforces this at import time by setting JAX_PLATFORM_NAME=cpu.

5.1.1. NVIDIA GPU (CUDA) — Linux

Requirements:

  • Linux operating system

  • NVIDIA GPU with CUDA support (Maxwell or newer, SM ≥ 5.2)

  • CUDA 12.x or 13.x drivers installed

  • cuDNN (bundled with JAX)

Installation:

# Recommended: use the Makefile target (auto-detects CUDA version)
make install-jax-gpu

# Or manually for CUDA 12
pip install --upgrade "jax[cuda12-local]"

# Or for CUDA 13
pip install --upgrade "jax[cuda13-local]"

Verify installation:

import jax

print(f"JAX version: {jax.__version__}")
print(f"Devices: {jax.devices()}")
print(f"Default backend: {jax.default_backend()}")

Expected output:

JAX version: 0.9.0
Devices: [CudaDevice(id=0)]
Default backend: gpu

5.1.2. AMD GPU (ROCm) — Linux

Requirements:

  • Linux operating system

  • AMD GPU with ROCm support

  • ROCm 5.x+ installed

Installation:

pip install --upgrade "jax[rocm]" -f https://storage.googleapis.com/jax-releases/jax_releases.html

5.1.3. macOS and Windows (CPU Only)

NLSQ enforces CPU-only mode on macOS and Windows at import time. No additional configuration is needed — the following environment variables are set automatically:

  • NLSQ_FORCE_CPU=1

  • JAX_PLATFORM_NAME=cpu

  • JAX_PLATFORMS=cpu

On macOS, additional guards prevent SIGBUS crashes from Metal/OpenGL/XLA conflicts (XLA_FLAGS, OMP_NUM_THREADS, MPLBACKEND, etc.).

# Just install JAX (CPU backend is automatic)
pip install jax jaxlib

5.1.4. Docker Setup

For containerized environments:

FROM nvidia/cuda:12.1-runtime-ubuntu22.04

RUN pip install jax[cuda12_pip] nlsq

Run with GPU access:

docker run --gpus all my-nlsq-container

5.1.5. Troubleshooting Installation

CUDA not found:

# Check CUDA installation
nvidia-smi
nvcc --version

# CUDA path may need to be set
export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH

JAX falls back to CPU:

import jax

if jax.default_backend() == "cpu":
    print("GPU not detected!")
    # Check CUDA drivers
    # Reinstall JAX with CUDA support

Out of memory errors:

# Limit GPU memory
import os

os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"
os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"] = "0.5"

Multiple GPUs not detected:

# Check all GPUs visible
nvidia-smi

# Set visible devices
export CUDA_VISIBLE_DEVICES=0,1

5.1.6. Verifying NLSQ GPU Usage

from nlsq import get_device

device = get_device()
print(f"NLSQ device: {device}")

# Check if GPU is available
import jax

if jax.default_backend() == "gpu":
    print("GPU acceleration enabled!")
else:
    print("Running on CPU")

5.1.7. Next Steps