LoRA Fine-Tuning Workflow in Needle: A Complete Guide to Tool-Calling Adaptation

Needle implements a complete LoRA (Low-Rank Adaptation) fine-tuning pipeline that adapts pre-trained language models to new tool-calling tasks through three phases: data generation, low-rank adapter training, and quantized model export.

The cactus-compute/needle repository provides a lightweight, end-to-end workflow for fine-tuning models on custom tool-use datasets. This guide walks through the complete LoRA fine-tuning workflow in Needle, from synthetic data generation to deploying a quantized .cact model.

Overview of the LoRA Fine-Tuning Pipeline

The Needle LoRA fine-tuning workflow consists of three distinct phases implemented across needle/model/finetune.py and needle/model/export.py:

1. Data Preparation: Generate or augment JSONL training examples using generate_main and augment_jsonl
2. LoRA Training: Train low-rank adapters on attention projections using finetune_local with cosine decay scheduling
3. Model Export: Merge adapters into the base model and quantize to .cact format via build_main and write_export

Phase 1: Data Preparation with generate-data

Defining Tool Schemas

Before generating training data, define your tool specifications in a JSON file. Needle expects an array of tool objects containing name, description, and parameters:

[
  { "name": "search", "description": "Run a web search", "parameters": {...} },
  { "name": "calc", "description": "Perform a calculation", "parameters": {...} }
]

Generating Synthetic Training Data

Needle can synthesize realistic training examples by prompting OpenRouter. The generate_main function orchestrates this process, utilizing _openrouter (lines 56‑66) in needle/model/finetune.py to send requests:

needle generate-data \
  --tools tools.json \
  --num-samples 200 \
  --output synthetic.jsonl

Internally, generate_examples constructs a system prompt using _GEN_SYSTEM and _GEN_TEMPLATE (lines 31‑53), then parses the LLM response via _parse_array (lines 69‑77). Each generated example is enriched with the original tool schemas (line 90) to ensure consistency.

Augmenting Existing Datasets

If you have existing hand-crafted examples, expand them with synthetic data using augment_jsonl:

needle generate-data \
  --augment existing.jsonl \
  --tools tools.json \
  --num-samples 100 \
  --output expanded.jsonl

The augment_jsonl function (lines 58‑71) reads the original file, extracts tool definitions via _collect_tools (lines 47‑55), generates additional samples, and writes a combined JSONL file suitable for training.

Phase 2: LoRA Training with finetune_local

The core training logic resides in finetune_local (lines 94‑100 and 292‑400) within needle/model/finetune.py. This function implements efficient low-rank adaptation targeting specific attention projection layers.

CLI Training Command

Execute the fine-tuning workflow via the Needle CLI:

needle finetune training.jsonl \
  --checkpoint checkpoints/needle2.pkl \
  --epochs 5 \
  --batch-size 32 \
  --lr 1e-4 \
  --lora-rank 16 \
  --lora-alpha 32 \
  --max-len 1024 \
  --out lora_adapter.pkl

Training Pipeline Implementation

fit_max_len (lines 20‑35) automatically determines the optimal sequence length by scanning the dataset, counting tokens for prompt-plus-target pairs, and rounding to the nearest power-of-two bucket for computational efficiency.

load_jsonl (lines 38‑51) encodes examples using _encode (lines 9‑17), converting text into token IDs and attention masks. The function handles the tool-calling format specific to Needle's architecture.

lora_target_paths identifies trainable parameters by walking the parameter tree and selecting kernels matching LORA_TARGETS (line 23), which includes attention projections: q_proj, k_proj, v_proj, gate_proj, and out_proj.

init_lora (lines 67‑82) initializes low-rank matrices A and B for each target weight. The rank is controlled by --lora-rank, while --lora-alpha determines the scaling factor applied during forward passes.

Optimization Configuration

The training loop utilizes Optax for optimization:

• Schedule: warmup_cosine_decay_schedule (lines 47‑49) provides learning rate warmup followed by cosine annealing
• Optimizer: adamw (line 50) with configurable learning rate and weight decay
• Loss: Cross-entropy computed on shifted targets (lines 54‑58), standard for autoregressive language modeling

After training, the adapter is serialized to a pickle file containing:

{
  "lora": { "layer/…/kernel": {"A": np.ndarray, "B": np.ndarray}, … },
  "scale": <float>,
  "base": "<path-to-base-checkpoint>",
  "rank": <int>
}

Phase 3: Model Export and Merging

Merging LoRA Weights with build_main

To create a deployment-ready model, merge the trained LoRA adapter into the base checkpoint using build_main (lines 4‑34) in needle/model/finetune.py:

needle build checkpoints/needle2.pkl \
  --lora lora_adapter.pkl \
  --bits 4 \
  --out needle_finetuned.cact

The build_main function loads the base checkpoint via load_checkpoint (imported from needle/model/run.py), applies the low-rank updates through merge_lora (lines 85‑90), and prepares the model for serialization.

Quantized Export to .cact Format

The write_export function in needle/model/export.py handles the final quantization and binary serialization. This produces a compact .cact file optimized for Needle's inference engine. The --bits parameter controls quantization precision (e.g., 4-bit for maximum compression).

The final output includes instructions for instantiating the model via the Needle class with the newly exported weights.

Programmatic API Usage

For integration into Python workflows, import functions directly from needle/model/finetune.py:

from needle.model.finetune import (
    generate_dataset, augment_jsonl, finetune_local, build_main
)

# 1. Data generation

tools = [{"name": "search", "description": "Web search", "parameters": {}}]
examples = generate_dataset(tools, 200)

# 2. Fine-tuning

class Args:
    jsonl_path = "training.jsonl"
    checkpoint = "checkpoints/needle2.pkl"
    epochs = 4
    batch_size = 32
    lr = 1e-4
    lora_rank = 16
    lora_alpha = 32
    max_len = 1024
    out = "lora.pkl"

finetune_local(Args())

# 3. Export

class BuildArgs:
    checkpoint = "checkpoints/needle2.pkl"
    lora = "lora.pkl"
    bits = "4"
    out = "final.cact"

build_main(BuildArgs())

Summary

Needle's LoRA fine-tuning workflow provides a complete solution for adapting language models to tool-calling tasks:

• Synthetic data generation via generate_main and augment_jsonl creates training corpora from tool schemas using OpenRouter
• Efficient fine-tuning through finetune_local targets attention projections (q_proj, k_proj, v_proj, gate_proj, out_proj) with configurable rank and alpha parameters
• Production deployment via build_main and write_export merges adapters and quantizes to the .cact format optimized for Needle's inference engine
• Full CLI and Python API support enables both interactive experimentation and automated training pipelines

Frequently Asked Questions

What LoRA rank and alpha values should I use for Needle fine-tuning?

Start with rank 16 and alpha 32 (alpha = 2×rank) for most tool-calling tasks. According to the implementation in needle/model/finetune.py, the lora_rank parameter controls the inner dimension of low-rank matrices A and B, while lora_alpha scales the adapter output. Higher ranks (32-64) capture more complex tool interactions but increase memory usage and checkpoint size.

Can I use the LoRA adapter without merging it into the base model?

The current Needle workflow requires merging via build_main before deployment. The finetune_local function saves only the low-rank matrices and scaling factors to a pickle file, but the inference engine expects consolidated weights. The merge_lora function (lines 85‑90) applies the low-rank update $W = W_0 + \alpha/r \cdot BA$ before quantization and export.

What data format does Needle expect for fine-tuning?

Needle requires a JSONL file where each line contains tool-call examples with specific fields. The load_jsonl function (lines 38‑51) expects entries containing tool definitions, user queries, and assistant responses. Use generate-data to create properly formatted synthetic examples, or ensure your manual data follows the schema extracted by _encode (lines 9‑17) in needle/model/finetune.py.

How does Needle handle sequence length during training?

fit_max_len automatically calculates the optimal context window by analyzing token counts in your dataset. Located at lines 20‑35 in needle/model/finetune.py, this function bins sequences to the nearest power-of-two (e.g., 512, 1024, 2048) to maximize training throughput while accommodating your longest examples. Set --max-len to cap this at a specific value if memory constraints arise.

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