5. GPU Acceleration

NLSQ uses JAX for GPU acceleration, providing significant speedups for large datasets with no code changes required.

• 5.1. GPU Setup
• 5.2. GPU Usage
• 5.3. Multi-GPU

5.4. Chapter Overview

GPU Setup (10 min)

Install JAX with GPU support and verify configuration.

GPU Usage (5 min)

Automatic GPU detection and usage patterns.

Multi-GPU (5 min)

Using multiple GPUs for parallel processing.

5.5. Quick Start

# Install JAX with CUDA support
pip install --upgrade "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

from nlsq import fit

# GPU is used automatically if available
popt, pcov = fit(model, x, y, p0=[...])

5.6. Performance Benefits

GPU acceleration provides significant speedups:

Dataset Size	CPU Time	GPU Time	Speedup
10K points	~1.0s	~0.5s	2x
100K points	~5.0s	~0.8s	6x
1M points	~30s	~2s	15x
10M points	~5min	~15s	20x

First fit includes JIT compilation overhead. Subsequent fits are faster.