# JAX Backend Options for Training Needle on CPU, GPU, TPU, and Apple Silicon

> Discover JAX backend options for Needle training across CPU, GPU, TPU, and Apple Silicon. Needle auto-detects and configures your hardware for optimal performance.

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

---

**Needle automatically detects and runs on five JAX backends—CPU, NVIDIA GPU (CUDA), TPU, Metal (Apple Silicon), or a manually forced platform—by checking `jax.default_backend()` at runtime and configuring environment variables before imports.**

The **Needle** training framework, developed by **cactus-compute/needle**, delegates all numerical computation to **JAX**. This design choice means Needle inherits JAX's hardware portability: the same Python code trains on laptops, servers, or cloud pods without modification. Understanding how Needle selects and configures JAX backends helps you optimize performance and troubleshoot deployment issues across different hardware.

## Supported JAX Backends in Needle

Needle works with any backend JAX supports. The framework detects hardware automatically, with special handling for Apple Silicon Macs.

### CPU Backend

The **CPU backend** runs on any machine without specialized accelerators. Needle falls back to this when `jax.default_backend()` returns `"cpu"`.

This backend requires no configuration. Performance is sufficient for small models or debugging, though large transformer training becomes prohibitively slow.

### NVIDIA GPU (CUDA) Backend

The **GPU backend** activates on machines with CUDA-capable NVIDIA GPUs and properly installed CUDA/cuDNN libraries.

In [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py) at lines 315-321, Needle checks the backend:

```python
backend = jax.default_backend().lower()
print(f"Using JAX backend: {backend}")

# Prints "gpu" when CUDA is available

```

The framework then uses `jax.device_put()` to move parameters to the GPU and `@jax.jit` to compile CUDA kernels. No manual intervention is required if your environment has working JAX GPU support.

### TPU Backend

The **TPU backend** runs on Google Cloud TPU pods or TPU-enabled environments. JAX reports `"tpu"` from `jax.default_backend()`, and Needle uses it automatically.

TPU training requires `jax[tpu]` installation and appropriate Cloud TPU access configuration. Needle's backend-agnostic code executes identically—only the underlying XLA compilation target changes to generate TPU-specific instructions.

### Metal Backend for Apple Silicon

The **Metal backend** enables GPU acceleration on Apple Silicon Macs (M1/M2/M3). This requires special handling in Needle because of JAX's PJRT plugin architecture.

In [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py) at lines 9-14, Needle forces Metal compatibility before any JAX import:

```python
import sys, os
if sys.platform == "darwin":
    os.environ["ENABLE_PJRT_COMPATIBILITY"] = "1"
    # Must be set BEFORE importing jax

```

After this environment setup, `jax.default_backend()` returns `"gpu"` (Metal is exposed as a GPU backend to JAX). This early configuration is mandatory—the newer PJRT API rejects the Metal plugin if JAX is imported first.

## How Needle Configures the JAX Backend

Needle's backend handling follows a three-stage pattern across its training pipeline.

### Stage 1: Environment Preparation

For macOS specifically, Needle sets `ENABLE_PJRT_COMPATIBILITY=1` in [`finetune.py`](https://github.com/cactus-compute/needle/blob/main/finetune.py) lines 9-14. This stage executes before any `import jax` statement.

### Stage 2: Backend Detection

After imports complete, Needle calls `jax.default_backend().lower()` at line 315 of [`finetune.py`](https://github.com/cactus-compute/needle/blob/main/finetune.py). This value is stored and logged:

```python
backend = jax.default_backend().lower()
print(f"Using JAX backend: {backend}")

```

The detected backend drives subsequent decisions about device placement and kernel selection.

### Stage 3: Device Placement and JIT Compilation

Throughout the codebase, Needle uses JAX primitives that respect the selected backend:

- **`jax.device_put()`**: In [`needle/model/run.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/run.py) at line 129, parameters are moved to the default device automatically
- **`@jax.jit`**: At line 103 of [`run.py`](https://github.com/cactus-compute/needle/blob/main/run.py), functions are JIT-compiled for the active backend
- **Backend-specific optimizations**: In [`needle/model/decode.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/decode.py) lines 93-95 and [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) lines 246-248, Needle selects optimized attention implementations when `jax.default_backend() == "gpu"`

```python

# From run.py - automatic device placement

params = jax.device_put(params)  # Routes to GPU/TPU/CPU as appropriate

# JIT compilation targets the selected backend

@jax.jit
def train_step(params, batch):
    # Compiles to CUDA kernels, Metal shaders, or CPU code

    ...

```

## Manually Forcing a Specific Backend

You can override Needle's automatic detection by setting `JAX_PLATFORM_NAME` **before** importing JAX. This is useful for reproducibility testing or avoiding GPU memory contention.

```python
import os
os.environ["JAX_PLATFORM_NAME"] = "cpu"  # Force CPU backend

# os.environ["JAX_PLATFORM_NAME"] = "tpu"  # Force TPU backend

import jax, jax.numpy as jnp
print("Backend:", jax.default_backend())  # → cpu

```

This environment variable takes precedence over hardware detection. Needle will use whatever backend JAX initializes, making it compatible with custom XLA backends or containerized environments.

## Key Files Controlling Backend Behavior

| File | Lines | Purpose |
|------|-------|---------|
| [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py) | 9-14 | Sets `ENABLE_PJRT_COMPATIBILITY` for macOS Metal |
| [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py) | 315-321 | Detects and logs `jax.default_backend()` |
| [`needle/model/run.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/run.py) | 103-129 | Uses `jax.jit` and `jax.device_put` for device placement |
| [`needle/model/decode.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/decode.py) | 93-95 | Selects cuDNN attention for GPU backend |
| [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) | 246-248 | Adapts attention implementation per backend |

These files demonstrate how Needle maintains hardware portability: environment configuration happens early, backend detection occurs once, and all computation uses JAX's backend-agnostic primitives.

## Summary

- **Five backends supported**: CPU, NVIDIA GPU (CUDA), TPU, Metal (Apple Silicon), and manually forced platforms
- **Automatic detection**: Needle uses `jax.default_backend()` checked at `finetune.py:315`
- **macOS special case**: `ENABLE_PJRT_COMPATIBILITY=1` must be set before JAX import at `finetune.py:9-14`
- **Portable primitives**: `jax.device_put()` and `@jax.jit` route computation to the detected hardware
- **Manual override**: Set `JAX_PLATFORM_NAME` before importing JAX to force any supported backend

## Frequently Asked Questions

### How do I verify which JAX backend Needle is using?

Check the training log for the line printed from [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py) at line 315-321, or run this code before training:

```python
import jax
print(jax.default_backend())  # cpu, gpu, or tpu

```

### Why does Needle fail with PJRT errors on my Mac?

The Metal plugin requires `ENABLE_PJRT_COMPATIBILITY=1` set **before** JAX imports. Needle handles this automatically in [`finetune.py`](https://github.com/cactus-compute/needle/blob/main/finetune.py) lines 9-14, but if you're running Needle components in a custom script, ensure this environment variable is set first.

### Can I train Needle on multiple GPUs?

JAX supports multi-GPU training through `pmap` or `shard_map`. Needle's current codebase uses `jax.device_put()` with the default backend, which targets a single device. For multi-GPU training, you would need to modify the device placement logic in [`needle/model/run.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/run.py) to distribute parameters across devices.

### Does Needle work with AMD GPUs?

JAX's ROCm support for AMD GPUs is experimental. If `jax.default_backend()` detects and reports an AMD GPU, Needle should work, but this depends on your JAX installation having working ROCm support—not all features are guaranteed.