Le principe de LoRA, écrit à la main

La leçon précédente disait : quand un fine-tuning est vraiment nécessaire, LoRA est le premier choix. Cette leçon explique ce que c'est et l'implémente de vos mains sur le petit GPT entraîné au module 09.

python lora_from_scratch.py

Le problème du fine-tuning complet

Fine-tuner, c'est poursuivre l'entraînement d'un modèle déjà entraîné avec de nouvelles données. La méthode la plus directe fait participer tous les paramètres : c'est le fine-tuning complet.

Pour un grand modèle, c'est coûteux. À l'entraînement, chaque paramètre exige, en plus de lui-même, le stockage de son gradient, et l'optimiseur AdamW stocke encore deux nombres supplémentaires par paramètre (la leçon 5 du module 08 l'a évoqué : il garde le « moment » de chaque paramètre). Pour un modèle de 7 milliards de paramètres, cela représente à soi seul plus de cent Go de mémoire graphique.

De plus, chaque fine-tuning produit un modèle complet de plus. Pour dix clients, dix versions fine-tunées, soit dix fichiers de plusieurs dizaines de Go à stocker.

L'idée de LoRA

L'idée de LoRA (Low-Rank Adaptation, adaptation de rang faible) : lors d'un fine-tuning, la modification nécessaire des poids n'a pas besoin d'être si « complexe ».

Les poids d'une couche linéaire forment une matrice W, par exemple 128×384. Le fine-tuning complet apprend une modification ΔW de même taille, qu'il ajoute à W. LoRA n'apprend pas directement ΔW, mais le décompose en produit de deux matrices très étroites :

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

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

Le 8 du milieu s'appelle le rang (rank, r). A et B réunis ne comptent que 4096 nombres, un douzième de ΔW. Plus r est petit, moins il y a de paramètres à apprendre, mais plus la modification exprimable est limitée.

À l'entraînement, le W d'origine est gelé, et seuls A et B sont entraînés :

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

Implémentation

Envelopper un nn.Linear déjà entraîné ne demande qu'une dizaine de lignes :

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

Deux détails.

B est initialisé à 0. Ainsi, A×B vaut 0 au départ, la dérivation n'a aucun effet, et le modèle est identique à l'original. L'entraînement part du modèle d'origine, et non d'un modèle perturbé par une dérivation aléatoire. A ne peut pas être aussi à 0 : sinon, les gradients de A et de B seraient tous deux nuls, et rien n'apprendrait jamais.

Le facteur d'échelle alpha / r. Il règle l'influence de la dérivation. Ainsi, en changeant r, on n'a pas besoin de réajuster le taux d'apprentissage. r=8 et alpha=16 est un réglage très courant.

On équipe ensuite de LoRA les qkv et proj de chaque couche d'attention du modèle :

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)

La tâche : ne plus écrire que des jueju à cinq caractères

Le petit GPT entraîné au module 09 écrit des poèmes de tous formats. Fine-tunons-le uniquement sur des jueju à cinq caractères par vers, pour voir s'il se met à n'écrire plus que cela.

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

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

Avant le fine-tuning, 67 poèmes sur 200 sont des jueju à cinq caractères. C'est bien plus que les 10,4 % des données d'entraînement, sans doute parce que c'est la forme la plus courte et la plus facile à terminer.

Une fois LoRA en place, il n'y a que 24 576 paramètres à entraîner, soit 1,5 % du modèle entier. 4 couches, deux LoRA par couche : pour qkv, A et B font 128×8 et 8×384, pour proj 128×8 et 8×128, soit 6144 par couche et 24 576 pour 4 couches.

« La sortie est-elle identique juste après l'ajout de LoRA ? True » confirme l'effet de l'initialisation de B à 0.

Entraînement

Avec pour seules données des jueju à cinq caractères, 300 pas :

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

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

191 poèmes sur 200 sont des jueju à cinq caractères, soit une hausse de 33,5 % à 95,5 %. Seuls 1,5 % des paramètres ont été touchés, pour 23 secondes d'entraînement.

Fait intéressant, la perte n'a presque pas changé, de 4,36 à 4,21. Cela montre que le « niveau » poétique du modèle (la justesse de sa prédiction de chaque caractère) n'a pas progressé ; ce qui a changé, c'est son « habitude » : après quatre vers de cinq caractères, il a appris à produire un retour à la ligne et à terminer le poème. C'est précisément ce genre de changement que LoRA réussit : ajuster format, style et comportement, et non enseigner au modèle beaucoup de connaissances nouvelles. « Le fine-tuning change le comportement, les connaissances viennent du RAG », disait la leçon précédente ; on le voit très clairement ici.

Enregistrer et fusionner

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

Premier avantage de LoRA : il suffit d'enregistrer A et B. Ici 96 Ko, contre 6,3 Mo pour le modèle lui-même. Sur un grand modèle, l'écart est encore plus criant : pour un modèle de plusieurs dizaines de Go, les poids LoRA ne font peut-être que quelques dizaines de Mo. Un fine-tuning par client pour dix clients, ce sont un modèle de base et dix petits fichiers.

Deuxième avantage : on peut fusionner. À la fin de l'entraînement, on ajoute directement A×B×facteur d'échelle au W d'origine, et l'on obtient une couche linéaire ordinaire :

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

Après fusion, structure et nombre de paramètres du modèle sont exactement ceux d'origine, et l'inférence n'est pas ralentie. L'écart de sortie avec l'état avant fusion est de 5,7e-06, une simple erreur de virgule flottante.

Pourquoi LoRA fonctionne

Une explication intuitive : le modèle pré-entraîné a déjà appris l'essentiel, et le fine-tuning n'y apporte que de « petits ajustements ». L'article de 2021 qui a introduit LoRA (Hu et al., « LoRA: Low-Rank Adaptation of Large Language Models ») part précisément de là : la modification nécessaire à cet ajustement est elle-même « de rang faible » et s'exprime avec un très petit r.

Notre expérience va dans ce sens : faire « terminer après quatre vers » est un changement de comportement très simple, et 1,5 % des paramètres suffisent amplement. Mais pour enseigner au modèle beaucoup de connaissances entièrement nouvelles, par exemple tout un nouveau domaine, LoRA n'égale pas forcément le fine-tuning complet.

LoRA a une variante courante, QLoRA : le modèle d'origine gelé est quantifié en 4 bits (la quantification est vue à la leçon 4) pour économiser la mémoire graphique, tandis que les A et B de LoRA sont entraînés en précision normale. Cela rend possible le fine-tuning de modèles de plusieurs milliards de paramètres sur une seule carte graphique grand public.

Exercices

1. Passez rank à 1, 2 et 32. Quelle est la proportion de jueju à cinq caractères après fine-tuning dans chaque cas ? Et le nombre de paramètres ?
2. Sans LoRA, faites un fine-tuning complet (tous les paramètres entraînés, taux d'apprentissage à 1e-4), également sur 300 pas, et comparez la proportion de jueju et le temps d'entraînement.
3. N'ajoutez LoRA qu'au réseau à propagation avant (block.mlp[0] et block.mlp[2]), pas à l'attention. En quoi le résultat diffère-t-il ?

Auto-test

1. Quels paramètres LoRA entraîne-t-il ? Que deviennent les paramètres d'origine ?

Les poids d'origine sont tous gelés et ne participent pas à l'entraînement. LoRA ajoute deux petites matrices A et B à côté de couches linéaires choisies, et n'entraîne qu'elles. La sortie est la sortie d'origine plus x×A×B×facteur d'échelle.

2. Pourquoi B est-il initialisé à 0 ? A peut-il l'être aussi ?

Avec B à 0, A×B vaut 0, et le modèle juste équipé de LoRA est identique à l'original : l'entraînement part du modèle d'origine. A ne peut pas être aussi à 0 : si A et B valent tous deux 0, leurs gradients sont nuls et les paramètres ne sont jamais mis à jour.

3. Après un entraînement LoRA, l'inférence est-elle ralentie ?

Elle peut ne pas l'être. Après l'entraînement, on ajoute A×B×facteur d'échelle aux poids d'origine pour fusionner ; le modèle obtenu a la même structure et le même nombre de paramètres que l'original, et la même vitesse d'inférence. Sans fusion, le calcul d'une dérivation s'ajoute, ce qui est un peu plus lent, mais permet de changer de LoRA à tout moment.