PythiloLesson 8 of 9 · progress kept in this browser

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.

This lesson uses torch, which Python in this browser does not include. Read it here, and run and check it in the Pythilo app on iPhone or iPad.