Xiaohongshu dévoile BigMac : mettre fin au compromis entre mémoire et vitesse dans l'entraînement multimodal

Les grands modèles linguistiques multimodaux s'imposent comme la pierre angulaire de l'IA de nouvelle génération, mais leur infrastructure d'entraînement se heurte depuis longtemps à un dilemme récurrent : les systèmes optimisés pour la vitesse sont souvent limités par des contraintes de mémoire, tandis que ceux qui privilégient l'efficacité mémoire ont tendance à être lents. L’équipe Dots Infra de Xiaohongshu a identifié ce goulot d’étranglement comme la frontière de Pareto de l’entraînement des pipelines multimodaux et vient de le surmonter. Le 22 juillet, l’équipe a mis en open source BigMac, un nouveau paradigme d’entraînement parallèle en pipeline spécialement conçu pour les scénarios multimodaux natifs. Le projet est désormais disponible sur GitHub dans le dépôt Dots-Infra.

Contrairement aux transformateurs standard, les modèles multimodaux sont intrinsèquement complexes. Un grand modèle de langage multimodal (MLLM) typique comprend trois composants distincts : des encodeurs modaux qui convertissent les images et l’audio en représentations, une architecture LLM pour l’inférence, et un générateur qui reconvertit les sorties du LLM en modalités cibles telles que des images ou de la parole. Les différences structurelles significatives entre ces composants posent des défis de taille lors de leur intégration au sein d’un pipeline d’entraînement unique.

Les solutions actuelles du secteur se répartissent généralement en deux catégories. La première approche privilégie l’efficacité computationnelle en découplant les encodeurs et les générateurs du pipeline LLM, pour les exécuter séparément. Cela empêche les fluctuations de la durée des modules modaux de créer des « bulles » dans le pipeline LLM, mais entraîne une augmentation de la mémoire d’activation proportionnelle au nombre de micro-lots, ce qui se traduit par des coûts élevés à grande échelle. La deuxième approche met l’accent sur l’efficacité mémoire en conservant tous les modules au sein du même pipeline, ce qui raccourcit les cycles de vie des activations et réduit l’utilisation de la mémoire. Cependant, si un encodeur ou un générateur est lent, l’ensemble du pipeline LLM se bloque, ce qui entraîne l’apparition de « bulles de queue ». À mesure que l’échelle du modèle augmente, ces deux conceptions révèlent des goulots d’étranglement critiques.

BigMac relève ce défi grâce à un principe simple mais essentiel : le pipeline LLM principal reste au centre des préoccupations. L’entraînement des LLM à grande échelle s’appuie déjà sur des stratégies de planification éprouvées telles que 1F1B ou 1F1B entrelacé, qui sont profondément intégrées dans les piles d’entraînement de niveau production. Plutôt que de remplacer ces méthodes établies, BigMac utilise la planification du LLM comme ligne de temps sous-jacente, en insérant stratégiquement les calculs de l’encodeur et du générateur lorsque les entrées sont prêtes, sans perturber l’ordre d’exécution du LLM. L’équipe qualifie cette architecture de « pipeline imbriqué quasi-sûr en termes de dépendances ».

Cette conception offre deux avantages clés. Premièrement, les fluctuations de durée de l’encodeur et du générateur ne se répercutent plus sur le pipeline du LLM, ce qui permet à ce dernier de conserver son rythme optimal. Deuxièmement, les besoins en mémoire pour les activations modales sont réduits à O(1) au niveau algorithmique. Les encodeurs n’ont plus besoin de conserver les activations de tous les micro-lots jusqu’à la fin du pipeline, et les générateurs évitent de prolonger les longues queues d’activation. En substance, BigMac ne sacrifie pas la mémoire au profit de pics de performance ; il restructure plutôt la logique de planification afin de rendre ces deux objectifs compatibles plutôt que mutuellement exclusifs.

Combler le fossé entre la conception de la planification et l’exécution effective implique d’importants défis techniques, notamment l’intégration du système, la définition des interfaces et le débogage des performances. BigMac est conçu pour relever ces défis spécifiques en offrant trois capacités principales. Tout d’abord, il rend la planification globale transparente : le planificateur génère un tableau d’opérateurs complet couvrant tous les rangs du pipeline, les micro-lots et les types de modules lors de l’exécution. L’exécuteur les répartit ensuite en séquences locales pour chaque rang, s’intégrant de manière transparente aux backends LLM tels que Megatron Core, aux environnements d’exécution modaux et aux couches de communication. Cette visibilité facilite la vérification et l’optimisation du backend.

Deuxièmement, BigMac propose une interface de parallélisme de pipeline qui reste invisible pour les ingénieurs en algorithmes. Les utilisateurs n’ont qu’à définir ce que chaque module produit et consomme ; le système gère en interne la division en étapes, l’activation et le transfert de gradients, ainsi que la communication inter-périphériques. Cela permet aux expériences multimodales vérifiées sur un seul GPU de s’adapter naturellement au parallélisme de pipeline. Troisièmement, il fournit une suite d’outils de profilage, de simulation et de visualisation. Ces outils décomposent les itérations d’entraînement jusqu’au niveau des opérateurs, révélant exactement ce que fait chaque rang à chaque pas de temps, identifiant les périodes d’inactivité et mettant en évidence les goulots d’étranglement liés aux dépendances. Le simulateur permet également aux équipes de tester diverses configurations de PP et combinaisons de micro-lots avant de se lancer dans un apprentissage à grande échelle, en estimant leur impact sur les bulles et le débit.

L’efficacité de BigMac a été validée sur deux charges de travail représentatives. Dans les tâches MLLM-Understanding, en utilisant Qwen3-30B-A3B comme architecture de base et un encodeur ViT de 1,3 milliard de paramètres, BigMac atteint un gain de vitesse de 1,08x à 1,1x par rapport à la référence Optimus, efficace sur le plan computationnel, et un gain de 1,6x à 1,9x par rapport à la référence Megatron-DistTrain, efficace en termes de mémoire. Point crucial : l’utilisation de la mémoire reste stable à mesure que la taille des lots par GPU augmente, tandis qu’Optimus subit des pics de consommation de mémoire qui finissent par entraîner des erreurs « Out-of-Memory » (OOM) avec des lots plus volumineux. Les tâches de génération MLLM sont plus complexes, impliquant un générateur MMDiT de 20 milliards, où l’écart de performances se creuse encore davantage : Optimus échoue pour toutes les tailles de lots testées en raison d’erreurs OOM, tandis que BigMac évite ce problème en exécutant efficacement le générateur à rebours et en libérant rapidement les activations. Par rapport à Megatron-DistTrain, BigMac offre toujours un gain de vitesse de 1,5 à 1,9 fois supérieur avec une utilisation stable de la mémoire. Cela met en évidence la valeur principale des pipelines imbriqués : prendre en charge à la fois les dépendances de l’encodeur et celles du générateur sans rétention excessive des activations ni temps d’inactivité significatif.

BigMac est désormais un composant essentiel de l’infrastructure d’entraînement des modèles multimodaux de Dots et fonctionne au sein de l’environnement de production de Xiaohongshu. L’article qui l’accompagne, intitulé « BigMac : Breaking the Pareto Frontier of Compute and Memory in Multimodal LLM Training », est disponible sur arXiv, ainsi que des exemples interactifs de traces générées par PP Profiler. Pour les chercheurs et les ingénieurs confrontés au compromis entre mémoire et vitesse dans l’entraînement multimodal, ce projet open source validé en production offre une solution pratique permettant un gain de temps et d’efforts considérable.