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

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.

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

{"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【/needle/model/finetune.py#L80-L90】 to create diverse examples from your tool schemas.

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【/needle/cli.py#L13-L18】. This function:

• Loads the base checkpoint via load_checkpoint in 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

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【/needle/cli.py#L44-L48】) merges your LoRA adapter into the base checkpoint and optionally re-quantizes to 2-bit or 4-bit precision.

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【/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

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)

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

Fine-Tune

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

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

Run Inference with 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	LoRA adapter creation, data synthesis (generate_dataset), training loop (finetune_local), merging logic (merge_lora)【/needle/model/finetune.py】
needle/model/run.py	Checkpoint loading (load_checkpoint) and inference utilities【/needle/model/run.py#L61-L68】
needle/model/architecture.py	SimpleAttentionNetwork definition used during fine-tuning
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. 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 with underlying logic in 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.

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