How do you implement LoRA in a coding interview?

In this article

1. Before you start
2. Why interviewers ask this
3. LoRA from first principles
4. Worked example
5. Trade-offs and edge cases
6. How to answer it in an interview
7. Common mistakes
8. Follow-up questions
9. Sources

Before you start

This page belongs to the machine learning (ML) coding round guide. It assumes you can read a short PyTorch module.

Why interviewers ask this

"Implement LoRA" is reported by candidates in our question bank of ML coding questions, with no public source naming a company, so treat it as a plausible question, not a certain one. It is short enough for one round and tests several things at once: shapes, initialisation, what is frozen, and what happens at inference. It sits alongside the other tasks in AI engineer coding interview questions.

A weak answer writes two matrices and adds them to the output without saying which is trained or how they start. A strong answer freezes the base layer, explains why one matrix starts at zero, scales the update, and states that the update can be merged into the base weight so inference costs nothing extra. A Staff-level answer adds when LoRA is the wrong tool and how quantized variants fit.

LoRA from first principles

LoRA

LoRA is a fine-tuning method that freezes a pretrained model's weights and trains a pair of small low-rank matrices per adapted layer, whose scaled product is added to that layer's output instead of changing its weights.

Rank

The rank r of a LoRA update is the shared inner size of its two matrices; it caps how many independent directions the update can change, and it sets the number of trainable parameters.

In plain words

Imagine a printed textbook you may not edit. Instead of rewriting pages, you attach a thin sheet of sticky notes that changes how a few pages read. The book stays the same; the notes carry everything new. The analogy breaks in one place: LoRA's notes can be folded into the book afterwards, so the finished model has no extra sheet.

Here is the same idea precisely. LoRA writes the update to a weight matrix as a product. The product is a d by r matrix B times an r by k matrix A. The size r is much smaller than d and k.

The original paper states it in its abstract. It freezes the pretrained weights. It injects trainable rank decomposition matrices into each layer of the Transformer.

Step 1: Freeze the base layer

Start from an existing linear layer with weight W0 of shape d by k. Mark every one of its parameters as not trainable, so the optimiser never changes them. Input: the pretrained layer. Output: the same layer, frozen.

This is the step that makes LoRA cheap. Gradients and optimiser state, such as the running averages Adam keeps for every trainable number, are only needed for parameters that train, so a frozen layer needs neither. In code it is a loop over the layer's parameters setting requires_grad to false. You now know what stays fixed; the next step creates what changes.

Step 2: Initialise B at zero and A at random

Create A with shape r by k and B with shape d by r. Fill A with small random numbers and B with zeros. Input: the three sizes d, k and r. Output: an update whose product is exactly zero.

Hu et al., §4.1 use a random Gaussian initialisation for A and zero for B, so the product of B and A is zero at the beginning of training and the adapted model begins identical to the base model. The reason is the gradient. With B at zero, the gradient reaching B is non-zero while the gradient reaching A is zero, so B moves first and then A follows.

In a toy check with illustrative shapes, the gradient of B at the start is non-zero and the gradient of A is exactly 0.0.¹ If both started at zero, both gradients would be zero and nothing would train. If both were random, the model would start with a noisy change to its outputs.

Step 3: Add the scaled low-rank path in the forward pass

Compute the frozen output as usual, then add the low-rank path: multiply the input by A, then by B, then by a scale of alpha divided by r. Alpha is a fixed number you choose, and it controls how strongly the update counts. Input: a batch of vectors x. Output: $h=W_{0}x+\frac{\alpha }{r}BAx$.

Multiply in that order: A first. That keeps the intermediate result at r numbers per example instead of building a d by k matrix. The scale keeps the update's size roughly steady when you change r. Hu et al., §4.1 set alpha to the first r they try and do not tune it, noting that the scaling "helps to reduce the need to retune hyperparameters when we vary r". With r set to 8 and alpha to 16 the scale is 2.0; with r at 16 and the same alpha it is 1.0 (illustrative).¹

Step 4: Train A and B, then merge

Train with a normal loop. Only A and B receive gradients. After training you can merge: add the scaled product of B and A into W0, and delete A and B. Input: trained A and B. Output: one ordinary weight matrix.

The paper reports no additional inference latency, unlike adapters, because B times A can be folded into W0. In the toy check, the merged and unmerged outputs agree to within a difference below 1e-12.¹ Merge only once, and only when you no longer need to switch adapters. If you keep A and B separate, you add two small matrix products per layer. You can then swap adapters for different tasks over one shared base.

1 · Freeze2 · Initialise3 · Add paths4 · MergeReplay

Step-by-step description

1. The base weight stays fixed.
2. A starts random, B starts at zero.
3. Frozen and low-rank outputs are summed.
4. B times A folds into W0.

• The base weight W0 is frozen.
• Matrix A starts random and matrix B starts at zero.
• The frozen output and the scaled low-rank output are added.
• After training, B times A is merged into W0.

Here is the whole module. Read it against the four steps.

import math
import torch
import torch.nn as nn

class LoRALinear(nn.Module):
    def __init__(self, base: nn.Linear, r: int, alpha: float):
        super().__init__()
        self.base = base
        for p in self.base.parameters():      # step 1: freeze
            p.requires_grad = False
        self.scale = alpha / r
        self.A = nn.Parameter(torch.empty(r, base.in_features))
        self.B = nn.Parameter(torch.zeros(base.out_features, r))
        nn.init.normal_(self.A, std=1.0 / math.sqrt(base.in_features))  # step 2
        self.merged = False

    def forward(self, x):
        out = self.base(x)
        if not self.merged:                   # step 3: low-rank path
            out = out + self.scale * ((x @ self.A.T) @ self.B.T)
        return out

    @torch.no_grad()
    def merge(self):                          # step 4: fold into W0
        if not self.merged:
            self.base.weight += self.scale * (self.B @ self.A)
            self.merged = True

The merged flag stops the update being counted twice. The std on A is a design choice for this sketch, so say so if asked.

Apply it to a model and check what trains

Applying LoRA to a real model means replacing chosen linear layers with LoRALinear and then checking how many parameters train. Freeze everything first, walk model.named_modules(), and use setattr on each match's parent module. Projection names differ between models, so the names below are examples.

def add_lora(model, names=("q_proj", "v_proj"), r=8, alpha=16):
    for p in model.parameters():
        p.requires_grad = False
    for name, module in list(model.named_modules()):
        if isinstance(module, nn.Linear) and name.split(".")[-1] in names:
            parent = model.get_submodule(name.rpartition(".")[0])
            setattr(parent, name.rpartition(".")[2], LoRALinear(module, r, alpha))
    return model

model = add_lora(model)  # model is any pretrained torch module you loaded earlier
trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
total = sum(p.numel() for p in model.parameters())
print(f"trainable {trainable:,} of {total:,} ({100 * trainable / total:.2f}%)")

Only the new A and B tensors train, so the printed share should be small. Depending on the rank and on which layers you adapt, it is typically a fraction of one percent up to a few percent of the total. A common addition is dropout on the input of the low-rank path only; many LoRA libraries expose it as a configuration setting.

Worked example

Take one hypothetical matrix of 4,096 by 4,096, repeated across a hypothetical 32 layers (illustrative sizes, not a claim about any named model). Full fine-tuning trains all 16,777,216 numbers of it. LoRA at rank 8 trains r * (d + k), which is 8 times 8,192, or 65,536 numbers, about 0.39% of the matrix.¹ Across 32 layers with two adapted matrices each, that is 4,194,304 trainable numbers against 1,073,741,824 for the same matrices trained fully (illustrative layout).¹

All figures in the table are computed from r * (d + k) against d * k.¹ For scale, the paper reports that on Generative Pre-trained Transformer 3 (generative pre-trained transformer (GPT)-3) 175B, LoRA can cut trainable parameters by 10,000 times and graphics processing unit (GPU) memory by 3 times compared with fine-tuning with Adam (Hu et al., 2021). That is a reported result on one model, not a figure for yours.

Text description of the diagram

• At rank 8, LoRA trains 65,536 numbers, about 0.39% of the 16,777,216 in the matrix.
• At rank 64, LoRA trains 524,288 numbers, 3.125% of the matrix.

Trade-offs and edge cases

LoRA trades some capacity for memory and speed. Because the update has rank at most r, it cannot represent every change full fine-tuning can. The paper reports that LoRA performs on par or better than fine-tuning on the models it tested, which is evidence on those models and tasks, not a rule. In the runs of Biderman et al. (2024), LoRA underperforms full fine-tuning on code and maths, and it better preserves the base model's performance on tasks outside the target domain.

Edge cases to name. Which layers you adapt is a choice: the paper adapts only the attention projections (Hu et al., §4.2 and §7), and adding more layers raises capacity and cost. Alpha is held fixed while r changes: the alpha over r scale keeps the update's size roughly steady across ranks, so tuning alpha then acts much like tuning the learning rate. An unmerged adapter adds two small products per layer at inference. After a merge, you cannot cleanly remove the adapter unless you keep a copy of W0.

Quantized LoRA (QLoRA) goes further. Dettmers et al. (2023) backpropagate through a frozen 4-bit quantized model into LoRA adapters, and report finetuning a 65B parameter model on a single 48 GB GPU. Sharding is a different tool: fully sharded data parallel splits full-model state across GPUs.

Text description of the diagram

• Full fine-tuning retrains all model parameters.
• LoRA freezes the base weights and trains low-rank matrices.
• QLoRA backpropagates through a frozen 4-bit quantized model into LoRA adapters.

How to answer it in an interview

The spoken structure: freeze, initialise and say why, give the forward formula, say what trains, then merge and the trade-off. Write the code in that order too; the GPT from scratch walkthrough shows the same habit of building a module piece by piece.

Common mistakes

• "I initialise A and B both to zero." (wrong) Both gradients would be zero, so nothing trains. (right) "B is zero, A is random, so the start equals the base model and B learns first."
• "LoRA always adds inference latency." (wrong) Only while the adapter is unmerged. (right) "After I merge the update into the weight, inference cost is the same as the base model."
• "I train the original weight too." (wrong) That removes the saving and defeats the method. (right) "The base layer is frozen; only A and B train."
• "I drop the alpha over r scale." (wrong) The update's strength then changes with r, so every new rank needs retuning. (right) "I keep the scale, set alpha once at the first rank I try, and treat it like a learning-rate knob."
• "I build the full d by k update first." (wrong) It wastes the memory LoRA saves. (right) "I multiply x by A first, then by B."

Follow-up questions

How do you count LoRA's trainable parameters?

For each adapted d by k matrix, LoRA trains r times (d plus k) numbers against d times k for full tuning; rank 8 on a 4,096 by 4,096 matrix is 65,536.¹ Printing trainable against total parameters, as above, confirms it. Attention layers like those in multi-head attention from scratch are where you would apply it.

Why is matrix B initialised to zero in LoRA?

A zero B makes the update zero at the start, so the adapted model equals the base model, and B still receives gradient because A is non-zero. A zero A and a zero B together give no gradient at all.

What is the difference between r and alpha in LoRA?

Rank r sets the capacity of the update and the number of trainable parameters. Alpha is a constant: the scale alpha over r keeps the update's size roughly steady as r changes, so you can fix alpha at the first rank you try, as Hu et al. do, and tuning it then acts like tuning the learning rate.

Does LoRA add inference latency?

Not once merged: the paper reports no additional inference latency. If you keep the adapter separate to switch tasks, you pay two small matrix products per adapted layer.

How is QLoRA different from LoRA?

QLoRA keeps the frozen base model in 4-bit form and trains LoRA adapters on top of it, which lowers memory further (Dettmers et al., 2023). The adapter maths is the same.

Which layers should you adapt?

State it as a choice with a cost. Adapting more matrices raises capacity and trainable parameters; the attention projections are the usual starting point, as in Hu et al., §4.2.

Sources

1. LoRA: Low-Rank Adaptation of Large Language Models (arXiv 2106.09685 v2) — arXiv (Hu et al.), 16 Oct 2021
2. QLoRA: Efficient Finetuning of Quantized LLMs (arXiv 2305.14314) — arXiv (Dettmers et al.), 23 May 2023
3. LoRA Learns Less and Forgets Less (arXiv 2405.09673) — arXiv (Biderman et al.), 15 May 2024