File size: 5,786 Bytes
2538ecc e0e8238 2538ecc e0e8238 2538ecc e0e8238 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 | #!/usr/bin/env python3
"""Generate LoRA notebook."""
import nbformat as nbf
nb = nbf.v4.new_notebook()
nb.metadata = {"kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"},"language_info": {"name": "python", "version": "3.12.0"}}
cells = []
def md(s): cells.append(nbf.v4.new_markdown_cell(s))
def code(s): cells.append(nbf.v4.new_code_cell(s))
md("# LoRA: Low-Rank Adaptation\n\nParameter-efficient fine-tuning of GPT via low-rank weight updates.")
md("""## 背景
LoRA（Hu et al. 2021）解决大模型全量微调成本过高的问题。核心思想：
冻结预训练权重 $W_0 \\in \\mathbb{R}^{d \\times k}$，学习一对低秩矩阵 $B \\in \\mathbb{R}^{d \\times r}, A \\in \\mathbb{R}^{r \\times k}$，
其中 $r \\ll \\min(d, k)$。
$$h = W_0 x + \\Delta W x = W_0 x + \\frac{\\alpha}{r} BAx$$
| 方法 | 可训练参数量 | 对 GPT (5.7M) |
|---|---|---|
| Full fine-tune | $d \\times k$ | 5,693,952 |
| LoRA r=8 | $(d + k) \\times r$ | **32,772** (0.58%) |
本项目在已训练好的本地 GPT（text8）上注入 LoRA，仅微调注意力层的 Q/K/V/O 投影矩阵。
""")
md("""## 数学原理
### 低秩分解
$$\\Delta W = BA, \\quad B \\in \\mathbb{R}^{d \\times r}, A \\in \\mathbb{R}^{r \\times k}$$
推理时将 $\\Delta W$ 合并回原始权重，零额外开销：
$$W_{\\text{merged}} = W_0 + \\frac{\\alpha}{r} BA$$
### 参数效率
$$\\text{Ratio} = \\frac{r(d + k)}{d \\cdot k} \\approx \\frac{r}{k} + \\frac{r}{d}$$
当 $r=8, d=k=256$ 时：
$$\\text{Ratio} = \\frac{8 \\times 512}{256 \\times 256} = \\frac{4096}{65536} \\approx 6.25\\%$$
""")
code("""\
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from datasets import load_dataset
from nlp.gpt.model import GPT
from nlp.gpt.tokenizer import WordTokenizer
from nlp.lora.model import LoRALayer, inject_lora, freeze_all_except_lora, lora_params_count
from utils.config import load_config
from utils.seed import set_seed
from utils.device import get_device
device = get_device()
print(f"Device: {device}")
""")
code("""\
# Load pretrained GPT
checkpoint = "nlp/gpt/gpt_text8.pt"
model = torch.load(checkpoint, map_location="cpu", weights_only=False).to(device)
full_params = sum(p.numel() for p in model.parameters())
print(f"Full model: {full_params:,} params")
# Compare different ranks
print(f"\\n{'r':>3} {'LoRA params':>12} {'Ratio':>8}")
print("-" * 25)
for r in [1, 4, 8, 16]:
# Quick count: each of 4 Linear × 4 layers × (d*r + r*d) = 4×4×2×d×r
d = 256
lora_count = 4 * 4 * 2 * d * r
print(f"{r:3d} {lora_count:>10,d} {lora_count/full_params:>7.2%}")
""")
code("""\
# Inject LoRA with r=8
freeze_all_except_lora(model)
model = inject_lora(model, r=8, alpha=16)
lora_count = lora_params_count(model)
trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"LoRA params: {lora_count:,} ({lora_count/full_params:.2%})")
print(f"Trainable: {trainable:,}")
""")
md("""## 训练
> ⏱ 预估耗时：**10 epoch × ~60s/epoch ≈ 10 分钟**（M4 Max, batch_size=32）
""")
code("""\
# Load text8 subset
ds = load_dataset("afmck/text8", split="train")
raw = ds[0]["text"]
words = raw.lower().split()
raw_chunks = [" ".join(words[i:i + 32]) for i in range(0, len(words), 32)]
tokenizer = WordTokenizer(vocab_size=5000)
tokenizer.build_vocab(raw_chunks)
sentences = [s.strip() for s in raw_chunks if len(s.strip()) > 5][:5000]
class TextDataset(torch.utils.data.Dataset):
def __init__(self, texts, tokenizer, max_len=64):
self.examples = []
for text in texts:
tokens, mask = tokenizer.encode(text, max_len)
self.examples.append({"input_ids": torch.tensor(tokens, dtype=torch.long)})
def __len__(self):
return len(self.examples)
def __getitem__(self, idx):
return self.examples[idx]
dataset = TextDataset(sentences, tokenizer, max_len=64)
loader = DataLoader(dataset, batch_size=32, shuffle=True)
print(f"Training chunks: {len(sentences):,}")
""")
code("""\
NUM_EPOCHS = 10
LR = 0.001
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=LR)
loss_hist, ppl_hist = [], []
for epoch in range(1, NUM_EPOCHS + 1):
model.train()
total_loss = 0.0
for batch in loader:
input_ids = batch["input_ids"].to(device)
labels = input_ids[:, 1:].contiguous()
inputs = input_ids[:, :-1].contiguous()
optimizer.zero_grad()
logits, _ = model(inputs)
loss = criterion(logits.view(-1, logits.size(-1)), labels.view(-1))
loss.backward(); optimizer.step()
total_loss += loss.item()
avg = total_loss / len(loader)
ppl = __import__('math').exp(avg)
loss_hist.append(avg)
ppl_hist.append(ppl)
print(f"Epoch [{epoch:2d}/{NUM_EPOCHS}] Loss: {avg:.4f} PPL: {ppl:.2f}")
""")
md("""## Loss 曲线""")
code("""\
import matplotlib.pyplot as plt
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))
ax1.plot(loss_hist, marker='o'); ax1.set_xlabel("Epoch"); ax1.set_ylabel("Loss"); ax1.grid(True)
ax2.plot(ppl_hist, marker='o', color='orange'); ax2.set_xlabel("Epoch"); ax2.set_ylabel("Perplexity"); ax2.grid(True)
plt.tight_layout(); plt.show()
""")
md("""\
## 思考题
1. LoRA 的 rank $r$ 越大越好还是越小越好？分析过拟合和表达能力之间的权衡。
2. 为什么只修改注意力层的 Q/K/V/O 而不修改 FFN？
3. 推理时如何将 LoRA 权重合并回原模型，实现零额外开销？
4. 尝试将 $\\alpha$ 从 16 改到 4 或 64，观察训练速度和 loss 的变化。
""")
nb.cells = cells
with open("nlp/lora/lora.ipynb", "w") as f:
nbf.write(nb, f)
print("Generated nlp/lora/lora.ipynb")
|