Fine-tuning avec LoRA

Tu as un modèle. Il fonctionne, plus ou moins. Tu veux le spécialiser : meilleur en code, dans un dialecte précis, ou pour produire du JSON. L’approche naïve consiste à continuer l’entraînement sur les nouvelles données en mettant à jour tous les paramètres.

Ça marche, mais c’est cher. Fine-tuner un modèle d’un milliard de paramètres pour une petite tâche coûte beaucoup de GPU et produit un checkpoint énorme par tâche. La plupart des poids changent à peine.

LoRA rend le fine-tuning bon marché. Au lieu de modifier W, on le fige. On ajoute une petite correction basse-rang A · B et on n’entraîne que cette correction :

W est [d × d], A est [d × r], B est [r × d], avec r ≪ d. Les paramètres entraînables passent de d² à 2 · d · r — un ratio d’économie de d / (2r). Le gain s’amplifie spectaculairement quand les modèles grossissent :

• Ton modèle du ch.12 (n_embd = 128, r = 8) : chaque nn.Linear carré passe de 16 384 → 2 048 — soit 8× moins.
• GPT-2 small (d = 1024, r = 8) : 1 048 576 → 16 384 — soit 64× moins.
• Échelle GPT-3 (d = 4 096, r = 8) : 16 777 216 → 65 536 — soit 256× moins.

LoRA prend tout son sens quand les modèles sont gros. Sur ton modèle local de 14M paramètres, envelopper chaque nn.Linear ramène quand même le compte entraînable à environ 65k — soit ~200× moins de paramètres entraînables que le modèle complet. Le gain relatif ne fait que croître quand tu montes en taille.

1. La passe avant LoRA

Écris la formule. La cellule reçoit une matrice W figée, des matrices basses-rang A et B, et un facteur alpha.

// 'x' is the input vector (length D).
// 'W' is the FROZEN pretrained matrix (D x D).
// 'A' is [D x R], 'B' is [R x D]. These are the only trainable params.
// 'alpha' is the LoRA scaling factor.
//
// LoRA forward:  y = x · (W + alpha · A · B)
//
// Compute it. (The order matters for performance — multiply x by W
// directly, and x by (A·B) separately. Same result, less memory.)
return new Array(W.length).fill(0);

La sortie LoRA peut imiter un full fine-tune si la mise à jour complète est approximativement basse-rang. Le papier LoRA montre que, pour beaucoup de tâches d’adaptation, cette hypothèse tient très bien.

2. Pourquoi ça marche

Ce qu’on adapte est souvent un déplacement étroit : parler dans un autre registre, suivre une instruction, citer une source, respecter un format. Ça ne demande pas forcément des millions de dimensions nouvelles. Quelques dimensions de “delta” suffisent.

Vois W comme la compétence de base du modèle, et α · A · B comme un overlay par tâche. Quand tu changes de tâche, tu changes d’overlay ; le modèle de base reste intact.

3. La version Python

Le pattern PyTorch consiste à envelopper les nn.Linear existants dans une version LoraLinear. Sauvegarde llm/lora.py :

"""llm/lora.py — LoRA adapter wrapper for nn.Linear."""
import math
import torch
import torch.nn as nn
 
class LoraLinear(nn.Module):
    """Wraps an existing nn.Linear with frozen weights + a trainable low-rank update."""
    def __init__(self, base: nn.Linear, r: int = 8, alpha: float = 16.0):
        super().__init__()
        self.base = base
        for p in self.base.parameters():
            p.requires_grad = False  # freeze
        self.r = r
        self.alpha = alpha
        d_in = base.in_features
        d_out = base.out_features
        self.A = nn.Parameter(torch.zeros(d_in, r))
        nn.init.kaiming_uniform_(self.A, a=math.sqrt(5))
        self.B = nn.Parameter(torch.zeros(r, d_out))  # zero init: model starts at base behavior
 
    def forward(self, x):
        return self.base(x) + self.alpha * (x @ self.A @ self.B)
 
def apply_lora(model, r=8, alpha=16):
    """Replace every nn.Linear inside the model with a LoraLinear wrapper."""
    for name, module in model.named_children():
        if isinstance(module, nn.Linear):
            setattr(model, name, LoraLinear(module, r=r, alpha=alpha))
        else:
            apply_lora(module, r=r, alpha=alpha)
    return model

Lis ce wrapper comme “comportement de base + petite correction apprise” :

• self.base = base garde la couche originale.
• requires_grad = False fige ses poids.
• A projette vers le rang r.
• B reprojette vers la sortie.
• x @ self.A @ self.B est le delta basse-rang.
• apply_lora parcourt le modèle et remplace les nn.Linear.

Deux détails :

• B initialisé à zéro : au départ, le modèle produit exactement la sortie du modèle de base.
• A initialisé avec Kaiming : petite init aléatoire prête pour les gradients.

Pour utiliser LoRA sur ton modèle du chapitre 13, modifie le setup d’entraînement :

from llm.model import GPT, GPTConfig
from llm.lora import apply_lora
 
# [1]
model = GPT(GPTConfig())
# [2]
model.load_state_dict(torch.load("checkpoints/model.pt"))
# [3]
model = apply_lora(model, r=8, alpha=16)
 
# only A and B are trainable now
# [4]
trainable = [p for p in model.parameters() if p.requires_grad]
print(f"trainable params: {sum(p.numel() for p in trainable):,}")
# ... rest of the training loop, using `trainable` for the optimizer

• [1] recrée l’architecture.
• [2] charge le checkpoint de base.
• [3] ajoute les adaptateurs.
• [4] donne à l’optimiseur seulement les paramètres entraînables.

3b. Qu’est-ce que le rang t’achète vraiment ?

Le rang est le seul bouton que tu vas régler, et son coût se calcule avant de dépenser quoi que ce soit. Calcule le nombre de paramètres entraînables et la taille du checkpoint pour une gamme de rangs sur GPT-2 small.

// GPT-2 small: 12 blocks, n_embd = 768, 124M parameters.
// You adapt two matrices per block (the qkv projection and the output
// projection). Each adapted matrix adds A [nEmbd x rank] and B [rank x nEmbd].
//
// For each rank in [1, 2, 4, 8, 16, 32, 64], return an object:
//   { rank, trainable, checkpointMb, percentOfFull }
// checkpointMb assumes float32 (4 bytes per parameter).
const nLayer = 12;
const nEmbd = 768;
const fullParams = 124000000;

return [];

Deux choses doivent sauter aux yeux. D’abord, la relation est linéaire en rang : doubler r double l’adaptateur, et rien d’autre ne change. Ensuite, même le rang 64 reste une erreur d’arrondi à côté du modèle de base gelé. Cette asymétrie est tout l’argument de LoRA : tu peux garder une douzaine de spécialisations sur disque pour moins cher qu’une seule copie complète.

4. Régime checkpoint

Un LoRA d’un modèle 1B avec r = 8 produit un checkpoint de quelques dizaines de Mo au lieu de plusieurs Go. Tu peux servir un modèle de base + des dizaines d’adaptateurs.

C’est pourquoi LoRA domine le fine-tuning open-source : on partage les deltas, pas tout le modèle.

Pour aller plus loin

• Hu et al., “LoRA”.
• Hugging Face PEFT.
• QLoRA.

Prochaine étape : quantification — LoRA réduit le coût d’adaptation ; la quantification réduit le coût d’inférence.