NVIDIA acelera o pré-treinamento de LLM: NVFP4 no Blackwell em conjunto com JAX e MaxText

O pré-treinamento de LLMs de fronteira esbarra nos limites de throughput dos sistemas computacionais. A NVIDIA demonstrou como a combinação de JAX, MaxText e o novo formato NVFP4 em chips Blackwell permite acelerar significativamente esse processo sem perda de qualidade.

Por Que Cada Porcento É Importante

Quando o treinamento acontece em trilhões de tokens através de milhares de aceleradores, economizar até um por cento do tempo em cada etapa resulta em vários dias de tempo de calendário real. Na escala do pré-treinamento de fronteira, isso é uma tradução direta em milhões de dólares em despesas de computação. NVFP4 — um formato de ponto flutuante de quatro bits, que estreou na arquitetura Blackwell — tornou-se uma das ferramentas-chave para acelerar operações matriciais.

Comparado ao FP8, ele compacta números duas vezes mais densamente, o que reduz a carga de memória e aumenta o throughput efetivo dos núcleos de tensor. O principal desafio: a grade numérica de quatro bits é esparsa. Com configuração inadequada, os gradientes facilmente ultrapassam seus limites — isso leva à divergência do treinamento.

A NVIDIA e a equipe MaxText resolveram isso através de esquemas de escala personalizados e escala de perda dinâmica.

Como o Treinamento Mixed-Precision Funciona com NVFP4

O treinamento mixed-precision não é uma abordagem nova: FP8 e BF16 já se tornaram um padrão industrial. NVFP4 vai um passo além, permitindo pesos de 4 bits nas multiplicações matriciais mais computacionalmente intensivas, mantendo precisão mais alta onde realmente importa.

• NVFP4 é aplicado a pesos e ativações em operações GEMM
• BF16 ou FP32 permanecem para acumuladores e normalização
• MaxText roteia automaticamente operações para o formato apropriado
• JAX compila o gráfico computacional através de XLA, otimizando núcleos para Blackwell
• Resultado — crescimento de throughput com consumo de energia comparável ou menor

O Stack e O Que Mudar no Código

MaxText é um framework de treinamento de alto desempenho de código aberto baseado em JAX, desenvolvido pelo Google. Foi originalmente criado para TPU, mas está sendo ativamente adaptado para clusters de GPU, e a parceria com a NVIDIA é natural nesse contexto. A NVIDIA incluiu núcleos NVFP4 de baixo nível como parte de cuBLAS e cuDNN, e JAX/XLA recebeu suporte para essas operações através de adaptadores especiais. Os desenvolvedores não precisam reescrever o código de treinamento manualmente — é suficiente ativar os sinalizadores necessários nas configurações de MaxText e garantir que o cluster tenha chips Blackwell instalados (B100, B200, GB200).

"Precisão numérica é um dos parâmetros mais relevantes, mas pré-treinamento mixed-precision de baixo bit é difícil de implementar corretamente," observa a equipe do NVIDIA

Developer Blog.

O Que Isso Significa

Para equipes envolvidas no pré-treinamento de modelos de fronteira, NVFP4 no Blackwell é praticamente aceleração gratuita: o stack existente em JAX e MaxText requer mudanças mínimas de configuração. Na escala de centenas e milhares de GPUs, até mesmo ganhos de 10–15% em throughput reduzem diretamente o tempo até checkpoint e o orçamento geral de computação. A corrida pela eficiência do pré-treinamento está entrando na fase de batalha pela precisão numérica.