Técnicas mantêm treinamento de MoE gigantes dentro da memória fixa de GPUs

O problema dos picos de memória no treinamento de MoE

Um artigo pré-publicado no arXiv (2609.14306) por Shrey Pandit, Xuan-Phi Nguyen, Yiran Zhao e Shafiq Joty ataca um gargalo do treinamento de modelos Mixture-of-Experts (MoE) em contextos longos ou com batches grandes: o treino falha assim que qualquer componente excede a memória da GPU em seu pico de alocação, de modo que o alvo precisa ser todos os picos ao mesmo tempo, não o consumo médio.

Os autores identificam quatro componentes que ficam sem limites nos planos de paralelismo de uso comum, cada um crescendo de forma diferente: o despacho de experts, que depende da matriz de roteamento; a projeção de vocabulário, que escala com tokens vezes vocabulário; os limites de gradient checkpointing, que crescem com profundidade vezes comprimento de sequência; e o estado do otimizador, proporcional ao número de parâmetros. Como o componente que estoura primeiro varia conforme o modelo, o comprimento do contexto e a quantidade de dispositivos, reduzir apenas o maior pico apenas expõe o próximo.

Quatro técnicas com conjunto de trabalho fixo

O artigo propõe limitar os quatro picos com escalonamentos cujo conjunto de trabalho da GPU é fixado no lançamento:

• PipelinedLLEP: estende o paralelismo de experts por menor carga (least-loaded) com um teto sobre os tokens que cada origem contribui para um bloco de despacho;
• Ring-DTP: circula ativações ou fragmentos de pesos em anel na projeção de vocabulário e incorpora cada bloco de logits em um log-sum-exp online;
• Selective checkpoint offload (SCO): mantém na memória da CPU o único tensor de longa duração de cada limite de checkpoint;
• OffloadStreamAdamW: transforma a atualização serial de Adam, executada na CPU no offload do otimizador, em um pipeline por buckets.

Segundo os autores, todas as quatro técnicas alteram apenas a ordem e a granularidade de computação e movimento de dados, de modo que a perda e os gradientes permanecem exatos.

Resultados reportados

Em testes pareados por componente, as técnicas reduzem o pico de despacho do MoE em até 59,3% sem perda de throughput, o pico da projeção de vocabulário em 86,6%, e aceleram em 2,05 vezes a etapa do otimizador com offload.

Compostas em modelos MoE de 120B a 667B parâmetros, as técnicas permitem treinar com comprimento de contexto de 1M de tokens - alcance 8 a 32 vezes maior que o de uma linha de base FSDP2 ajustada - com throughput de até 10,4 vezes o dessa linha de base.

O trabalho é classificado como preprint na área de Computação Distribuída, Paralela e de Clusters, e os números apresentados vêm dos próprios autores do artigo.