Day 65: LoRA Engineering: Weight Freezing, Gradients & Optimizer Collapse

import torch
import torch.nn as nn

class TinyLoRA(nn.Module):
def __init__(self, d_in=8, d_out=4, r=2, alpha=4):
super().__init__()
self.linear = nn.Linear(d_in, d_out, bias=False)
self.linear.weight.requires_grad_(False) # freeze W0
self.lora_A = nn.Parameter(torch.randn(r, d_in) * 0.01)
self.lora_B = nn.Parameter(torch.zeros(d_out, r)) # B = 0, so the adapter starts as a no-op
self.scale = alpha / r

def forward(self, x):
return self.linear(x) + self.scale * (x @ self.lora_A.T @ self.lora_B.T)

torch.manual_seed(0)
model = TinyLoRA()
x, target = torch.randn(5, 8), torch.randn(5, 4)
loss = nn.functional.mse_loss(model(x), target)
loss.backward()

assert model.linear.weight.grad is None # no .grad buffer was ever allocated
assert model.lora_A.grad is not None
assert model.lora_B.grad is not None

# Surprise worth recording: at step 0, dL/dA = scale * B.T @ dL/dy @ x is exactly zero because
# B = 0, while dL/dB = scale * dL/dy.T @ (x @ A.T) is not. A receives signal only after B moves.
assert torch.count_nonzero(model.lora_A.grad) == 0
assert torch.count_nonzero(model.lora_B.grad) > 0

trainable = [n for n, p in model.named_parameters() if p.requires_grad]
frozen = sum(p.numel() for p in model.parameters() if not p.requires_grad)
print(trainable, frozen) # ['lora_A', 'lora_B'] 32
# Saving: 32 frozen weights x 4 bytes = 128 bytes of fp32 gradient that is never allocated;
# for a 7B base in bf16 the same rule saves 7e9 x 2 bytes = 14 GB.

Breakdown per trainable parameter: 4 B fp32 master weight + 4 B first moment + 4 B second moment + 4 B gradient = 16 bytes.

Full fine-tuning AdamW VRAM = 1,000,000,000 * 16 bytes = 16,000,000,000 bytes = 16.0 GB (14.90 GiB).
LoRA AdamW VRAM = 10,000,000 * 16 bytes = 160,000,000 bytes = 160.0 MB (152.59 MiB).
Reduction = 1 - 160 MB / 16,000 MB = 99.0% reduction in optimizer memory.

(b) The LoRA run still holds the frozen base. In bf16 that is 1,000,000,000 * 2 bytes = 2.0 GB, with no gradient and no optimizer state. Total static memory: 2.0 GB + 0.16 GB = 2.16 GB, against 16.0 GB for full fine-tuning, which is 86.5% less. The saving is large but smaller than the optimizer-only 99.0%, because the base weights do not disappear.

(c) 7B, all 7 projections, r=16: 39,976,960 * 16 bytes = 639,631,360 bytes = 639.6 MB = 610.0 MiB. Against full fine-tuning's 112 GB this is the ~99.4% reduction. Adding the frozen 7B base in bf16 (about 14.0 GB) gives about 14.64 GB of static memory before activations.

Units: 1 GB = 10^9 bytes and 1 GiB = 2^30 bytes, so 16.0 GB is 14.90 GiB. Tools such as nvidia-smi report binary units, so always say which one you mean.

Common errors: counting only the two moments (8 bytes); forgetting the frozen base in the total; mixing GB and GiB in one comparison.

import torch
import torch.nn as nn

def count_target_params(model, targets, rank):
count = 0
for name, module in model.named_modules():
# Match the last name component exactly: 'layers.0.self_attn.q_proj' -> 'q_proj'.
# A substring test would also match names such as 'q_proj_norm'.
if isinstance(module, nn.Linear) and name.split('.')[-1] in targets:
count += rank * (module.in_features + module.out_features) # A is r x in, B is out x r
return count

class Layer(nn.Module): # the projection shapes of one LLaMA-7B decoder layer
def __init__(self, d=4096, mlp=11008):
super().__init__()
for n in ('q_proj', 'k_proj', 'v_proj', 'o_proj'):
setattr(self, n, nn.Linear(d, d, bias=False, device='meta'))
self.gate_proj = nn.Linear(d, mlp, bias=False, device='meta')
self.up_proj = nn.Linear(d, mlp, bias=False, device='meta')
self.down_proj = nn.Linear(mlp, d, bias=False, device='meta')

model = nn.ModuleList([Layer() for _ in range(32)]) # device='meta' allocates no memory
qv = count_target_params(model, {'q_proj', 'v_proj'}, 16)
all7 = count_target_params(model, {'q_proj', 'k_proj', 'v_proj', 'o_proj', 'gate_proj', 'up_proj', 'down_proj'}, 16)
print(qv, all7, round(all7 / qv, 2)) # 8388608 39976960 4.77

assert qv == 32 * 2 * 16 * (4096 + 4096) # 8,388,608
assert all7 == 32 * (4 * 16 * 8192 + 3 * 16 * (4096 + 11008)) # 39,976,960
assert count_target_params(model, {'q_proj', 'v_proj'}, 32) == 2 * qv # linear in the rank
# Per layer: 4 attention projections x 131,072 + 3 MLP projections x 241,664 = 1,249,280.
# Against a 6.74B base, all7 is 0.59% of the weights: still far below 1%.

import torch
import torch.nn as nn

class TinyLM(nn.Module):
# Frozen embedding + frozen W0 + trainable LoRA A, B: the smallest model with a real
# next-token objective. Activity 1's TinyLoRA uses the same attribute names.
def __init__(self, vocab=16, d=8, r=2, alpha=4):
super().__init__()
self.emb = nn.Embedding(vocab, d)
self.linear = nn.Linear(d, vocab, bias=False)
self.lora_A = nn.Parameter(torch.randn(r, d) * 0.01)
self.lora_B = nn.Parameter(torch.zeros(vocab, r))
self.scale = alpha / r
for p in (*self.emb.parameters(), *self.linear.parameters()):
p.requires_grad_(False)

def forward(self, idx):
h = self.emb(idx)
return self.linear(h) + self.scale * (h @ self.lora_A.T @ self.lora_B.T)

torch.manual_seed(0)
model = TinyLM()
tokens = torch.randint(0, 16, (4, 9))
x, y = tokens[:, :-1], tokens[:, 1:] # next-token pairs, Day 49

W0_initial = model.linear.weight.detach().clone()
lora_A_initial = model.lora_A.detach().clone()
lora_B_initial = model.lora_B.detach().clone()
optimizer = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=1e-1)

loss_history = []
for _ in range(10):
loss = nn.functional.cross_entropy(model(x).reshape(-1, 16), y.reshape(-1))
optimizer.zero_grad()
loss.backward()
optimizer.step()
loss_history.append(loss.item())

assert model.linear.weight.grad is None
assert torch.equal(model.linear.weight, W0_initial) # bit-for-bit, not allclose
assert not torch.equal(model.lora_A, lora_A_initial)
assert not torch.equal(model.lora_B, lora_B_initial)
assert loss_history[-1] < loss_history[0]
trainable = {id(p) for group in optimizer.param_groups for p in group['params']}
assert trainable == {id(model.lora_A), id(model.lora_B)}
print([round(v, 3) for v in loss_history])

# Mutation check: run the suite once with the freeze line deleted. The first and second asserts
# must fail; a test that still passes then is not testing the freeze.