Résumé Technique : Briser la Bulle avec PACI

Énoncé du Problème

Le parallélisme de pipeline est une technique critique pour l'entraînement de réseaux de neurones à grande échelle, particulièrement les transformers, en partitionnant les modèles entre les dispositifs et en chevauchant le calcul. Cependant, les stratégies d'ordonnancement existantes font face à un compromis persistant entre l'utilisation du matériel, l'efficacité de la mémoire et la cohérence de l'optimisation :

• Ordonnancements Synchrones (ex: 1F1B-Flush) : Ceux-ci préservent la cohérence des poids forward/backward (garantissant qu'un micro-batch utilise les mêmes paramètres lors de sa passe arrière que lors de sa passe avant) mais souffrent de « bulles de pipeline » (temps d'inactivité), réduisant ainsi le débit.
• Ordonnancements Asynchrones : Ces derniers éliminent les bulles pour atteindre une pleine utilisation du matériel mais introduisent une incohérence de version des poids forward/backward. Dans ces ordonnancements, la passe arrière d'un micro-batch peut s'exécuter en utilisant une version de paramètre qui a considérablement évolué depuis sa passe avant.
• Mitigations Existantes : Les méthodes asynchrones actuelles tentent de résoudre cette incohérence par le stockage de poids (weight stashing) (stockage de multiples copies de paramètres), la prédiction (approximation des poids futurs) ou la synchronisation globale. Ces approches entraînent un surcoût de mémoire important, un calcul supplémentaire, une complexité système accrue ou des sémantiques d'optimisation non standard.

Les auteurs se demandent si l'incohérence forward/backward peut être rendue suffisamment faible pour être tolérée sans être totalement éliminée, permettant ainsi de conserver l'efficacité de l'exécution asynchrone tout en évitant les coûts des stratégies de mitigation existantes.

Méthodologie : PACI

Le papier introduit PACI (Pipeline Asynchronous training with Controlled Inconsistency), un ordonnancement de pipeline 1F1B asynchrone qui limite la dérive de version sans stockage de poids, sans prédiction, ni synchronisation globale.

Mécanisme Central : L'Accumulation de Gradients comme Contrôle de Version

PACI exploite l'observation selon laquelle l'incohérence n'est pas régie par l'asynchronisme seul, mais par le nombre de mises à jour de l'optimiseur qui se produisent entre la passe avant et la passe arrière d'un micro-batch.

1. Découplage des Mises à Jour et des Passes Arrières : Dans un 1F1B asynchrone standard, les paramètres sont mis à jour après chaque passe arrière. PACI découple cela : les étapes exécutent les opérations backward immédiatement dès l'arrivée de l'entrée, mais accumulent les gradients localement sur une fenêtre de $a$ passes arrière avant d'appliquer une seule étape d'optimiseur.
2. Limitation de la Dérive de Version : En ralentissant l'évolution des versions de paramètres par rapport au délai du pipeline, PACI limite le nombre de mises à jour traversées par un micro-batch. L'incohérence maximale ($\Delta_{max}$) est bornée par la profondeur du pipeline ($N$) et le facteur d'accumulation ($a$) :
   $$\Delta_{max} \le \left\lfloor \frac{N-1}{a} \right\rfloor$$
3. Contrôle de Flux Local : Pour gérer les temps de stage non uniformes, PACI emploie une règle de contre-pression (backpressure) locale. Chaque étape suit le nombre de passes avant non résolues ($q_i$) et n'autorise de nouvelles passes avant que si $q_i < N + 1 - i$. Cela empêche les étapes amont de progresser arbitrairement loin devant les étapes aval sans introduire de barrières de synchronisation globales.
4. Efficacité des Ressources : PACI nécessite zéro mémoire de poids supplémentaire. Il utilise une copie unique de paramètre par étape et réutilise les tampons de gradient standards déjà présents dans l'optimiseur. Aucun stockage de poids, tampon de prédiction ou synchronisation globale n'est nécessaire.

Distinction par rapport au Micro-batching

Dans les ordonnancements synchrones, l'augmentation du nombre de micro-batches ($m$) est principalement utilisée pour amortir le surcoût de la bulle. Dans PACI, le micro-batching (via le facteur d'accumulation $a$) est principalement utilisé comme un curseur de cohérence pour contrôler la dérive de version. Cela permet à PACI de séparer l'objectif du système (haute utilisation) de l'objectif d'optimisation (faible incohérence).

Contributions Clés

1. Incohérence Bornée sans Stockage : Le papier démontre que l'accumulation locale de gradients peut servir de mécanisme de contrôle de version ajustable, limitant l'incohérence forward/backward sans mémoire supplémentaire, prédiction ou synchronisation.
2. Amélioration du Temps d'Atteinte de la Précision (Time-to-Accuracy) : PACI permet un pré-entraînement stable avec une perplexité finale comparable au 1F1B-flush synchrone, tout en réduisant le temps d'exécution de bout en bout jusqu'à 1,69× (pour une taille de batch de 128) et 1,41× (pour une taille de batch de 256) par rapport aux bases de référence flush les plus rapides.
3. Analyse de Débit en Accord avec la Théorie : Les auteurs montrent que l'écart de débit entre le 1F1B-flush synchrone et PACI est entièrement expliqué par l'efficacité de la bulle de pipeline. PACI atteint le débit théorique pleinement utilisé du pipeline tout en conservant la même empreinte mémoire de pointe que le flushing synchrone.

Résultats Expérimentaux

Les auteurs ont évalué PACI sur le pré-entraînement de GPT-2 Medium à partir de zéro sur OpenWebText en utilisant un pipeline de 8 étapes sur 8 GPU NVIDIA RTX PRO 6000.

• Stabilité (Q1) : Pour une incohérence bornée de $\Delta_{max} \le 2$ (atteinte avec un facteur d'accumulation $a=4$ ou $a=8$), PACI correspond étroitement à la dynamique de perte du 1F1B-flush synchrone. Il présente une perte finale comparable et, notablement, une plus faible variabilité d'un essai à l'autre (écart type RMS de la perte d'entraînement).
• Temps d'Atteinte de la Précision (Q2) : PACI atteint les niveaux de perplexité cibles nettement plus tôt que les bases de référence synchrones.
  • À une taille de batch de 128, PACI est 1,69× plus rapide que la meilleure base de référence 1F1B-flush optimisée.
  • Comparé à une configuration de micro-batch correspondante, PACI atteint jusqu'à 2,04× d'accélération.
• Débit et Mémoire (Q3) :
  • PACI atteint un débit pleinement utilisé, alors que le débit du 1F1B-flush évolue avec l'efficacité de la bulle (s'améliorant à mesure que les micro-batches augmentent, mais n'atteignant jamais une utilisation complète).
  • PACI correspond à l'empreinte mémoire de pointe du 1F1B-flush, tandis que d'autres méthodes asynchrones (comme PipeDream ou PipeMare) nécessitent beaucoup plus de mémoire pour le stockage de poids ou les états auxiliaires.
  • Des extrapolations à grande échelle (jusqu'à 28,3 milliards de paramètres) suggèrent que PACI peut atteindre le débit des ordonnancements Zero-Bubble (ZB-2p) tout en conservant l'efficacité mémoire du 1F1B-flush.

Signification et Revendications

L'article soutient que l'incohérence forward/backward n'a pas besoin d'être éliminée pour obtenir un entraînement stable. Au contraire, lorsqu'elle est explicitement bornée via l'accumulation de gradients, elle peut être échangée contre des gains d'efficacité substantiels.

• Point de Fonctionnement : PACI établit un nouveau point de fonctionnement dans le parallélisme de pipeline : zéro bulle de pipeline, zéro mémoire de poids supplémentaire et une faible incohérence bornée.
• Changement de Paradigme : Ce travail remet en question la nécessité du stockage de poids ou de la prédiction pour l'entraînement asynchrone. Il suggère que le contrôle du taux d'évolution des paramètres par rapport au délai du pipeline est suffisant pour maintenir la stabilité de l'optimisation.
• Impact Pratique : En supprimant l'overhead de la « bulle » sans engendrer de pénalités de mémoire, PACI offre une voie pratique pour accélérer le temps d'atteinte de la précision des grands modèles de langage, particulièrement dans les régimes où les ordonnancements synchrones sont limités par le temps d'inactivité.

Les auteurs restent modestes quant aux limites, notant que leur évaluation est actuellement limitée à GPT-2 Medium, des pipelines à 8 étapes et des configurations d'entraînement spécifiques. Ils reconnaissent que l'intégration du checkpointing d'activation et la gestion des gradients invalides (NaNs) via le rollback restent des travaux futurs.