NVIDIA Megatron ganha suporte a Muon e outros otimizadores emergentes pra treino de LLM em escala

A NVIDIA publicou um deep dive sobre como o Megatron Core agora suporta o Muon (MomentUm Orthogonalized by Newton-Schulz) e outros otimizadores de ordem superior. Muon já foi usado pra treinar modelos open source de peso, incluindo Kimi K2 e GLM-5, e a promessa é acelerar a convergência sem estourar custo computacional.

Os números no GB300 NVL72

Treino medido com NeMo Megatron Bridge 26.02. Kimi K2 em 256 GPUs NVIDIA GB300 com paralelismo PP4DP64EP64. Qwen3 30B-A3B em 8 GPUs GB300 com DP8EP8. Throughput em MXFP8:

• Kimi K2: AdamW 1.051 TFLOPs/s/GPU vs Muon 1.080 (sendo 1.029 do modelo + 51 do Muon)
• Qwen3 30B-A3B: AdamW 713 vs Muon 721 (686 modelo + 35 Muon)

Ou seja, Muon empata ou passa o AdamW em MFU (Model FLOPs Utilization) quando você conta os FLOPs das multiplicações de matriz nas iterações de Newton-Schulz. Na prática, o custo extra do otimizador some no ruído, o que é a barreira principal pra adoção em produção.

Por que Muon é chato de distribuir

O passo de preconditioning (via iteração Newton-Schulz ou decomposição de autovalores) aumenta custo de compute e memória. Soma isso a instabilidade numérica em treino mixed precision e ao gargalo de comunicação pra sincronizar updates ortogonalizados em milhares de GPUs, e você entende por que até hoje o AdamW dominava.

A NVIDIA ataca dois problemas:

1. Layer-wise distributed optimizer

Otimizador distribuído tradicional (element-wise, que funciona bem com AdamW) fatia os estados do otimizador byte a byte entre GPUs. Isso quebra com Muon: ele precisa do gradiente da camada INTEIRA pra calcular o update dos pesos daquela camada. Se peso e estado estão picadinhos entre DP ranks, cada GPU não tem dado suficiente pra calcular nada sozinha.

Solução: distribuir por camada. Cada GPU recebe camadas inteiras, permitindo calcular o preconditioner localmente. O preço é comunicação de tamanho variável (camadas têm tamanhos diferentes), resolvida com all_gatherv. Implementação tá em layer_wise_optimizer.py no Megatron Core.

2. Newton-Schulz distribuído (3 modos)

Quando você usa tensor parallelism (TP), a matriz de pesos fica fatiada entre GPUs, e o passo de ortogonalização do Muon precisa da matriz de momentum inteira. Três estratégias no TensorParallelMuon:

• Duplicated mode: faz all-gather do momentum no domínio TP uma vez, cada GPU roda todas as iterações NS localmente. Otimizado pra latência de rede (uma comunicação só, independente de quantas iterações NS).
• Distributed mode: distribui o compute das iterações NS entre GPUs do TP. Cada iteração faz 3 matmuls e precisa de um all-reduce depois do primeiro. Otimizado pra compute, mais chatinho na rede.
• Blockwise mode: ortogonaliza só com momentum e pesos locais, sem comunicação. Mais barato, mas não é matematicamente equivalente à ortogonalização da matriz inteira (é ortogonalização em blocos).

A escolha depende do gargalo: se rede tá apertada, duplicated; se compute satura, distributed; se você topa aproximação por economia, blockwise.

O que ainda tá no forno

A NVIDIA lista otimizações em andamento no projeto Emerging Optimizers:

• Communication hiding: adiar o gather de parâmetros pro forward do próximo batch e sobrepor com compute, do jeito que já se faz em otimizadores element-wise.
• Load balancing: hoje as camadas vão pras GPUs em round-robin, o que desbalanceia porque MLPs e attention têm tamanhos bem diferentes. Dá pra melhorar estimando custo por camada.
• SYRK + all-reduce fundidos: duas das três matmuls de cada iteração NS podem virar SYRK (SYmmetric Rank-K update), cortando quase metade dos FLOPs. Kernels em Triton já foram liberados, versão CuTe DSL com comunicação fundida tá planejada.

Outros otimizadores suportados

Além do Muon, o repo NVIDIA-NeMo/Emerging-Optimizers tem:

• MOP (Momentum Orthogonalized by Polar decomposition): versão "final" do otimizador ortogonalizado.
• REKLS: variante avançada do SOAP (ShampoO with Adam in the Preconditioner's eigenbasis) que atualiza a base de autovalores a cada passo com decomposição + correção KL.

Pra reproduzir

O repo NVIDIA-NeMo/Megatron-Bridge tem as performance recipes. Pra rodar Kimi em 256 GB300, o exemplo é direto:

CONTAINER="nvcr.io/nvidia/nemo:26.04"
python scripts/performance/setup_experiment.py \
  -m kimi -mr kimi_k2 \
  --num_gpus 256 --gpus_per_node 4 \
  -g gb300 -c fp8_mx

Pra Muon, o flag use_distributed_optimizer despacha automaticamente pro wrapper layer-wise em vez do element-wise.

Pra quem treina modelo próprio em escala, vale olhar: Muon entregando paridade com AdamW em MFU abre espaço pra ganho de convergência de graça. O teto é quem tem GB300 na mão, mas o código roda em qualquer GPU com suporte a MXFP8 ou precisão menor.