# LoRA Adapter Merging at Export Time in Needle 2: How and When It Works

> Learn when and how LoRA adapter merging occurs at export time in Needle 2. Discover how the --lora flag integrates LoRA weights into base checkpoints for efficient export.

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

---

**LoRA adapter merging in Needle 2 happens only during the `needle build` command when the `--lora` flag is provided, merging low-rank adaptation weights into the base checkpoint before exporting to a `.cact` file.**

Needle 2, the JAX-based inference engine from Cactus Compute, supports efficient fine-tuning through Low-Rank Adaptation (LoRA). Understanding when and how LoRA adapter merging occurs at export time is essential for production deployments. This article explains the exact mechanism, source code locations, and practical usage based on the official Needle 2 repository.

## When LoRA Adapter Merging Occurs

LoRA adapter merging in Needle 2 is **explicitly opt-in** through the CLI. The merge operation is skipped entirely unless you request it.

The trigger condition is straightforward:

- **With `--lora`**: The adapter merges into the base checkpoint, producing a standalone `.cact` file
- **Without `--lora`**: The original checkpoint exports unchanged, with no LoRA weights applied

This design keeps base model exports fast and avoids accidental merging of experimental adapters.

## How the Merge Process Works

The merging pipeline follows three sequential steps in [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py), within the `build_main` function (lines 12–30).

### Step 1: Adapter Loading

The adapter file specified via `--lora` is unpickled and reformatted for JAX consumption.

```python

# Lines 13-17 of needle/model/finetune.py

with open(args.lora, "rb") as f:
    adapter = pickle.load(f)
lora = {
    tuple(k.split("/")): {"A": jnp.asarray(v["A"]), "B": jnp.asarray(v["B"])}
    for k, v in adapter["lora"].items()
}

```

Key transformations performed:

- String paths like `"layer00/self_attn/q_proj"` become tuple keys: `("layer00", "self_attn", "q_proj")`
- Stored NumPy arrays convert to JAX tensors via `jnp.asarray()`

### Step 2: Weight Merging

The `merge_lora` function applies the low-rank update to each target weight tensor.

```python

# Lines 85-92 of needle/model/finetune.py

def merge_lora(params, lora, scale):
    for path, mats in lora.items():
        A, B = mats["A"], mats["B"]
        delta = scale * (A @ B.T if A.shape[0] < B.shape[0] else B @ A.T)
        node = params
        for key in path[:-1]:
            node = node[key]
        node[path[-1]] = node[path[-1]] + delta
    return params

```

The mathematical operation performed is:

```

W_merged = W_base + scale × (A @ B)

```

Where **scale equals α / r** (the LoRA scaling factor stored in `adapter["scale"]`). The function automatically handles both orientations of the low-rank matrices based on their dimensions.

### Step 3: Export to .cact Format

The merged parameters pass to `write_export` for final serialization.

```python

# Lines 26-30 of needle/model/finetune.py

write_export(
    params,
    config,
    args.output,
    bits=args.bits,
    tokenizer=tokenizer,
    kv_window=effective_kv_window(config),
)

```

The resulting `.cact` file contains the fully merged weights—no LoRA metadata remains, enabling efficient inference without runtime adapter overhead.

## Command-Line Usage

Export a merged model using the `needle build` command with the `--lora` argument:

```bash
needle build path/to/base.ckpt --lora path/to/adapter.pkl --out model.cact

```

Optional parameters control the export:

| Flag | Purpose | Default |
|------|---------|---------|
| `--bits` | Quantization bits (4, 8, or 16) | 4 |
| `--lora` | Path to pickled LoRA adapter | None |
| `--out` | Output `.cact` file path | Required |

If `--lora` is omitted, the checkpoint exports without merging.

## Programmatic Merging

For custom workflows, replicate the CLI logic directly in Python:

```python
import pickle
import jax.numpy as jnp
from needle.model.finetune import merge_lora
from needle.model.export import write_export
from needle.model.run import load_checkpoint
from needle.model.tokenizer import get_tokenizer
from needle.model.architecture import effective_kv_window

# Load base checkpoint

params, config, _ = load_checkpoint("base.ckpt", return_run=True)

# Load and reformat LoRA adapter

with open("adapter.pkl", "rb") as f:
    adapter = pickle.load(f)

lora = {
    tuple(k.split("/")): {"A": jnp.asarray(v["A"]), "B": jnp.asarray(v["B"])}
    for k, v in adapter["lora"].items()
}

# Merge: W_base + scale × (A @ B)

params = merge_lora(params, lora, adapter["scale"])

# Export merged model

write_export(
    params,
    config,
    "merged.cact",
    bits=4,
    tokenizer=get_tokenizer(config.vocab_size),
    kv_window=effective_kv_window(config),
)

```

This approach enables batch merging, adapter composition, or custom scaling factors beyond standard usage.

## Source Code Reference

The core merging logic resides in these locations:

- **`needle/model/finetune.py:85-92`** — `merge_lora()` implementation
- **`needle/model/finetune.py:12-30`** — `build_main()` orchestrating load, merge, and export
- **`needle/cli.py:62-68`** — `--lora` argument definition for `needle build`
- **`needle/model/export.py:87-93`** — `write_export()` final serialization

## Summary

- **LoRA adapter merging in Needle 2 is export-time only**, triggered by `needle build --lora`
- The `merge_lora` function in [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py) applies scaled low-rank updates (`A @ B`) to base weights
- Merged models export as standalone `.cact` files with no runtime LoRA overhead
- The process is fully reproducible via both CLI and Python API

## Frequently Asked Questions

### What happens if I export without the `--lora` flag?

The base checkpoint exports unchanged. Any previously trained LoRA adapter remains separate and must be loaded at runtime if you want its effects—you lose the performance benefits of a merged model.

### Can I merge multiple LoRA adapters at once?

The current `needle build` command accepts only a single `--lora` argument. For multi-adapter merging, use the programmatic approach: load multiple adapters, combine their deltas manually, then call `merge_lora` once with the composite update.

### Does merging affect quantization?

No. Quantization (controlled by `--bits`) applies **after** merging. The merged weights are quantized during `write_export`, so the LoRA deltas benefit from the same precision reduction as base weights.

### Where does the scaling factor `scale = α / r` come from?

The scale value is stored in `adapter["scale"]` during `needle finetune` training, computed from your `--lora-alpha` and `--lora-rank` hyperparameters. It travels with the adapter pickle file and applies automatically during `merge_lora`.