# How Needle 2 Supports Fine-Tuning Using LoRA Adapters: A Complete Guide

> Discover how Needle 2 efficiently supports LoRA fine-tuning by optimizing adapter matrices with minimal overhead. Learn this complete guide to parameter-efficient fine-tuning.

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

---

**Needle 2 implements LoRA fine-tuning by freezing base model weights and optimizing low-rank adapter matrices that are merged on-the-fly during forward passes, enabling efficient parameter-efficient fine-tuning with minimal computational overhead.**

Needle 2 provides a complete LoRA (Low-Rank Adaptation) workflow that allows you to fine-tune large attention networks without modifying the original checkpoint. The implementation in [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py) orchestrates the entire pipeline—from identifying eligible weight matrices to serializing compact adapter files that can be distributed and merged later.

## What Are LoRA Adapters in Needle 2?

**LoRA adapters** are small, trainable matrices that approximate weight updates using a low-rank decomposition. Instead of modifying the base model parameters directly, Needle 2 learns two matrices—**A** and **B**—for each target layer. During inference, the product **A × B** (scaled by `lora_alpha / lora_rank`) is added to the original weights, producing the adapted output while keeping the base checkpoint immutable.

## The LoRA Fine-Tuning Pipeline in Needle 2

The fine-tuning process follows a strict six-step workflow defined in [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py), ensuring that only adapter parameters receive gradient updates.

### Step 1: Identifying LoRA-Eligible Weights

The helper function `lora_target_paths` scans the flattened parameter tree to locate attention projection kernels. It selects weights whose names match target strings including `"q_proj"`, `"k_proj"`, `"v_proj"`, `"gate_proj"`, and `"out_proj"` (defined in the `LORA_TARGETS` constant at line 23). Only kernels with non-trivial norms are retained for adaptation.

```python

# From needle/model/finetune.py lines 54-63

def lora_target_paths(params, targets=LORA_TARGETS):
    # Walks parameter tree and returns paths matching target projection layers

    return [path for path, value in tree_flatten_with_path(params) 
            if any(t in str(path) for t in targets) and jnp.linalg.norm(value) > 1e-8]

```

### Step 2: Initializing Low-Rank Adapter Matrices

For each selected weight, `init_lora` (lines 67-82) creates matrix **A** with shape *(in_dim, rank)* and matrix **B** with shape *(rank, out_dim)*. Matrix **A** initializes from a normal distribution scaled by the rank, while **B** starts as zeros. This initialization ensures stable training from the first forward pass.

```python

# From needle/model/finetune.py lines 67-82

def init_lora(weight, rank=8):
    in_dim, out_dim = weight.shape
    A = jax.random.normal(key, (in_dim, rank)) * 0.01
    B = jnp.zeros((rank, out_dim))
    return {"A": A, "B": B}

```

### Step 3: Merging Adapters On-the-Fly

During each forward pass, `merge_lora` (lines 85-92) computes the low-rank update and adds it to the base kernel. The merged parameters are supplied to the model without ever writing back to the original checkpoint, preserving the frozen base weights.

```python

# From needle/model/finetune.py lines 85-92

def merge_lora(params, lora_params, scale):
    # scale = lora_alpha / lora_rank

    merged = {}
    for path in lora_params.keys():
        merged[path] = params[path] + scale * (lora_params[path]["A"] @ lora_params[path]["B"])
    return merge_dicts(params, merged)

```

### Step 4: Training with Frozen Base Weights

The `finetune_local` function (starting at line 94) loads the base checkpoint, builds a `SimpleAttentionNetwork`, and runs a standard JAX/Optax training loop. Crucially, gradients are computed on the merged model but applied only to the LoRA matrices, leaving base weights untouched.

```python

# Training loop applies gradients only to LoRA parameters

merged_params = merge_lora(params, lora_params, scale)
loss, grads = jax.value_and_grad(loss_fn)(lora_params, merged_params, batch)
updates, opt_state = optimizer.update(grads, opt_state, lora_params)
lora_params = jax.tree_map(lambda p, u: p + u, lora_params, updates)

```

## Saving and Loading LoRA Adapters

### Persisting Adapter Checkpoints

After training, Needle 2 serializes the adapter dictionaries (containing **A**, **B**, the scaling factor, and a reference to the base checkpoint) using `pickle`. The resulting compact `.pkl` file contains only the low-rank matrices, typically reducing storage from gigabytes to megabytes.

```python

# From needle/model/finetune.py lines 91-99

import pickle
adapter_bundle = {
    "lora_params": lora_params,
    "scale": lora_alpha / lora_rank,
    "base_checkpoint": checkpoint_path
}
with open("adapters/needle2_lora.pkl", "wb") as f:
    pickle.dump(adapter_bundle, f)

```

### Loading for Inference

The `build_main` command accepts a `--lora` flag to merge adapters before export. Needle 2 loads the adapter, calls `merge_lora` to combine the matrices with base weights, and proceeds to quantize and export the fully merged model.

```bash
needle build \
  --checkpoint checkpoints/needle2.pkl \
  --lora adapters/needle2_lora.pkl \
  --bits 4 \
  --out needle2_quantized.cact

```

## Practical Examples: Fine-Tuning with LoRA in Needle 2

### Fine-Tuning via Command Line

Use the `needle finetune` command to train LoRA adapters on your dataset. The `--lora_rank` controls the rank of the decomposition (typical values: 4-32), while `--lora_alpha` adjusts the scaling factor.

```bash
needle finetune \
  --checkpoint checkpoints/needle2.pkl \
  --jsonl_path data/needle_data.jsonl \
  --lora_rank 8 \
  --lora_alpha 16 \
  --batch_size 32 \
  --epochs 4 \
  --lr 1e-4 \
  --out adapters/needle2_lora.pkl

```

**Key parameters:**
- `--lora_rank`: Sets the rank **R** of matrices **A** and **B** (parameter count scales as `2 × (in_dim × R + R × out_dim)`)
- `--lora_alpha`: Multiplies the adapter contribution, enabling interpolation between frozen and adapted behavior
- `--jsonl_path`: Points to training data generated by `generate_dataset`

### Merging Adapters for Deployment

To create a standalone model with merged weights for production inference:

```bash
needle build \
  --checkpoint checkpoints/needle2.pkl \
  --lora adapters/needle2_lora.pkl \
  --bits 4 \
  --out needle2_quantized.cact

```

This command loads both the base checkpoint and LoRA adapter, merges them using `merge_lora`, then quantizes to 4-bit precision for efficient deployment.

### Python API Inference

Load a LoRA-augmented model directly in Python for interactive inference:

```python
from needle import Needle

# Load base weights with LoRA adapter attached

model = Needle(
    weights="checkpoints/needle2.pkl", 
    lora="adapters/needle2_lora.pkl", 
    tools=[...]
)

prompt = "Summarize the latest research on LoRA adapters."
response = model(prompt)
print(response)

```

## Summary

- **Needle 2 implements LoRA in [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py)** through three core functions: `lora_target_paths` for weight discovery, `init_lora` for matrix initialization, and `merge_lora` for on-the-fly merging.
- **Only attention projection layers are adapted**, specifically `"q_proj"`, `"k_proj"`, `"v_proj"`, `"gate_proj"`, and `"out_proj"` kernels, minimizing parameter overhead while preserving model capacity.
- **Training freezes base weights** by computing gradients on merged parameters but applying updates solely to the low-rank **A** and **B** matrices.
- **Adapters are portable** as pickled `.pkl` files containing only the low-rank matrices and scaling metadata, enabling distribution without sharing full model weights.
- **JIT compatibility** ensures that all LoRA operations compile through XLA, adding negligible runtime overhead to the forward pass.

## Frequently Asked Questions

### What is the difference between lora_rank and lora_alpha in Needle 2?

**`lora_rank`** determines the dimensionality of the low-rank decomposition, controlling how many parameters are trained (lower rank = fewer parameters). **`lora_alpha`** scales the contribution of the adapter during merging (calculated as `scale = lora_alpha / lora_rank`). Higher alpha values increase the influence of the trained adapter relative to the base model, allowing you to interpolate between frozen and fine-tuned behavior without changing the rank.

### Which model layers does Needle 2 adapt with LoRA?

Needle 2 targets only **attention projection layers** as defined in the `LORA_TARGETS` constant at line 23 of [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py). Specifically, it adapts kernels named `"q_proj"`, `"k_proj"`, `"v_proj"`, `"gate_proj"`, and `"out_proj"` within the `SimpleAttentionNetwork` architecture defined in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py). All other layers—including feed-forward networks not matching these patterns—remain frozen during fine-tuning.

### Can I use multiple LoRA adapters with the same base model in Needle 2?

The current implementation supports loading a single LoRA adapter via the `--lora` flag in `build_main`. While you could theoretically merge multiple adapters by loading their **A** and **B** matrices and summing their contributions before calling `merge_lora`, the standard CLI workflow expects a single adapter file. For complex multi-adapter scenarios, you would need to manually compose the adapter parameters using the Python API before instantiation.

### How does Needle 2 keep the base model weights unchanged during LoRA training?

Needle 2 maintains immutability through a **functional merging strategy**. The `merge_lora` function computes `base_weight + scale * (A @ B)` during the forward pass but never writes this result back to the original parameter dictionary. The `finetune_local` training loop (starting at line 94) explicitly passes the merged parameters to `model.apply()` for loss computation, while the optimizer updates apply only to the separate `lora_params` dictionary containing matrices **A** and **B**. This ensures that the base checkpoint loaded from `checkpoints/needle2.pkl` remains identical before and after training.