Équilibrage de charge par lot global : presque un déjeuner gratuit pour améliorer l’entraînement des LLM MoE

L'architecture Mixture-of-Experts (MoE) est une technique populaire pour augmenter le nombre de paramètres d'un modèle, mais les frameworks d'entraînement existants comme Megatron-core utilisent une perte d'équilibrage au niveau du micro-lot, ce qui peut nuire aux performances lorsque les micro-lots contiennent des données homogènes (par exemple, uniquement du code). Alibaba Qwen introduit une perte d'équilibrage au niveau du lot global, qui synchronise la fréquence de sélection des experts sur tous les groupes parallèles et calcule la perte de manière globale. Des expériences avec trois configurations MoE (3,4B/0,6B activés, 15B/2,54B activés, 43B/6,6B activés) et jusqu'à 400 milliards de tokens montrent que l'équilibrage global surpasse systématiquement l'équilibrage par micro-lot sur diverses tâches. Il conduit également à une spécialisation significative des experts par domaine, absente dans les modèles équilibrés par micro-lot. L'ajout d'une petite perte d'équilibrage par micro-lot (poids 0,01) en plus de la perte globale améliore la vitesse d'entraînement sans affecter l'efficacité. La méthode est quasiment gratuite car la synchronisation ne nécessite qu'un vecteur de taille égale au nombre d'experts. Ce travail est publié dans la prépublication arXiv 2501.11873.