LRMoE poda a passagem reversa do MoE e acelera o treinamento em 21–45%

Uma equipe de autores do meio acadêmico propôs o LRMoE, um sistema para treinar modelos MoE em nós nos quais a memória das GPUs não é suficiente. MoE (Mixture-of-Experts, «mistura de especialistas») é uma arquitetura em que a camada é formada por dezenas ou centenas de subredes pequenas, os especialistas, e o roteador ativa apenas algumas delas para cada token. Cada uma dessas escolhas «token — especialista» é chamada pelos autores de par.

O principal número do trabalho: nos modelos Qwen3-30B-A3B e GLM4.5-106B-A12B, o tempo acumulado de treinamento das camadas MoE até um valor previamente definido da função de perda (loss) caiu 21,2–44,8% em comparação com o SYMI. Esse é o sistema anterior, que separa o armazenamento dos parâmetros do modelo e dos estados do otimizador para posicionar os especialistas com mais flexibilidade.

Por que o desbalanceamento de carga não se resolve com cópias de especialistas

Os especialistas ficam distribuídos entre as GPUs, e os tokens viajam até as placas certas por meio do all-to-all, uma troca de «todos com todos». Os especialistas mais populares recebem mais tokens, suas GPUs viram gargalo e as demais ficam esperando. O passo síncrono avança na velocidade do participante mais sobrecarregado.

As respostas usuais a esse problema são reposicionar os especialistas ou criar cópias dos mais populares. É assim que funcionam FasterMoE, FlexMoE, SmartMoE, LAER-MoE, UltraEP e TopoEP. Mas as cópias têm um custo:

• a réplica ocupa memória que seria necessária para os buffers de tokens e as ativações salvas;
• as cópias precisam ser sincronizadas, e os gradientes, somados;
• o volume total de trabalho não diminui, apenas se desloca para outras placas.

Os autores mostram que, com um limite rígido de memória na GPU, mesmo uma distribuição perfeitamente uniforme pode deixar a carga total acima do que a memória e a rede conseguem atender. Desequilíbrio e sobrecarga geral são problemas diferentes, e o balanceamento resolve só o primeiro.

O que o LRMoE faz: forward completo, backward podado

A ideia é tratar a carga da passagem reversa (backward) não como consequência inevitável do roteamento, mas como uma grandeza ajustável. O trabalho segue em duas etapas.

• A passagem direta (forward) continua completa. Todos os pares token–especialista são calculados como de costume, porque nessa fase ainda não há gradientes e não existe base para avaliar o valor de cada par.
• Cotas antes do cálculo. O controlador vê quantos tokens chegaram a cada especialista e, levando em conta a memória e a rede, decide quantos pares de cada especialista manter na passagem reversa. Ele também planeja o posicionamento das réplicas. A cota fica fixa durante toda a chamada da camada.
• Seleção pela métrica LRM. Quando os gradientes chegam à camada, estima-se, para cada par, a contribuição para a atualização dos pesos do especialista. A contribuição de um par é o produto externo entre o gradiente de entrada e a entrada salva da projeção. Por isso, o quadrado de sua norma se decompõe no produto de três fatores: o quadrado do peso de roteamento, o quadrado da norma do gradiente e o quadrado da norma da entrada. Não é preciso calcular as matrizes completas de gradientes para isso. Dentro da cota, ficam os pares com maior pontuação.
• Redistribuição. Os pares restantes são distribuídos, por um algoritmo de programação inteira, entre as cópias permitidas do seu especialista. O objetivo é minimizar o pico de carga em uma única GPU. Nenhuma réplica nova é criada na passagem reversa.
• Compactação física. As linhas restantes de gradientes e ativações são reunidas em blocos densos por destinatário e por especialista. A troca all-to-all fica mais compacta, e as multiplicações de matrizes rodam como grouped GEMM de tamanho variável, sem preenchimento com zeros. Depois, o resultado volta ao layout original.

Um detalhe importante: simplesmente mascarar as linhas excedentes em tensores de tamanho completo não traz nenhum ganho, nem de tempo, nem de memória. O benefício só aparece quando os tamanhos reais das transferências e das matrizes diminuem.

O que mostraram os experimentos e onde está o porém

Além do número principal, de 21,2–44,8%, os autores fizeram uma reprodução separada em GPU e obtiveram redução de 57,8% na memória total usada para o posicionamento dos especialistas, as réplicas e os buffers de tokens, em comparação com a propriedade fixa dos especialistas. Essa é uma medição de memória, não do tempo de treinamento.

O porém é que a passagem reversa se torna aproximada: parte das contribuições ao gradiente é simplesmente descartada. A métrica LRM é uma estimativa no nível de uma única projeção e não oferece garantias de convergência. Por isso, os autores comparam os sistemas pelo tempo até um loss fixo, e não pela velocidade do passo: o ganho só vale se o modelo chegou à mesma qualidade. Os custos extras de calcular as pontuações, empacotar e rotear estão incluídos nas medições.

Para os engenheiros, a conclusão é a seguinte. Se o treinamento de MoE esbarra na memória da GPU e na rede, e não há mais espaço para adicionar réplicas, vale olhar para a redução do próprio trabalho da passagem reversa, e não apenas para sua redistribuição. As ressalvas também importam. Os resultados vêm de dois modelos e foram comparados com um único sistema de referência, então a transferência para outras arquiteturas e para execuções longas precisa ser verificada à parte. O custo dos gradientes aproximados pode aparecer em metas de qualidade mais exigentes.