Resumo Técnico: DisagMoE

1. Declaração do Problema

Arquiteturas de Mistura de Especialistas (MoE) em grande escala permitem o treinamento de Modelos de Linguagem de Grande Escala (LLMs) com trilhões de parâmetros, ativando esparsamente um subconjunto de especialistas. No entanto, o treinamento desses modelos enfrenta gargalos severos devido à comunicação all-to-all exigida pelo paralelismo de especialistas (EP). À medida que os modelos escalam em múltiplos nós de GPU, a largura de banda limitada da rede inter-nós (por exemplo, InfiniBand) exacerba a latência de comunicação, que pode consumir até 50% do tempo de uma etapa de treinamento.

Estratégias de otimização existentes focam principalmente em sobrepor comunicação e computação no nível do operador (por exemplo, fragmentando cálculos de FFN) ou usando pipeline de micro-lotes. No entanto, essas abordagens sofrem de duas limitações fundamentais:

1. Paradas Residuais: Desequilíbrios inerentes nas razões de computação-comunicação entre as camadas de atenção e FFN impedem uma sobreposição perfeita. A computação de atenção escala quadraticamente com o comprimento da sequência, enquanto FFN e comunicação escalam linearmente. Isso cria "caudas de comunicação" que não podem ser ocultas, especialmente à medida que os comprimentos das sequências aumentam.
2. Ineficiência de Recursos: Abordagens desagregadas anteriores (separando atenção e FFN em hardware diferente) frequentemente exigem que cada grupo mantenha parâmetros completos do modelo, levando a erros de Memória Insuficiente (OOM) em escala, ou falham em adaptar a alocação de recursos às características específicas de computação-comunicação de cada componente.

2. Metodologia: DisagMoE

DisagMoE é um sistema de treinamento projetado para mitigar gargalos all-to-all, otimizando conjuntamente o posicionamento do modelo e o agendamento através de uma arquitetura desagregada.

2.1 Posicionamento de Componentes Desagregados

Diferentemente do Paralelismo de Especialistas tradicional, onde camadas de atenção e FFN coexistem no mesmo grupo de GPU, o DisagMoE particiona o modelo por tipo de componente:

• Grupos de Atenção (A-Workers): Estes grupos contêm camadas de auto-atenção de múltiplas camadas de transformador. Eles replicam componentes densos e distribuem dados de treinamento via Paralelismo de Dados (DP).
• Grupos de FFN (F-Workers): Estes grupos contêm camadas FFN (especialistas) de múltiplas camadas de transformador. Os especialistas são particionados (sharded) entre GPUs dentro do grupo via Paralelismo de Especialistas (EP).
• Atribuição de Camadas: As camadas são atribuídas aos grupos de forma intercalada (por exemplo, o Grupo 0 lida com as camadas 0, $p$, $2p$; o Grupo 1 lida com as camadas 1, $p+1$, etc.), maximizando a utilização do dispositivo sob restrições de memória.

2.2 AF-Pipe: Pipeline Multi-Estágio com Sobreposição

DisagMoE introduz o AF-Pipe, um agendamento de pipeline que trata a comunicação como um estágio de primeira classe, e não como um efeito colateral.

• Comunicação Muitos-para-Muitos (M2N): Em vez de trocas all-to-all padrão dentro de um único grupo de nós, o AF-Pipe estabelece trocas de tokens unidirecionais, muitos-para-muitos, entre os grupos de Atenção e FFN.
• Alinhamento de Estágios: O pipeline alinha limites entre grupos, fundindo as operações tradicionais de "dispatch" (envio) e "combine" (combinação) em um único estágio de comunicação unificado. Isso reduz a sobrecarga de comunicação em aproximadamente $1/k$ (onde $k$ é o número de top-$k$ especialistas) e elimina transferências de dados redundantes.
• Execução Assíncrona: O sistema utiliza três fluxos coordenados (forward, backward e comunicação) dentro de cada worker. Isso permite que o grupo FFN envie gradientes (N2M) enquanto simultaneamente computa a passagem forward nos dados disponíveis, garantindo execução não bloqueante e utilização quase contínua da GPU.

2.3 Alocação Adaptativa de Workers

Para abordar o desequilíbrio entre atenção (limitada por computação) e FFN (limitada por comunicação), o DisagMoE emprega um alojador adaptativo guiado por um Modelo de Teto Computação-Comunicação.

• Modelagem: O sistema modela o ponto de virada ($\hat{I}$) onde um componente muda de limitado pela rede para limitado por computação. Camadas de atenção atingem o teto de computação mais rapidamente do que camadas FFN à medida que o comprimento da sequência aumenta.
• Otimização: O sistema formula a partição de GPU e Placa de Interface de Rede (NIC) como um problema de Programação Linear Inteira Mista (MILP). Ele aloca dinamicamente mais largura de banda de rede para grupos FFN (para elevar seu teto de comunicação) e mais recursos de computação para grupos de Atenção, equilibrando os estágios do pipeline para minimizar bolhas.
• Estratégia de Busca: Um algoritmo usa o modelo de teto para gerar uma configuração inicial (seed), que é então refinada via busca local usando tempos medidos por etapa para contabilizar sobrecargas não modeladas.

3. Contribuições Principais

1. Posicionamento de Componentes Desagregados: Uma arquitetura inovadora que separa camadas de atenção e FFN em grupos de workers disjuntos, permitindo escalabilidade independente e alocação de recursos.
2. Agendamento AF-Pipe: Um pipeline multi-estágio que substitui o all-to-all convencional por trocas M2N inter-grupo, sobrepondo sistematicamente a comunicação com as computações de atenção e FFN.
3. Alocação Adaptativa de Recursos: Um mecanismo guiado por modelo de teto que dimensiona conjuntamente recursos de GPU e NIC para grupos de atenção e FFN, equilibrando efetivamente a razão computação-comunicação através de comprimentos de sequência e configurações de modelo variáveis.

4. Resultados Experimentais

Os autores avaliaram o DisagMoE em clusters de 16 nós (128× GPUs H800) usando três modelos MoE: DeepSeek-MoE, GPT-OSS-120B e Qwen3, com comprimentos de sequência variando de 4K a 32K.

• Melhoria de Throughput: DisagMoE superou consistentemente bases fortes, incluindo Megatron-LM (1F1B intercalado), Tutel, Comet e DualPipe.
  • Alcançou até 1,81× de aceleração sobre Megatron-1F1B.
  • Alcançou até 1,34× de aceleração sobre sistemas de sobreposição MoE mais avançados.
• Redução de Comunicação: Comparado às bases, o DisagMoE reduziu o tempo de comunicação não sobreposta em até 88% (vs. Tutel), 75% (vs. Comet) e 45% (vs. DualPipe).
• Escalabilidade: O sistema demonstrou eficácia através de diferentes valores de top-$k$ e tamanhos de Paralelismo de Especialistas (EP).
• Sensibilidade a Recursos: Experimentos confirmaram que as proporções ótimas de GPU entre grupos de atenção e FFN variam com o comprimento da sequência (por exemplo, sequências mais longas favorecem mais recursos de atenção), validando a necessidade de alocação adaptativa.

5. Significado e Alegações

O artigo afirma que o DisagMoE aborda uma limitação fundamental no treinamento atual de MoE: a incapacidade de sobrepor perfeitamente a comunicação devido ao desequilíbrio estrutural entre camadas de atenção e FFN. Ao desagregar esses componentes e alocar recursos adaptativamente com base em um modelo de teto, o sistema expande a janela de sobreposição efetiva além do que a provisão simétrica de largura de banda pode alcançar.

Os autores enfatizam que sua abordagem é particularmente significativa para cargas de trabalho de pré-treinamento em grande escala com comprimentos de sequência fixos. Eles observam que, embora seu design atual assuma formas estáveis (comprimento de sequência fixo e micro-lote), a ideia central — de que hardware e software futuros devem visar razões favoráveis de computação-comunicação entre componentes, em vez de simplesmente escalar a largura de banda bruta — oferece uma nova direção para escalar o treinamento de MoE com eficiência. O sistema é implementado no Megatron-LM e demonstrou ser uma solução viável para treinar modelos com centenas de bilhões de parâmetros sem as restrições de memória que afligiram tentativas anteriores de treinamento desagregado.