NVIDIA relata treinamento de modelo JAX até 73% mais rápido em GPUs Blackwell com formato NVFP4

A NVIDIA divulgou melhorias de desempenho para o treinamento de grandes modelos de linguagem com um novo formato numérico de baixa precisão chamado NVFP4 em suas mais recentes GPUs com arquitetura Blackwell. O anúncio, baseado em experimentos com o framework JAX do Google e a biblioteca de treinamento MaxText, reflete o esforço contínuo da indústria para reduzir o custo e o tempo necessários para o treinamento de modelos de inteligência artificial.

Segundo uma publicação em blog para desenvolvedores, a NVIDIA obteve acelerações que variaram de 1,31× a 1,73× em relação a uma base FP8 ao treinar o modelo Llama 3 8B em GPUs Blackwell usando o formato NVFP4. Isso representa uma melhoria de desempenho de até 73%, e a empresa informou que esses ganhos ocorreram sem perda mensurável de precisão. A companhia afirmou ter mantido uma curva de perda de treinamento semelhante ao longo de 10.000 etapas de pré-treinamento.

Equilíbrio entre precisão numérica e eficiência de treinamento

A precisão numérica no treinamento de modelos de IA envolve um equilíbrio entre velocidade de computação, uso de memória e qualidade final do modelo. Tradicionalmente, o formato de ponto flutuante de 32 bits (FP32) era o padrão, mas, nos últimos anos, a indústria passou a adotar formatos de 16 bits (FP16), Brain Float 16 (BF16) e 8 bits (FP8). Cada etapa reduziu a precisão em troca de maior throughput computacional e menores requisitos de largura de banda de memória.

O NVFP4 amplia essa tendência com um formato de ponto flutuante de 4 bits. Em teoria, um formato de 4 bits pode reduzir pela metade o uso de memória e aumentar o throughput em comparação com formatos de 8 bits. No entanto, na prática, a faixa numérica representável e a precisão são limitadas, o que pode gerar instabilidade numérica ou problemas de convergência durante o treinamento.

Os resultados da NVIDIA são relevantes porque sugerem que o NVFP4 pode funcionar em treinamento real de grandes modelos de linguagem sem perda de precisão, apesar dessas preocupações teóricas. A empresa informou que curvas de perda de treinamento semelhantes foram mantidas ao longo de 10.000 etapas de pré-treinamento, indicando que o modelo aprendeu em um padrão comparável ao do FP8.

O papel da arquitetura Blackwell

Esses ganhos de desempenho estão diretamente ligados ao projeto de hardware das GPUs Blackwell. Blackwell é a mais recente arquitetura de GPUs para data center da NVIDIA, incorporando aceleradores de hardware dedicados para aritmética de baixa precisão. O formato NVFP4 foi projetado para usar esses recursos de hardware, combinando otimização de software com suporte de hardware.

MaxText é uma biblioteca de treinamento de alto desempenho baseada em JAX, desenvolvida pelo Google, que fornece implementações para treinamento de grandes modelos de linguagem. A ênfase da NVIDIA na integração com o MaxText destaca a colaboração dentro do ecossistema JAX e sugere que os recursos do Blackwell podem ser aproveitados em frameworks além de PyTorch ou TensorFlow.

Contexto da indústria e cenário competitivo

Este anúncio faz parte de um esforço mais amplo da indústria para reduzir os custos de treinamento de IA. O treinamento de grandes modelos de linguagem pode exigir despesas computacionais substanciais, com tempos de treinamento que variam de semanas a meses. Uma aceleração de 73% tem potencial para reduzir esses custos e prazos, tornando o treinamento de modelos em larga escala mais acessível para mais organizações.

Os concorrentes seguem em direções semelhantes. A AMD está desenvolvendo seus próprios formatos de baixa precisão, as TPUs do Google são otimizadas em torno de formatos Brain Float, e a Intel e outros novos participantes buscam posições no mercado de aceleradores de IA. O anúncio do NVFP4 da NVIDIA pode ser visto nesse contexto competitivo.

Considerações práticas e limitações

No entanto, aplicar esses resultados em ambientes de produção envolve várias considerações. Primeiro, os resultados divulgados pela NVIDIA se baseiam em um modelo específico (Llama 3 8B) e em uma configuração específica de treinamento (receita MaxText). Se resultados semelhantes ocorrerão com diferentes arquiteturas de modelo, conjuntos de dados ou hiperparâmetros de treinamento requer validação adicional.

Segundo, 10.000 etapas de pré-treinamento podem representar apenas uma parte do processo completo de treinamento. Modelos grandes passam por centenas de milhares a milhões de etapas de treinamento, e erros numéricos podem se acumular ao longo de períodos prolongados. Não está claro se a NVIDIA confirmou a mesma manutenção de precisão em execuções de treinamento mais longas.

Terceiro, o NVFP4 é um formato específico da arquitetura Blackwell, de modo que seu uso exige a atualização para o hardware mais recente. Organizações que utilizam GPUs das gerações Hopper ou Ampere não podem se beneficiar imediatamente desses recursos.

Perspectivas futuras

Os avanços em treinamento de baixa precisão são importantes à medida que a escala e a complexidade dos modelos de IA continuam a aumentar. A indústria já discute modelos com trilhões de parâmetros, e os recursos computacionais necessários para treiná-los continuam a crescer. Tecnologias como o NVFP4 podem ajudar a moderar esse crescimento e permitir um treinamento mais eficiente.

Além disso, formatos de baixa precisão podem desempenhar um papel importante na etapa de inferência. Ao implantar modelos treinados em ambientes de produção, menor precisão pode significar tempos de resposta mais rápidos e custos operacionais mais baixos. Se o mesmo formato de baixa precisão puder ser usado tanto no treinamento quanto na inferência, a eficiência de toda a cadeia de IA poderá melhorar.

O anúncio da NVIDIA mostra como a colaboração entre fabricantes de hardware, desenvolvedores de frameworks e pesquisadores de modelos pode resultar em melhorias práticas de desempenho. A velocidade com que as comunidades JAX e MaxText adotarem o NVFP4, e se resultados semelhantes poderão ser reproduzidos com outros modelos e tarefas, ajudará a determinar o impacto de longo prazo dessa tecnologia.

A adoção de formatos de baixa precisão também tem implicações econômicas e ambientais. A redução do tempo de treinamento pode diminuir o consumo de energia, o que pode reduzir tanto os custos operacionais de data centers quanto as emissões de carbono. À medida que a indústria de IA enfrenta pressão por sustentabilidade, tecnologias de treinamento eficientes oferecem uma forma de abordar simultaneamente considerações ambientais e econômicas.