# JAX Training Infrastructure for Needle 2: GPU and Apple Silicon Support Explained

> Explore JAX training infrastructure for Needle 2. Discover seamless GPU and Apple Silicon support with automatic CUDA and Metal detection for optimized performance.

- Repository: [Cactus Compute, Inc./needle](https://github.com/cactus-compute/needle)
- Tags: deep-dive
- Published: 2026-08-25

---

**Needle 2 uses JAX as its unified training backend, automatically detecting NVIDIA GPUs via CUDA or Apple Silicon GPUs via Metal and applying platform-specific optimizations through conditional configuration in [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py).**

The Needle 2 project from cactus-compute/needle provides a device-agnostic fine-tuning pipeline built entirely on JAX. Whether you're training on an NVIDIA data center GPU or a MacBook with Apple Silicon, the same Python codebase adapts transparently—only environment variables and installation extras change. This article breaks down how the JAX training infrastructure for Needle 2 handles backend detection, Metal-specific workarounds, and the GPU path.

## Backend Detection and Configuration

When training starts, the `finetune_local` function (line 315) calls `jax.default_backend()` to determine which accelerator is available and applies the appropriate configuration.

```python

# From needle/model/finetune.py, lines 315-322

backend = jax.default_backend().lower()
if backend == "metal":
    config.flash = False
    config.remat = False
    config.scan_unroll = config.num_layers

```

This single conditional branch is the only platform-specific logic in the training loop. The rest of the pipeline—checkpoint loading, LoRA initialization, and the JIT-compiled training step—remains identical across hardware.

## Metal-Specific Optimizations for Apple Silicon

Apple Silicon support requires pre-flight environment configuration before JAX initializes. Lines 9-14 of [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py) handle this automatically:

```python
import sys, os
if sys.platform == "darwin":
    os.environ.setdefault("ENABLE_PJRT_COMPATIBILITY", "1")

```

This sets `ENABLE_PJRT_COMPATIBILITY=1` on macOS to ensure JAX's PJRT runtime works correctly with Metal. Once the backend resolves to `"metal"`, three training optimizations are disabled:

- **Flash attention** (`config.flash = False`) — Metal does not yet support the fused attention kernels
- **Gradient rematerialization** (`config.remat = False`) — Memory-saving recomputation is skipped
- **Layer stack unrolling** (`config.scan_unroll = config.num_layers`) — The full layer stack is unrolled rather than scanned

These trade-offs prioritize stability and compatibility over peak performance on Apple hardware.

## GPU (CUDA) Path for NVIDIA Accelerators

The NVIDIA GPU path requires no code changes. Users install the GPU extra and JAX automatically selects the CUDA backend:

```bash
pip install "cactus-needle[gpu]"
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) (lines 46-51), the `[gpu]` extra pulls in `jax[cuda12]` or equivalent platform-specific CUDA wheels. The training script proceeds without the Metal compatibility flags, keeping flash attention and rematerialization enabled for maximum throughput.

## Core Training Pipeline Components

The unified training flow in [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py) consists of four stages:

1. **Checkpoint loading** — `load_checkpoint` restores base model weights from a `.pkl` file
2. **Device placement** — `jax.device_put` moves parameters to the detected device in `float32` precision
3. **LoRA adapter initialization** — `init_lora` creates low-rank adapters targeting projection matrices defined in `LORA_TARGETS`
4. **Compiled training loop** — `train_step` is JIT-compiled with `jax.jit` and uses Optax's warm-up-cosine learning rate schedule

Because every operation uses JAX primitives—`jax.jit`, `jax.random`, `jax.tree.map`—the identical Python script executes correctly on both backends.

## Installation and Usage Examples

### Training on NVIDIA GPU (Linux/Windows)

```bash

# Install GPU variant

pip install "cactus-needle[gpu]"

# Fine-tune and merge adapter

needle finetune data.jsonl --epochs 10 --out adapter.pkl
needle build checkpoints/needle2.pkl --lora adapter.pkl --out tuned.cact

```

### Training on Apple Silicon (macOS)

```bash

# Install Metal variant

pip install "cactus-needle[metal]"

# Same commands—backend detection is automatic

needle finetune data.jsonl --epochs 10 --out adapter.pkl
needle build checkpoints/needle2.pkl --lora adapter.pkl --out tuned.cact

```

### Minimal Backend-Aware Configuration

This snippet reproduces the core detection logic from `finetune_local`:

```python
import sys, os

# Metal compatibility must be set before importing JAX

if sys.platform == "darwin":
    os.environ.setdefault("ENABLE_PJRT_COMPATIBILITY", "1")

import jax

backend = jax.default_backend().lower()
if backend == "metal":
    config.flash = False
    config.remat = False
    config.scan_unroll = config.num_layers

# Parameters automatically move to the selected accelerator

params = jax.device_put(params)

```

## Key Source Files

| File | Purpose |
|------|---------|
| [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py) | Core training implementation with backend detection and LoRA logic |
| [`doc/finetuning.md`](https://github.com/cactus-compute/needle/blob/main/doc/finetuning.md) | User-facing documentation for installation and training workflows |
| [`needle/model/run.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/run.py) | Checkpoint utilities and inference scaffolding |
| [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) | `SimpleAttentionNetwork` definition that the trainer instantiates |

## Summary

- **JAX provides the abstraction layer** — The same training script runs on CUDA and Metal without modification
- **Metal requires pre-configuration** — `ENABLE_PJRT_COMPATIBILITY=1` and three disabled optimizations
- **GPU path is zero-code-change** — Install `[gpu]` extra and JAX handles the rest
- **Device placement is explicit** — `jax.device_put` moves parameters after backend detection
- **LoRA and Optax schedules are backend-agnostic** — JIT compilation ensures performance portability

## Frequently Asked Questions

### What JAX version does Needle 2 require?

The repository pins compatible JAX versions through its extras. Install `cactus-needle[gpu]` for CUDA 12 support or `cactus-needle[metal]` for Apple Silicon, and pip resolves the correct JAX wheel with matching XLA backend.

### Why does Metal disable flash attention?

Apple's Metal Performance Shaders do not yet expose the fused attention operations that JAX's flash attention implementation requires. The fallback standard attention is functionally equivalent but uses more memory and compute.

### Can I force a specific backend instead of auto-detection?

JAX respects `JAX_PLATFORMS` and `JAX_BACKEND_TARGET` environment variables. Set `JAX_PLATFORMS=cpu` to disable accelerators entirely, though this is not recommended for training due to speed.

### How does checkpoint portability work between GPU and Metal?

Checkpoints store NumPy arrays on disk. When `load_checkpoint` runs, `jax.device_put` transfers parameters to whichever backend is active—training on Metal, then moving the `.pkl` to a GPU server works seamlessly.