# How to Fine-Tune Needle with LoRA Using JSONL Data Format

> Learn to fine-tune Needle with LoRA using JSONL data. Follow simple steps to prepare your data and train a low-rank adapter efficiently. Get started now!

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

---

**To fine-tune Needle with LoRA, prepare a JSONL file with `query` and `answers` fields, then run `needle finetune data.jsonl --lora-rank 16 --lora-alpha 32` to train a low-rank adapter and save it as a `.pkl` file for later merging.**

Fine-tuning Needle with **LoRA (Low-Rank Adaptation)** enables efficient adapter training on custom tool-calling datasets without modifying the full base weights. This guide covers the complete workflow—from JSONL preparation through adapter serialization—based directly on the Needle source code in `cactus-compute/needle`.

## Preparing Your JSONL Dataset

Each line in your JSONL file must be a valid JSON object containing at least two fields: `query` (the input prompt) and `answers` or `function_calls` (the expected tool outputs). An optional `tools` array can describe the tool schemas used for that example.

```json
{"query":"Add a reminder to my calendar for tomorrow at 9 am","answers":[{"name":"create_event","arguments":{"title":"reminder","datetime":"2023-05-01T09:00:00"}}],"tools":[{"name":"create_event","description":"Create a calendar event","parameters":{"type":"object","properties":{"title":{"type":"string"},"datetime":{"type":"string"}}}}]}

```

The `load_jsonl` function in [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py) (lines 38–52) parses this format, tokenizes each example, and pads or truncates to your specified `max_len`.

## Running the Fine-Tune Command

Needle exposes fine-tuning through the `finetune` sub-command registered in [`needle/cli.py`](https://github.com/cactus-compute/needle/blob/main/needle/cli.py) (lines 31–40). This invokes `finetune_local` in [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py), which orchestrates the full training pipeline.

```bash
needle finetune data.jsonl \
  --checkpoint checkpoints/needle2.pkl \
  --epochs 5 \
  --batch-size 32 \
  --lora-rank 16 \
  --lora-alpha 32 \
  --max-len 1024 \
  --out my_adapter.pkl

```

### Key Hyper-Parameters

- **`--lora-rank`** — Rank of the low-rank matrices `A` and `B`. Higher ranks increase expressiveness but also parameter count.
- **`--lora-alpha`** — Scaling factor applied to the LoRA update, where `scale = alpha / rank`. Typical setting is `alpha = 2 * rank`.
- **`--max-len`** — Maximum sequence length for padding and truncation.
- **`--generate N`** — Optional synthetic data augmentation using OpenRouter before training begins.

## How LoRA Training Works Internally

Understanding the internal mechanics helps you debug and optimize your fine-tuning runs.

### LoRA Target Discovery

The `lora_target_paths` function (lines 54–62 in [`finetune.py`](https://github.com/cactus-compute/needle/blob/main/finetune.py)) automatically identifies which weight matrices to adapt. It walks the checkpoint's parameter tree and selects all kernels belonging to attention layers—specifically `q_proj`, `k_proj`, `v_proj`, and `o_proj` projections.

### Adapter Initialization

`init_lora` (lines 66–82) creates the low-rank decomposition for each target:

- Matrix `A` is initialized with random normal values
- Matrix `B` is initialized to zeros
- Both are scaled by the user-provided rank

This ensures the adapter starts with zero effect on the base model and learns incrementally.

### Training Loop

Each epoch shuffles the dataset and processes batches through `train_step`, which is JIT-compiled for performance. The `merge_lora` function temporarily combines LoRA weights with base weights for the forward pass. Optional **CQ-STE quantization** can be applied during this merge.

Loss computation uses masked cross-entropy over the tokenized target sequence, ignoring padding tokens.

### Validation and Checkpointing

If you specify `--val-split > 0`, a hold-out set is evaluated after each epoch to report validation loss. After training completes, the adapter dictionary containing `lora` weights, `scale` values, and metadata is pickled to your output path (default: `<checkpoint_dir>/needle_lora.pkl`).

## Optional Data Augmentation

When training data is scarce, the `--generate N` flag triggers `augment_jsonl` in the pipeline. This process:

1. Extracts existing tool schemas via `_collect_tools`
2. Uses OpenRouter to synthesize N additional examples with similar structure
3. Adds the synthetic data to your training set before the main loop begins

This is particularly useful for bootstrapping tool-calling datasets in new domains.

## Merging and Deploying Your Adapter

The saved `.pkl` adapter is not directly deployable—it must be merged with a base checkpoint. Use the `build` command to produce a final `.cact` archive:

```bash
needle build checkpoints/needle2.pkl \
  --lora my_adapter.pkl \
  --out needle_finetuned.cact

```

The resulting archive loads normally for inference:

```python
from needle import Needle

model = Needle(
    weights='needle_finetuned.cact',
    tools=[...]  # your tool schemas

)

```

## Summary

- **JSONL format**: Each line needs `query` + `answers`/`function_calls`; optional `tools` array for schema context
- **CLI entry point**: `needle finetune` defined in [`needle/cli.py`](https://github.com/cactus-compute/needle/blob/main/needle/cli.py), implemented in [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py)
- **LoRA mechanics**: `lora_target_paths` discovers attention weights, `init_lora` creates `A`/`B` matrices, `train_step` handles JIT-compiled training with optional quantization
- **Output**: Pickled adapter saved to `.pkl`, merged via `needle build` into deployable `.cact` format

## Frequently Asked Questions

### What JSONL fields are required for Needle fine-tuning?

You need at least `query` (string) and `answers` or `function_calls` (array of tool call objects). The `tools` array describing schemas is optional but recommended for consistent training. Each line must be valid, compact JSON with no trailing commas.

### How do I choose LoRA rank and alpha values?

Start with `--lora-rank 8` or `16` and `--lora-alpha` equal to `2 * rank`. Higher ranks capture more complex adaptations but increase memory usage and training time. For simple tool-calling tasks, rank 8 often suffices; complex multi-step reasoning may need 32 or 64.

### Can I fine-tune without a GPU?

The JAX-based training in [`finetune.py`](https://github.com/cactus-compute/needle/blob/main/finetune.py) supports TPU and CPU backends, though CPU training will be significantly slower. The LoRA approach itself reduces memory requirements compared to full fine-tuning, making it feasible on modest hardware with small batch sizes.

### What happens if my examples exceed `--max-len`?

The tokenizer in [`needle/model/tokenizer.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/tokenizer.py) truncates sequences to `max_len` and applies appropriate padding. Long examples are not dropped—they are silently truncated from the right, which may cut off target outputs if your max length is too conservative for your data.