PyTorch atinge desempenho SOTA em normalização com torch.compile

Métodos de normalização como LayerNorm e RMSNorm são fundamentais em aprendizado profundo, pois padronizam os valores de entrada para garantir um treinamento mais suave dos modelos. Recentemente, a equipe do PyTorch avaliou e otimizou o desempenho do torch.compile para LayerNorm e RMSNorm em GPUs NVIDIA H100 e B200, alcançando desempenho próximo ao estado da arte (SOTA) em nível de kernel, além de ganhos adicionais por meio de fusões automáticas. Essa evolução representa um marco significativo para desenvolvedores que buscam maximizar a eficiência de seus modelos de machine learning, especialmente em tarefas que exigem alta performance computacional. A otimização não se limita apenas ao hardware de ponta, mas também melhora a experiência de uso para uma ampla gama de aplicações, desde treinamento de modelos de linguagem até processamento de dados em larga escala.

O LayerNorm, introduzido em 2016, normaliza as entradas calculando a média e a variância dos valores ao longo de uma dimensão específica, além de aplicar escalonamento por meio de parâmetros treináveis, denominados gamma (peso) e beta (viés). Essa abordagem é amplamente utilizada em arquiteturas como Transformers, onde a estabilidade do treinamento é crítica. Já o RMSNorm, apresentado como uma alternativa ao LayerNorm em 2019, simplifica o processo ao substituir o cálculo da média pelo da raiz quadrática média (RMS), eliminando a necessidade do parâmetro beta. Essa diferença reduz a complexidade computacional, tornando o RMSNorm uma opção atraente para modelos que priorizam eficiência sem sacrificar desempenho.

A passagem direta (forward pass) de ambos os métodos é semelhante: envolve reduções ao longo de uma dimensão contígua e operações pontuais adicionais. O RMSNorm, no entanto, apresenta vantagem em eficiência, pois requer menos operações de ponto flutuante (flops) e não utiliza o viés (bias). Para este estudo, os resultados de benchmark foram apresentados de forma intercambiável entre LayerNorm e RMSNorm, dada a similaridade de seus kernels. Essa decisão reflete a flexibilidade do PyTorch em adaptar-se a diferentes abordagens de normalização, permitindo que os desenvolvedores escolham a que melhor se adequa ao seu caso de uso específico.

O Quack, uma biblioteca de kernels otimizados em CuteDSL desenvolvida por Tri Dao, serviu como referência SOTA para comparação. Os benchmarks iniciais mostraram que o torch.compile apresentava desempenho aproximadamente 50% inferior ao Quack em kernels de redução para GPUs H100. Essa lacuna motivou a equipe a investigar e implementar melhorias significativas no compilador, culminando em um desempenho competitivo. A análise detalhada dos kernels revelou oportunidades de otimização, especialmente em relação ao gerenciamento de reduções e fusão de operações, que são cruciais para o desempenho em hardware de alta gama.

O torch.compile utiliza o Inductor, um compilador JIT (Just-In-Time) que converte código PyTorch em kernels otimizados para GPU. Para LayerNorm e RMSNorm, o Inductor gera kernels que processam reduções de forma eficiente, aproveitando a contiguidade das dimensões de entrada. Quando a dimensão de redução (rnumel) é pequena (geralmente menor que 1024), o Inductor opta por uma redução persistente, eliminando loops desnecessários e acessando diretamente a média. Essa estratégia reduz a sobrecarga de controle de fluxo e melhora o desempenho em cenários onde a dimensão de redução é compacta, como em modelos de atenção em Transformers.

Os resultados dos benchmarks comparando o torch.compile versão 2.11 com o Quack (março de 2026) em formas comuns de uso revelaram que o PyTorch alcançou desempenho equivalente ao Quack em grande parte dos casos. No entanto, duas classes de regressões foram identificadas: shapes com dimensões muito grandes ou padrões incomuns que ainda não haviam sido otimizados. Essas exceções destacam a importância do autotuning contínuo e da adaptação do compilador a diferentes cenários de uso, garantindo que o desempenho permaneça competitivo em uma ampla variedade de aplicações.

A passagem reversa (backward pass) para LayerNorm e RMSNorm é mais complexa, pois envolve o cálculo de pelo menos dois gradientes: dX para a entrada, dW para os pesos e, no caso do LayerNorm, dB para o viés. Para otimizar o desempenho, essas operações de redução são realizadas em núcleos separados, o que pode duplicar a leitura de dados (dY) e aumentar a latência, especialmente em GPUs com alta largura de banda de memória. Essa limitação é particularmente crítica em modelos onde a normalização é aplicada em grandes lotes ou sequências longas, como em tarefas de processamento de linguagem natural.

A solução desenvolvida pela equipe do PyTorch foi a implementação de kernels fusionados, que combinam as reduções para dW, dB e dX em um único kernel. Essa abordagem reduz a duplicação de leitura de dados e melhora a eficiência, desde que a forma (shape) do tensor não exceda o limite de memória compartilhada ou registradores em um bloco de threads. O conceito de fusão de reduções não é novo, sendo explorado em projetos como o Liger, da Meta, e no Quack. O Inductor do PyTorch adotou essa estratégia, representando reduções com tipos distintos e permitindo a fusão de kernels para diferentes ordens de redução, um avanço chamado de MixOrderReduction.

Para shapes onde a dimensão xnumel (geralmente o lote ou batch) é muito maior que rnumel, o Inductor utiliza uma técnica de redução dividida (split reduction). Nessa abordagem, somas parciais são calculadas em paralelo e, em seguida, reduzidas em um kernel secundário. Essa estratégia melhora o paralelismo, mas requer um tensor de espaço de trabalho (workspace) para armazenar as somas intermediárias. O Inductor implementa essa técnica sem o uso de atomicidades, garantindo que cada bloco de threads (CTA) processe múltiplas linhas, evitando gargalos de sincronização e inconsistências numéricas.

A otimização MixOrderReduction combina fusão e redução dividida, permitindo que o Inductor gere kernels backward de normalização de alto desempenho de forma automática. Além disso, a geração desses kernels possibilita autotuning e fusões automáticas com operações circundantes, como multiplicações matriciais (GEMM) ou atenção, que frequentemente acompanham normalizações em modelos modernos. O desafio principal dessa otimização é fundir reduções com a mesma entrada, mas em ordens diferentes, um problema resolvido pelo MixOrderReduction, que introduz uma camada adicional de redução para consolidar os resultados parciais no tensor de espaço de trabalho.

A escolha do SPLIT_SIZE é crítica para o desempenho dos kernels de MixOrderReduction. Benchmarks em uma GPU H100 com dtype bfloat16 mostraram que reduzir o SPLIT_SIZE em 32 vezes aumentou a vazão (throughput) de 0,417 TB/s para 1,912 TB/s em um shape específico. Essa descoberta levou à implementação de heurísticas próprias para o SPLIT_SIZE, substituindo as escolhas padrão do Inductor, que eram baseadas em critérios menos relevantes para esse tipo de kernel. A autotuning tornou-se, assim, uma ferramenta essencial para maximizar o desempenho em diferentes cenários.

Outra descoberta importante foi a implementação de software pipelining (ou prefetching de cargas) nos kernels backward. Embora essa técnica seja tradicionalmente associada a workloads computacionalmente intensivos, como GEMM e atenção, a equipe observou que kernels de normalização também se beneficiavam dela, especialmente em shapes com grande M e pequeno N. A adição do parâmetro num_stages para autotuning permitiu ganhos de até 20% em algumas formas, demonstrando que até mesmo kernels de redução podem se beneficiar de técnicas avançadas de otimização de memória.

Os benchmarks finais comparando o MixOrderReduction com versões anteriores do PyTorch, além de baselines como Quack e Liger, foram realizados em uma máquina B200 com 750W de potência, utilizando CUDA 12.9 no final de 2025. Os resultados revelaram que o torch.compile sem MixOrderReduction já apresentava um desempenho significativamente superior ao PyTorch eager, mas com a nova otimização, o ganho quase dobrou em relação à versão anterior do compilador. Essa melhora aproximou o desempenho do PyTorch ao pico de largura de banda de memória da GPU, um feito notável considerando a complexidade das operações envolvidas.

Os testes também foram estendidos ao LayerNorm, cujos resultados foram semelhantes aos do RMSNorm, reforçando a eficácia da abordagem. A equipe observou que o novo kernel gerado pelo Inductor não apenas alcançou desempenho SOTA, mas também permitiu fusões automáticas com operações circundantes, melhorando ainda mais o desempenho final. Essa capacidade de integração automática é um diferencial do torch.compile, que simplifica o desenvolvimento de pipelines de machine learning otimizados sem a necessidade de intervenção manual extensiva.

A evolução do PyTorch neste domínio reflete um compromisso contínuo com a otimização de desempenho em aprendizado profundo. Ao atingir desempenho próximo ao estado da arte em kernels de normalização, a equipe não apenas superou limitações técnicas, mas também abriu caminho para o desenvolvimento de modelos mais eficientes e escaláveis. Essa conquista é particularmente relevante para a comunidade brasileira de IA, que pode se beneficiar diretamente dessas melhorias em aplicações como processamento de linguagem natural, visão computacional e análise preditiva.

Para desenvolvedores interessados em explorar essas otimizações, o PyTorch oferece documentação abrangente e tutoriais detalhados, que cobrem desde conceitos básicos até técnicas avançadas de otimização. A comunidade open-source também desempenha um papel crucial nesse ecossistema, contribuindo com feedback e melhorias contínuas. Com essas ferramentas, é possível construir modelos de IA mais rápidos e eficientes, reduzindo custos computacionais e acelerando a inovação em diversas áreas.

A implementação do MixOrderReduction e das demais otimizações no torch.compile representa um avanço significativo para o ecossistema PyTorch, solidificando sua posição como uma das principais plataformas para desenvolvimento de modelos de aprendizado profundo. À medida que o hardware de GPU evolui, com GPUs como a NVIDIA B200 oferecendo desempenho sem precedentes, ferramentas como o Inductor e o torch.compile tornam-se essenciais para aproveitar todo o potencial dessas máquinas. Para a comunidade brasileira, que cada vez mais adota tecnologias de IA em setores como saúde, finanças e manufatura, essas melhorias são um passo fundamental rumo a soluções mais acessíveis e eficientes.

O sucesso dessas otimizações também destaca a importância da colaboração entre desenvolvedores, pesquisadores e empresas. Projetos como o Quack e o Liger, que serviram de inspiração para as melhorias no PyTorch, demonstram como o compartilhamento de conhecimento e a inovação aberta podem acelerar o progresso tecnológico. No Brasil, iniciativas semelhantes vêm ganhando tração, com comunidades locais contribuindo para o desenvolvimento de soluções de IA adaptadas às necessidades regionais.

À medida que o PyTorch continua a evoluir, espera-se que novas otimizações sejam introduzidas, especialmente em áreas como fusão de operações, autotuning e suporte a novos hardwares. Para desenvolvedores brasileiros, acompanhar essas atualizações e participar ativamente da comunidade pode ser um diferencial competitivo, permitindo a criação de modelos de IA mais robustos e escaláveis. Com ferramentas como o torch.compile e o Inductor, o futuro do desenvolvimento de IA no Brasil e no mundo parece promissor, com possibilidades praticamente ilimitadas.

Em resumo, a otimização do torch.compile para LayerNorm e RMSNorm em GPUs H100 e B200 marca um avanço significativo no campo do aprendizado profundo. Ao alcançar desempenho próximo ao estado da arte e introduzir técnicas inovadoras como o MixOrderReduction e o software pipelining, a equipe do PyTorch não apenas melhorou a eficiência de seus kernels, mas também estabeleceu um novo padrão para o desenvolvimento de modelos de IA. Para a comunidade brasileira, essas melhorias representam uma oportunidade valiosa de impulsionar a inovação local, aproveitando tecnologias de ponta para resolver desafios reais em diversos setores.

À medida que mais desenvolvedores e empresas adotam essas tecnologias, espera-se que o ecossistema de IA no Brasil cresça ainda mais, com soluções cada vez mais avançadas e acessíveis. O PyTorch, com suas ferramentas poderosas e comunidade ativa, continuará a desempenhar um papel central nesse processo, capacitando desenvolvedores a construir o futuro da inteligência artificial. Com otimizações como as apresentadas neste artigo, o caminho para modelos de IA mais rápidos, eficientes e escaláveis está mais claro do que nunca.

Para aqueles que desejam se aprofundar nessas tecnologias, o PyTorch oferece recursos valiosos, desde documentação técnica até tutoriais práticos. A comunidade open-source também é um excelente ponto de partida para aprender e contribuir, compartilhando conhecimento e colaborando em projetos inovadores. Com o compromisso contínuo de otimização e inovação, o futuro do aprendizado profundo no Brasil e no mundo está repleto de possibilidades empolgantes.