un — Développer un Modèle de Langage : Backprop en CUDA Personnalisé [DESCRIPTION /]

un

invité

1 / ?

retour aux leçons

Les Gradients Locaux se Multiplient

Forward & Backward Kernels

Le Passage Forward

Le passage forward d'ANDREA-120M fait passer l'entrée à travers une séquence d'opérations :

x = embed(token_ids)         # embeddings des tokens
for layer in 12_layers:
x = x + attn(LN(x))      # sous-couche d'attention
x = x + mlp(LN(x))       # sous-couche MLP
logits = LN(x) @ embed.T     # projection de sortie liée
loss   = cross_entropy(logits, targets)

Chaque opération lit des tenseurs d'entrée et produit des tenseurs de sortie. Le passage avant se termine par un scalaire unique : la perte d'entropie croisée pour ce lot.

Le Passage Arrière

L'entraînement met à jour les poids dans la direction qui diminue la perte. Pour obtenir les directions de mise à jour, le moteur a besoin de :

dL/dW pour chaque W apprenable dans le modèle

La règle de la chaîne le fournit. Pour une chaîne loss = f(g(h(x))) :

dL/dx = (dL/df) * (df/dg) * (dg/dh) * (dh/dx)

Chaque facteur est un gradient local : la façon dont la sortie d'une opération change lorsque son entrée change d'une petite quantité. Multiplier les gradients locaux vers l'arrière à travers le graphe propage le signal de perte vers chaque poids.

Différentiation en Mode Inverse

Backprop calcule les gradients dans l'ordre inverse : en commençant par dL/dlogits = 1, puis en remontant à travers l'entropie croisée, puis la projection de sortie, puis la normalisation de couche, puis douze blocs transformeurs, puis les embeddings. À chaque étape, multiplier le gradient entrant par le Jacobien local.

Le mode inverse est efficace quand la sortie est un scalaire unique (la perte) & qu'il y a de nombreuses entrées (les poids). Un passage arrière produit les gradients pour chaque poids du modèle. Le mode avant nécessiterait un passage par poids ; pour ANDREA-120M avec ~120M poids, le mode avant est infaisable.

Pourquoi le Mode Inverse

ANDREA-120M a ~120M poids et produit une perte scalaire unique par étape d'entraînement. Comparez la différentiation automatique en mode reverse par rapport au mode forward. Indiquez (1) quel mode produit tous les gradients des poids en une seule passe arrière ; (2) combien de passes forward-mode seraient nécessaires pour calculer tous les 120M gradients des poids ; (3) quel mode ANDREA utilise et pourquoi.

Chaque Opération Forward Obtient un Jumeau Backward

La Discipline de Paires

microgpt_cuda.cu fournit deux kernels CUDA pour chaque opération : un qui calcule la sortie forward, un qui calcule les gradients d'entrée donnés les gradients de sortie. L'appariement est un-à-un :

Noyau forward	Noyau backward	Opération
`k_embed_fwd`	`k_embed_bwd`	Recherche d'embedding de token
`k_layernorm_fwd`	`k_layernorm_bwd`	Normalisation de couche
`k_attn_qkv_fwd`	`k_attn_qkv_bwd`	Projections Q, K, V
`k_attn_fwd`	`k_attn_bwd`	Attention produit scalaire normalisé
`k_attn_out_fwd`	`k_attn_out_bwd`	Projection de sortie W_O
`k_mlp_fwd`	`k_mlp_bwd`	MLP (avec GELU)
`k_residual_add`	`k_residual_add_bwd`	Connexion résiduelle
`k_loss_fwd`	`k_loss_bwd`	Perte d'entropie croisée

Huit paires d'opérations couvrent le transformeur complet. Plus quelques noyaux utilitaires : k_grad_norm_partial, k_grad_norm_final, k_grad_scale pour le clipping des gradients (voir activité 75).

Le rôle d'un noyau arrière

Étant donné le gradient provenant des couches ultérieures (grad_output), un noyau arrière calcule :

1. grad_input : le gradient par rapport au tenseur d'entrée de l'opération. Celui-ci est transmis plus loin en arrière.

2. grad_weight : le gradient par rapport aux paramètres apprenables dans l'opération. Celui-ci est transmis à l'état de l'optimiseur.

Les deux sont calculés dans un seul lancement de noyau. Les threads CUDA coopèrent sur des tuiles du tenseur de gradient en parallèle.

Tenseurs Sauvegardés

Le calcul arrière nécessite souvent des valeurs du passage avant. Par exemple, k_layernorm_bwd a besoin de la moyenne et de la variance calculées pendant l'avant ; k_mlp_bwd a besoin de la pré-activation GELU. Le moteur d'entraînement les stocke dans des tampons dédiés pendant l'avant, puis les lit pendant l'arrière.

Coût mémoire : à peu près la même forme que la sortie forward pour chaque tenseur sauvegardé. Pour ANDREA-120M avec batch=8, seq=1024, d_model=768, un tenseur sauvegardé fait 8 × 1024 × 768 × 4 bytes = 25 MB. Sur 12 couches & plusieurs tenseurs sauvegardés par couche, les activations dominent la VRAM pendant l'entraînement (~5-10 GB sur une carte 24 GB).

Traçage d'un pas Backward

ANDREA-120M termine un forward pass à travers un bloc transformer. Tracez ce qui se passe pendant le backward pass à travers ce même bloc (structure pre-norm : `x = x + Attention(LN(x))` puis `x = x + MLP(LN(x))`). Nommez les kernels backward dans l'ordre où ils s'exécutent, & indiquez avec quel kernel forward chacun est apparié. Couvrez au moins 4 kernels.

Où les Gradients Résident en Mémoire

Un Tenseur de Gradient par Tenseur de Poids

Chaque tenseur de poids apprenable dans ANDREA-120M possède un tenseur de gradient correspondant de forme identique. Pour chaque bloc :

W_Q       [768, 768]     ↔   grad_W_Q       [768, 768]
W_K       [768, 768]     ↔   grad_W_K       [768, 768]
W_V       [768, 768]     ↔   grad_W_V       [768, 768]
W_O       [768, 768]     ↔   grad_W_O       [768, 768]
W_1       [768, 3072]    ↔   grad_W_1       [768, 3072]
W_2       [3072, 768]    ↔   grad_W_2       [3072, 768]
LN1.gamma [768]          ↔   grad_LN1.gamma [768]
LN1.beta  [768]          ↔   grad_LN1.beta  [768]
LN2.gamma [768]          ↔   grad_LN2.gamma [768]
LN2.beta  [768]          ↔   grad_LN2.beta  [768]

Plus les embeddings de tokens, les embeddings de position, & une normalisation de couche finale. La mémoire totale du buffer de gradients correspond à la mémoire des poids : ~120M flottants, ~480 MB en FP32, ~240 MB en FP16.

Accumulation Sur Plusieurs Micro-lots

Le batch_size = 8 d'ANDREA tient en VRAM en FP16. Des lots effectifs plus grands nécessitent l'accumulation de gradients : exécuter plusieurs passes forward+backward sur de petits lots, sommer les gradients dans le même buffer, puis effectuer une étape d'optimiseur.

for microbatch in range(n_microbatches):
forward(microbatch)
backward()           # AJOUTE aux tampons de gradients, n'écrase pas
scale_grads(1.0 / n_microbatches)  # moyenne sur les microbatches
optimizer_step()
zero_grads()             # réinitialise pour l'étape d'entraînement suivante

Les noyaux de rétropropagation utilisent une sémantique +=, pas =. Chaque appel ajoute les contributions de gradient au tampon existant ; le tampon conserve la somme courante jusqu'à ce que zero_grads() le vide.

L'État de l'Optimiseur

AdamW (activité 73) contient deux buffers supplémentaires par poids : premier moment m & second moment v. Mémoire totale pendant l'entraînement :

poids :     1× nombre de poids
gradients : 1× nombre de poids
Adam m :    1× nombre de poids
Adam v:     1× nombre de poids
acts sauvés: ~2-4× selon les couches & le lot
──────────────────────────────────────────
total:      ~6-8× nombre de poids

ANDREA-120M en FP16: ~240 MB × 4 buffers (poids, grad, m, v) + ~5-10 GB activations = ~10-12 GB total. Confortablement sous le plafond de 24 GB de la RTX 4090. ANDREA-12M entraîné en 1,4 GB ; l'échelle de paramètres 10× apporte ~10× mémoire.

Dimensionnement des Buffers de Gradient

ANDREA-120M contient ~120 000 000 poids & utilise l'accumulation de gradients sur 4 micro-batchs par étape d'entraînement. Calcul : (a) taille du buffer de gradients en MB à FP16 ; (b) mémoire totale pour poids + gradients + Adam m + Adam v à FP16 ; (c) combien d'appels séparés `forward()` + `backward()` sont exécutés par étape d'entraînement. Montrez vos calculs.

Contrôle total de la mémoire & de la précision

Ce que coûtent les frameworks génériques

PyTorch & JAX rendent l'autograd pratique : écrivez du code Python, obtenez les gradients automatiquement. Le coût : une couche de dispatch générique entre votre code & CUDA. Chaque opération passe par la surcharge de l'interpréteur Python, la gestion du framework, & la sélection dynamique de kernels. Pour entraîner un petit modèle de langage sur un GPU, cette surcharge compte.

Coûts concrets évités par ANDREA :

1. Latence de l'interpréteur Python. Chaque opération PyTorch franchit la frontière Python/C++. Pour ~100 lancements de noyaux par étape d'entraînement à ~9 étapes/min, cela fait ~900 franchissements de frontière par minute. Le dispatch au niveau C élimine cela.

2. Imprévisibilité de l'alloueur du framework. L'alloueur de mise en cache de PyTorch offre un bon débit en moyenne mais une mémoire de pointe imprévisible. Le moteur d'entraînement d'ANDREA pré-alloue chaque tampon au démarrage ; pas de réallocation pendant l'entraînement, pas de fragmentation, pas d'OOM surprises à l'étape 100K.

3. Sélection de noyaux génériques. PyTorch choisit les noyaux à l'exécution via des heuristiques. ANDREA choisit les noyaux au moment de la compilation, adaptés aux tailles de tuiles des cœurs tenseurs RTX 4090.

4. Plomberie en précision mixte. Le chemin cuBLAS FP16 d'ANDREA-120M et les expériences tensor core FP8 E4M3 d'ANDREA nécessitent un contrôle précis sur les tenseurs qui vivent à quelle précision. Les frameworks génériques exposent ce contrôle via des API en couches ; les écritures CUDA personnalisées l'écrivent directement.

Le compromis

Coûts du CUDA personnalisé : plus de code à écrire, plus de bugs à trouver, pas d'écosystème communautaire. Le fichier microgpt_cuda.cu d'ANDREA fait ~6000 lignes de CUDA écrites à la main et a pris des mois à déboguer. Chaque nouvelle opération nécessite d'écrire un kernel forward, un kernel backward, & des tests.

Ce qu'ANDREA gagne :

- Reproductibilité totale. Le pipeline d'entraînement est un binaire C unique plus un proxy Python. Pas de dérive de version entre les releases PyTorch, pas d'incompatibilités de version CUDA avec les wheels du framework.

- Reprises bit-exactes. SIGTERM déclenche l'écriture d'un checkpoint qui capture chaque tenseur exactement comme le voit le GPU. La reprise reprend la même trajectoire de perte que celle sur laquelle était l'exécution.

- Mémoire prévisible. ANDREA-120M entraîné pendant 200K étapes sans OOM. La mémoire a été comptabilisée au démarrage du moteur.

- Accès direct au matériel. Tailles de tuiles des cœurs tenseurs, paramètres FP8 E4M3, copies mémoire asynchrones : tout directement adressable en CUDA, opaque dans les frameworks génériques.

La Reproductibilité Comme Mission

La section 9 du whitepaper ANDREA liste la pile complète de reproductibilité :

Moteur d'entraînement : microgpt/microgpt_cuda.cu
Proxy d'entraînement : microgpt/training_proxy.py
Configurations d'expériences : experiments/ANDREA-*-TRAIN.json
Pipeline de données : scripts/pull-hermes3.py, scripts/prep-megachat.py
Tableau de bord : scripts/live-loss-dashboard.html
Spécification Bandit : docs/FIREHOSE-BANDIT.md
Documentation du modèle : docs/ANDREA.md

Exigence matérielle : une carte NVIDIA GPU avec ≥8 GB VRAM (RTX 3060 ou mieux). N'importe qui peut reproduire ANDREA-12M à partir de ces artefacts. Le chemin CUDA personnalisé en fait partie : pas de versions de framework figées, pas de surprises de dépendances dans cinq ans.

Signaux & Points de contrôle

La boucle d'entraînement CUDA répond à deux signaux POSIX :

- SIGTERM : écrit un point de contrôle immédiat, puis sort. Utilisé pour arrêter l'entraînement proprement. ```

- SIGUSR1 : écrit un point de contrôle immédiat, continue l'entraînement. Utilisé pendant le pivot de polissage en v3 pour capturer l'état sans interrompre l'exécution.

Format du point de contrôle : [int32 step][int32 n_params][n_params × float32 weights][n_params × float32 m][n_params × float32 v]. Compteur d'étapes, nombre de poids, puis poids suivis des moments Adam. Reprend bit-exactement. Le proxy archive .samples.json & .state.json séparément sur polish ; .loss.json n'est jamais archivé (il accumule l'historique complet de l'entraînement).

Pourquoi pas PyTorch

ANDREA aurait pu utiliser l'autograd de PyTorch au lieu d'écrire `microgpt_cuda.cu` à la main. Donnez deux raisons d'ingénierie distinctes pour lesquelles ANDREA a choisi un CUDA personnalisé. Une raison doit faire référence au contrôle de la mémoire ou de la précision ; l'autre doit faire référence à la reproductibilité, aux dépendances de framework ou à la maintenance à long terme.