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

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, 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.

# 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.

# 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.

# 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:

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:

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 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.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}