LoRA Adapter - PyTorch Basics

LoRA Adapter

What you'll build: lora_forward(x, W, A, B, alpha, r), the forward pass of a linear layer with a frozen weight W plus a trainable low-rank update A @ B.

Background

Fine-tuning every weight of a large model is expensive in memory: Adam stores two extra tensors per parameter. LoRA (Hu et al., 2021) freezes the pretrained W and learns a low-rank correction ΔW = (alpha / r) * A @ B with a small rank r. For a 4096x4096 layer and r = 8 that is 2 * 4096 * 8 = 65,536 trainable parameters instead of 16,777,216, about 99.6% fewer. After training, ΔW can be merged into W, so inference costs nothing extra.

Formula

output = x @ W + (alpha / r) * (x @ A) @ B

Symbol	Shape	Role
x	(batch, in_features)	input
W	(in_features, out_features)	frozen pretrained weight
A	(in_features, r)	down-projection (trainable, random init)
B	(r, out_features)	up-projection (trainable, zero init, so training starts exactly at the pretrained model)
alpha / r	scalar	scaling, so changing r doesn't require retuning the learning rate

Expected Return Format

An array of shape (batch, out_features).

Edge Cases and Conventions

• Compute (x @ A) @ B, not x @ (A @ B): the result is the same, but the first costs O(batch * r * (in + out)) while forming A @ B costs O(in * out * r) and a full-size matrix.
• alpha = 0 gives the base output x @ W.

Constraints

• Plain matrix products; no loops.
• PyTorch track: the hidden tests check that, with W frozen, autograd gives gradients to A and B only.

Follow-up

• Which weight matrices of a transformer do people usually adapt with LoRA, and why?
• QLoRA stores W in 4 bits. Which parts still need full precision?

Input: lora_forward(np.array([[1.0, 2.0]]), np.eye(2), np.array([[1.0], [1.0]]), np.array([[1.0, 0.0]]), alpha=2, r=1)
Output: array([[7., 2.]])

Input: x (2, 768), W (768, 768), A (768, 8), B (8, 768), alpha=16, r=8
Output: shape (2, 768)

import numpy as np

def lora_forward(x, W, A, B, alpha, r):
    """
    LoRA forward pass: x @ W + (alpha / r) * (x @ A) @ B.

    Args:
        x: Input (batch, in_features)
        W: Frozen weights (in_features, out_features)
        A: Down-projection (in_features, r)
        B: Up-projection (r, out_features)
        alpha: Scaling factor
        r: Rank

    Returns:
        Output (batch, out_features)
    """
    # Your code here
    pass