# How to Fine-Tune Needle 2 for Custom Tasks: A Complete LoRA Training Guide

> Learn to fine-tune Needle 2 for custom tasks with our comprehensive LoRA training guide. Train lightweight adapters, merge, and quantize for efficient deployment.

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

---

**Fine-tune Needle 2 using lightweight LoRA adapters on its frozen 45M-parameter base model, then merge and quantize to a single `.cact` file for deployment.**

Needle 2 ships as a compact, frozen base model optimized for tool-calling. You can adapt it to custom domains without retraining the full network by attaching small **LoRA (Low-Rank Adaptation)** adapters. This guide walks through the complete workflow—from data preparation to building a quantized, self-contained engine.

## Overview of the Fine-Tuning Workflow

The fine-tuning pipeline in `cactus-compute/needle` follows three stages:

1. **Prepare training data** in JSON-L format with queries, available tools, and correct tool call answers
2. **Train LoRA adapters** on the frozen base weights using the `needle finetune` CLI command
3. **Merge and quantize** into a `.cact` file via `needle build` for standalone inference

All stages are implemented in Python modules under `needle/model/` and exposed through the CLI in [`needle/cli.py`](https://github.com/cactus-compute/needle/blob/main/needle/cli.py).

## Step 1: Prepare Your Training Data

Each line in your JSON-L file must contain a complete training example with three fields: **query**, **tools**, and **answers**. The exact format is documented in the README at `data format`【/README.md#L81-L86】.

```json
{"query": "What's the weather in Tokyo?", "tools": [...], "answers": [...]}

```

You have two options for populating this dataset:

- **Manual authoring**: Write examples by hand following the schema
- **Synthetic generation**: Use the built-in generator that calls OpenRouter

### Generating Synthetic Data

The `needle generate-data` command invokes `generate_dataset` in [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py)【/needle/model/finetune.py#L80-L90】 to create diverse examples from your tool schemas.

```bash
export OPENROUTER_API_KEY=sk-xxxx
needle generate-data \
    --tools my_tools.json \
    --num-samples 500 \
    --output data.jsonl

```

This requires an `OPENROUTER_API_KEY` environment variable. The generator contacts OpenRouter and returns JSON examples conforming to your tool schemas.

## Step 2: Run LoRA Fine-Tuning

The `needle finetune` command executes `finetune_local` in [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py)【/needle/cli.py#L13-L18】. This function:

- Loads the base checkpoint via `load_checkpoint` in [`needle/model/run.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/run.py)【/needle/model/run.py#L61-L68】
- Computes maximum sequence length with `fit_max_len`【/needle/model/finetune.py#L20-L35】
- Instantiates LoRA adapters through `init_lora`【/needle/model/finetune.py#L67-L83】
- Trains with **AdamW** and **cosine decay schedule** from `optax.warmup_cosine_decay_schedule`【/needle/model/finetune.py#L46-L50】

### Fine-Tuning Command

```bash
needle finetune data.jsonl \
    --epochs 10 \
    --lora-rank 16 \
    --lora-alpha 32 \
    --batch-size 16 \
    --max-len 1024 \
    --val-split 0.1

```

Key parameters:

| Parameter | Purpose |
|-----------|---------|
| `--lora-rank` | Rank of the low-rank matrices (smaller = fewer parameters) |
| `--lora-alpha` | Scaling factor applied to LoRA weights |
| `--max-len` | Maximum sequence length; computed automatically if omitted |
| `--val-split` | Fraction of data held out for validation metrics |

During training, loss and validation metrics print periodically【/needle/model/finetune.py#L70-L86】. Upon completion, the adapter saves to `checkpoints/needle_lora.pkl`【/needle/model/finetune.py#L89-L99】.

## Step 3: Build the Tuned `.cact` File

The `needle build` command (registered in [`needle/cli.py`](https://github.com/cactus-compute/needle/blob/main/needle/cli.py)【/needle/cli.py#L44-L48】) merges your LoRA adapter into the base checkpoint and optionally re-quantizes to 2-bit or 4-bit precision.

```bash
needle build checkpoints/needle2.pkl \
    --lora checkpoints/needle_lora.pkl \
    --out my_needle.cact \
    --bits 2

```

The build pipeline executes `build_main` → `merge_lora` → `quantize`, with merging logic implemented in [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py)【/needle/model/finetune.py#L85-L91】. The resulting `.cact` contains the complete engine and loads directly by the runtime with no further compilation【/README.md#L9-L12】.

## Complete Working Example

### Define Your Tools

```bash
cat > my_tools.json <<'EOF'
[
  {
    "name": "set_lights",
    "parameters": {
      "type": "object",
      "properties": {
        "room": {"type": "string"},
        "brightness": {"type": "integer"}
      },
      "required": ["room"]
    }
  }
]
EOF

```

### Generate Synthetic Data (Optional)

```bash
export OPENROUTER_API_KEY=sk-xxxx
needle generate-data --tools my_tools.json --num-samples 500 --output data.jsonl

```

### Fine-Tune

```bash
needle finetune data.jsonl \
    --epochs 10 \
    --lora-rank 16 \
    --lora-alpha 32 \
    --batch-size 16 \
    --max-len 1024 \
    --val-split 0.1

```

### Build and Deploy

```bash
needle build checkpoints/needle2.pkl \
    --lora checkpoints/needle_lora.pkl \
    --out my_needle.cact \
    --bits 2

```

### Run Inference with Python

```python
import needle

@needle.tool
def set_lights(room: str, brightness: int):
    """Set the brightness of a given room."""
    return {"room": room, "brightness": brightness}

agent = needle.Needle(weights="my_needle.cact", tools=[set_lights])
resp = agent.run("Dim the kitchen lights to 10%")
print(resp["results"])

# => [{'room': 'kitchen', 'brightness': 10}]

```

## Key Source Files and Architecture

| File | Role |
|------|------|
| [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py) | LoRA adapter creation, data synthesis (`generate_dataset`), training loop (`finetune_local`), merging logic (`merge_lora`)【/needle/model/finetune.py】 |
| [`needle/model/run.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/run.py) | Checkpoint loading (`load_checkpoint`) and inference utilities【/needle/model/run.py#L61-L68】 |
| [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) | `SimpleAttentionNetwork` definition used during fine-tuning |
| [`needle/cli.py`](https://github.com/cactus-compute/needle/blob/main/needle/cli.py) | CLI entry points for `finetune`, `generate-data`, `build` sub-commands【/needle/cli.py#L13-L18】【/needle/cli.py#L44-L48】 |

The frozen base model uses a **SimpleAttentionNetwork** from [`architecture.py`](https://github.com/cactus-compute/needle/blob/main/architecture.py). Only the LoRA adapters—small low-rank matrices injected into attention layers—are trained, keeping memory requirements minimal.

## Summary

- **Needle 2 fine-tuning** uses LoRA adapters on a frozen 45M-parameter base, avoiding full model retraining
- **Data preparation** requires JSON-L format with query/tools/answers; synthetic generation available via OpenRouter
- **Training** runs through `needle finetune`, implemented in `finetune_local` with AdamW and cosine scheduling
- **Deployment** produces a quantized `.cact` file via `needle build`, containing the merged engine ready for inference
- **All workflows** are accessible through the CLI in [`needle/cli.py`](https://github.com/cactus-compute/needle/blob/main/needle/cli.py) with underlying logic in [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py)

## Frequently Asked Questions

### What hardware is required for fine-tuning Needle 2?

LoRA fine-tuning runs efficiently on consumer GPUs and Apple Silicon due to the small adapter size. The frozen 45M-parameter base model stays in memory while only the low-rank adapters (typically rank 8-32) receive gradients, keeping VRAM requirements modest compared to full fine-tuning.

### Can I fine-tune without using OpenRouter for data generation?

Yes. The `needle generate-data` command is optional. You can manually author all training examples in the JSON-L format specified in the README【/README.md#L81-L86】. OpenRouter synthesis simply accelerates dataset creation for complex tool schemas.

### How does the LoRA rank affect model quality and size?

Higher `--lora-rank` values increase adapter capacity but also parameter count and training time. The default rank of 16 with alpha 32 offers a practical balance. Lower ranks (8) work for simple tasks; higher ranks (32-64) may help for complex multi-tool reasoning. The adapter file remains small regardless—typically under 10MB even at rank 64.

### What quantization options are available when building the final model?

The `--bits` parameter accepts `2` or `4` for aggressive or moderate compression. Two-bit quantization minimizes file size for edge deployment; four-bit preserves slightly more precision. The base checkpoint is re-quantized after LoRA merging, with logic handled in `quantize` within [`finetune.py`](https://github.com/cactus-compute/needle/blob/main/finetune.py).