Federação de especialistas: inferência distribuída eficiente em comunicação para grandes modelos de linguagem

Resumo

A mistura de especialistas emergiu como o principal mecanismo para tornar os Modelos de Linguagem de Grande Porte (LLMs) computacionalmente eficientes. No entanto, em ambientes distribuídos, a comunicação de embeddings de tokens entre especialistas constitui um gargalo significativo. Apresentamos a nova arquitetura Federação de Especialistas (FoE). A FoE reestrutura o bloco MoE de uma camada de transformador em múltiplos clusters MoE. Cada cluster é responsável por apenas uma das cabeças KV, e o paralelismo de especialistas é aplicado entre esses especialistas. Entre os clusters, uma soma sincroniza os resíduos pós-atenção, os quais, por sua vez, orientam o roteamento e a distribuição para o próximo bloco MoE. Em um ambiente de nó único, a FoE elimina completamente a comunicação todos-para-todos, pois todos os especialistas dentro de um grupo estão contidos na mesma GPU. Em ambientes multinó, a FoE confina a comunicação todos-para-todos à infraestrutura intra-nó, reduzindo assim significativamente a sobrecarga de comunicação. Uma implementação da FoE demonstra que, no LongBench, a FoE melhora significativamente a taxa de transferência de inferência e a latência tanto em ambientes de nó único quanto multinó, reduzindo a latência de passagem direta de ponta a ponta em até 5,2 vezes, o TTFT em 3,62 vezes e o TBT em 1,95 vezes. Isso é alcançado mantendo uma qualidade de geração comparável à de um modelo de mistura de especialistas do mesmo tamanho e configuração de treinamento.

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.