# Needle Training Hardware Requirements: GPU and Apple Silicon Support Explained

> Discover Needle training hardware requirements. Learn how to enable GPU and Apple Silicon support by installing optional JAX dependencies for accelerated training.

- Repository: [Cactus Compute, Inc./needle](https://github.com/cactus-compute/needle)
- Tags: getting-started
- Published: 2026-09-06

---

**You can train Needle on any accelerator that JAX supports—NVIDIA GPU, Apple Silicon GPU, or CPU—by installing the appropriate optional dependency.**

Needle's training pipeline is built entirely on **plain JAX**, which abstracts hardware acceleration behind a unified API. This means the same training scripts run unchanged across platforms, with JAX automatically detecting and using the best available backend. The only configuration required is selecting the correct `pip` extra during installation.

## Platform-Specific Installation Requirements

Each hardware target requires a specific JAX build provided through Needle's optional extras. Choose the command that matches your machine.

### CPU (Default)

No additional dependencies are needed. JAX automatically uses the CPU backend when no GPU extras are installed.

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

```

This installs the core training stack including `jax`, `flax`, and `optax` from [[`requirements-train.txt`](https://github.com/cactus-compute/needle/blob/main/requirements-train.txt)](https://github.com/cactus-compute/needle/blob/main/requirements-train.txt). Training will function but run considerably slower than on GPU accelerators.

### NVIDIA GPU Support

For CUDA-capable hardware, install the **`gpu`** extra. This pulls the correct `jax[cuda]` wheels and enables JAX's GPU backend.

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

```

According to the [main README.md](https://github.com/cactus-compute/needle/blob/main/README.md), this is the official method for enabling NVIDIA acceleration: "runs on any accelerator jax supports" with the specific install command `pip install "cactus-needle[train,gpu]"`.

### Apple Silicon GPU Support

For M-series Macs, install the **`metal`** extra. This includes JAX's Metal plugin for Apple Silicon GPU access.

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

```

**Important limitation:** The Metal backend pins an older JAX version because Apple Metal support stopped functioning after JAX 0.4.38. This constraint is enforced by the extra dependency specification.

## Running Hardware-Agnostic Training

Once installed, the same commands work regardless of backend. JAX selects the appropriate hardware at import time based on available drivers and installed packages.

### Basic Finetuning Command

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

```

As documented in [[`doc/finetuning.md`](https://github.com/cactus-compute/needle/blob/main/doc/finetuning.md)](https://github.com/cactus-compute/needle/blob/main/doc/finetuning.md), this command runs unchanged on CPU, NVIDIA GPU, or Apple Silicon GPU. The underlying [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py) script relies on JAX's runtime backend selection, requiring no code modification for different hardware targets.

### Building a Tuned Model

```bash
needle build checkpoints/needle2.pkl --lora adapter.pkl --out tuned.cact

```

This post-training step remains hardware-agnostic and functions identically across all platforms.

## How JAX Backend Selection Works

Needle's architecture delegates all hardware abstraction to JAX. When you import `jax` in [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py), the library probes for available accelerators in priority order: GPU (CUDA or Metal) → CPU. The selected backend handles:

- **Device memory management** (GPU/Metal VRAM or system RAM)
- **Kernel compilation and execution** (cuDNN, Metal Performance Shaders, or XLA CPU)
- **Distributed and batched operations** across compute units

This design eliminates hardware-specific code paths in Needle itself.

## Version Constraints and Compatibility

| Platform | JAX Version | Notes |
|----------|-------------|-------|
| CPU | Latest stable | No version restrictions |
| NVIDIA GPU | Latest stable CUDA build | Requires compatible CUDA drivers |
| Apple Silicon | Capped at ≤0.4.38 | Metal plugin incompatibility above this version |

The `metal` extra enforces this cap automatically through its dependency specification.

## Summary

- **Needle training runs on any JAX-supported accelerator** without code changes
- **Install `cactus-needle[train,gpu]`** for NVIDIA GPU support
- **Install `cactus-needle[train,metal]`** for Apple Silicon GPU support (pins JAX ≤0.4.38)
- **Base install `cactus-needle[train]`** functions on CPU with no extra dependencies
- **Source references:** Commands are documented in [README.md](https://github.com/cactus-compute/needle/blob/main/README.md) and [doc/finetuning.md](https://github.com/cactus-compute/needle/blob/main/doc/finetuning.md)

## Frequently Asked Questions

### Can I train Needle without a GPU?

Yes. The base installation `pip install "cactus-needle[train]"` runs entirely on CPU. Performance will be significantly slower for large models, but all functionality remains available.

### Why does Apple Silicon require an older JAX version?

Apple's Metal backend for JAX stopped functioning after release 0.4.38 due to breaking changes in JAX's device handling. The `metal` extra pins to a compatible version automatically.

### Do I need to modify training scripts for different hardware?

No. Needle's training scripts in [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py) use standard JAX operations that dispatch to the best available backend at runtime. The same `needle finetune` command works on CPU, NVIDIA GPU, and Apple Silicon without modification.

### Can I switch between CPU and GPU training on the same machine?

Yes, though this requires reinstalling with different extras. JAX loads its backend at import time based on installed packages, so you cannot dynamically switch without changing your environment. For GPU development with CPU fallback testing, maintain separate virtual environments with `[train]` and `[train,gpu]` respectively.