# How Needle's LoRA Fine-Tuning Merges Adapters: Complete Technical Guide

> Learn how Needle merges LoRA adapters by adding a scaled low-rank update to original weights. This technical guide explains the process before inference or export.

- Repository: [Cactus Compute, Inc./needle](https://github.com/cactus-compute/needle)
- Tags: deep-dive
- Published: 2026-08-16

---

**Needle merges LoRA adapters by adding a scaled low-rank update (A @ B) to the original weight tensors, performed by the `merge_lora` function in [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py) before inference or export.**

Needle implements **Low-Rank Adaptation (LoRA)** fine-tuning as a parameter-efficient method where adapter weights remain separate from the base model during training. The critical merge operation—combining these adapters back into the original weights—enables efficient deployment without runtime overhead. This article explains exactly how Needle's LoRA fine-tuning merge adapters step-by-step, with complete source code references.

## How LoRA Adapters Work in Needle

LoRA fine-tuning in Needle follows a two-phase approach: separate training of low-rank matrices, then merge for use.

### The Adapter Structure

For each target weight matrix, Needle maintains:

- **Matrix A**: Shape `(in_dim, rank)` — initialized randomly
- **Matrix B**: Shape `(rank, out_dim)` — initialized to zero
- **Scale factor**: Computed as `alpha / rank`

During training, the base model parameters in `params` remain frozen. Only the `lora` dictionary containing these low-rank matrices is updated.

### When Merging Occurs

According to the source code in [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py), merging happens in two contexts:

1. **During training** — The forward pass computes loss using merged parameters: `model.apply({"params": merge_lora(params, lora, scale)}, ids)` [lines 55-56](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py#L55-L56)
2. **At export time** — `build_main` loads saved adapters and merges before checkpoint export [lines 12-19](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py#L12-L19)

## The `merge_lora` Function: Step-by-Step Implementation

The core merge logic resides in `merge_lora` at [lines 85-92 of [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py)](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py#L85-L92).

### Step 1: Flatten the Parameter Tree

Needle uses JAX/Flax utilities to create a flat view of nested parameters:

```python
flat_params = flax.traverse_util.flatten_dict(params)

```

This converts the hierarchical `params` dictionary into a flat structure where each weight tensor is addressable by its path key (e.g., `"transformer/gate/dense/kernel"`).

### Step 2: Iterate LoRA Adapters

For every entry in the `lora` dictionary, the function identifies the corresponding original weight tensor by matching path keys.

### Step 3: Compute the Low-Rank Update

The update follows the standard LoRA formula:

```python
update = scale * (A @ B)

```

Where:
- `A` has shape `(in_dim, rank)`
- `B` has shape `(rank, out_dim)`
- `scale` is passed as a parameter (typically `alpha / rank`)

### Step 4: Add Update to Original Weight

The computed update is added to the original weight tensor, with explicit dtype preservation to prevent precision issues.

### Step 5: Un-Flatten the Structure

After all updates are applied:

```python
merged_params = flax.traverse_util.unflatten_dict(flat_params)

```

This restores the nested structure expected by the model's `apply` method.

## Complete Merge Code Examples

### Basic Adapter Merge

```python
from needle.model.finetune import merge_lora

# Load original model parameters (frozen during training)

params, config = load_checkpoint("checkpoints/needle2.pkl")

# Load trained LoRA adapter from .pkl file

lora_adapter = load_adapter("adapters/lora_rank8.pkl")

# Compute scale from hyperparameters

scale = config.lora_alpha / config.lora_rank  # e.g., 16 / 8 = 2.0

# Perform the merge

merged_params = merge_lora(params, lora_adapter, scale)

```

### Inference with Merged Weights

```python
from needle.model.run import load_checkpoint
from needle.model.architecture import SimpleAttentionNetwork
from needle.model.finetune import merge_lora

# Load base checkpoint

params, config = load_checkpoint("checkpoints/needle2.pkl")

# Load and merge LoRA adapter

lora_adapter = load_adapter("fine_tuned_adapter.pkl")
scale = config.lora_alpha / config.lora_rank
params = merge_lora(params, lora_adapter, scale)

# Initialize model with merged weights

model = SimpleAttentionNetwork(config)

# Standard forward pass—no adapter overhead at runtime

logits = model.apply({"params": params}, input_ids)

```

## Key Supporting Components

Several utilities work alongside `merge_lora` to enable the complete workflow:

| Component | Location | Purpose |
|-----------|----------|---------|
| `lora_target_paths` | [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py) [lines 54-62](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py#L54-L62) | Selects which weight matrices are eligible for LoRA adaptation |
| `init_lora` | [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py) | Creates initial A/B matrices for selected weight paths |
| `build_main` | [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py) [lines 12-19](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py#L12-L19) | Handles adapter loading and merge at export time |
| CLI flags | [`needle/cli.py`](https://github.com/cactus-compute/needle/blob/main/needle/cli.py) [lines 127-153](https://github.com/cactus-compute/needle/blob/main/needle/cli.py#L127-L153) | Exposes `--lora`, `--lora-rank`, `--lora-alpha` parameters |

## Why Merge Instead of Runtime Composition?

Needle's design choice to **merge rather than compose** at runtime provides several advantages:

- **No inference overhead** — After merging, the model runs at full speed without computing `W + scale * (A @ B)` on every forward pass
- **Standard deployment** — Merged checkpoints are compatible with any JAX/Flax inference pipeline
- **Composable adapters** — Multiple LoRA adapters can be merged sequentially before final export

The tradeoff is minimal: a single merge operation at load time versus repeated low-rank matrix multiplications during every forward pass.

## Summary

- **Needle's LoRA merge** is implemented in `merge_lora` within [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py)
- The merge computes `scale * (A @ B)` for each adapter and adds it to the original weight tensor
- Flattening/un-flattening via `flax.traverse_util` enables clean path-based parameter updates
- Merging occurs both during training (for loss computation) and at export time (for deployment)
- Resulting merged parameters eliminate runtime adapter overhead entirely

## Frequently Asked Questions

### What file contains Needle's LoRA merge implementation?

The `merge_lora` function is defined in [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py) at [lines 85-92](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py#L85-L92). This file also contains related utilities `init_lora` and `lora_target_paths`.

### Does merging modify the original base model weights?

No—the merge creates a **new parameter dictionary** with updated weights. The original `params` and `lora` adapter remain unchanged, allowing multiple merges with different adapters or rollback to base weights.

### What is the scale factor in Needle's LoRA implementation?

The scale factor is computed as `alpha / rank`, where `alpha` and `rank` are hyperparameters set via CLI flags (`--lora-alpha`, `--lora-rank`). This scaling controls the magnitude of the low-rank adaptation applied during merging.

### Can I merge multiple LoRA adapters with Needle?

Yes—since `merge_lora` returns a new parameter dictionary, you can chain merges: apply the first adapter's merged result as the `params` input to a second `merge_lora` call with a different adapter.