Autores: Abhishek Patnia* , Aman Gupta* e Shao Tang*.
*estes autores contribuíram igualmente
Introdução
No Nubank, treinamos modelos Transformer em sequências de transações financeiras para aprender padrões de comportamento financeiro. Para melhorar esses modelos, pesquisadores executam experimentos constantemente para explorar mudanças de arquitetura ou de dados [8]. Cada experimento nos ajuda a avaliar uma nova ideia de modelagem, mas também consome tempo de GPU. Melhorar a eficiência do treinamento encurta esse ciclo de feedback e nos permite testar mais ideias com o mesmo orçamento de hardware.
Este artigo descreve como aceleramos o treinamento de um modelo híbrido que combina camadas Gated DeltaNet (GDN) com atenção softmax na proporção de 3:1 [4]. Escolhemos essa arquitetura porque ela entrega melhor qualidade de modelo para a nossa tarefa e treina mais rápido do que um modelo com atenção em todas as camadas. Para estudar sua eficiência de treinamento, usamos sequências de 16.384 tokens em um cluster de B200 e analisamos duas configurações de otimizador: AdamW e Muon.
Check our job opportunies
Medindo a eficiência do treinamento
Medimos a eficiência de execução usando a utilização de FLOPs do modelo (MFU), que estima quanto da capacidade teórica de computação do hardware é destinada aos passes forward e backward exigidos pelo modelo [1].
MFU (%) = 100 × (tokens/s × FLOPs/token) / (pico de FLOPs/s)
FLOPs/token estima o trabalho de forward e backward do modelo por token, excluindo a recomputação de ativações. A taxa de tokens medida e a capacidade de computação de pico referem-se à mesma alocação de GPUs, com a capacidade de pico considerada na precisão de treinamento.
Nosso modelo híbrido GDN-Atenção, em contextos mais longos como 16.384 tokens exige menos FLOPs por token do que um modelo só de atenção, ao mesmo tempo em que entrega melhor qualidade de modelo nos nossos experimentos. Por isso, ele consegue processar mais tokens por segundo por GPU mesmo quando seu MFU é menor. Para ilustrar essa vantagem, também reportamos o MFU que um modelo só de atenção precisaria atingir para igualar a taxa de tokens medida do modelo híbrido. Chamamos isso de MFU equivalente all-SDPA. Por exemplo, nossa baseline híbrida com AdamW registra 21,1% de MFU, enquanto um modelo só de atenção processando o mesmo número de tokens por segundo por GPU precisaria de 51,4% de MFU. Ao longo deste artigo, apresentamos as duas métricas para mostrar a eficiência de execução do modelo híbrido junto com a vantagem de velocidade de treinamento da sua arquitetura.
Partindo dessa baseline híbrida, melhoramos a implementação do treinamento um passo de cada vez. Dentro da série de cada otimizador, mantemos fixos a configuração do modelo, o batch size, a receita de treinamento e a alocação de hardware. As duas séries seguem a mesma progressão, passando pela seleção do kernel de atenção e pela integração com o compilador, sendo que o Muon adiciona um passo final que agrupa em lotes suas atualizações de pesos. Juntas, essas melhorias elevam o MFU do modelo híbrido de 21,1% para 39,3% com AdamW e de 20,2% para 39,2% com Muon. A Figura 1 mostra como as melhorias se acumulam, usando tanto o MFU do modelo híbrido quanto o MFU equivalente all-SDPA.
Figura 1. MFU em quatro estágios cumulativos para AdamW e cinco para Muon. Os valores representam o último intervalo registrado no passo 999 do otimizador. O MFU equivalente all-SDPA altera a contabilização de FLOPs, não a arquitetura executada.
Escolhendo o kernel de atenção certo
A interface de scaled dot-product attention (SDPA) do PyTorch suporta várias implementações, então o mesmo código de modelo pode executar kernels diferentes por baixo. Na prática, o PyTorch seleciona o kernel por meio de um conjunto de regras de backend, e o kernel selecionado nem sempre é o mais rápido disponível. Nosso primeiro passo foi verificar qual implementação estava de fato sendo executada [2].
Em experimentos anteriores, corrigir nossa configuração de software habilitou os kernels de atenção otimizados cuDNN da NVIDIA, com desempenho comparável ao da implementação de terceiros que estávamos testando. Usar o cuDNN pela interface existente do PyTorch manteve a mudança pequena e evitou a manutenção de uma integração de atenção separada.
Aplicamos essa lição aos nossos experimentos com o híbrido GDN, em que o cuDNN cuida das camadas de atenção softmax enquanto os kernels GDN permanecem inalterados. As duas configurações de otimizador se beneficiaram, com o modelo e a carga de trabalho mantidos fixos em cada série. Nossa primeira melhoria veio da seleção do backend certo, sem alterar a arquitetura do modelo.
MFU após habilitar o cuDNN
AdamW: 21,1% → 28,6% +7,5 p.p.
Muon: 20,2% → 27,1% +6,9 p.p.
Comparado com a baseline de cada otimizador; p.p. = pontos percentuais.
Aproveitando o Torch Compile
Além do kernel de atenção, um passo de treinamento também inclui muitas operações de tensores ao redor. O torch.compile do PyTorch pode otimizar esse trabalho combinando operações e reduzindo o tráfego de memória intermediário. Para obter esses benefícios, o código compilado precisa continuar reutilizável à medida que o treinamento avança.
Em trabalhos anteriores de treinamento, um peso de loss variável armazenado como número Python disparava recompilações repetidas. Representar esse valor como tensor permitiu que o código compilado lidasse com pesos variáveis sem gerar uma nova versão para cada valor [3]. Habilitar a compilação, portanto, foi só parte do trabalho; também precisávamos resolver os problemas que impediam a reutilização do código compilado.
Com o cuDNN já habilitado, a compilação melhorou o MFU para os dois otimizadores. Em seguida, atacamos outro obstáculo: a integração do GDN dividia a computação em regiões compiladas separadas, limitando o que o compilador conseguia otimizar entre elas.
MFU após habilitar o torch.compile
AdamW: 28,6% → 32,0% +3,4 p.p.
Muon: 27,1% → 30,1% +3,0 p.p.
Comparado com o estágio com cuDNN habilitado; p.p. = pontos percentuais.
Reempacotando os kernels FLA para o GDN
Nossas camadas GDN já usavam kernels do Flash Linear Attention (FLA) [5]. O problema restante era como o código ao redor desses kernels interagia com o torch.compile. Esse código dividia a computação em regiões compiladas separadas, limitando a capacidade do compilador de otimizar as operações ao redor.
Resolvemos isso em dois passos. Primeiro, expusemos operações-chave do FLA por meio de interfaces amigáveis ao compilador, com regras de gradiente explícitas. Depois, tratamos uma propriedade fixa do hardware da GPU como constante, eliminando outra fonte de graph breaks. Essas mudanças preservaram a arquitetura do modelo, as equações do GDN e os kernels principais do FLA.
Figura 2. Corrigir a compilação do GDN reduziu o número de grafos capturados em um diagnóstico isolado e melhorou o MFU de treinamento para as duas configurações de otimizador. Um grafo capturado não implica um único kernel de GPU nem um único grafo para todo o job de treinamento.
Com o cuDNN e o torch.compile já habilitados, o conjunto completo de correções de compilação do GDN melhorou o MFU tanto para AdamW quanto para Muon nos nossos experimentos em B200 com carga de trabalho fixa. A melhoria veio de tornar a computação existente do GDN mais amigável ao compilador, sem alterar o modelo nem substituir seus kernels principais. Isso mostrou por que o código ao redor de um kernel otimizado merece tanta atenção quanto o próprio kernel.
MFU após as correções de compilação do GDN
AdamW: 32,0% → 39,3% +7,3 p.p.
Muon: 30,1% → 36,7% +6,6 p.p.
Comparado com o estágio compilado anterior; p.p. = pontos percentuais.
Mantendo baixo o overhead do otimizador
À medida que os passes forward e backward do modelo ficam mais rápidos, a execução do otimizador pode ocupar uma fatia maior de cada passo de treinamento. O Muon prepara as atualizações de pesos por meio de operações matriciais repetidas, e processar cada matriz separadamente gera muitas operações pequenas na GPU. Focamos em executar esse trabalho de forma mais eficiente, mantendo a regra de atualização usada na nossa pesquisa anterior com Muon [6].
Nossa implementação MuonMegaBatch agrupa atualizações compatíveis com o mesmo formato de matriz e as processa juntas. Ela também agrupa em lotes as atualizações de momentum e de parâmetros ao redor, reduzindo lançamentos repetidos por matriz. Adaptamos essa estratégia de execução da implementação Dion da Microsoft [7]. Aqui, batching se refere às operações do otimizador — o batch size de treinamento permanece inalterado.
Figura 3. O MuonMegaBatch agrupa atualizações de matrizes compatíveis preservando a atualização independente de cada matriz. Este esquema ilustra a execução, não uma contagem literal de kernels de GPU.
Cada matriz continua recebendo sua própria atualização independente. O MuonMegaBatch preserva a regra de atualização pretendida do Muon, embora a aritmética em lote possa introduzir diferenças de ponto flutuante. Mantendo todas as melhorias anteriores de kernel e compilação, essa otimização final, exclusiva do Muon, aumentou o MFU na nossa comparação em B200 com carga de trabalho fixa. A melhoria veio de mudar a forma como as atualizações do otimizador são executadas, sem alterar o modelo nem o batch size de treinamento.
MFU após habilitar o MuonMegaBatch
Muon: 36,7% → 39,2% +2,5 p.p.
Comparado com o Muon padrão após as correções de compilação do GDN; p.p. = pontos percentuais.
Conclusão
Juntas, essas mudanças elevaram o MFU de cerca de 20% para quase 40% com AdamW e de 20,2% para 39,2% com Muon nas nossas comparações em B200 com carga de trabalho fixa. Alcançamos esses ganhos sem alterar a arquitetura do modelo nem o batch size de treinamento. Os ganhos vieram da seleção do kernel de atenção, de código compilado reutilizável, de correções na compilação do GDN e, no caso do Muon, de atualizações do otimizador em lote.
A lição mais ampla é que a eficiência do treinamento depende de todo o caminho de execução. Kernels rápidos importam, mas também importam o código ao redor deles e o trabalho realizado entre os passes do modelo. Manter a carga de trabalho fixa e medir cada mudança nos ajudou a identificar quais melhorias se traduziram em melhor uso do hardware.
Referências
[1] Chowdhery, A., et al. (2022). PaLM: Scaling Language Modeling with Pathways. arXiv:2204.02311, Section 4.1 (Training Efficiency).
https://arxiv.org/html/2204.02311v5#S4.SS1
[2] PyTorch contributors. (n.d.). torch.nn.functional.scaled_dot_product_attention. PyTorch 2.9 documentation.
https://docs.pytorch.org/docs/2.9/generated/torch.nn.functional.scaled_dot_product_attention.html
[3] PyTorch contributors. (n.d.). torch.compile Troubleshooting: Wrapping constants with tensors. PyTorch 2.9 documentation. https://docs.pytorch.org/docs/2.9/torch.compiler_troubleshooting.html#wrapping-constants-with-tensors
[4] Yang, S., Kautz, J., & Hatamizadeh, A. (2025). Gated Delta Networks: Improving Mamba2 with Delta Rule. International Conference on Learning Representations.
https://arxiv.org/abs/2412.06464
[5] Flash Linear Attention contributors. (n.d.). Flash Linear Attention [Computer software]. GitHub.
https://github.com/fla-org/flash-linear-attention
[6] Braithwaite, D., Yoshida, A., Celente, R., & Gupta, A. (2025, September 15). Muon for Improved Foundation Model Pretraining Data Efficiency. Building Nubank.
https://building.nu.com/muon-for-improved-foundation-model-pretraining-data-efficiency/
[7] Microsoft. (n.d.). Dion [Computer software]. GitHub. Prior work for the same-shape batching strategy.
https://github.com/microsoft/dion
[8] Yoshida, A., Braithwaite, D., Buga, M., Foust, T. (2026, September 03). How we built a benchmarking framework to horizontally accelerate transaction model research.
Check our job opportunies
Quick Navigation