Beyond Task-Agnostic: Task-Aware Grouping for Communication-Efficient Multi-Task MoE Inference

Resumo Técnico: Além do Agnóstico à Tarefa: Agrupamento Consciente da Tarefa para Inferência de MoE Multitarefa com Eficiência de Comunicação

1. Definição do Problema

Modelos de Mistura de Especialistas (MoE) de ativação esparsa escalam a capacidade de forma eficiente via computação condicional, mas enfrentam desafios sistêmicos significativos durante a inferência distribuída, especificamente o excesso de overhead de comunicação entre GPUs e um desequilíbrio de carga severo. As estratégias existentes de posicionamento de especialistas tentam mitigar esses custos ao co-localizar especialistas frequentemente co-ativados, mas tipicamente derivam um plano de implantação global único a partir de traços de roteamento agnósticos à tarefa (agregados globalmente).

Os autores identificam uma falha crítica nesta abordagem: a co-ativação de especialistas é fortemente condicionada à tarefa. Conforme demonstrado na análise do artigo (Figura 1), pares de especialistas fortemente acoplados em uma família de tarefas (ex: geração de código) são frequentemente não correlacionados ou até mesmo concorrentes em outra (ex: raciocínio matemático). A média desses padrões heterogêneos em um único grafo de afinidade colapsa os grupos de especialistas específicos de cada tarefa, deixando a localidade de comunicação inexplorada e falhando em se adaptar a mudanças de distribuição em cargas de trabalho online.

2. Metodologia

O framework proposto, Task-Aware Coactivation Grouping (TACG) em conjunto com o Generic Expert Shared Replication (GESR), aborda esses problemas através de um processo de implantação em dois estágios:

A. Task-Aware Coactivation Grouping (TACG)
O TACG converte estatísticas de co-ativação condicionadas à tarefa em um plano de posicionamento viável em termos de capacidade:

1. Pontuação de Preferência de Família de Tarefas: Usando traços de calibração, o sistema computa vetores de uso e matrizes de co-ativação para cada família de tarefas. Ele calcula uma "pontuação de preferência suave" ($s^{(f)}_{l,e}$) para cada especialista em relação a famílias de tarefas específicas, normalizando essas pontuações para criar uma distribuição de probabilidade ($p^{(f)}_{l,e}$) que representa a afinidade de um especialista por uma família.
2. Grafo de Co-ativação Modulado pela Tarefa: Em vez de usar um grafo de co-ativação agrupado, o TACG constrói um grafo modulado pela tarefa ($\tilde{A}_l$). Este grafo repondera as arestas de co-ativação padrão pelo produto das preferências de família dos especialistas ($p^{(f)}_{l,e} p^{(f)}_{l,e'}$). Isso garante que as arestas entre especialistas que co-ativam e compartilham uma forte preferência de família sejam amplificadas, enquanto arestas entre diferentes famílias sejam atenuadas.
3. Agrupamento Viável em Capacidade: O framework particiona os especialistas em $M$ grupos (um por GPU) maximizando a afinidade intra-grupo no grafo modulado pela tarefa sob restrições estritas de capacidade. Isso é resolvido via uma abordagem de agrupamento espectral de dois estágios seguida por um passo de "reparo de capacidade" local para garantir que os limites exatos de carga da GPU sejam atendidos.

B. Generic Expert Shared Replication (GESR)
Para lidar com o desvio de carga causado por especialistas "genéricos" (aqueles ativos em todas as famílias e propensos a se tornarem hotspots), o GESR introduz um mecanismo companheiro leve:

1. Identificação: Ele identifica especialistas genéricos usando uma pontuação que recompensa a centralidade global e a consistência entre famílias, enquanto penaliza desvios fortes específicos de tarefas.
2. Replicação: Esses especialistas são replicados em um pequeno conjunto de GPUs secundárias escolhidas com base na compatibilidade estrutural (afinidade de co-ativação).
3. Guarda Online: No momento do serviço, uma regra de seleção consciente de localidade e carga escolhe o melhor candidato de GPU para um especialista genérico, usando sinais de carga com decaimento exponencial para evitar hotspots sem alterar o roteador do MoE.

3. Principais Contribuições

• Insight: O artigo estabelece que a co-ativação de especialistas em serviço MoE multitarefa é fortemente condicionada à tarefa, argumentando que o posicionamento eficaz deve ser impulsionado pela co-ativação consciente da tarefa, em vez de um único grafo de afinidade média global.
• TACG: Uma estratégia consciente da tarefa que repondera o grafo de co-ativação usando as preferências especialista-família, produzindo um posicionamento que generaliza estritamente o agrupamento agnóstico à tarefa.
• GESR: Um módulo de replicação leve que replica seletivamente especialistas genéricos entre famílias e realiza a seleção online consciente da carga para manter o equilíbrio sem sacrificar a eficiência de comunicação.
• Validação Empírica: Avaliação extensiva em três modelos MoE de código aberto representativos (DeepSeek-MoE-16B, Qwen1.5-MoE-A2.7B, Moonlight-16B-A3B) demonstrando melhorias consistentes sobre bases fortes.

4. Resultados Experimentais

O framework foi avaliado em uma configuração de paralelismo de especialistas de 16 GPUs em três modelos com quatro famílias de tarefas (Código, Consulta, Matemática, Raciocínio Complexo).

• Redução de Comunicação: O framework completo (TACG + GESR) reduziu o custo médio de comunicação por token em 31,39% em comparação com o baseline, superando a estratégia de posicionamento anterior mais forte (ExFlow) em 12,58%.
• Equilíbrio de Carga: O método manteve um alto índice de equidade de Jain médio de 0,9975 e uma baixa violação de carga máxima média (MaxVio) de 0,0736, igualando ou excedendo o equilíbrio de carga dos baselines enquanto reduzia significativamente a comunicação.
• Robustez: Sob severas mudanças de distribuição de tarefas (cargas de trabalho de tarefa única fora da distribuição), a abordagem consciente da tarefa manteve reduções de comunicação estáveis, enquanto variantes agnósticas à tarefa sofreram grandes oscilações de desempenho e formação de hotspots (MaxVio subindo ~7x em tarefas de Código).
• Ablação: A remoção do sinal de co-ativação modulado pela tarefa reduziu os ganhos de comunicação de 31,39% para 16,51%, confirmando que o condicionamento à tarefa é o fator dominante. O GESR mostrou-se essencial para controlar as violações de carga (MaxVio) sem sacrificar a eficiência de comunicação.

5. Significância e Alegações

O artigo afirma que sua significância reside na mudança do paradigma de posicionamento de especialistas MoE de uma visão estática e agnóstica à tarefa para uma perspectiva dinâmica e consciente da tarefa. Ao modelar explicitamente a heterogeneidade da co-ativação de especialistas entre diferentes famílias de tarefas, o framework desbloqueia a localidade de comunicação que é invisível para métodos de agregação global.

Os autores afirmam que esta abordagem fornece uma melhoria estrita na fronteira de equilíbrio comunicação-carga. Diferente de métodos anteriores que frequentemente sacrificam o equilíbrio de carga em prol de ganhos de comunicação (ou vice-versa), o TACG + GESR alcança uma redução substancial de comunicação enquanto preserva um equilíbrio de carga quase perfeito. O trabalho destaca que o serviço multitarefa eficaz requer estratégias de implantação que respeitem a estrutura intrínseca da colaboração de especialistas específica de cada tarefa, em vez de apenas tirá-la pela média.

O artigo permanece modesto quanto ao seu escopo, observando que foca no posicionamento offline estático e deixa como trabalho futuro o reagrupamento dinâmico dentro de janelas de requisição, co-design de memória (pré-busca/cache) e especialistas especializados em modelos particionados por domínio.