Das Prinzip von LoRA, von Hand geschrieben

Die vorige Lektion sagte: Muss wirklich feingetunt werden, ist LoRA die erste Wahl. Diese Lektion klärt, was es ist, und setzt es selbst am in Modul 09 trainierten kleinen GPT um.

python lora_from_scratch.py

Das Problem mit vollständigem Feintuning

Feintuning heißt, ein trainiertes Modell mit neuen Daten weiterzutrainieren. Am direktesten ist es, alle Parameter mittrainieren zu lassen; das heißt vollständiges Feintuning.

Bei großen Modellen ist das teuer. Beim Training muss für jeden Parameter außer ihm selbst auch sein Gradient gespeichert werden, und der Optimierer AdamW speichert pro Parameter noch zwei weitere Zahlen (in Modul 08, Lektion 5 erwähnt: er hält das „Momentum“ jedes Parameters fest). Ein Modell mit 7 Milliarden Parametern braucht allein dafür über hundert GB Grafikspeicher.

Außerdem erhält man bei jedem Feintuning ein vollständiges neues Modell. Feintunt man für zehn Kunden je eine Version, muss man zehn Dateien zu Dutzenden GB speichern.

Die Idee von LoRA

LoRA (Low-Rank Adaptation) beruht auf der Idee: Die Änderung, die die Gewichte beim Feintuning brauchen, muss gar nicht so „komplex“ sein.

Die Gewichte einer linearen Schicht sind eine Matrix W, etwa 128×384. Vollständiges Feintuning lernt eine gleich große Änderung ΔW und addiert sie zu W. LoRA lernt ΔW nicht direkt, sondern zerlegt es in das Produkt zweier sehr schmaler Matrizen:

      ΔW       =       A       ×     B
  (128 × 384)      (128 × 8)     (8 × 384)

  49152 个数       1024 个数   +  3072 个数  =  4096 个数

Die 8 in der Mitte heißt Rang (rank, r). A und B haben zusammen nur 4096 Zahlen, ein Zwölftel von ΔW. Je kleiner r, desto weniger Parameter sind zu lernen, desto begrenzter aber auch die ausdrückbare Änderung.

Beim Training bleibt das ursprüngliche W eingefroren, trainiert werden nur A und B:

输出 = x × W  +  x × A × B × 缩放系数
       ─────     ─────────────────────
      原来的路        新加的支路

Umsetzung

Um ein trainiertes nn.Linear zu umhüllen, braucht es nur gut zehn Zeilen:

class LoRALinear(nn.Module):
    """包住一个已经训练好的线性层：原来的权重冻结不动，旁边加一条 A、B 两个小矩阵的"支路"。"""

    def __init__(self, base: nn.Linear, rank=8, alpha=16):
        super().__init__()
        self.base = base
        for p in self.base.parameters():
            p.requires_grad = False  # 原来的参数一个都不训练
        self.A = nn.Parameter(torch.randn(base.in_features, rank) * 0.01)  # 输入维度 × r
        self.B = nn.Parameter(torch.zeros(rank, base.out_features))  # r × 输出维度，初始为 0
        self.scale = alpha / rank

    def forward(self, x):
        # 原来的输出 + 支路的输出。B 一开始是 0，所以刚加上时模型的行为和原来完全一样
        return self.base(x) + (x @ self.A @ self.B) * self.scale

Zwei Details.

B wird mit 0 initialisiert. So ist A×B anfangs 0, der Nebenzweig hat keinerlei Wirkung, und das Modell ist genau wie vorher. Das Training beginnt beim ursprünglichen Modell, nicht bei einem Modell, das ein zufälliger Nebenzweig durcheinanderbringt. A darf nicht ebenfalls 0 sein, sonst sind die Gradienten von A und B beide 0 und es lernt nie etwas.

Der Skalierungsfaktor alpha / r. Er bestimmt, wie stark der Nebenzweig wirkt. So muss man beim Ändern von r die Lernrate nicht neu einstellen. r=8 und alpha=16 ist eine sehr übliche Einstellung.

Dann bekommen qkv und proj jeder Attention-Schicht im Modell LoRA:

for p in model.parameters():
    p.requires_grad = False
for block in model.blocks:
    block.attn.qkv = LoRALinear(block.attn.qkv)
    block.attn.proj = LoRALinear(block.attn.proj)

Die Aufgabe: nur noch Jueju mit fünf Zeichen pro Zeile

Das in Modul 09 trainierte kleine GPT schreibt Gedichte in jedem Format. Jetzt wird es nur mit Jueju zu fünf Zeichen pro Zeile (五言绝句) feingetunt, um zu sehen, ob es danach nur noch diese schreibt.

原来的训练数据里，五言绝句占 10.4%（3538 首）
微调前：生成 200 首，其中五言绝句 67 首

模型共 1,639,296 个参数，要训练的只有 24,576 个（1.50%）
刚装上 LoRA 时，输出和原模型一样吗？True

Vor dem Feintuning sind 67 von 200 Gedichten Jueju mit fünf Zeichen pro Zeile. Das ist deutlich mehr als die 10,4 % in den Trainingsdaten, vermutlich weil diese Form die kürzeste ist und sich am leichtesten vollständig schreiben lässt.

Mit LoRA sind nur 24576 Parameter zu trainieren, 1,5 % des ganzen Modells. 4 Schichten mit je zwei LoRA: A und B für qkv sind 128×8 und 8×384, für proj 128×8 und 8×128; pro Schicht 6144, bei 4 Schichten 24576.

„Ist die Ausgabe direkt nach dem Einbau von LoRA gleich? True“ bestätigt die Wirkung der Initialisierung von B mit 0.

Training

Nur Jueju mit fünf Zeichen pro Zeile als Daten, 300 Schritte:

  第 1 步  损失 4.356
  第 100 步  损失 4.294
  第 200 步  损失 4.335
  第 300 步  损失 4.213
训练 300 步用了 23 秒

微调后：生成 200 首，其中五言绝句 191 首。前 5 首：
  白头还作雨，一鬓自如霜。何时之饮在，试死势悠扬。
  白发勤鸿急，秋贫梦不闲。夜寒休绕郡，秋杀杜陵陂。
  才子俭高人，何言亦重重。言知不见人，辄空一一言。
  一丛生碧霄，危棹度清湍。沙渚穿花发，苍苍隔鹤闲。
  九陌两金阁，一枝归草间。何当重枕簟，归复独裴回。

191 von 200 Gedichten sind jetzt Jueju mit fünf Zeichen pro Zeile, von 33,5 % auf 95,5 % gestiegen. Nur 1,5 % der Parameter wurden angefasst, das Training dauerte 23 Sekunden.

Interessant ist, dass sich der Verlust fast nicht ändert, von 4,36 auf 4,21. Das „Niveau“ des Modells beim Dichten (wie genau es jedes Zeichen vorhersagt) ist also nicht gestiegen; geändert hat sich seine „Gewohnheit“: Nach vier Zeilen mit fünf Zeichen hat es gelernt, einen Zeilenumbruch auszugeben und das Gedicht zu beenden. Genau solche Änderungen beherrscht LoRA: Format, Stil und Verhalten anpassen, nicht dem Modell viel neues Wissen beibringen. „Feintuning ändert Verhalten, Wissen kommt über RAG“ aus der vorigen Lektion sieht man hier sehr deutlich.

Speichern und zusammenführen

LoRA 的权重单独存下来只有 96 KB，整个模型是 6.3 MB
把 LoRA 合并回原来的权重之后，模型又变回 1,614,720 个参数，输出和合并前最大差别 5.7e-06

Der erste Vorteil von LoRA: Man muss nur A und B speichern. Hier sind das 96 KB, das Modell selbst hat 6,3 MB. Bei großen Modellen ist der Unterschied noch krasser: Ein Modell von Dutzenden GB hat womöglich LoRA-Gewichte von nur Dutzenden MB. Feintunt man für zehn Kunden je einmal, hat man ein Basismodell plus zehn kleine Dateien.

Der zweite Vorteil: Man kann zusammenführen. Nach dem Training addiert man A×B×Skalierungsfaktor direkt zum ursprünglichen W und erhält eine gewöhnliche lineare Schicht:

def merged(self):
    """把支路合并回原来的权重，得到一个普通的线性层：推理时没有任何额外开销。"""
    layer = nn.Linear(self.base.in_features, self.base.out_features)
    with torch.no_grad():
        # nn.Linear 的权重形状是 (输出, 输入)，所以要转置
        layer.weight.copy_(self.base.weight + (self.A @ self.B).T * self.scale)
        layer.bias.copy_(self.base.bias)
    return layer

Nach dem Zusammenführen sind Struktur und Parameterzahl des Modells genau wie vorher, und die Inferenz wird nicht langsamer. Die Ausgabe weicht um 5,7e-06 von der vor dem Zusammenführen ab, nur Gleitkommafehler.

Warum LoRA funktioniert

Eine intuitive Erklärung: Das vortrainierte Modell hat das meiste schon gelernt, Feintuning nimmt darauf nur einige „kleine Anpassungen“ vor. Das Paper, das LoRA 2021 vorstellte (Hu et al., „LoRA: Low-Rank Adaptation of Large Language Models“), geht genau davon aus: Die für diese Anpassung nötige Änderung ist selbst von „niedrigem Rang“ und lässt sich mit sehr kleinem r ausdrücken.

Unser Experiment passt zu dieser Erklärung: Das Modell „nach vier Zeilen aufhören“ zu lassen, ist eine sehr einfache Verhaltensänderung, und 1,5 % der Parameter reichen völlig. Will man dem Modell aber viel ganz neues Wissen beibringen, etwa den Inhalt eines ganzen neuen Fachgebiets, ist LoRA dem vollständigen Feintuning nicht unbedingt ebenbürtig.

Eine verbreitete Variante von LoRA heißt QLoRA: Das eingefrorene ursprüngliche Modell wird auf 4 Bit quantisiert (Quantisierung kommt in Lektion 4), um Grafikspeicher zu sparen, während A und B von LoRA weiter in normaler Genauigkeit trainiert werden. Das macht es möglich, Modelle mit Milliarden Parametern auf einer einzigen Consumer-Grafikkarte feinzutunen.

Übungen

1. Setz rank auf 1, 2 und 32. Wie hoch ist jeweils der Anteil der Jueju mit fünf Zeichen pro Zeile nach dem Feintuning? Wie viele Parameter sind es jeweils?
2. Ohne LoRA vollständig feintunen (alle Parameter trainieren, Lernrate 1e-4), ebenfalls 300 Schritte, und Anteil der Jueju sowie Trainingszeit vergleichen.
3. LoRA nur beim Feedforward-Netz einbauen (block.mlp[0] und block.mlp[2]), nicht bei der Attention. Wie unterscheidet sich die Wirkung?

Selbsttest

1. Welche Parameter trainiert LoRA? Was geschieht mit den ursprünglichen Parametern?

Die ursprünglichen Gewichte werden vollständig eingefroren und nicht trainiert. LoRA fügt neben ausgewählten linearen Schichten zwei kleine Matrizen A und B hinzu und trainiert nur diese. Die Ausgabe ist die ursprüngliche Ausgabe plus x×A×B×Skalierungsfaktor.

2. Warum wird B mit 0 initialisiert? Kann auch A mit 0 initialisiert werden?

Ist B 0, ist A×B 0, und das Modell direkt nach dem Einbau von LoRA ist mit dem ursprünglichen völlig identisch; das Training beginnt beim ursprünglichen Modell. A darf nicht ebenfalls 0 sein: Sind A und B beide 0, sind auch beide Gradienten 0, und die Parameter werden nie aktualisiert.

3. Wird die Inferenz nach dem LoRA-Training langsamer?

Nicht unbedingt. Nach dem Training addiert man A×B×Skalierungsfaktor zu den ursprünglichen Gewichten und führt zusammen; das entstehende Modell hat dieselbe Struktur und Parameterzahl wie das ursprüngliche und ist gleich schnell. Ohne Zusammenführen kommt die Berechnung eines Nebenzweigs hinzu, was etwas langsamer ist, dafür kann man jederzeit ein anderes LoRA einsetzen.