Needle Training Hardware Requirements: GPU and Apple Silicon Support Explained

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.

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). 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.

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

According to the 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.

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

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), this command runs unchanged on CPU, NVIDIA GPU, or Apple Silicon GPU. The underlying needle/model/finetune.py script relies on JAX's runtime backend selection, requiring no code modification for different hardware targets.

Building a Tuned Model

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, 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 and 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 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.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}