🤖 machine learning

FORGE: Fused On-Register Gradient Elimination for Memory-Efficient LLM Training

FORGE est une méthode d'entraînement de LLM économisant la mémoire qui fusionne l'étape de l'optimiseur dans la passe arrière pour éliminer les gradients matérialisés en les traitant entièrement dans les registres, réduisant ainsi de moitié l'utilisation de la mémoire de l'optimiseur, accélérant l'entraînement et préservant la fidélité en pleine précision sans modifier la logique de mise à jour sous-jacente.

Auteurs originaux : Dikshant Kukreja, Kritarth Prasad, Avinash Anand, Zhengkui Wang, Erik Cambria, Timothy Liu, Aik Beng Ng, Simon See, Bapi Chatterjee

Publié 2026-06-23

📖 5 min de lecture🧠 Analyse approfondie

CC BY 4.0

Auteurs originaux : Dikshant Kukreja, Kritarth Prasad, Avinash Anand, Zhengkui Wang, Erik Cambria, Timothy Liu, Aik Beng Ng, Simon See, Bapi Chatterjee

Article original sous licence CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Ceci est une explication générée par l'IA de l'article ci-dessous. Elle n'a pas été rédigée ni approuvée par les auteurs. Pour une précision technique, consultez l'article original. Lire la clause de non-responsabilité complète

Le gros problème : L'« atelier encombré »

Imaginez que vous entraînez un modèle d'IA géant (comme un cerveau de robot) sur un ordinateur. Pour enseigner à ce robot, l'ordinateur doit effectuer une quantité massive de calculs mathématiques.

Dans la méthode standard pour faire cela (appelée « différenciation en mode inverse »), l'ordinateur suit un processus strict en deux étapes :

Calculer les erreurs : Il examine les données, détermine où le robot s'est trompé, et rédige une liste géante de « notes de correction » (gradients) pour chaque partie du cerveau du robot. Il écrit ces notes sur un immense tableau blanc (la mémoire de l'ordinateur).
Appliquer les corrections : Ce n'est qu'une fois que le tableau blanc est entièrement rempli que l'ordinateur prend une gomme et un marqueur pour réellement mettre à jour le cerveau du robot en fonction de ces notes.

Le goulot d'étranglement : Le problème est que l'ordinateur doit conserver ce gigantesque tableau blanc de notes dans sa mémoire pendant qu'il contient également le cerveau du robot et les notes pour l'étape suivante. Ce « tableau blanc » prend tellement de place qu'il remplit souvent la mémoire de l'ordinateur avant même que l'entraînement puisse commencer. C'est comme essayer de réparer une voiture dans un garage, mais devoir garder l'intégralité du plan de la voiture, les outils et le manuel de réparation étalés sur le sol en même temps. Si le garage est trop petit, vous ne pouvez pas faire entrer la voiture.

La solution : FORGE (La méthode de l'« instantané »)

Les auteurs de cet article, FORGE (Fused On-Register Gradient Elimination), disent : « Pourquoi écrire les notes sur le tableau blanc ? »

Ils soutiennent que la liste géante de notes n'est qu'un artefact de la manière dont nous faisons les mathématiques, et non quelque chose dont le processus d'apprentissage a réellement besoin.

Comment fonctionne FORGE :
Au lieu d'écrire les notes puis de les relire, FORGE effectue le calcul et la correction en même temps, dans la zone de stockage la plus rapide et la plus petite de l'ordinateur (appelée « registres »).

L'analogie : Imaginez un chef cuisinier préparant un plat complexe.
- L'ancienne méthode : Le chef coupe un oignon, écrit « oignon coupé » sur un bloc-notes, met l'oignon dans un bol, puis passe à l'ingrédient suivant. Une fois que tous les ingrédients sont découpés et listés, le chef lit son bloc-notes et mélange le tout. Le bloc-notes prend de la place sur le plan de travail.
- La méthode FORGE : Le chef coupe un oignon et le jette immédiatement dans la marmite. Ensuite, il coupe une carotte et la jette immédiatement. Il n'écrit jamais rien. Il n'a jamais besoin d'un bloc-notes. Le plan de travail reste dégagé.

Pourquoi c'est une avancée majeure

L'article revendique trois avantages majeurs découlant de cette approche « sans notes » :

1. Cela économise énormément d'espace (Mémoire)
Parce que l'ordinateur n'écrit jamais la liste géante de notes dans la mémoire principale, cela libère un espace considérable.

Le résultat : Sur une puce informatique standard (H200), ils ont été capables d'entraîner un modèle de 8 milliards de paramètres en utilisant 53 % de mémoire en moins.
L'analogie : C'est comme pouvoir faire tenir une table de salle à manger pour 10 personnes dans un studio parce que vous avez arrêté de stocker les chaises supplémentaires dans le couloir.

2. C'est réellement plus rapide
Parce que l'ordinateur n'a pas besoin de s'arrêter pour écrire des notes, attendre, puis les relire, l'ensemble du processus est accéléré.

Le résultat : Les étapes d'entraînement sont environ 1,5 fois plus rapides.
L'analogie : C'est la différence entre un livreur qui dépose un colis, revient en courant au camion pour chercher le suivant, et recommence (Ancienne méthode), contre un livreur qui dispose d'un tapis roulant qui charge directement le colis sur le camion pendant qu'il le saisit (FORGE).

3. Cela ne perd pas en précision
Généralement, lorsque l'on essaie d'économiser de l'espace en sautant des étapes, on perd en précision (le robot apprend un peu moins bien). Les auteurs prouvent que, comme ils effectuent les calculs dans les « registres » de plus haute qualité de l'ordinateur (pleine précision) avant de rejeter les données, l'apprentissage est mathématiquement identique à l'ancienne méthode lente.

Le résultat : Le robot apprend tout aussi bien, mais de manière beaucoup plus efficace.

L'astuce du « Tuile » (Tile)

Comment font-ils cela sans chaos ? Ils décomposent le problème mathématique géant en petits morceaux gérables appelés « tuiles » (comme des carrés de 128x128).

Ils traitent une tuile : Calculent l'erreur, corrigent le cerveau, et jettent l'erreur immédiatement.
Puis ils passent à la tuile suivante.
Comme ils procèdent ainsi, tuile par tuile, l'ordinateur n'a jamais besoin de contenir la liste entière des erreurs à la fois.

Ce que cela nous permet de faire

L'article montre qu'avec FORGE :

Vous pouvez entraîner des modèles plus grands sur le même ordinateur.
Vous pouvez utiliser des « lots » (batches) plus grands (enseigner plus d'exemples au robot à la fois) sans manquer de mémoire.
Dans un test spécifique, ils ont pu entraîner un modèle sur quatre GPU qui aurait nécessité huit GPU avec l'ancienne méthode.

Résumé

FORGE est une nouvelle façon d'entraîner l'IA qui empêche l'ordinateur d'encombrer sa mémoire avec des « notes de correction » temporaires. En calculant les erreurs et en corrigeant le modèle instantanément dans la partie la plus rapide de la puce, cela réduit l'utilisation de la mémoire de moitié et accélère l'entraînement, le tout sans rendre l'IA moins intelligente. Cela transforme un atelier encombré et lent en une ligne d'assemblage rationalisée et à haute vitesse.

Résumé Technique : FORGE (Fused On-Register Gradient Elimination)

1. Le Problème : Le Plancher de Gradient

Les méthodes actuelles d'entraînement de modèles de langage étendus (LLM) économes en mémoire ciblent principalement trois des quatre pools de mémoire sur un GPU : les états de l'optimiseur (via la compression), les activations (via la recomputation) et les poids du modèle. Cependant, elles laissent intact le quatrième pool — le gradient de poids matérialisé.

Dans une différenciation standard en mode inverse, la passe arrière calcule le tenseur de gradient complet ( $\nabla W$ ) pour chaque couche, l'écrit dans la mémoire à haute bande passante (HBM), puis l'optimiseur le relit pour mettre à jour les poids. Ce cycle à deux phases crée un « plancher de gradient » où, à la frontière entre la passe arrière et l'étape de l'optimiseur, les gradients de toutes les couches doivent coexister en mémoire. Pour un modèle de 8 milliards de paramètres, ce pool de gradient transitoire consomme à lui seul environ 16 Gio, dictant souvent si une configuration d'entraînement peut tenir sur un seul GPU. Les méthodes existantes qui fusionnent les étapes de rétropropagation et d'optimisation (par exemple, LOMO, FlashOptim) opèrent à la granularité de paramètres entiers ou de couches entières, ce qui signifie que le tenseur de gradient complet pour un paramètre est toujours matérialisé en mémoire avant que la mise à jour ne se produise.

2. Méthodologie : FORGE

FORGE (Fused On-Register Gradient Elimination) élimine entièrement le tenseur de gradient matérialisé en fusionnant l'étape de l'optimiseur directement dans le calcul du gradient de poids à une granularité de tuile (spécifiquement des tuiles de $128 \times 128$ ).

Mécanisme Central

Au lieu de calculer la matrice de gradient complète et de la stocker, FORGE traite la matrice de poids une tuile à la fois :

Accumulation de Gradient par Tuile : Pour une tuile de poids spécifique, le noyau accumule les produits partiels du gradient amont ( $\Delta Y$ ) et de l'activation d'entrée ( $X$ ) directement dans les registres FP32 du GPU.
Optimisation en Registre : Une fois que le gradient de la tuile est entièrement accumulé dans les registres, la mise à jour de l'optimiseur (par exemple, AdamW) est appliquée immédiatement à cette tuile, également au sein des registres.
Abandon Immédiat : La tuile de gradient est consommée dès l'instant où elle est produite. Elle n'est jamais écrite dans l'HBM sous forme de tenseur.
Réécriture (Write-Back) : Seules la tuile de poids mise à jour et les tuiles d'états de l'optimiseur (moments, optionnellement quantifiées) sont réécrites dans la mémoire globale.

Précision et Fidélité

Exactitude : Pour tout optimiseur élément par élément séparable (où la mise à jour d'un poids dépend uniquement de la coordonnée de ce poids, ex: AdamW, SGD, Lion), la mise à jour par tuile est mathématiquement identique à une étape de gradient complet standard. La seule déviation est l'ordre de la sommation en virgule flottante sur les tokens, ce qui est inhérent à toute implémentation de GEMM et ne constitue pas une approximation de la logique de l'optimiseur.
Fidélité en Basse Précision : Dans les régimes d'entraînement BF16 ou INT8, FORGE préserve la fidélité de pleine précision de deux manières :
- Le gradient n'est jamais tronqué en BF16 avant l'optimisation ; il est calculé et consommé dans des registres FP32.
- Pour éviter la « zone morte » du BF16 (où les petites mises à jour sont arrondies à zéro), FORGE utilise l'arrondi stochastique pour les réécritures de poids, assurant une accumulation non biaisée des mises à jour de sous-ULP.
- Pour les états d'optimiseur INT8, FORGE utilise des blocs plus petits (64 contre 2048 dans les optimiseurs 8-bit standards) pour minimiser l'erreur de quantification.

Compatibilité avec l'Entraînement Distribué

FORGE est prouvé exact sous Parallélisme de Tenseur (TP) et Parallélisme de Séquence (SP) car le gradient local pour un fragment est complet au moment de la rétropropagation. Cependant, il est incompatible avec le Parallélisme de Données (DP) ou le Parallélisme de Contexte dans sa forme actuelle, car ceux-ci nécessitent un all-reduce des gradients partiels avant que la mise à jour puisse être appliquée, ce qui nécessiterait de matérialiser le tenseur de gradient.

3. Contributions Clés

Élimination Structurelle, Non Compression : FORGE identifie le gradient matérialisé comme un artefact de l'ordonnancement de l'entraînement plutôt que comme une nécessité mathématique. Il élimine le terme de mémoire de gradient $2P$ (nombre de paramètres), réduisant le pool de gradient de $O(P)$ à $O(T_R T_C)$ (taille de la tuile), indépendamment de la taille du modèle.
Exactitude Prouvée : L'article prouve que pour les optimiseurs élément par élément, le schéma fusionné par tuile produit des mises à jour identiques à un AdamW standard sur des gradients FP32. Cette exactitude est maintenue sous parallélisme de tenseur et de séquence.
Fidélité de Pleine Précision à Basse Précision : En conservant les gradients dans les registres FP32 jusqu'à l'optimisation, FORGE évite le biais introduit par la troncature des gradients en BF16 avant l'optimisation. Combiné à l'arrondi stochastique, il égale la convergence du master FP32 même lors de l'utilisation d'états BF16 ou INT8.
Généralité : La méthode est agnostique à l'architecture pour les couches linéaires et supporte huit optimiseurs élément par élément entièrement fusionnés (AdamW, NAdam, RAdam, Lion, RMSprop, AdaGrad, SGD, SGD-momentum).

4. Résultats Expérimentaux

Évalué sur des GPU NVIDIA H200 et B200 pour des modèles allant de 0,6B à 14B de paramètres (Llama-3.1-8B, famille Qwen3) :

Réduction de Mémoire : FORGE réduit la mémoire de crête de 38 à 53 % par rapport à un AdamW fusionné. Pour Llama-3.1-8B, la mémoire de crête chute d'environ 75 Gio à 35,4 Gio (une réduction de 53 %).
Débit (Throughput) : FORGE atteint une accélération de 1,52× du temps de l'étape (99,1 ms contre 150,9 ms) en éliminant le lancement de noyau d'optimiseur séparé et l'aller-retour HBM pour les gradients.
Scalabilité :
- Sur un seul H200, FORGE entraîne Qwen3-14B (ce qui provoque des erreurs Out-Of-Memory avec l'AdamW fusionné standard) à 63,2 Gio.
- Dans une configuration 4-GPU A100-40GB, FORGE permet une taille de micro-batch de 4 (tenant dans 24,6 Gio), alors que l'AdamW fusionné standard ne permet qu'un micro-batch de 1. Cela permet à FORGE d'atteindre un débit 2,42× plus élevé que la base standard sur le même matériel.
Convergence : FORGE suit la convergence de l'AdamW fusionné à 0,001 nat près (écart moyen) sur des tâches de pré-entraînement continu, confirmant que les déviations de basse précision sont dans le bruit de mesure habituel.
Composabilité : FORGE se compose proprement avec la quantification FP8 et le checkpointing d'activation, absorbant les coûts de déquantification qui pénalisent habituellement les optimiseurs standards.

5. Signification et Revendications

L'article affirme que FORGE représente un changement structurel plutôt qu'une optimisation incrémentale. En supprimant le plancher de gradient, il convertit la surcharge mémoire de l'optimiseur en une marge utilisable pour des tailles de lots plus grandes, des séquences plus longues ou des modèles plus grands sur le matériel existant.

Limite de Capacité : Les auteurs affirment que l'élimination du plancher de gradient n'est pas seulement une amélioration d'un facteur constant, mais une limite de capacité. Par exemple, les configurations d'entraînement standard pour des modèles de 14B sur un seul H200 sont impossibles à cause du plancher de gradient, alors que FORGE les rend réalisables.
Orthogonalité : La méthode est orthogonale aux autres techniques d'économie de mémoire (comme la quantification d'état ou la recomputation d'activation) et peut être combinée avec elles pour réduire davantage l'empreinte mémoire.
Limitations : La méthode est actuellement limitée aux optimiseurs séparables élément par élément et ne supporte pas le parallélisme de données sans modification. Elle nécessite également que le gradient d'entrée soit lu avant la mise à jour du poids, excluant les embeddings liés ou les experts MoE avec plusieurs consommateurs de gradients dans une seule étape.

En résumé, FORGE démontre que le « plancher de gradient » est un artefact inutile des boucles d'entraînement actuelles. En fusionnant l'optimiseur dans la passe arrière au niveau de la tuile, il atteint des économies de mémoire et des accélérations significatives tout en maintenant l'exactitude mathématique pour une large classe d'optimiseurs.