🤖 machine learning

FOAM: Blocked State Folding for Memory-Efficient LLM Training

Auteurs originaux : Ziqing Wen, Jiahuan Wang, Ping Luo, Dongsheng Li, Tao Sun

Publié 2026-05-14

📖 5 min de lecture🧠 Analyse approfondie

Auteurs originaux : Ziqing Wen, Jiahuan Wang, Ping Luo, Dongsheng Li, Tao Sun

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

Imaginez que vous essayez d'enseigner à un robot géant et ultra-intelligent (un Modèle de Langage de Grande Taille) comment écrire, discuter et raisonner. Pour ce faire, vous lui montrez une bibliothèque massive de livres (données d'entraînement). Le robot apprend en commettant des erreurs, en vérifiant son travail et en ajustant ses « réglages internes du cerveau » (paramètres) pour s'améliorer.

La méthode standard pour effectuer cet ajustement est comparable à un comptable très méticuleux nommé Adam. Adam tient un registre détaillé pour chaque réglage du cerveau, suivant non seulement l'erreur actuelle, mais aussi l'historique des erreurs passées pour décider de l'ampleur des changements à apporter.

Le Problème : Le Comptable est Trop Lourd
Le problème est que ce « comptable » (l'optimiseur) est incroyablement lourd. Pour un robot comportant des milliards de réglages du cerveau, le registre du comptable occupe deux fois plus de mémoire que le cerveau du robot lui-même.

Analogie : Imaginez essayer de déménager une maison. Les meubles (le modèle) sont grands, mais le camion de déménagement (les états de l'optimiseur) est encore plus grand. Si vous n'avez qu'un petit camion (mémoire informatique limitée), vous ne pouvez pas déménager la maison du tout, peu importe la qualité de la maison. C'est le « goulot d'étranglement de la mémoire » qui empêche les chercheurs d'entraîner des robots plus grands et plus intelligents.

Les Anciennes Solutions : Ruser
Les tentatives précédentes pour réduire la taille du camion de déménagement présentaient des défauts :

Geler des parties de la maison : Vous arrêtez de déplacer certains meubles et n'ajustez que quelques petites boîtes. Cela économise de l'espace mais rend la maison moins flexible (performances réduites).
Prendre des photos floues : Au lieu de déplacer les meubles réels, vous prenez des photos basse résolution pour économiser de l'espace. Mais le calcul de ces photos demande beaucoup de temps et d'énergie (surcharge de calcul).
Partager les registres : Vous faites en sorte que différentes pièces partagent le même cahier. Cela économise de l'espace mais rend les ajustements moins précis.

La Nouvelle Solution : FOAM (La Méthode de Pliage Intelligente)
L'article présente FOAM (Folded Optimizer with Approximate Moment, soit Optimiseur Plié avec Moment Approximatif). Considérez FOAM comme un expert en emballage qui utilise une astucieuse technique de « pliage » pour faire tenir le camion de déménagement massif dans un petit fourgon, sans perdre aucun détail important.

Voici comment FOAM fonctionne, en utilisant des métaphores simples :

1. Le « Pliage par Blocs » (Compression)

Au lieu de suivre chaque réglage du cerveau individuellement, FOAM les regroupe en petits blocs (comme une grille de 8 ou 16 réglages).

La Métaphore : Imaginez que vous avez une file de 100 personnes et que vous devez retenir leur taille moyenne. Au lieu d'écrire 100 nombres séparés, vous les regroupez en 10 équipes de 10. Vous écrivez juste un nombre : la taille moyenne de chaque équipe.
Le Résultat : Cela réduit considérablement la mémoire nécessaire pour le « registre » (jusqu'à 90 % de moins).

2. La « Correction Résiduelle » (Le Filet de Sécurité)

Si vous utilisiez uniquement les moyennes, vous perdriez les détails uniques de chaque personne. Peut-être qu'une personne dans une équipe est très grande et une autre très petite. La moyenne masque cela.

La Métaphore : FOAM est intelligent. Après avoir noté la « moyenne de l'équipe », il calcule rapidement la différence (l'erreur) entre les vraies personnes et la moyenne. Il appelle cela le « Résidu ».
L'Astuce : Il ne conserve pas cette différence pour toujours. Il la calcule, l'utilise pour corriger la mise à jour à cet instant précis, puis la jette. C'est comme un chef qui goûte une soupe, ajuste le sel, puis oublie le goût exact de la cuillère, plutôt que de garder une trace de chaque cuillère jamais goûtée.
Pourquoi c'est important : Cela garantit que le robot apprend toujours les détails uniques de chaque réglage du cerveau, même si la mémoire est compressée.

3. Le « Dépliement » (Restauration)

Qu'il est temps de mettre à jour réellement le cerveau du robot, FOAM reprend ces « moyennes d'équipe » compressées et les réexpande à leur taille complète, en ajoutant les corrections « Résiduelles » par-dessus.

Le Résultat : Le robot reçoit une mise à jour précise, tout comme avec le comptable lourd, mais l'ordinateur n'a eu à transporter que la version pliée et minuscule.

Ce que l'Article Affirme (Les Résultats)
Les auteurs ont testé FOAM sur l'entraînement de divers modèles de robots (de petits modèles de 60 millions de paramètres à de grands modèles de 7 milliards de paramètres). Voici ce qu'ils ont découvert :

Économies Massives de Mémoire : FOAM réduit la mémoire nécessaire pour l'optimiseur jusqu'à 90 %. Au total, il réduit l'empreinte mémoire de l'entraînement d'environ 50 %. Cela signifie que vous pouvez entraîner des modèles plus grands sur le même matériel, ou entraîner les mêmes modèles beaucoup plus vite.
Aucune Perte de Performance : Malgré la compression intensive, les modèles entraînés avec FOAM fonctionnent aussi bien (et parfois même mieux) que les modèles entraînés avec le comptable lourd standard. Ils apprennent plus vite et atteignent une précision plus élevée.
Vitesse : Parce qu'il n'a pas à porter une charge lourde, le processus d'entraînement est plus rapide.
Polyvalence : Il fonctionne sur différents types d'architectures de robots (comme LLaMA, Qwen et autres) et peut même être combiné avec d'autres astuces d'économie de mémoire.

En Résumé
FOAM est comme une valise magique. Il vous permet de ranger une quantité massive d'informations (la mémoire de l'optimiseur) dans un tout petit espace en les pliant soigneusement et en ajoutant une rapide « note de correction » pour garantir que rien n'est perdu. Cela permet aux chercheurs de construire des modèles d'IA plus intelligents et plus grands sans avoir besoin d'ordinateurs massifs et extrêmement coûteux.

Résumé Technique : FOAM – Repliement d'État Bloqué pour l'Entraînement Économe en Mémoire des LLM

1. Énoncé du Problème

Les grands modèles de langage (LLM) ont atteint des performances remarquables grâce à l'augmentation du nombre de paramètres et des données d'entraînement, mais leur échelle introduit des goulots d'étranglement significatifs en mémoire durant l'entraînement. L'optimiseur de facto, Adam (et AdamW), nécessite le stockage des estimations des moments d'ordre un et deux pour chaque paramètre, entraînant une surcharge mémoire de l'état de l'optimiseur égale au double de la taille du modèle. Cela rend l'entraînement préalable et le fine-tuning de grands modèles (par exemple, 7 milliards de paramètres ou plus) limités par la mémoire, nécessitant souvent plusieurs centaines de gigaoctets de VRAM même avec de petites tailles de lot.

Les solutions existantes économes en mémoire font face à des compromis :

Gel des paramètres (par exemple, LoRA) : Restreint les mises à jour à des sous-espaces de faible rang, limitant potentiellement l'expressivité et les performances du modèle.
Projections de faible rang (par exemple, GaLore, APOLLO) : Compresse les états de l'optimiseur en utilisant la décomposition en valeurs singulières (SVD) ou des projections aléatoires. Bien qu'efficaces, la SVD introduit une surcharge computationnelle substantielle, et les projections aléatoires peuvent dégrader les performances sous des contraintes mémoire strictes.
Partage de blocs (par exemple, Adam-Mini) : Partage les estimations du moment d'ordre deux entre les blocs pour réduire le stockage, mais laisse souvent la surcharge du moment d'ordre un élevée ou plafonne la réduction mémoire à environ 50 %.

2. Méthodologie : FOAM

Les auteurs proposent FOAM (Folded Optimizer with Approximate Moment, Optimiseur Replié avec Moment Approximatif), un cadre structuré de repliement d'état conçu pour compresser les états de l'optimiseur sans sacrifier la fidélité de l'optimisation. FOAM fonctionne sur le principe de la moyenne bloquée combinée à une correction résiduelle.

Mécanisme Central

Étant donné une matrice de gradient $G_t \in \mathbb{R}^{m \times n}$ à l'étape de temps $t$ et un niveau de repliement $l$ :

Moyenne Bloquée (Repliement) : Le gradient est partitionné en blocs de $2^l$ éléments consécutifs. Un opérateur de repliement $A^{(l)}$ calcule la moyenne de chaque bloc, comprimant la dimension du gradient de $n$ à $n/2^l$ . Les moments d'ordre un et deux ( $\tilde{M}_t, \tilde{V}_t$ ) sont mis à jour sur la base de ces gradients compressés.
Dépliement : Pour appliquer les mises à jour à l'espace complet des paramètres, un opérateur de repliement $E^{(l)}$ réplique les entrées compressées vers les dimensions d'origine.
Correction Résiduelle : La réplication directe des mises à jour fait que tous les éléments d'un bloc partagent des valeurs identiques, dépendant uniquement de l'initialisation pour la différenciation. Pour récupérer les informations perdues, FOAM calcule un terme résiduel $R_t = G_t - \tilde{G}_t E^{(l)}$ $R_{t} = G_{t} - \tilde{G}_{t} E^{(l)}$ , représentant l'erreur entre le gradient original et la version repliée-dépliée.
- Le moment d'ordre un est reconstruit comme suit : $M_t = \tilde{M}_t E^{(l)} + R_t$ .
- Le moment d'ordre deux est reconstruit comme suit : $V_t = \tilde{V}_t E^{(l)} + R_t^2$ .
- Crucialement, $R_t$ est calculé à la volée et libéré immédiatement ; il n'est pas stocké en mémoire.

Application Hybride

FOAM est appliqué spécifiquement aux modules MLP et Attention, tandis qu'Adam standard est utilisé pour les couches d'embedding. Un coefficient de mise à l'échelle $\alpha$ ajuste le ratio du taux d'apprentissage entre les modules FOAM et Adam, une pratique courante dans les optimiseurs économes en mémoire.

3. Contributions Clés

Conception Innovante d'Optimiseur : FOAM atteint une réduction allant jusqu'à 90 % de la surcharge mémoire de l'état de l'optimiseur et une réduction d'environ 50 % de l'empreinte mémoire totale de bout en bout par rapport à Full-Adam. Contrairement aux méthodes basées sur la projection, il ne nécessite aucune matrice de projection et prend en charge l'optimisation de paramètres complets.
Garanties Théoriques : Sous des hypothèses standard d'optimisation non convexe (lissité L, variance bornée, gradients bornés), les auteurs prouvent que FOAM atteint le même taux de convergence de $O(\frac{\log T}{\sqrt{T}})$ que l'Adam vanilla. Le ratio d'erreur résiduelle $\delta_l$ est borné par 1 et n'affecte pas négativement le taux de convergence principal.
Évaluation Étendue : La méthode est validée sur différentes tailles de modèles (de 60 M à 7 B), architectures (LLaMA, Qwen, GPT, DeBERTa) et tâches (entraînement préalable sur C4, fine-tuning sur GLUE et MMLU).

4. Résultats Expérimentaux

L'article rapporte les constatations empiriques suivantes :

Performance d'Entraînement Préalable : FOAM surpasse systématiquement les bases de référence économes en mémoire (GaLore, APOLLO, Adam-Mini, Muon) en termes de perplexité de validation (PPL) tout en maintenant une utilisation de mémoire significativement plus faible. Par exemple, sur LLaMA-1,3B, FOAM-3 réduit la consommation mémoire totale d'environ 50 % par rapport à Adam tout en atteignant un PPL plus faible.
Vitesse de Convergence : FOAM démontre une convergence plus rapide, nécessitant souvent moins d'itérations d'entraînement pour atteindre un PPL cible par rapport à d'autres méthodes économes en mémoire.
Évolutivité : FOAM maintient des performances supérieures sur les modèles plus grands (LLaMA-3B, 7B) et dans des régimes d'entraînement de séquences longues (jusqu'à 2048 tokens).
Fine-tuning : Sur les benchmarks GLUE et MMLU, FOAM atteint des performances comparables ou supérieures à Full-Adam et aux autres bases de référence, démontrant son utilité au-delà de l'entraînement préalable.
Robustesse : Les études d'ablation confirment que le terme résiduel est critique pour les performances ; son retrait entraîne une convergence plus lente et un PPL final plus élevé. FOAM reste également robuste lorsqu'il est combiné avec une quantification sur 8 bits et lorsqu'il est intégré avec d'autres optimiseurs comme Adam-Mini et Muon.

5. Signification et Revendications

L'article revendique que FOAM offre une solution pratique, hautement efficace et agnostique à l'optimiseur aux goulots d'étranglement mémoire dans l'entraînement de modèles de langage à grande échelle. En éliminant le besoin de gel des poids ou de calculs de projection coûteux, FOAM permet l'optimisation de paramètres complets avec une surcharge mémoire drastiquement réduite. Les auteurs positionnent FOAM comme une alternative complémentaire aux techniques de projection de faible rang, fournissant une méthode qui correspond théoriquement à la convergence d'Adam tout en offrant empiriquement une efficacité mémoire supérieure et souvent de meilleures performances d'entraînement.

Ce travail suggère que FOAM peut permettre l'entraînement de modèles plus grands sur le matériel existant ou réduire les coûts computationnels de l'entraînement des modèles actuels de l'état de l'art, rendant l'entraînement de LLM haute performance plus accessible.