第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.

学習可能パラメータ 1 個あたりの内訳：fp32 マスター重み 4 B + 1 次モーメント 4 B + 2 次モーメント 4 B + 勾配 4 B = 16 バイト。

フルファインチューニングの 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）。
削減率 = 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 などのツールは 2 進単位で表示するので、どちらの単位かを必ず明記する。

よくある誤り：2 つのモーメントだけを数える（8 バイト）。合計に凍結ベースを入れ忘れる。1 つの比較の中で 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.