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

> Master the LoRA fine-tuning workflow in Needle. This guide covers data generation, adapter training, and model export for efficient tool-calling adaptation. Optimize your LLMs now.

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

---

**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](https://github.com/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`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py) and [`needle/model/export.py`](https://github.com/cactus-compute/needle/blob/main/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`:

```json
[
  { "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`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py) to send requests:

```bash
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`:

```bash
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`](https://github.com/cactus-compute/needle/blob/main/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:

```bash
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:

```python
{
  "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`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py):

```bash
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`](https://github.com/cactus-compute/needle/blob/main/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`](https://github.com/cactus-compute/needle/blob/main/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`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py):

```python
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`](https://github.com/cactus-compute/needle/blob/main/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`](https://github.com/cactus-compute/needle/blob/main/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`](https://github.com/cactus-compute/needle/blob/main/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.