NVIDIA acelera o pré-treinamento de LLMs com NVFP4 no Blackwell, JAX e MaxText

O pré-treinamento de LLMs frontier esbarra na capacidade de processamento dos sistemas de computação. A NVIDIA mostrou como a combinação de JAX, MaxText e do novo formato NVFP4 em chips Blackwell permite acelerar bastante esse processo sem perda de qualidade.

Por que cada ponto percentual importa

Quando o treinamento roda em trilhões de tokens por milhares de aceleradores, economizar até um por cento de tempo em cada passo se transforma em vários dias de tempo real de calendário. Na escala frontier, isso se traduz diretamente em milhões de dólares em custos de computação.

O NVFP4, formato de ponto flutuante de quatro bits que estreou na arquitetura Blackwell, virou uma das principais ferramentas para acelerar operações matriciais. Em comparação com o FP8, ele empacota os números com o dobro de densidade, o que reduz a carga na memória e aumenta a vazão efetiva dos núcleos tensoriais.

A principal dificuldade é que a grade numérica de quatro bits é esparsa. Com uma configuração incorreta, os gradientes saem facilmente do seu intervalo, o que leva à divergência do treinamento. A NVIDIA e a equipe do MaxText resolveram isso com esquemas de escalonamento personalizados e loss scaling dinâmico.

Como funciona a precisão mista com NVFP4

O treinamento de precisão mista não é uma abordagem nova: FP8 e BF16 já são padrões da indústria. O NVFP4 dá um passo adiante e permite usar pesos de 4 bits nas multiplicações matriciais mais pesadas, mantendo maior precisão onde ela realmente importa.

• O NVFP4 é aplicado a pesos e ativações nas operações GEMM
• BF16 ou FP32 continuam sendo usados em acumuladores e normalização
• O MaxText direciona automaticamente as operações para o formato adequado
• O JAX compila o grafo de computação via XLA, otimizando os kernels para o Blackwell
• O resultado é maior throughput com consumo de energia comparável ou menor

O stack e o que mudar no código

O MaxText é um framework de treinamento de código aberto e alto desempenho baseado em JAX, desenvolvido pelo Google. Foi criado originalmente para TPUs, mas vem sendo adaptado ativamente para clusters de GPUs, então a parceria com a NVIDIA é natural.

A NVIDIA incluiu kernels NVFP4 de baixo nível no cuBLAS e no cuDNN, e o JAX/XLA ganhou suporte a essas operações por meio de adaptadores específicos. Os desenvolvedores não precisam reescrever o código de treinamento manualmente: basta ativar as flags corretas na configuração do MaxText e garantir que o cluster tenha chips Blackwell (B100, B200, GB200).

"A precisão numérica é um dos parâmetros de maior alavancagem, mas o pré-treinamento de precisão mista com poucos bits é difícil de implementar corretamente", observa a equipe do NVIDIA

Developer Blog.

O que isso significa

Para as equipes que fazem pré-treinamento de modelos frontier, o NVFP4 no Blackwell é uma aceleração praticamente gratuita: o stack atual de JAX e MaxText exige mudanças mínimas de configuração. Na escala de centenas ou milhares de GPUs, mesmo um ganho de 10-15% em throughput reduz diretamente o time-to-checkpoint e o orçamento total de computação.

A corrida pela eficiência do pré-treinamento entra numa fase de disputa pela precisão numérica.