🤖 machine learning

DMuon: Efficient Distributed Muon Training with Near-Adam Overhead

O artigo apresenta o DMuon, uma implementação distribuída de código aberto do otimizador Muon que se integra perfeitamente em pipelines de treinamento existentes para reduzir drasticamente a latência do passo do otimizador e alcançar eficiência próxima à do AdamW, permitindo assim o uso escalável de otimização baseada em ortogonalização de matrizes em grandes modelos de fundação.

Autores originais: Vincent Chen, Starrick Liu, Regis Cheng, Dance Yang, Shalfun Li, Ryan Yu, Lucy Liang, Hang Su, Roy Gan, Hao Wang, Qian Wang

Publicado 2026-06-26

📖 5 min de leitura🧠 Leitura aprofundada

CC BY 4.0

Autores originais: Vincent Chen, Starrick Liu, Regis Cheng, Dance Yang, Shalfun Li, Ryan Yu, Lucy Liang, Hang Su, Roy Gan, Hao Wang, Qian Wang

Artigo original sob licença CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Esta é uma explicação gerada por IA do artigo abaixo. Não foi escrita nem endossada pelos autores. Para precisão técnica, consulte o artigo original. Ler aviso legal completo

Imagine que você está treinando o cérebro de um robô gigante (um Modelo de Linguagem de Grande Escala ou uma IA incorporada) para aprender. Para ensiná-lo, você precisa de um "otimizador" — um treinador que observa os erros do robô e ajusta suas conexões internas (pesos) para torná-lo mais inteligente.

Por anos, o treinador padrão foi o AdamW. É como uma equipe de milhares de pequenos trabalhadores, cada um consertando um pequeno parafuso do robô de cada vez. É rápido e eficiente, mas trata cada parafuso de forma independente.

Recentemente, um treinador mais inteligente chamado Muon foi descoberto. Em vez de consertar parafusos um por um, o Muon olha para painéis inteiros do robô (matrizes completas de dados) e os ajusta todos juntos. Essa abordagem de "terapia de grupo" ajuda o robô a aprender mais rápido e a atingir um desempenho mais alto. No entanto, há um porém: o Muon é incrivelmente lento para rodar em um sistema distribuído (um cluster de muitos computadores trabalhando juntos).

O Problema: O Gargalo da Reunião de "Todos à Mão"
Em uma configuração de treinamento distribuído, o cérebro do robô é dividido entre muitos computadores (GPUs).

AdamW funciona como uma equipe remota: o Computador A conserta seus parafusos, o Computador B conserta os dele, e eles não precisam conversar muito entre si.
Muon funciona como um comitê: Para consertar um painel, cada computador precisa ver o painel inteiro primeiro.

Em uma implementação ingênua do Muon, cada computador tem que parar o que está fazendo, baixar os dados completos de todos os outros, realizar a matemática complexa e então enviar os resultados de volta. Isso é como realizar uma reunião massiva onde todos leem um relatório de 1.000 páginas antes de tomar uma única decisão. Leva tanto tempo que a "reunião" (o passo do otimizador) demora mais do que o próprio trabalho (os passos de aprendizado). Na verdade, o artigo observa que isso pode custar 2x mais tempo do que o próprio aprendizado!

A Solução: DMuon (Distributed Muon)
A Equipe de Robôs X Square construiu o DMuon, um novo sistema que faz o Muon rodar quase tão rápido quanto o AdamW. Eles fizeram isso mudando como a equipe trabalha, não mudando a matemática do treinador.

Aqui está como eles resolveram o gargalo usando três truques principais:

1. A Estratégia do "Especialista Designado" (Execução Centrada no Proprietário)

Em vez de todos tentarem consertar todos os painéis, o DMuon atribui um computador específico para ser o "Proprietário" de cada painel.

Antes: Todos reuniam os dados, todos faziam a matemática, todos perdiam tempo.
Agora: Apenas o computador "Proprietário" reúne os dados e faz a matemática pesada. Os outros computadores apenas esperam ou fazem outro trabalho.
Analogia: Imagine um canteiro de obras. Em vez de cada trabalhador tentar construir a casa inteira, um especialista é designado para o telhado, outro para o encanamento. Os outros não tentam construir o telhado; eles apenas entregam as ferramentas que o especialista precisa. Isso impede que todos façam o mesmo trabalho duas vezes.

2. O Sistema de "Esteira de Produção" (Sobreposição de Comunicação)

Enquanto o "Proprietário" está fazendo a matemática, os outros computadores não ficam apenas parados. O DMuon cria um pipeline:

Passagem Direta (Forward Pass): Enquanto o robô está "pensando" (processando dados), o sistema envia silenciosamente o próximo conjunto de ferramentas para os especialistas em segundo plano.
Passagem Reversa (Backward Pass): Enquanto o robô está "aprendendo com os erros", o sistema já está reunindo o próximo lote de dados para os especialistas.
Analogia: É como a cozinha de um restaurante. O chef (o Proprietário) está cortando vegetais. Em vez de esperar o chef terminar antes de o garçom trazer o próximo pedido, o garçom traz o próximo pedido enquanto o chef ainda está cortando. O tempo de espera é escondido dentro do tempo de corte.

3. A "Linha de Montagem Inteligente" (Loteamento e Ajuste)

A matemática que o Muon faz é complexa. Às vezes, as peças são enormes (como uma parede gigante) e, às vezes, são minúsculas (como um pequeno azulejo).

O Problema: Se você tentar processar um azulejo minúsculo em uma máquina gigante, a máquina ficará ociosa. Se você processar uma parede gigante, levará uma eternidade.
A Correção: O DMuon agrupa muitos azulejos minúsculos em um único "lote" para que a máquina permaneça ocupada. Ele também usa um "ajustador" que encontra automaticamente a maneira mais rápida de cortar cada formato específico de azulejo.
Analogia: Em vez de enviar um caminhão de entrega para um único envelope, o DMuon espera até ter uma carga completa de envelopes e os envia todos de uma vez. Ele também altera a rota do caminhão dependendo se o destino é uma cidade ou uma fazenda.

Os Resultados

O artigo testou o DMuon em modelos do mundo real, incluindo um modelo de treinamento de robô (Wall-OSS) e um modelo de linguagem (Qwen2.5).

Velocidade: O DMuon tornou o "passo do otimizador" (a reunião) de 6 a 163 vezes mais rápido do que a antiga forma ingênua.
Eficiência: O tempo total para treinar um passo é agora apenas 2% mais lento do que o método padrão AdamW.
Conclusão: O DMuon permite que as equipes usem o estilo de aprendizado superior de "terapia de grupo" do Muon sem pagar a enorme penalidade de tempo. Ele transforma um processo lento e desajeitado em um processo suave e eficiente, tornando-o pronto para uso em produção em modelos de IA massivos.

Em resumo, o DMuon é o sistema de controle de tráfego que permite que o treinador mais inteligente (Muon) rode em uma rodovia de computadores sem causar um congestionamento, finalmente tornando prático o treinamento da próxima geração de IA.

Resumo Técnico: DMuon – Treinamento de Muon Distribuído Eficiente com Overhead Próximo ao Adam

1. Declaração do Problema

Otimizadores conscientes de matrizes, particularmente o Muon, demonstraram propriedades de convergência superiores em comparação com otimizadores elemento a elemento como o AdamW, ao aplicar iterações de Newton–Schulz (NS) aos gradientes agregados pelo momentum das matrizes de pesos. Esta abordagem direciona os valores singulares das atualizações para um espectro uniforme, oferecendo ganhos significativos de eficiência computacional no pré-treinamento de grandes modelos (ex: LLMs e modelos de fundação incorporados/embodied).

No entanto, implantar o Muon em ambientes de treinamento distribuído modernos enfrenta um gargalo crítico de sistemas. As pilhas de treinamento distribuído contemporâneas (ex: ZeRO, FSDP) operam sob um contrato implícito de otimização elemento a elemento, onde as atualizações são computadas independentemente nos fragmentos (shards) de parâmetros. O Muon viola este contrato porque sua iteração central de NS exige que a matriz de peso completa esteja presente em um único dispositivo para computar os produtos matriciais necessários (especificamente $XX^\top$ ).

Em uma implementação distribuída ingênua ("Gather-then-Compute"), cada rank deve:

Realizar o all-gather da matriz de gradiente reduzida completa de todos os shards.
Executar a iteração NS localmente.
Descartar os resultados para os shards que não possui.

Esta abordagem introduz dois custos principais:

Sobrecarga de Comunicação: Reconstruir matrizes completas a cada passo introduz comunicação coletiva que escala com o tamanho do modelo e a largura distribuída.
Computação Redundante: Cada rank realiza a mesma ortogonalização dispendiosa, multiplicando o custo computacional pelo número de ranks ( $D$ ).

Na prática, estes custos fazem com que o tempo do passo do otimizador no Muon distribuído rivalize ou até exceda o tempo de execução (wall-clock time) combinado das passagens forward e backward, erodindo as vantagens algorítmicas do Muon. Esta penalidade é particularmente severa para o treinamento de modelos incorporados (ex: VLA), onde as passagens forward-backward são computacionalmente mais leves em relação ao passo do otimizador comparado ao pré-treinamento padrão de LLMs.

2. Metodologia: Design de Sistema DMuon

DMuon é uma implementação distribuída de código aberto projetada para fechar a lacuna de desempenho entre o Muon e o AdamW. Ele alcança isso através de um design de sistema co-projetado que integra otimização de comunicação, balanceamento de carga e eficiência ao nível de kernel, sem modificar o framework de treinamento subjacente (ex: PyTorch/FSDP2).

2.1 Estratégia de Execução Centrada no Proprietário (Owner-Centric)

Em vez de replicar a computação em todos os ranks, o DMuon atribui cada parâmetro de matriz a um único rank proprietário (owner rank).

Papel do Proprietário (Owner): O proprietário mantém o estado autoritativo do otimizador e computa a atualização Muon (iteração Newton–Schulz) exatamente uma vez por matriz.
Papel do Não-Proprietário: Outros ranks materializam o parâmetro apenas quando necessário para a computação forward/backward e não realizam a ortogonalização.
Fluxo de Comunicação:
- Forward: Os proprietários publicam parâmetros atualizados para os não-proprietários de forma assíncrona, em pipeline com a execução da camada.
- Backward: Os gradientes são acumulados em buffers compactados e reduzidos ao proprietário.
- Sobreposição (Overlap): O DMuon emprega uma hierarquia de dois estágios (intra-nó e inter-nó) com layouts de pesos de grão fino para sobrepor a comunicação com a computação, evitando a contenção de coletivos.

2.2 Iteração Gram Newton–Schulz Eficiente

Para reduzir o custo computacional da atualização no lado do proprietário, o DMuon adota a formulação Gram Newton–Schulz:

Recorrência no Espaço Gram: Em vez de operar no espaço original $m \times n$ , a iteração é realizada no espaço Gram $m \times m$ ( $G = XX^\top$ ). Isso reduz a complexidade de $O(m^2n)$ para $O(m^3)$ quando $m < n$ .
Exploração de Simetria: Como $G$ é simétrica, o DMuon utiliza um caminho de execução do tipo SYRK para computar apenas a parte triangular inferior, reduzendo pela metade o trabalho aritmético dos produtos de matrizes dominantes.
Agrupamento (Batching): Para lidar com problemas de ocupação de GPU com matrizes pequenas, o DMuon agrupa matrizes de formatos semelhantes e executa a recorrência Gram NS conjuntamente, amortizando o overhead de despacho de kernel.

2.3 Balanceamento de Carga Consciente da Computação

Como os custos de atualização do Muon variam significamente com o formato da matriz, uma atribuição ingênua (ex: round-robin) leva a ranks retardatários (stragglers). O DMu do DMuon formula a atribuição de proprietários como um Programa Linear de Inteiros Mistos (MILP):

Perfilamento (Profiling): Durante a inicialização, o sistema realiza testes de desempenho (benchmarks) do tempo de execução de vários tamanhos de lote e cronogramas de kernel para cada formato de matriz distinto.
Otimização: O solver MILP minimiza o makespan previsto (carga máxima) entre os ranks proprietários, garantindo que as matrizes mais pesadas sejam distribuídas para equilibrar a carga de trabalho.
Alternativa (Fallback): Para espaços de busca extremamente grandes, um heurístico ganancioso (greedy) é utilizado para limitar o tempo de inicialização.

2.4 Autotuning de Kernel Adaptável ao Formato

O DMuon utiliza um fluxo de trabalho de autotuning orientado por DSL (usando as DSLs TileLang e CUTE) para gerar kernels de alto desempenho.

Ele busca sobre formas de tile, tamanhos de bloco e configurações de pipeline de software específicas para o hardware alvo e formatos de matriz.
A configuração ideal é armazenada em cache e reutilizada para passos subsequentes, garantindo que o cronograma do kernel esteja sempre próximo do ideal para as dimensões de problema específicas encontradas no treinamento de Transformers.

3. Principais Contribuições

Otimização de Comunicação de Grão Fino: Um tempo de execução centrado no proprietário que minimiza o overhead de comunicação via escalonamento hierárquico e layouts de pesos de grão fino, preservando a semântica exata do otimizador.
Stack de Execução Adaptável ao Formato: Um motor de alto desempenho que combina Newton–Schulz Gram em lote, kernels conscientes de simetria e autotuning para lidar eficientemente com cargas de trabalho de matrizes heterogêneas.
Balanceamento de Carga Consciente da Computação: Uma abordagem de otimização baseada em custo medido (MILP) que atribui proprietários de matrizes com base em perfilamento de tempo de execução real, em vez de contagens teóricas de FLOPs.
Integração Direta (Drop-in): Um módulo que se integra em pipelines de treinamento existentes (ex: FSDP2) com esforço mínimo, não exigindo modificações no nível do framework.
Validação em Escala de Produção: Implementação bem-sucedida em cargas de trabalho de robótica (modelos incorporados) e treinamento de modelos de linguagem.

4. Resultados Experimentais

Os autores avaliaram o DMuon em quatro cargas de trabalho de produção: Wall-OSS-0.5, Pi0, Wall-WM e Qwen2.5-7B, escalando de 8 a 256 GPUs A800.

Tempo de Passo do Otimizador: O DMuon alcança um aceleração (speedup) de 6,85× a 163,00× no tempo do passo do otimizador em comparação com uma implementação ingênua "Gather-then-Compute" do Muon distribuído.
Tempo de Passo End-to-End: O DMuon alcança uma aceleração de 1,48× a 3,01× no tempo total do passo.
Overhead vs AdamW: A latência por passo do DMuon é trazida a níveis próximos ao do AdamW, com um overhead médio de apenas +2% em relação ao AdamW em todas as cargas de trabalho avaliadas.
Escalabilidade: No Wall-OSS, o DMuon mantém o rendimento (throughput) próximo ao do AdamW conforme o número de GPUs aumenta, enquanto o Muon ingênuo sofreria de gargalos severos de comunicação.

Detalhamento da Aceleração:

Kernel Gram Simétrico: 48% da aceleração (redução pela metade da aritmética nos produtos dominantes).
Escalonamento de Proprietário e Balanceamento de Carga: 32% (remoção da computação redundante $D$ -vez).
Autotuning e Agrupamento (Batching): 16% (redução do overhead de lançamento de kernel para matrizes pequenas).

5. Significância e Alegações

O artigo afirma que o DMuon resolve com sucesso o descompasso de sistemas entre otimizadores conscientes de matriz e infraestruturas de treinamento distribuído. Ao eliminar a computação redundante e otimizar a comunicação, o DMuon permite que os benefícios de convergência do Muon (como maior eficiência computacional e melhor geometria de otimização) se traduzam em acelerações práticas de tempo de execução (wall-clock speedups).

Os autores posicionam o DMuon não como um novo algoritmo, mas como um habilitador de sistemas necessário que torna o Muon um substituto direto (drop-in replacement) viável para o AdamW em treinamentos de larga escala em produção. O trabalho demonstra que a "penalidade" da otimização consciente de matriz pode ser reduzida a níveis insignificantes, permitindo o escalonamento eficiente tanto para Large Language Models quanto para Modelos de Fundação Incorporados (VLA).

O artigo observa explicitamente que o DMuon não altera a regra de atualização ou as propriedades de convergência do Muon; ele aborda estritamente o overhead de sistema. Consequentemente, os benefícios são mais pronunciados em configurações distribuídas, enquanto o treinamento em GPU única vê uma melhoria mais modesta (aprox. 2× de aceleração no otimizador) devido à ausência de redundância distribuída para eliminar.