Fédération d'experts : Inférence distribuée efficace en communication pour les grands modèles de langage

Résumé

Le mélange d'experts est devenu le mécanisme principal pour rendre les grands modèles de langage (LLMs) efficaces sur le plan computationnel. Cependant, dans les environnements distribués, la communication des plongements de jetons entre experts constitue un goulot d'étranglement significatif. Nous présentons la nouvelle architecture de Fédération d'experts (FoE). FoE restructure le bloc MoE d'une couche de transformateur en plusieurs clusters MoE. Chaque cluster est responsable d'une seule des têtes KV et le parallélisme entre experts est appliqué entre ces experts. Entre les clusters, une somme synchronise les résidus post-attention, qui pilotent ensuite le routage et la distribution pour le bloc MoE suivant. Dans un environnement mono-nœud, FoE élimine complètement la communication tous-à-tous car tous les experts d'un groupe sont contenus sur le même GPU. Dans les environnements multi-nœuds, FoE confine la communication tous-à-tous au réseau intra-nœud, réduisant ainsi considérablement la surcharge de communication. Une implémentation de FoE montre que sur LongBench, FoE améliore significativement le débit d'inférence et la latence dans les environnements mono-nœud et multi-nœuds, réduisant la latence de passage avant de bout en bout jusqu'à 5,2x, le TTFT de 3,62x et le TBT de 1,95x. Ceci tout en atteignant une qualité de génération comparable à celle d'un modèle de mélange d'experts de même taille et configuration d'entraînement.

English

Mixture of experts has emerged as the primary mechanism for making Large Language Models (LLMs) computationally efficient. However, in distributed settings, communicating token embeddings between experts is a significant bottleneck. We present the novel Federation of Experts (FoE) architecture. FoE restructures the MoE block of a transformer layer into multiple MoE clusters. Each cluster is responsible for only one of the KV heads and expert parallelism is applied between those experts. Between clusters, a sum synchronizes the post-attention residuals, which then drives routing and dispatch for the next MoE block. In a single-node setting, FoE completely eliminates all-to-all communication as all experts within a group are contained on the same GPU. In multi-node settings, FoE confines all-to-all communication to the intra-node fabric, thus significantly reducing communication overhead. An implementation of FoE finds that on LongBench, FoE significantly improves inference throughput and latency in both single-node and multi-node settings, reducing the end-to-end forward-pass latency by up to 5.2x, TTFT by 3.62x, and TBT by 1.95x. It does so while achieving comparable generation quality to a mixture of experts model of the same size and training configuration.