Deep Optimizer States: Towards Scalable Training of Transformer Models Using Interleaved Offloading

Ce papier présente Deep Optimizer States, une nouvelle technique qui améliore l'efficacité de l'entraînement des grands modèles de langage en exploitant les fluctuations de l'utilisation de la mémoire GPU pour intercaler dynamiquement le transfert et la mise à jour des états de l'optimiseur entre le CPU et le GPU, permettant ainsi d'accélérer les itérations d'un facteur 2,5 par rapport aux approches de pointe.

L'essentiel

🚀 Le Problème : Le Mur de la Mémoire

Imaginez que vous essayez de construire un gratte-ciel gigantesque (un modèle d'intelligence artificielle comme ceux qui écrivent des poèmes ou codent des logiciels) dans un garage très petit (la mémoire de votre carte graphique, ou GPU).

Le problème, c'est que ce gratte-ciel est si grand qu'il ne rentre pas tout entier dans le garage. Pour continuer à travailler, les chercheurs ont dû inventer une astuce : ils gardent les plans du bâtiment dans le garage, mais ils stockent les "notes de calcul" (l'état de l'optimiseur) dans le sous-sol de la maison (la mémoire de l'ordinateur, ou CPU).

Mais il y a un gros hic :
Le garage (GPU) est un ouvrier ultra-rapide, capable de poser des briques à la vitesse de l'éclair. Le sous-sol (CPU) est un ouvrier beaucoup plus lent. De plus, il y a un escalier étroit (le câble PCIe) entre les deux.
Dans les méthodes actuelles, l'ouvrier rapide du garage attend patiemment que l'ouvrier lent du sous-sol lui apporte les notes, puis l'ouvrier lent attend que l'ouvrier rapide lui rende les plans. C'est comme une course de relais où le coureur rapide passe son temps à attendre le coureur lent. Résultat : le travail avance au rythme du plus lent, et le garage reste souvent vide (sous-utilisé).

💡 La Solution : "Deep Optimizer States" (Des États d'Optimiseur Profonds)

Les auteurs de cet article ont eu une idée brillante : pourquoi ne pas faire travailler les deux ouvriers en même temps, en se relayant intelligemment ?

Ils ont créé une nouvelle méthode appelée "Deep Optimizer States". Voici comment cela fonctionne avec une analogie simple :

1. Le Concept de "Relais Intercalé" (Interleaved Offloading)

Au lieu de dire "Tous les plans vont au sous-sol, puis tout revient", ils découpent le travail en petits paquets (des sous-groupes).

• Pendant que l'ouvrier lent du sous-sol travaille sur le paquet A, l'ouvrier rapide du garage travaille sur le paquet B.
• Pendant que le garage finit le paquet B, il envoie le paquet C au sous-sol, et le sous-sol renvoie le paquet A terminé.

C'est comme une chaîne de montage où l'on ne s'arrête jamais. On mélange les tâches pour que l'escalier (le câble) soit toujours rempli de boîtes qui montent et descendent, et que personne n'attende jamais les bras croisés.

2. L'Analogie du Camion de Déménagement

Imaginez que vous déménagez une maison avec un camion (le GPU) et un entrepôt (le CPU).

• Méthode ancienne : Vous chargez tout le camion, vous descendez tout à l'entrepôt, vous attendez que l'entrepôt range les meubles, puis vous remontez tout. Le camion reste à l'arrêt la moitié du temps.
• Méthode Deep Optimizer States : Vous chargez une partie du camion, vous la descendez, mais pendant que le camion revient chercher la partie suivante, l'entrepôt range déjà la première partie. Le camion ne s'arrête jamais, et l'entrepôt ne s'ennuie jamais.

3. L'Intelligence Artificielle du Système

Le système ne devine pas au hasard. Il utilise une formule mathématique (un modèle de performance) pour calculer exactement combien de travail il faut donner au garage et combien au sous-sol à chaque instant.

• Si le camion est très rapide mais l'escalier lent, il ajuste le rythme.
• Si le sous-sol a beaucoup de bras (beaucoup de cœurs CPU), il envoie plus de travail là-bas.
  C'est comme un chef d'orchestre qui ajuste la musique en temps réel pour que tout le monde joue parfaitement ensemble.

🏆 Les Résultats : Pourquoi c'est génial ?

Grâce à cette méthode, les chercheurs ont obtenu des résultats spectaculaires :

• Vitesse : L'entraînement des modèles est 2,5 fois plus rapide qu'avec les meilleures méthodes actuelles.
• Économie d'énergie : On utilise mieux les ressources. Le garage n'est plus en pause, et l'escalier est utilisé à plein régime.
• Accessibilité : Cela permet de faire tourner des modèles géants sur des machines plus petites et moins chères, car on ne gaspille plus de temps d'attente.

En Résumé

Imaginez que vous avez deux équipes pour construire un mur : une équipe de robots ultra-rapides et une équipe de humains plus lents, séparés par un couloir étroit.
Avant, les robots attendaient que les humains finissent leur part avant de reprendre le travail.
Avec Deep Optimizer States, on a appris aux robots et aux humains à travailler en parallèle : pendant que les humains travaillent sur une brique, les robots en préparent une autre, et les deux équipes s'échangent les briques sans jamais s'arrêter.

C'est une révolution pour l'intelligence artificielle : cela permet d'entraîner des "cerveaux" numériques beaucoup plus grands et plus intelligents, beaucoup plus vite, sans avoir besoin de construire des usines (des supercalculateurs) encore plus énormes.

Résumé technique

1. Problématique : Le « Mur de Mémoire » dans l'entraînement des LLM

L'entraînement des modèles de langage de grande taille (LLM) et des transformateurs se heurte à une limitation critique : le mur de mémoire.

• Contexte : Les modèles modernes comptent des centaines de milliards, voire des billions de paramètres. Même avec des techniques de parallélisme 3D (pipeline, tenseur, données), la mémoire GPU est insuffisante pour stocker simultanément les paramètres du modèle, les gradients, les activations et, surtout, les états de l'optimiseur (paramètres FP32, momentums, variances pour des optimiseurs comme Adam).
• Solution actuelle et ses limites : Les approches de pointe (comme DeepSpeed Offload, ZeRO-Offload) déchargent une partie ou la totalité de l'état de l'optimiseur vers la mémoire hôte (CPU). Cependant, cela introduit deux goulots d'étranglement majeurs :
  1. Bande passante I/O limitée : Le transfert de données entre le CPU et le GPU via PCIe (généralement 25-50 Go/s) est lent par rapport à la vitesse de calcul.
  2. Faiblesse de calcul du CPU : Les CPU sont plusieurs ordres de grandeur plus lents que les GPU pour les mises à jour de paramètres (dans les tests de l'article, les GPU mettent à jour ~100 milliards de paramètres/seconde contre ~8 milliards pour les CPU).
• Inefficacité actuelle : Les solutions hybrides statiques (ex: TwinFlow) allouent une fraction fixe de l'optimiseur au GPU et le reste au CPU. Cela laisse souvent la mémoire GPU sous-utilisée pendant les phases de calcul (forward/backward) et ne permet pas un chevauchement optimal entre les transferts de données et les calculs, entraînant des temps d'attente inutiles.

2. Méthodologie : Deep Optimizer States

Les auteurs proposent Deep Optimizer States, une technique middleware qui optimise dynamiquement l'entraînement hybride CPU-GPU en exploitant les fluctuations d'utilisation de la mémoire et de la bande passante PCIe.

Principes de conception clés :

1. Mise à jour entrelacée (Interleaved Updates) : Au lieu d'une répartition statique, le système divise l'état de l'optimiseur en sous-groupes (subgroups). Il planifie dynamiquement la mise à jour de certains sous-groupes sur le GPU et d'autres sur le CPU au sein d'une même itération.
2. Chevauchement computationnel et I/O : L'architecture permet de chevaucher les calculs de mise à jour sur le CPU avec les transferts de données (H2D et D2H) et les calculs sur le GPU. Pendant que le CPU calcule la mise à jour d'un sous-groupe, le GPU peut précharger le prochain sous-groupe ou effectuer sa propre mise à jour.
3. Gestion intelligente des gradients :
   • Les gradients générés en FP16 sur le GPU sont transférés vers le CPU.
   • Pour éviter les conversions coûteuses en temps réel (FP16 $\to$ FP32) sur le CPU, l'approche effectue la conversion sur le GPU (à très haute vitesse, ~1,2 To/s) avant le transfert vers la mémoire hôte, ou utilise des tampons épinglés (pinned buffers) pour des transferts directs.
4. Modèle de performance dynamique : Un modèle mathématique calcule le ratio optimal de sous-groupes à mettre à jour sur le GPU (le « pas de mise à jour » ou update stride, noté $k$). Ce modèle équilibre le temps de calcul CPU, le temps de calcul GPU et les temps de transfert PCIe pour maximiser le chevauchement.

Algorithme de planification :

L'algorithme (Algorithm 1) itère sur les sous-groupes de l'optimiseur. Pour chaque sous-groupe, il décide s'il doit être mis à jour sur le GPU (selon le ratio $k$) ou sur le CPU. Il orchestre des transferts asynchrones (via des streams CUDA) pour :

• Précharger (prefetch) les états de l'optimiseur (paramètres, momentums, variances) du CPU vers le GPU.
• Vider (flush) les paramètres mis à jour du GPU vers le CPU.
• Exécuter les mises à jour en parallèle sur les deux processeurs.

3. Contributions Clés

1. Analyse fine des caractéristiques système : Les auteurs ont démontré que l'utilisation de la mémoire GPU et de la bande passante PCIe fluctue considérablement entre les phases forward, backward et update, créant des opportunités d'optimisation souvent ignorées.
2. Conception d'un middleware hybride dynamique : Intégration de Deep Optimizer States dans DeepSpeed et Megatron-LM, permettant un entraînement évolutif sans nécessiter de modifications majeures du code utilisateur.
3. Modèle de performance et algorithme de planification : Développement d'une formule pour déterminer le ratio optimal CPU/GPU en fonction des spécificités matérielles (débit PCIe, vitesse de calcul CPU/GPU).
4. Optimisation des transferts de précision : Utilisation de la conversion FP16 $\to$ FP32 sur le GPU pour éviter les goulots d'étranglement de la mémoire hôte et des transferts asynchrones pour masquer la latence.

4. Résultats Expérimentaux

Les expériences ont été menées sur un nœud équipé de 4 GPU NVIDIA H100 (80 Go) et de 192 cœurs CPU, avec des modèles allant de 7B à 20B de paramètres.

• Accélération des itérations : Deep Optimizer States est 2,5 fois plus rapide que l'approche de référence (DeepSpeed ZeRO-3 avec déchargement complet) et 1,7 à 2,3 fois plus rapide que les solutions hybrides statiques (DeepSpeed TwinFlow).
• Débit de mise à jour : L'approche atteint un débit de mise à jour 70 % supérieur à celui de ZeRO-3, en exploitant efficacement la puissance de calcul du GPU pour une partie des mises à jour tout en masquant les transferts I/O.
• Utilisation des ressources :
  • Le GPU atteint une utilisation proche de 100 % pendant la phase de mise à jour (contre ~8 % avec ZeRO-3 standard).
  • La bande passante PCIe est utilisée à ~40 % de son pic, contre <10 % dans les approches statiques.
• Évolutivité : La méthode maintient son avantage (jusqu'à 2,5x) même avec un parallélisme de données accru et sur des configurations avec un nombre variable de cœurs CPU par GPU.
• Robustesse : Le modèle de performance prédit correctement le ratio optimal ($k=2$, soit une mise à jour GPU tous les deux sous-groupes) sur différentes architectures matérielles.

5. Signification et Impact

• Démocratisation de l'entraînement de LLM : Cette technique permet d'entraîner des modèles de taille moyenne (jusqu'à 20B de paramètres) sur des nœuds uniques avec des ressources limitées, ce qui est crucial pour le fine-tuning rapide de modèles spécialisés sans accès à des supercalculateurs massifs.
• Efficacité des ressources existantes : Elle maximise l'utilisation du matériel hétérogène (CPU + GPU) déjà disponible, réduisant le coût et l'empreinte énergétique de l'entraînement.
• Préparation pour l'avenir : L'approche est particulièrement pertinente pour les futures architectures (comme les systèmes Grace Hopper) qui offrent des interconnexions CPU-GPU à très haut débit, où le déchargement entrelacé dynamique deviendra la norme pour exploiter pleinement la bande passante.

En résumé, Deep Optimizer States résout le problème de la mémoire wall non pas en ajoutant plus de mémoire, mais en réorganisant intelligemment le flux de données et de calculs pour transformer les temps d'attente I/O en temps de calcul utile, offrant ainsi une accélération significative pour l'entraînement des modèles d'IA à grande échelle.