Pretraining Large Language Models with NVFP4

Este artigo apresenta uma abordagem inovadora para o pré-treinamento estável e preciso de modelos de linguagem grandes no formato NVFP4, que combina transformadas de Hadamard aleatórias, esquemas de quantização bidimensionais e arredondamento estocástico para alcançar desempenho comparável ao FP8 em um modelo de 12 bilhões de parâmetros treinado com 10 trilhões de tokens.

O essencial

Imagine que você está construindo um arranha-céu gigante (um Modelo de Linguagem Grande, ou LLM) para ensinar uma máquina a pensar como um humano. Para fazer isso, você precisa de milhões de tijolos (dados) e uma equipe de engenheiros (o computador) trabalhando 24 horas por dia.

O problema é que, até agora, os tijolos eram feitos de ouro maciço (precisão de 8 bits ou FP8). São tijolos perfeitos, super resistentes, mas pesados e caros. Construir um prédio inteiro com ouro gasta uma fortuna de energia e tempo, e o computador fica cansado.

A NVIDIA, nesta nova pesquisa, descobriu como construir o mesmo prédio usando tijolos de plástico leve (precisão de 4 bits, chamados NVFP4). Isso tornaria a construção duas ou três vezes mais rápida e usaria metade da energia.

Mas há um problema: Tijolos de plástico são frágeis. Se você tentar construir um arranha-céu gigante só com eles, o prédio pode desmoronar ou ficar torto porque o plástico não aguenta certas pressões (os "picos" de dados).

Aqui está a solução mágica que a NVIDIA encontrou, explicada com analogias simples:

1. O Formato NVFP4: O "Bloco Inteligente"

Antes, os tijolos de plástico eram todos iguais e rígidos. O novo formato NVFP4 é como um bloco de construção inteligente.

• Em vez de tratar 32 tijolos como um grupo grande e desajeitado, o NVFP4 agrupa apenas 16 tijolos de cada vez.
• Para cada grupo de 16, ele usa uma "régua de medição" (escala) muito mais precisa.
• A analogia: Imagine que você está empilhando caixas. O método antigo usava uma régua gigante para medir 32 caixas de uma vez, o que deixava algumas caixas pequenas sem espaço e outras grandes esmagadas. O NVFP4 usa uma régua menor e mais precisa para grupos menores, garantindo que nada seja esmagado e nada fique solto.

2. O Problema dos "Gigantes" (Outliers)

Em qualquer pilha de dados, existem alguns números que são gigantes (valores extremos) e outros que são minúsculos. Em 4 bits, esses "gigantes" podem quebrar o sistema.

• A Solução (Transformada de Hadamard): Imagine que você tem uma sala cheia de pessoas, e algumas estão gritando muito alto (os outliers), atrapalhando todos. O método NVFP4 usa um "espelho mágico" (a Transformada de Hadamard) que espalha o grito de uma pessoa por toda a sala, transformando um grito estridente em um murmúrio suave que todos podem ouvir sem se assustar. Isso evita que o prédio desabe por causa de um único valor extremo.

3. A Regra do "Espelho Quebrado" (Escalonamento 2D)

Quando você constrói algo, você precisa olhar para ele de frente (para a frente) e de trás (para trás) para garantir que está tudo certo.

• O problema é que, ao usar tijolos de plástico, a forma como você mede a frente é diferente da forma como mede o fundo. É como se você olhasse para um objeto e ele mudasse de forma dependendo do ângulo. Isso confunde o cérebro da máquina.
• A Solução (Escalonamento 2D): A NVIDIA criou uma técnica onde eles garantem que, mesmo que a régua mude de ângulo, a "forma" do tijolo permaneça a mesma. É como se eles usassem um molde especial que garante que o tijolo seja idêntico, seja visto de frente ou de trás. Isso evita que a matemática da construção fique confusa.

4. A Estratégia de "Segurança" (Precisão Mista)

Mesmo com tijolos de plástico inteligentes, alguns cantos do prédio são tão críticos que não podem correr riscos.

• A Solução: Eles deixaram os últimos andares do prédio (as camadas finais do modelo) feitos de ouro (precisão alta, BF16).
• A analogia: A maior parte do prédio é feita de plástico leve e rápido, mas os alicerces e o topo (onde a decisão final é tomada) são reforçados com ouro. Isso garante que, se algo der errado lá embaixo, o prédio não desabe.

5. O "Jogo de Sorte" (Arredondamento Estocástico)

Quando você corta um tijolo de plástico para caber em um espaço, você precisa arredondar. Se você sempre arredonda para cima ou sempre para baixo, acumula um erro (viés).

• A Solução: Eles introduziram um pouco de sorte. Às vezes, eles arredondam para cima, às vezes para baixo, baseado em uma probabilidade.
• A analogia: É como jogar uma moeda para decidir se você ganha ou perde um centavo. A longo prazo, isso cancela os erros e mantém a conta justa, evitando que o prédio fique torto por causa de pequenos desvios acumulados.

O Resultado Final?

A NVIDIA testou isso construindo um modelo gigante de 12 bilhões de parâmetros com 10 trilhões de palavras (tokens).

• O Milagre: O prédio feito com "tijolos de plástico" (NVFP4) ficou tão forte e preciso quanto o prédio feito de "ouro" (FP8).
• Eles conseguiram a mesma inteligência, o mesmo raciocínio e a mesma capacidade de responder perguntas, mas usando metade da memória e até 3 vezes mais velocidade.

Resumo da Ópera:
A NVIDIA aprendeu a usar materiais mais leves e baratos para construir supercomputadores, mas criou um conjunto de regras de engenharia (os truques matemáticos acima) para garantir que a estrutura não desmorone. Isso significa que, no futuro, poderemos treinar modelos de IA muito mais inteligentes e rápidos, gastando menos energia e dinheiro. É como passar de construir com blocos de ouro para construir com blocos de Lego super-resistentes.

Resumo técnico

Resumo Técnico: Pré-treinamento de Modelos de Linguagem Grandes com NVFP4

1. O Problema

O treinamento de Modelos de Linguagem Grandes (LLMs) de ponta exige investimentos massivos de computação, tempo e energia (na ordem de dezenas a centenas de yottaflops). Embora o treinamento em precisão de ponto flutuante de 8 bits (FP8) tenha sido amplamente adotado para reduzir custos, a transição para precisões ainda mais estreitas, como 4 bits (FP4), promete dobrar ou triplicar o desempenho aritmético e reduzir o uso de memória pela metade.

No entanto, o treinamento em FP4 apresenta desafios significativos:

• Instabilidade de Treinamento: A quantização agressiva pode levar à divergência do modelo, especialmente em grandes escalas e horizontes longos de tokens.
• Problemas de Convergência: A presença de "outliers" (valores extremos) e erros de viés na estimativa de gradientes comprometem a precisão final.
• Limitações de Formatos Existentes: Formatos de microescala existentes (como MXFP4) utilizam blocos grandes (32 elementos) e fatores de escala baseados em potências de dois, o que pode desperdiçar dinâmicas de representação e causar saturação ou perda de precisão em valores críticos.

2. Metodologia Proposta

A NVIDIA introduz o formato NVFP4 e uma metodologia de treinamento integrada para superar as limitações do FP4. O núcleo da solução não é apenas o formato de dados, mas um conjunto de técnicas algorítmicas:

• Formato NVFP4:

  • Bloco Menor: Reduz o tamanho do bloco de 32 para 16 elementos, capturando melhor a faixa dinâmica local.
  • Fatores de Escala Precisos: Utiliza um formato de escala E4M3 (em vez de UE8M0 usado no MXFP4), trocando parte do alcance do expoente por mais bits de mantissa para maior precisão.
  • Escala em Dois Níveis: Combina um fator de escala FP8 por bloco com um fator de escala FP32 por tensor. Isso permite que pelo menos 6,25% dos valores em cada bloco (os de maior magnitude) sejam representados com precisão próxima ao FP8, enquanto o resto permanece em FP4.

• Técnicas de Treinamento Integradas:

  1. Precisão Mista Seletiva: Mantém camadas linearmente sensíveis (especificamente as últimas camadas da rede, cerca de 15%) em precisão mais alta (BF16) para garantir a convergência, enquanto o restante é treinado em NVFP4.
  2. Transformadas de Hadamard Aleatórias (RHT): Aplicadas especificamente nas entradas dos gradientes de peso (Wgrad). Isso redistribui os outliers para uma distribuição aproximadamente Gaussiana, facilitando sua representação no formato FP4 sem causar saturação.
  3. Escala Bidimensional (2D) para Pesos: Garante que a representação quantizada dos pesos seja consistente entre as passagens direta (forward) e inversa (backward). Como a passagem inversa transpõe os tensores, uma escala 1D tradicional quebraria a regra da cadeia. A escala 2D (blocos de 16x16) resolve isso.
  4. Arredondamento Estocástico (Stochastic Rounding): Aplicado exclusivamente aos gradientes para eliminar o viés sistemático introduzido pelo arredondamento determinístico (como "round-to-nearest-even"), essencial para a estabilidade em horizontes longos.

3. Contribuições Principais

• Primeira Demonstração Pública em Grande Escala: O artigo relata o pré-treinamento bem-sucedido de um modelo de 12 bilhões de parâmetros em 10 trilhões de tokens usando precisão de 4 bits, o maior treinamento documentado publicamente em FP4 até a data.
• Validação de NVFP4 vs. MXFP4: Demonstra que o formato NVFP4 supera o MXFP4, exigindo menos tokens para atingir a mesma perda (perplexidade), devido à sua melhor gestão de outliers e precisão de escala.
• Análise de Ablação Detalhada: Identifica que cada componente da metodologia (RHT, escala 2D, arredondamento estocástico e camadas mistas) é crítico para a estabilidade, especialmente em modelos grandes. A remoção de qualquer um leva à divergência ou degradação significativa.
• Suporte de Hardware: Valida o uso dos Tensor Cores da arquitetura NVIDIA Blackwell, que oferecem suporte nativo para NVFP4, proporcionando aceleração de 2x a 3x em relação ao FP8.

4. Resultados

O modelo de 12B parâmetros treinado com NVFP4 foi comparado a uma linha de base treinada em FP8:

• Perda de Validação: A curva de perda do NVFP4 rastreia quase perfeitamente a do FP8. Durante a fase estável, o erro relativo permanece abaixo de 1%, expandindo-se levemente para 1,5% no final do treinamento (devido ao decaimento da taxa de aprendizado).
• Precisão em Tarefas (Downstream): O desempenho em benchmarks é comparável ao FP8:
  • MMLU-Pro: 62,58% (NVFP4) vs. 62,62% (FP8).
  • MATH: 81,48% vs. 83,32%.
  • GSM8k: 92,27% vs. 89,08% (NVFP4 superou o FP8 em algumas tarefas matemáticas).
  • Código (MBPP+): Houve uma pequena queda (55,91% vs. 59,11%), possivelmente devido a ruídos na avaliação ou necessidade de ajuste fino no final do treinamento.
• Eficiência de Tokens: Ao comparar NVFP4 e MXFP4 no mesmo modelo (8B), o MXFP4 precisou de 36% mais tokens (1,36T vs 1T) para atingir a mesma perda que o NVFP4, demonstrando a superioridade do formato NVFP4 em eficiência de dados.

5. Significado e Conclusão

Este trabalho representa um marco fundamental na evolução do treinamento de LLMs. Ele prova que é possível treinar modelos de ponta com precisão de 4 bits sem sacrificar a qualidade do modelo, desde que se utilize uma metodologia algorítmica robusta.

• Impacto Econômico e Ambiental: A redução de 50% no uso de memória e o aumento de 2x-3x no throughput computacional prometem reduzir drasticamente o custo e a pegada de carbono do treinamento de futuros modelos de fronteira.
• Futuro do Treinamento: O estudo sugere que o futuro do treinamento de LLMs pode ser dominado por formatos de baixa precisão (FP4), desde que combinados com técnicas como escala 2D e transformadas de Hadamard.
• Disponibilidade: O suporte para treinamento NVFP4 já está disponível no Transformer Engine da NVIDIA, facilitando a adoção imediata pela comunidade de pesquisa e indústria.

Em resumo, a NVIDIA demonstrou que a barreira da precisão de 4 bits foi superada, abrindo caminho para modelos maiores e mais eficientes treinados em hardware Blackwell.