# Needle Finetune Training Configuration: Complete CLI Options Guide

> Explore needle finetune CLI options for model loading, LoRA, optimizers, and data augmentation. Master your training configuration with this comprehensive guide.

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

---

**The `needle finetune` command accepts 15+ command-line arguments defined in [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py) that control model loading, LoRA adapter dimensions, optimizer scheduling, and data augmentation.**

The `cactus-compute/needle` repository provides a JAX-based fine-tuning pipeline for attention-based language models. All training configuration options are parsed in the `finetune_local` function within **[`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py)**, where they directly initialize the LoRA adapters, cosine learning-rate schedule, and data loaders.

## Model and Data Path Configuration

The foundation of any **needle finetune training configuration** starts with specifying the base model and dataset locations.

- **`--checkpoint`** – Path to the base model weights (default: `checkpoints/needle2.pkl`). This loads the frozen parameters that will be adapted via LoRA. Referenced at line 106 in [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py).

- **`--jsonl_path`** – Path to the training dataset in JSON-L format. This is the primary data source consumed by the pipeline at line 107.

- **`--checkpoint_dir`** – Output directory where the final LoRA adapter (`needle_lora.pkl`) is serialized (line 148).

- **`--out`** – Optional explicit filename for the adapter; if omitted, the filename is derived from `checkpoint_dir` (line 149).

## LoRA Adapter Parameters

Low-Rank Adaptation (LoRA) settings determine the capacity and scaling of the trainable adapters injected into the frozen base model.

- **`--lora_rank`** – Defines the rank of the LoRA decomposition matrices. Lower values reduce memory usage and trainable parameters; higher values increase model capacity. Implemented at line 131.

- **`--lora_alpha`** – Scaling factor for the LoRA updates. The actual scale applied during the forward pass is calculated as `alpha / rank`. This controls the strength of the adaptation relative to the base weights, also at line 131.

## Optimization and Scheduling

The **needle finetune** implementation uses `optax.warmup_cosine_decay_schedule` to manage learning rate dynamics across training steps.

- **`--lr`** – Peak learning rate for the cosine decay schedule (line 147). The schedule warms up from zero to this peak, then decays following a cosine curve for the remainder of training.

- **`--batch_size`** – Number of examples processed per gradient step (line 143). This directly impacts GPU memory consumption and throughput.

- **`--epochs`** – Total number of complete passes over the training dataset (line 145).

## Data Generation and Validation

Advanced configuration options control data augmentation and validation splits.

- **`--generate`** – When provided, triggers on-the-fly dataset augmentation via `augment_jsonl` before training begins (lines 108–110). The integer value specifies how many additional examples to synthesize.

- **`--workers`** – Number of parallel threads used during the data generation phase (line 99). Increasing this speeds up preprocessing for large datasets.

- **`--val_split`** – Fraction of data reserved for validation monitoring (default 0.1) at line 136. This split enables early stopping and perplexity reporting without contaminating the training set.

- **`--max_len`** – Maximum token sequence length after tokenization (line 123). Sequences exceeding this length are truncated; shorter sequences are padded to this limit.

## Implementation in [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py)

The **`finetune_local`** function orchestrates the entire pipeline. It first parses the CLI arguments, then:

1. Loads the base checkpoint using utilities from [`needle/model/run.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/run.py).
2. Initializes LoRA parameters via `lora_rank` and `lora_alpha`.
3. Constructs the `optax.warmup_cosine_decay_schedule` using the `--lr` peak value.
4. Instantiates the tokenizer from [`needle/model/tokenizer.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/tokenizer.py) to encode JSON-L examples.
5. Executes the training loop for the specified `--epochs`.
6. Exports the adapter to `--checkpoint_dir` or the path specified by `--out`.

Supporting modules include [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) (defining `SimpleAttentionNetwork`) and [`needle/model/export.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/export.py) (handling final model serialization).

## Practical Usage Examples

Basic fine-tuning with explicit LoRA dimensions:

```bash
needle finetune \
  --checkpoint checkpoints/needle2.pkl \
  --jsonl_path data/needle_data.jsonl \
  --batch_size 32 \
  --epochs 10 \
  --lr 3e-4 \
  --lora_rank 4 \
  --lora_alpha 32 \
  --max_len 1024 \
  --checkpoint_dir ./lora_adapters

```

On-the-fly data augmentation with increased adapter capacity:

```bash
needle finetune \
  --jsonl_path data/needle_data.jsonl \
  --generate 5000 \
  --batch_size 64 \
  --epochs 5 \
  --lr 2e-4 \
  --lora_rank 8 \
  --lora_alpha 64

```

Custom validation split with multi-worker preprocessing:

```bash
needle finetune \
  --jsonl_path data/needle_data.jsonl \
  --val_split 0.2 \
  --workers 12 \
  --batch_size 16 \
  --epochs 20 \
  --lr 1e-4

```

## Summary

- The **`needle finetune`** command exposes 15+ configuration options defined in [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py).
- **LoRA parameters** (`--lora_rank`, `--lora_alpha`) control adapter capacity and scaling at line 131.
- The optimizer uses a **cosine decay schedule** with warmup, parameterized by `--lr` at line 147.
- **Data augmentation** via `--generate` and parallel processing via `--workers` happen before the main training loop.
- Validation splits (`--val_split`, default 0.1) are handled at line 136 for monitoring generalization.

## Frequently Asked Questions

### What is the default learning rate schedule in needle finetune?

According to the source code in [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py) at line 147, the learning rate follows a cosine decay schedule with linear warmup. The schedule warms up from zero to the peak value specified by `--lr`, then decays following a cosine curve for the remainder of the training steps calculated from `--epochs` and `--batch_size`.

### How does the `--generate` flag work in needle finetune?

When the `--generate` flag is passed with an integer value, the pipeline invokes `augment_jsonl` at lines 108–110 of [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py) to synthesize additional training examples before the main loop begins. The `--workers` argument at line 99 controls the thread parallelism during this generation phase.

### What is the relationship between `--lora_rank` and `--lora_alpha`?

In [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py) at line 131, both parameters initialize the LoRA adapter. The rank determines the dimensionality of the trainable low-rank matrices, while alpha provides a scaling factor. The effective adaptation strength is proportional to `alpha / rank`, allowing you to decouple the parameter count (rank) from the update magnitude (alpha).

### Where does needle finetune save the trained adapter?

By default, the adapter is saved as `needle_lora.pkl` inside the directory specified by `--checkpoint_dir` (line 148). If you provide the `--out` argument at line 149, the adapter is written to that specific path instead, overriding the default naming convention.