# Needle 2 Fine-Tuning Workflow: A Complete Guide to LoRA Adapter Training

> Master the Needle 2 fine-tuning workflow with this guide on LoRA adapter training. Efficiently customize your model using lightweight adapter weights and tool-specific data without full retraining.

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

---

**Needle 2 uses a lightweight LoRA-based fine-tuning strategy that freezes the base model and trains small adapter weights on your tool-specific data, enabling efficient customization without full model retraining.**

The `cactus-compute/needle` repository implements a JAX-based parameter-efficient fine-tuning pipeline designed specifically for tool-calling agents. This workflow keeps the base checkpoint (`checkpoints/needle2.pkl`) frozen while optimizing low-rank adapters that modify attention projections, allowing you to specialize the model for your specific API schemas with minimal compute resources.

## Preparing Your Training Data

The first step in the Needle 2 fine-tuning workflow involves creating a line-delimited JSONL file where each record contains a `query`, optional `tools` array, a `reasoning` line, and the exact `answers` representing tool calls.

According to the documentation in [`doc/finetuning.md`](https://github.com/cactus-compute/needle/blob/main/doc/finetuning.md), the required format includes the natural language query, available tool schemas, the model’s reasoning trace, and a JSON array of tool calls with precise arguments. The following Python example creates a minimal training example:

```python

# Write a minimal JSONL training file (data.jsonl)

# Each line is a JSON object; see the doc for details.

with open("data.jsonl", "w") as f:
    f.write(
        '{"query":"Summarize the weather forecast for tomorrow.",'
        '"tools":[{"name":"weather_summary","parameters":{"type":"object","properties":{"location":{"type":"string"}}}}],'
        '"answers":[{"name":"weather_summary","arguments":{"location":"San Francisco"}}]}\n'
    )

```

## Augmenting the Dataset (Optional)

Before training, you can expand your dataset using the built-in generator that synthesizes additional examples from your tool schemas via the OpenRouter API.

The data augmentation logic resides in [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py) at lines 58-66, where the generator creates diverse variations of tool calls based on your schema definitions. Run the following command to produce 1,000 augmented samples:

```bash

# Generate more examples

needle generate-data --augment data.jsonl --num-samples 1000

# This creates data.augmented.jsonl

```

## Training the LoRA Adapter

The core training stage invokes `needle finetune`, which initializes a LoRA adapter with **rank 16 and alpha 32** and optimizes it using JAX/Optax while keeping the base model parameters frozen.

In [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py), the `finetune_local` function (lines 94-106) implements the training loop that targets specific attention projections: `q_proj`, `k_proj`, `v_proj`, `gate_proj`, and `out_proj` (defined as `LORA_TARGETS` in the same file). The loss function specifically covers the grounding portion—reasoning lines plus JSON tool calls—rather than the full sequence, which means training typically starts near a loss of 1.0 rather than random initialization.

Execute the fine-tuning command with default 10 epochs:

```bash

# Fine-tune a LoRA adapter (default 10 epochs, rank 16)

needle finetune data.jsonl --epochs 10 --out adapter.pkl

# Output ends with: "adapter <path>/adapter.pkl"

```

## Exporting and Merging Weights

After training, you must merge the adapter into the base checkpoint and export a portable `.cact` archive using the `needle build` command.

The export routine in [`needle/model/export.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/export.py) handles the merge operation and version-locks the output file to ensure compatibility. Note that `.cact` files created with older package versions cannot be loaded by newer code, requiring re-export after upgrades. Run the merge and export:

```bash

# Merge the adapter into the base checkpoint and export a .cact archive

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

# After export you'll see: "wrote ... tuned.cact"

```

For distribution, you can upload directly to Hugging Face by setting the `NEEDLE_HF_REPO` environment variable and adding the `--upload` flag:

```bash
export NEEDLE_HF_REPO=myusername/needle-tuned
needle build checkpoints/needle2.pkl --lora adapter.pkl --out tuned.cact --upload

```

## Deploying Fine-Tuned Models

Load the merged archive using the public `Needle` class exposed in [`needle/__init__.py`](https://github.com/cactus-compute/needle/blob/main/needle/__init__.py), passing the `.cact` file path to the `weights` parameter.

When using fine-tuned weights, the confidence head is automatically disabled, meaning the `confidence` attribute in responses will be `None`. The model retains the base tokenizer and inherits the frozen layers while applying your trained adapter modifications to the attention projections.

```python

# Load the tuned model in Python

import needle

agent = needle.Needle(
    tools=[...],                # your tool definitions

    weights="tuned.cact"       # the merged archive

)

# Now the agent can call your tools with the fine-tuned behavior

response = agent("Extract the title from the PDF at https://example.com/report.pdf")
print(response)

```

## Summary

- **LoRA Configuration**: Needle 2 uses rank-16, alpha-32 adapters targeting only attention projections (`q_proj`, `k_proj`, `v_proj`, `gate_proj`, `out_proj`) while freezing all base parameters.
- **Data Requirements**: A few hundred examples improve tool selection accuracy, but thousands of examples are necessary for precise argument grounding in complex function calls.
- **Training Scope**: The loss function optimizes only the grounding portion (reasoning plus JSON calls), leveraging the base model's prior knowledge to accelerate convergence.
- **Export Constraints**: Merged `.cact` archives are version-locked and must be re-exported when upgrading the Needle package to ensure binary compatibility.
- **Inference Differences**: Fine-tuned models disable the confidence calibration head, returning `None` for confidence scores while maintaining full tool-calling functionality.

## Frequently Asked Questions

### What is LoRA and why does Needle 2 use it for fine-tuning?

**Low-Rank Adaptation (LoRA)** is a parameter-efficient fine-tuning method that injects trainable rank decomposition matrices into frozen neural network layers. Needle 2 implements LoRA in [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py) to avoid storing full model copies, reducing memory overhead by 99%+ while allowing specialized adapters for different tool schemas. This approach enables rapid iteration on custom datasets without modifying the base JAX checkpoint.

### How much training data is required for effective fine-tuning?

**Dataset size depends on your optimization target**: a few hundred examples (200-500) typically suffice to improve tool selection accuracy, while thousands of diverse examples are required to correct argument grounding errors. If your model selects the correct tools but populates parameters incorrectly, you should increase your training set size and use the data augmentation feature in [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py) (lines 58-66) to generate additional variations.

### Why is the confidence head disabled in fine-tuned models?

The confidence calibration head remains tied to the base model's distribution and is automatically disabled when loading fine-tuned weights via [`needle/__init__.py`](https://github.com/cactus-compute/needle/blob/main/needle/__init__.py). Since the LoRA adapters modify the attention mechanisms without retraining the confidence estimation layers, enabling confidence scores on adapted weights would produce miscalibrated predictions. The `Needle` class sets `confidence=None` for all responses from `.cact` archives to prevent misleading uncertainty estimates.

### Can I share my fine-tuned Needle 2 models on Hugging Face?

Yes, the `needle build` command supports direct uploads via the `--upload` flag when you set the `NEEDLE_HF_REPO` environment variable to your Hugging Face repository path. The exported `.cact` file contains the merged base weights and LoRA adapters in a portable format, though recipients must use the same Needle package version that created the archive to load the weights successfully.