LoRA From Scratch in 19 Lines: It Matches PEFT Exactly, and Saved Less Memory Than We Predicted

LoRA From Scratch in 19 Lines: It Matches PEFT Exactly, and Saved Less Memory Than We Predicted

Before running anything, we wrote down that LoRA would cut peak training memory on Qwen2.5-0.5B to less than half of full fine-tuning, in fp32 at batch 4 and 512 tokens. It cut it to 66%: 10.31 GiB against 15.56 GiB. At 4,096 tokens, the second registered condition, neither run fit on an 80 GB A100.

The predictions about matching PEFT all held. Our 19-line LoRA class gave exactly the same logits as the PEFT library, and when both trained on the same batch for 30 steps, every loss matched to the last bit. Merging the trained adapter into the weights missed its prediction, though in fp32 not one top-1 token changed. In bf16 the adapter model and the merged model came apart, and the adapter model was the one further from the fp32 result.

This post is a tutorial. You build the layer, put it into a real model, check it against PEFT, merge it, and measure the GPU memory yourself. It is Part 3 of the LoRA variants series. Part 2 put five variants on one table; if LoRA itself is new to you, start with Part 1.

What you need

We ran everything with PyTorch 2.7.0, Transformers 5.20.2 and PEFT 0.21.2 on one A100 80 GB. The model is Qwen/Qwen2.5-0.5B-Instruct: 24 layers, hidden size 896, and two key-value heads of 64, so v_proj maps 896 numbers to 128. Steps 1 to 4 run on any GPU with a few GB free; only the long inputs in step 5 need the big card.

Step 1: the layer

A linear layer computes h = W x. LoRA keeps W frozen and adds a low-rank path: h = W x + (α/r) · B A x. A is r × in and starts random, B is out × r and starts at zero, so at step 0 the layer gives exactly what the original gave.

import math
import torch
import torch.nn as nn

class LoRALinear(nn.Module):
    """h = W x + (alpha / r) * B A x ; W frozen, only A and B train."""
    def __init__(self, base: nn.Linear, r: int, alpha: float):
        super().__init__()
        self.base, self.scale = base, alpha / r
        for p in self.base.parameters():
            p.requires_grad_(False)
        w = base.weight   # A and B live on W's device and dtype
        self.A = nn.Parameter(torch.empty(r, base.in_features, device=w.device, dtype=w.dtype))
        self.B = nn.Parameter(torch.zeros(base.out_features, r, device=w.device, dtype=w.dtype))
        nn.init.kaiming_uniform_(self.A, a=math.sqrt(5))   # PEFT's default; the paper says Gaussian

    def forward(self, x):
        return self.base(x) + self.scale * ((x @ self.A.T) @ self.B.T)

    def merge(self) -> nn.Linear:
        with torch.no_grad():
            self.base.weight += (self.scale * (self.B @ self.A)).to(self.base.weight.dtype)
        return self.base

The class is 19 lines. Three choices in it matter, each for a different reason:

• The scale alpha / r is what makes the output match PEFT. With r = 8 and alpha = 16 it is 2. Use another scale and the layer computes something else.
• A starts from Kaiming-uniform with a = sqrt(5), PEFT's default. The LoRA paper says Gaussian. This does not decide the match in step 3, where we copy A into PEFT anyway; it means a fresh layer starts the way a PEFT layer would.
• The order (x @ A.T) @ B.T is about cost. It squeezes each token down to 8 numbers and opens it back up. Computing B @ A first would build a full out × in matrix on every call, which the low-rank path is there to avoid.

The animation follows one v_proj call through those steps.

Step 2: put it into the model

Freeze everything, then swap chosen linear layers for LoRALinear. We adapt q_proj and v_proj, as the LoRA paper's main experiments did.

def add_lora(model, r=8, alpha=16, targets=("q_proj", "v_proj")):
    for p in model.parameters():
        p.requires_grad_(False)
    for module in list(model.modules()):
        for t in targets:
            child = getattr(module, t, None)
            if isinstance(child, nn.Linear):
                setattr(module, t, LoRALinear(child, r, alpha))
    return model

For each adapted layer, A has 8 × in numbers and B has out × 8.

The calculation gives 540,672, and so does counting requires_grad parameters in our model and in PEFT's: 0.11% of the 494,032,768 numbers in the frozen model.

Step 3: check it against PEFT

Load the model twice. Wrap one copy with add_lora, the other with PEFT's LoraConfig(r=8, lora_alpha=16, lora_dropout=0.0, target_modules=["q_proj", "v_proj"]). Then copy our A and B into PEFT's layers so both start from the same numbers:

peft_layers = dict(peft_model.base_model.model.named_modules())
with torch.no_grad():
    for name, ours in mine.named_modules():
        if isinstance(ours, LoRALinear):
            theirs = peft_layers[name]
            theirs.lora_A["default"].weight.copy_(ours.A)
            theirs.lora_B["default"].weight.copy_(ours.B)

We loaded the model in fp32 and turned off TF32 (torch.backends.cuda.matmul.allow_tf32 = False), so both models do plain fp32 matrix multiplications, and used eager attention. We did not test whether the match holds without these settings.

With B at zero, both models equal the untouched base model, which tests nothing. So we filled B with small random numbers (standard deviation 0.02, seed 1), copied them across, and compared logits on four sentences, two English and two Korean. The largest difference was 0.0, against logits as large as 25.5.

Then we reset B to zero, gave each model its own AdamW at learning rate 1e-3, and trained both on the same batch for 30 steps.

Every one of the 30 losses was identical, and so were A and B after the last step. The loss falling to 0.009 only means the model memorized four sentences; this check is about the two implementations agreeing, not about learning anything useful.

Step 4: merge the adapter into W

After training, you can fold the adapter into the weight, W' = W + (α/r) · B A, and throw away A and B. The model is the original shape again, with no extra work at inference.

def merge_lora(model):
    for module in list(model.modules()):
        for name, child in list(module.named_children()):
            if isinstance(child, LoRALinear):
                setattr(module, name, child.merge())
    return model

We had predicted the merged model would match the adapter model within 1e-4. After 30 training steps the largest logit difference, the measure we registered, was 7.3e-4. That count includes the padding positions; on the 82 real token positions alone it was 1.1e-4 (added after), still above the line. The two are the same in exact arithmetic; the computer adds in a different order, and each order rounds differently. The logits moved that little and no further: the top-1 token matched on all 82 positions, and the loss was 0.0079393 both ways.

Merging into bf16 weights, which we tested after the predictions, is different. We put the same trained A and B on a bf16 copy of the model and merged. Of the 22,020,085 nonzero entries of 2 · B A, 3.2% vanished: added to a bf16 weight, they were too small to change it. The adapter and merged models then disagreed on the top-1 token at 2 of 82 positions, and their logits differed by up to 15.1 on the real tokens, where the largest logit was 35.75.

To see which of the two moved, we compared each with the fp32 adapter model. The bf16 adapter model was the one further off: logits up to 14.6 away and loss 0.137, against up to 5.2 away and loss 0.0096 for the merged model, and 0.0079 in fp32. Both picked the same top-1 token as fp32 at 81 of 82 positions. Keeping A, B and the low-rank product in fp32 on top of the bf16 base did not help (up to 14.5 away, loss 0.143), so the gap does not come from the precision of A, B and their product. In that run the final addition was still in bf16. We have not found where it does come from, and this is one overfit batch of four sentences.

Step 5: measure the memory yourself

Peak memory is what decides whether a fine-tune fits on your card. This is the loop we used, once with every weight trainable and once after add_lora. Load the model with attn_implementation="sdpa" and the dtype you want to test, and move it to the GPU with .to("cuda").

model.train()
opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=1e-5)
ids = torch.randint(0, model.config.vocab_size, (4, 512), device="cuda")   # batch 4, 512 tokens
for step in range(3):
    torch.cuda.reset_peak_memory_stats()
    loss = model(input_ids=ids, labels=ids).loss
    loss.backward()
    opt.step(); opt.zero_grad(set_to_none=True)
    print(step, f"{torch.cuda.max_memory_allocated() / 2**30:.2f} GiB peak")

Read the last step. AdamW creates its two state tensors at the first opt.step(), so step 0 does not yet carry them.

Where the memory went

The registered prediction was fp32, batch 4: LoRA below half of full fine-tuning at 512 tokens.

The fixed part came out as calculated. At the end of the backward pass, full fine-tuning holds weights, gradients and two Adam states: 7.40 GiB, close to the 7.9 GB (7.4 GiB) we had worked out. LoRA holds 1.87 GiB, almost all of it the frozen weights. That difference in the fixed part, about 5.5 GiB, is the same at every length. The gap between the peaks still grows, from 5.26 to 9.96 GiB down the first three rows, because of the part that scales with tokens.

That part is what P5 missed. Every token keeps activations in each of the 24 layers for the backward pass, and every token produces 151,936 logits, one per vocabulary entry. Even at 512 tokens it was not small: in the bf16 split further down (added after), the 24 layers kept 2.02 GiB and the logits and loss 1.74 GiB, at batch 4.

fp32 adds a second cost. We measured at 1,024 and 2,048 tokens (added after): doubling the length multiplied the layers' activations by 2.8 in fp32 with SDPA, the same as eager attention, which keeps the full tokens-by-tokens attention matrix; in bf16 with SDPA it was 2.0. We did not check which kernel ran. If that matrix is kept, our count puts it at about 21 GiB at batch 4 and 2,048 tokens, and at about 84 GiB at 4,096 tokens, more than the whole card by itself.

In bf16

So we reran the comparison in bf16, the setting we would use to plan a real run. All of it was added after the registered runs.

LoRA saved 2.5 GiB at 512 tokens and 7.5 GiB at 4,096. At the end of the forward pass it held about 18% less activation memory than full fine-tuning at every length; our reading is that frozen layers do not need to keep their inputs for a weight gradient. Total memory grows faster than that saving, so the share LoRA saves shrinks from 28% to 14%.

These are pure bf16 runs: weights, gradients and Adam states all in bf16. Mixed-precision training that keeps fp32 master weights would make full fine-tuning heavier and LoRA's fixed saving larger.

In bf16 at 4,096 tokens, checkpointing saved more than LoRA

To see what LoRA leaves behind, we split (added after) the memory kept for the backward pass into the 24 layers and the logits plus loss, then turned on gradient checkpointing, which throws away the layers' activations and recomputes them during the backward pass.

Checkpointing took the layers from 15.5 GiB to 0.74 GiB and the peak down by 14.8 GiB; LoRA had saved 7.5 GiB against full fine-tuning at the same length. The logits did not move: 4 × 4,096 tokens produce 2.49 billion of them, 4.64 GiB in bf16, plus 9.27 GiB for the fp32 copy our script makes to compute the loss. Neither LoRA nor checkpointing touches that 13.9 GiB.

The breakdown comes from a separate script that computes the loss by hand, so its peaks (48.9 and 34.1 GiB) sit a little apart from the 44.3 GiB in the chart above, and the two savings come from different scripts.

LoRA still decided one case. In fp32 at 4,096 tokens and batch 2, full fine-tuning ran out of memory on the 80 GB card and LoRA ran at 72.4 GiB.

Our predictions, scored

The predictions file names an earlier version of the layer than the one we ran. The version we ran changed A's initialization from a Gaussian scaled by 1/√in to Kaiming-uniform and added device and dtype handling and merge(). P1 to P3 copy A into PEFT, so the change cannot move those results, but the file we registered is not the file we ran. P6 stays untested: the batch 1 run above is a different condition.

What we did not check

• One model, one rank, two targets. Everything here is Qwen2.5-0.5B at rank 8 on q_proj and v_proj. A 7B model has a much larger fixed part, so we expect LoRA's share of the saving to be larger there, but we did not measure it.
• Memory is PyTorch's peak allocated memory, not what nvidia-smi shows, which adds the CUDA context and the allocator's cache.
• The training check uses one batch of four sentences. It shows the two implementations compute the same thing, not anything about fine-tuning quality.
• Everything marked "added after" was chosen to explain a registered result: the extra fp32 rows, the attention probe, the bf16 chart, the breakdown, the real-token merge measure and both bf16 merge checks.
• We did not try a chunked loss, which computes the logits a slice at a time and is the usual answer to the 13.9 GiB above.

Next

The next post changes this layer into the four variants from Part 2. Each one is a few lines of difference: freeze A for LoRA-FA, share A and B and train two vectors for VeRA, give B its own learning rate for LoRA+, and add the change in B A to W for Delta-LoRA. We will count what each one trains and measure its memory with the loop from step 5. The newsletter form below sends it when it is out.