Fine-tuning with LoRA
Fine-tuning adapts a trained model to a new task. Updating every weight is costly, so LoRA freezes the original weights and learns a small correction beside each layer: two thin matrices whose product has the layer's shape.
import torch
out_f, in_f, r = 6, 4, 2
A = torch.randn(r, in_f) * 0.01
B = torch.zeros(out_f, r)
print((B @ A).shape) # torch.Size([6, 4])
print(r * (in_f + out_f), in_f * out_f)B starts at zero, so at first the adapted layer behaves exactly like the original. Training changes only A and B, a fraction of the weights, and the frozen base can be shared by many adapters.
Your task
Write LoRALinear(base, rank, alpha=1.0): keep the nn.Linear base and freeze its parameters, add A (rank by in_features, random times 0.01) and B (out_features by rank, zeros) as nn.Parameters. forward(x) returns base(x) + (x @ A.T @ B.T) * (alpha / rank). Write trainable(model) returning how many parameter values require gradients.
The explanation opens here once every check passes.