Fine-tuning Needle

Needle was designed to be customised. Its capacity is a ladder, and a subnetwork as small as two layers, fine-tuned on one product's tools, runs on devices far smaller than the full model needs. Constraining a narrow, well-defined task is what lets it reach frontier accuracy there: fine-tuning on DroidCall lifts every subnetwork by 18 to 36 points, and from four layers up the tuned subnetwork passes DeepSeek V4 Flash.

Base and tuned are both scored with forced calls; DeepSeek V4 Flash runs through its cloud API. This guide is the local path in the Python package. The Cactus Platform runs the same pipeline with the 2-bit post-training and quantisation behind the shipped model, on our infrastructure, and adds evaluation tracking and dataset management.

What a fine-tune is

LoRA adapters of rank 16 on the five attention projections of every layer, trained on your JSONL with the base frozen, then merged into the weights at export. Training runs at the full 20 layers through the same 4-bit quantisation-aware numerics the export uses, so the adapter matches the archive needle build writes. The engine, the tokenizer and the confidence head are untouched, and the output is a single .cact you pass as weights=.

Install

The runtime package does not carry the training stack. Add the train extra, and gpu or metal to train on an NVIDIA GPU or Apple Silicon:

pip install "cactus-needle[train]"
pip install "cactus-needle[train,gpu]"
pip install "cactus-needle[train,metal]"

Training is plain JAX. On NVIDIA the same command trains on the GPU with nothing else changed. Apple GPUs go through the jax-metal plugin, which does not work past jax 0.4.38, so the metal extra pins an older stack and Needle adapts to it (manual attention, no rematerialisation, unrolled layers). On an M5 Max a step takes 0.71 seconds against 2.90 on CPU at the same shape, with a one-time compile of about 23 seconds. Training is float32 on every backend.

Data

One JSON object per line. query and tools describe the turn, answers lists the exact calls the model should emit, and reasoning is one short line deriving each argument from its span in the query:

{"query": "dim the kitchen to 10", "tools": [{"name": "set_lights", "parameters": {"type": "object", "properties": {"room": {"type": "string"}, "brightness": {"type": "integer"}}, "required": ["room"]}}], "answers": [{"name": "set_lights", "arguments": {"room": "kitchen", "brightness": 10}}], "reasoning": "'kitchen' -> room; 'to 10' -> brightness"}

reasoning is optional but include it: the model produces the derivation before the call, and examples that show where each value comes from teach grounding, not just tool selection. The rules that matter:

1. Arguments contain only values present in the query. Omit optional fields with no evidence; never fill them with placeholders or empty strings.
2. Include off-topic examples with "answers": []. The built-in generator produces about one in eight. Without them the tuned model calls a tool on everything.
3. When the catalogue has similar tools, include ambiguous queries resolved to the correct one.
4. An optional "system" field per example becomes a system turn, matching Needle(system=...) at inference.
5. Each rendered example must fit within --max-len (default 1024) tokens; longer ones are silently truncated. Padding rounds up to the longest example, so a short dataset trains fast regardless of the cap: 245-token examples train six times faster padded to 256 than to 1024 on CPU.

Extraction uses the same format with the record as the tool and the passage as the query. To grow a small hand-written set, seed the generator with it (needs OPENROUTER_API_KEY; OPENROUTER_URL points it at another OpenAI-compatible gateway):

export OPENROUTER_API_KEY=sk-or-...
needle generate-data --tools my_tools.json --num-samples 500 --output data.jsonl
needle generate-data --augment data.jsonl --num-samples 1000

Train, build, run

needle finetune data.jsonl --epochs 10 --out adapter.safetensors
needle build --lora adapter.safetensors --out tuned.cact
needle build --lora adapter.safetensors --layers 8 --out tuned_8l.cact
needle build --lora adapter.safetensors --platform linux-arm64 --layers 2 --out ./device

The base checkpoint downloads from Hugging Face on the first run. needle build merges the adapter into the full base, slices the subnetwork you asked for, quantises to 4 bits and packages the tokenizer; with --platform it also downloads that platform's engine and header, so --out becomes a folder ready to copy onto the device. Defaults: batch 16, learning rate 1e-4 with warmup and cosine decay, gradient clipping at norm 1, rank 16, alpha 32, max length 1024, validation split 0.1, seed 0. --seed reproduces or deliberately varies the LoRA initialisation, the validation selection and the epoch shuffle.

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

The engine is weights-agnostic, so a tuned archive runs on it directly. Set NEEDLE_HF_REPO=<you>/<model> and pass --upload to publish it; needle download <you>/<model>/tuned.cact pulls it on any machine.

Reading the loss

The loss covers only the target: the reasoning line plus the JSON call. Much of the call is boilerplate the base already predicts (the name, the braces, the field names), so a run starts near 1.0 rather than near random. Judge it by the trend, not the level.

Step count is what small datasets get wrong. 200 examples at batch 16 is 13 steps per epoch, and three epochs is 39 steps, which barely moves a rank-16 adapter at the default learning rate. For a few hundred examples run 10 to 30 epochs and expect a clear downward trend. If the curve sits at its starting value after a few hundred steps, raise the epochs first, then the learning rate. A validation loss prints at each epoch end; when it rises while the training loss keeps falling, the run is overfitting: stop there, or add data.

Sizing the dataset

Tool selection moves first: a few hundred clean examples measurably improve which tool gets picked. Argument grounding moves later and needs more, on the order of thousands of examples with reasoning lines and varied phrasings and values. If evaluation shows correct tools with wrong argument values, the dataset is too small or too uniform, not mislabelled. For grounding-heavy tasks --lora-rank 32 doubles adapter capacity and the adapter stays tiny.

For a large catalogue, consider two passes at inference instead of more training: one turn against the full catalogue to pick the tool, then one turn declaring only that tool, which constrains the grammar to exactly that call.

What a fine-tune does not change

The confidence head. Its calibration holds for the base model, and fine-tuning does not update it, so Needle(weights=...) reports confidence as None and warns once. Route tuned models on your own validation instead of the score, or keep the base model for the decision and the tuned one for the call.

The tokenizer. Non-English text fragments into roughly 1.7 times more tokens (measured on Spanish), which taxes both quality and the context budget.

The bits. Local fine-tuning trains and exports at 4 bits. The 2-bit post-training and quantisation behind the shipped model, enriched with Cactus datasets, run on the Cactus Platform: upload your tools and examples there to get a customised Needle at the shipped size and accuracy.

Troubleshooting

1. failed to load weights: the .cact format is tied to the engine version, so an archive exported by an older package will not load. Rebuild it with the current package.
2. Loss hovers at its starting value: the run is undertrained, not broken. See reading the loss.
3. The tuned model calls a tool on everything: the dataset has no [] examples. Add them.
4. Correct tools, wrong values: more examples with reasoning lines, more varied values, then rank 32.