第 65 天：LoRA 工程實現：基座權重凍結、梯度路由與顯存雪崩

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.

每個可訓練參數的拆解：4 B fp32 主權重 + 4 B 一階動量 + 4 B 二階動量 + 4 B 梯度 = 16 位元組。

全量微調 AdamW 顯存 = 10^9 * 16 位元組 = 16,000,000,000 位元組 = 16.0 GB（14.90 GiB）。
LoRA 微調 AdamW 顯存 = 10^7 * 16 位元組 = 160,000,000 位元組 = 160.0 MB（152.59 MiB）。
顯存削減幅度 = 1 - 160 MB / 16,000 MB = 99.0%。

(b) LoRA 執行仍然持有凍結基座。bf16 下為 1,000,000,000 * 2 位元組 = 2.0 GB，沒有梯度也沒有優化器狀態。整體靜態顯存：2.0 GB + 0.16 GB = 2.16 GB，而全量微調為 16.0 GB，少 86.5%。節省很大，但由於基座權重不會消失，比只看優化器時的 99.0% 要小。

(c) 7B、全部 7 個投影、r=16：39,976,960 * 16 位元組 = 639,631,360 位元組 = 639.6 MB = 610.0 MiB。與全量微調的 112 GB 相比，這就是約 99.4% 的削減。加上 bf16 的凍結 7B 基座（約 14.0 GB），不含激活的靜態顯存約為 14.64 GB。

單位：1 GB = 10^9 位元組，1 GiB = 2^30 位元組，所以 16.0 GB 是 14.90 GiB。nvidia-smi 等工具使用二進位單位，所以一定要說明用的是哪種。

常見錯誤：只算兩個動量（8 位元組）；總量中忘記凍結基座；在同一個比較中混用 GB 和 GiB。

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.