How to Fine-Tune a Needle Model Using the `needle finetune` CLI: Complete Guide

Fine-tune a Needle model by training a lightweight LoRA adapter on JSON-L data using needle finetune, then merge it into a production-ready checkpoint with needle build.

The needle finetune command provides a streamlined, CPU-friendly way to adapt Needle's base language model to domain-specific tool-calling tasks. This guide walks through the complete workflow—from data preparation to deployment—based on the cactus-compute/needle source code.

---

What the needle finetune CLI Does

Under the hood, needle finetune invokes finetune_local() in needle/model/finetune.py (lines 94-100). The process:

1. Loads a frozen base checkpoint (default: checkpoints/needle2.pkl)—see lines 24-27
2. Tokenizes training examples and computes optimal sequence length (fit_max_len)—lines 20-35
3. Identifies target weight groups for LoRA adaptation (lora_target_paths)—lines 54-62
4. Creates a rank-R adapter (init_lora)—lines 67-82
5. Trains with cosine decay using Optax (schedule)—lines 46-51
6. Serializes the adapter to a .pkl file at your specified --out path

The training loop in train_step (lines 54-65) computes cross-entropy loss on concatenated prompt+target sequences, masks the prompt portion from loss calculation, and updates only LoRA parameters while keeping the base model frozen.

---

Step 1: Prepare Your Training Data

Needle requires JSON-L format with a specific schema. Each line contains:

• query: The user instruction
• tools (optional): Tool definitions in JSON Schema format
• answers: List of exact tool calls the model should emit
• reasoning (optional): Explanation mapping arguments to query spans

cat > data.jsonl <<'EOF'
{"query":"Extract title and author from the citation","tools":[{"name":"extract_citation","parameters":{"type":"object","properties":{"title":{"type":"string"},"author":{"type":"string"}},"required":["title","author"]}}],"answers":[{"name":"extract_citation","arguments":{"title":"My Paper","author":"Doe, J."}}],"reasoning":"title follows the word 'title'; author follows the word 'by'"}
EOF

For full schema documentation, see doc/finetuning.md in the repository.

---

Step 2: Run needle finetune

The CLI is defined in needle/cli.py (lines 31-41). Key parameters:

Parameter	Default	Description
--epochs	3	Training epochs
--batch-size	32	Batch size
--lr	2e-4	Peak learning rate
--lora-rank	16	LoRA rank (R)
--lora-alpha	32	LoRA scaling factor
--max-len	Auto-detected	Maximum sequence length
--out	Required	Output adapter path

needle finetune data.jsonl \
      --epochs 5 \
      --batch-size 32 \
      --lr 2e-4 \
      --lora-rank 32 \
      --lora-alpha 64 \
      --max-len 1024 \
      --out my_adapter.pkl

Data Augmentation with --generate

To synthesize additional training examples via OpenRouter, add the --generate N flag. Requires OPENROUTER_API_KEY in your environment.

needle finetune data.jsonl --generate 200 --out my_adapter.pkl

---

Step 3: Export a Production Checkpoint

After training, use needle build to merge the LoRA adapter into the base weights (merge_lora functionality in needle/model/finetune.py, lines 12-18). This produces a self-contained .cact archive.

needle build checkpoints/needle2.pkl \
      --lora my_adapter.pkl \
      --out tuned.cact

---

Loading Your Fine-Tuned Model

Use the merged checkpoint directly in Python:

import needle

agent = needle.Needle(
    tools=[...],
    weights="tuned.cact"
)

---

Key Implementation Files

Understanding these files helps with debugging and customization:

• needle/cli.py — CLI argument parsing and command dispatch
• needle/model/finetune.py — Dataset loading, LoRA initialization, training loop, and adapter serialization
• needle/model/architecture.py — SimpleAttentionNetwork architecture used during training
• needle/model/run.py — Checkpoint loading for base model and LoRA merging
• doc/finetuning.md — Official documentation for data format and best practices

---

Summary

• needle finetune trains a LoRA adapter on frozen base weights using Optax cosine-decay scheduling
• Input data must be JSON-L with query, answers, and optional tools/reasoning fields
• Key hyperparameters: --lora-rank, --lora-alpha, --epochs, --lr control adaptation quality and compute cost
• needle build merges adapters into self-contained .cact checkpoints for production
• Only LoRA parameters are trained—base weights remain frozen, minimizing memory and compute requirements

---

Frequently Asked Questions

What hardware do I need to run needle finetune?

The LoRA-based approach is designed for CPU-only training. Because only low-rank adapter matrices are updated while the base model stays frozen, memory requirements stay modest—typically under 8GB RAM for rank-32 adapters.

How do I choose LoRA rank and alpha values?

Start with --lora-rank 16 --lora-alpha 32 for baseline adaptation. Increase to rank 32-64 with alpha 2x the rank for complex domains with many tool variants. Higher ranks capture more patterns but increase training time and adapter size. According to needle/model/finetune.py, these values directly parameterize the init_lora function (lines 67-82).

Can I resume interrupted training or train multiple epochs?

The current needle finetune implementation trains from scratch for the specified --epochs. There is no built-in checkpoint resumption—ensure your environment has stable resources for the full run, or reduce epochs and iterate on data quality instead.

Why does my model output format look wrong after fine-tuning?

Verify your answers field contains exact, valid tool calls with correct argument types. The training loss masks the prompt portion—if answers are malformed, the model learns incorrect patterns. Use the reasoning field to help the model learn argument grounding, especially for extractive tasks.

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