FR
Un système d’entraînement adapte la mémoire aux pics de charge des experts
BaGuaLu-2: Taming Dynamic Memory Fluctuation for Efficient MoE Model Training over 10,000 GPUs · arXiv:2610.23064
En termes simples
Certains grands modèles de langage dirigent chaque jeton vers certains réseaux d’experts. Le nombre de jetons envoyé à chaque expert peut devenir soudainement déséquilibré, ce qui augmente la mémoire temporaire nécessaire pendant l’entraînement.
Une stratégie mémoire fixe peut gaspiller des ressources pendant les lots ordinaires ou provoquer des plantages lors d’un pic. BaGuaLu-2 est un système d’entraînement qui adapte les mesures d’économie de mémoire à chaque petit lot et superpose une partie de ce travail à d’autres tâches.
Lors d’une exécution sur 512 GPU, sa mesure publiée de l’utilisation des capacités de calcul est restée stable pendant des pics qui ont fait baisser celle de deux systèmes de comparaison.
La recherche
L’entraînement à grande échelle des modèles à mélange d’experts répartit les jetons entre de nombreux réseaux d’experts. De rares pics de routage peuvent augmenter les besoins en mémoire d’activations ; une stratégie fixe doit alors soit ralentir le travail ordinaire, soit exposer l’exécution à des plantages dus à un manque de mémoire.
Comment l’entraînement peut-il réagir à ces variations brèves sans payer le même coût pour chaque lot ?
BaGuaLu-2 choisit des mesures d’économie de mémoire pour chaque petit lot d’entraînement, puis superpose ce travail d’un lot à l’autre selon les variations de pression mémoire.
BaGuaLu-2 estime les besoins en mémoire d’activations et choisit, pour chaque micro-lot, de recalculer des valeurs, de les déplacer hors du GPU, de réduire leur précision ou de ne rien changer.
Il superpose ce travail aux opérations de lots voisins pour réduire le temps consacré à la gestion de la mémoire.
Pour les configurations de modèles et de grappes évaluées, de 128 à 12 288 GPU, les auteurs rapportent un débit moyen de bout en bout égal à 1.14 fois celui de Megatron-Kuai et à 1.13 fois celui d’AdaPipe.
Le débit correspond ici à la moyenne, pour chaque itération de l’ensemble de l’exécution, de l’utilisation des FLOPS du modèle (MFU), une mesure de la part de capacité de calcul disponible des GPU effectivement utilisée. Le calcul inclut le temps de remise en route des machines après les plantages. Ces résultats étayent l’observation d’un débit global supérieur dans les tests réalisés.
La comparaison moyenne porte sur les configurations évaluées, qui ont commencé avec le même degré de parallélisme choisi hors ligne ; l’article ne précise pas si les points de départ des modèles étaient identiques. Le calcul du coût des redémarrages inclut le temps de remise en route des machines, mais pas celui que les personnes consacrent au diagnostic et aux interventions : il sous-estime donc cette partie du coût.
Ce que montrent les preuves
- De rares pics de routage
Les modèles à mélange d’experts entraînés à grande échelle répartissent les jetons entre des réseaux d’experts. La charge de ces experts peut varier d’une itération d’entraînement à l’autre, ce qui change les besoins en mémoire temporaire pour les activations.
Les auteurs ont mesuré un déséquilibre supérieur à 1.2 fois la moyenne dans 2.9% des itérations, avec un taux maximal de 2.6 fois la moyenne. La question centrale est donc de savoir comment gérer ces pics occasionnels de mémoire sans recourir en permanence à une stratégie fixe coûteuse.
- Adapter la mémoire à chaque lot
BaGuaLu-2 est un système d’entraînement qui estime les besoins en mémoire de chaque micro-lot, puis choisit quelles opérations recalculer, déplacer hors du GPU, exécuter avec une précision réduite ou laisser inchangées. Son budget mémoire s’adapte aux conditions actuelles du pipeline d’entraînement.
Il entrelace aussi le travail de micro-lots voisins pour masquer une partie du temps supplémentaire consacré à la gestion de la mémoire.
- Comparer les entraînements complets
Les auteurs ont comparé BaGuaLu-2 à Megatron-Kuai et AdaPipe pour des configurations allant de 128 à 12 288 GPU. Les systèmes ont commencé avec la même configuration de parallélisme, choisie hors ligne ; les systèmes de comparaison utilisaient des stratégies mémoire statiques et pouvaient être redémarrés manuellement après un plantage.
Le débit global correspond à la moyenne, pour chaque itération de l’ensemble de l’exécution, de l’utilisation des FLOPS du modèle (MFU), une mesure de la part de capacité de calcul disponible des GPU effectivement utilisée. Le calcul inclut le temps de remise en route des machines après les redémarrages.
La moyenne de BaGuaLu-2 était égale à 1.14 fois celle de Megatron-Kuai et à 1.13 fois celle d’AdaPipe.
- Résultats à grande échelle et apport des composants
La plus grande exécution évaluée a entraîné un modèle de 1 040 milliards de paramètres sur 12 288 GPU et rapporté un débit total de 2.012 exaFLOPS. Une comparaison progressive des composants a également associé l’ajout successif d’une atténuation statique plus précise, d’une atténuation dynamique et d’une planification qui superpose les tâches à des gains de débit par rapport à la configuration précédente.
Ces résultats décrivent l’exécution à grande échelle et l’étude d’ablation rapportées ; ils ne permettent pas d’isoler l’effet de chaque composant pour toutes les charges de travail.
- Les performances pendant les pics
Lors d’une exécution représentative d’un modèle de 423 milliards de paramètres sur 512 GPU, la MFU des systèmes de comparaison a diminué pendant les pics de déséquilibre, tandis que celle de BaGuaLu-2 est restée stable. Pendant cette phase de pics, le débit moyen de BaGuaLu-2 était égal à 1.28 fois celui des systèmes de comparaison.
Le résultat porte sur la situation ciblée par le système, mais provient d’un seul modèle et d’une seule configuration de grappe.
Détails
- Auteurs de l’article original
- Mingshu Zhai, Minkun Guo, Hongrui Lian, Shuangyu Li, Qi Zhang
- Soumis à arXiv
- Licence
- CC BY 4.0