# How to Enable GPU Acceleration for Needle Training: Complete Setup Guide

> Learn how to enable GPU acceleration for Needle training. Install the CUDA-compatible JAX build via the gpu extra for automatic NVIDIA GPU utilization. Accelerate your training now.

- Repository: [Cactus Compute, Inc./needle](https://github.com/cactus-compute/needle)
- Tags: how-to-guide
- Published: 2026-09-04

---

**GPU acceleration in Needle is enabled by installing the CUDA-compatible JAX build via the `gpu` extra, allowing the training pipeline to automatically detect and utilize NVIDIA GPUs without requiring code changes.**

The `cactus-compute/needle` repository implements its fine-tuning pipeline on JAX, which abstracts hardware acceleration through backend detection. To enable GPU acceleration for Needle training, you must install the optional dependencies that include the CUDA-enabled JAX wheel rather than the default CPU-only version.

## Install GPU-Enabled Dependencies

Needle distributes hardware-specific builds as Python extras defined in the package configuration. According to the project README (lines 26-33), you must specify both the `train` and `gpu` extras when installing `cactus-needle` to pull in the NVIDIA CUDA dependencies.

Run the following command to install the GPU-accelerated training stack:

```bash
pip install "cactus-needle[train,gpu]"

```

This command installs the `jax[cuda]` build specified in [`requirements-train.txt`](https://github.com/cactus-compute/needle/blob/main/requirements-train.txt), replacing the CPU-only JAX installation with a version that compiles against local CUDA libraries. No manual CUDA toolkit installation is required beyond having compatible drivers, as JAX bundles the necessary CUDA runtime libraries.

## Verify JAX Device Detection

Before running training jobs, confirm that JAX can enumerate your GPU devices. Create a short verification script to probe the available accelerators:

```python
import jax

devices = jax.devices()
print(f"Detected devices: {devices}")

```

Execute this check in your terminal:

```bash
python -c "import jax; print(jax.devices())"

```

The output should list one or more `cuda` or `gpu` entries. If you see only `cpu` devices, the GPU installation failed or your CUDA drivers are not properly configured. Successful detection means the `needle` CLI will automatically route computations to the GPU.

## Run Training Commands on GPU

Once JAX detects the GPU, the standard Needle fine-tuning command executes on the accelerator without additional flags. The training loop implemented in [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py) delegates device placement to JAX, which defaults to the fastest available backend.

Execute training as documented in [`doc/finetuning.md`](https://github.com/cactus-compute/needle/blob/main/doc/finetuning.md) (lines 57-63):

```bash
needle finetune data.jsonl --epochs 10 --out adapter.pkl

```

The same command that runs on CPU will now utilize the GPU for forward passes, backward propagation, and optimizer steps. Batch sizes and learning rates may need adjustment to leverage the increased memory and throughput of GPU hardware.

## Enable Metal Acceleration on Apple Silicon

For macOS systems with Apple Silicon GPUs (M1/M2/M3), install the `metal` extra instead of `gpu` to use the JAX-Metal plugin. As detailed in [`doc/finetuning.md`](https://github.com/cactus-compute/needle/blob/main/doc/finetuning.md) (lines 63-66), this variant targets Apple’s Metal Performance Shaders rather than CUDA.

Install the Metal-enabled build:

```bash
pip install "cactus-needle[train,metal]"

```

After installation, run the same JAX device verification and `needle finetune` commands. JAX will automatically route operations to the `metal` backend when available.

## Summary

- **Install GPU extras**: Use `pip install "cactus-needle[train,gpu]"` to pull in the CUDA-enabled JAX build referenced in the README (lines 26-33).
- **Verify detection**: Run `jax.devices()` to confirm GPU availability before training.
- **No code changes**: The same `needle finetune` CLI works immediately after GPU installation, as the JAX backend in [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py) handles device placement automatically.
- **Apple Silicon**: Substitute `gpu` with `metal` in the pip command for macOS GPU acceleration.

## Frequently Asked Questions

### Do I need to set environment variables to force GPU usage in Needle?

No. Needle relies on JAX’s default device selection logic, which automatically prioritizes GPUs when available. As long as you install the correct extras from [`doc/finetuning.md`](https://github.com/cactus-compute/needle/blob/main/doc/finetuning.md), the training loop in [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py) executes on the detected accelerator without requiring `CUDA_VISIBLE_DEVICES` or `JAX_PLATFORM_NAME` overrides.

### Can I install both GPU and CPU versions simultaneously?

You cannot use both builds in the same Python environment. The `cactus-needle[train]` extra pulls in `jax` (CPU), while `[train,gpu]` upgrades it to `jax[cuda]`. Installing the GPU extra replaces the CPU-only JAX wheel; attempting to install both causes package conflicts. Choose the extra matching your hardware.

### Why does JAX still show only CPU devices after installing the GPU extra?

This typically indicates a CUDA driver mismatch. The `jax[cuda]` package requires NVIDIA drivers compatible with the bundled CUDA version. Run `nvidia-smi` to verify driver installation, and ensure your CUDA version matches the JAX wheel requirements listed in [`requirements-train.txt`](https://github.com/cactus-compute/needle/blob/main/requirements-train.txt). Reinstall the `cactus-needle[train,gpu]` package after updating drivers if necessary.

### Is mixed-precision training available on GPU in Needle?

Needle delegates all numerical operations to JAX, which supports mixed-precision via `jax.numpy` and `jax.lax` APIs. While the `needle finetune` CLI uses standard precision by default, the underlying [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py) implementation runs on JAX’s XLA compiler, which automatically optimizes GPU tensor cores when available. Explicit mixed-precision policies require modifying the training script rather than CLI flags.