Resumo Técnico: TetraJet-v2

1. Declaração do Problema

O pré-treinamento de Modelos de Linguagem de Grande Escala (LLM) está cada vez mais proibitivamente caro, impulsionando a necessidade de Treinamento Totalmente Quantizado (FQT) de baixa precisão. Embora formatos de 4 bits, como NVFP4, ofereçam ganhos significativos de eficiência em relação a FP8 e BF16, alcançar treinamento estável e quase sem perdas nessa precisão permanece um desafio crítico. Tentativas anteriores de treinamento totalmente em 4 bits sofreram com degradação severa de desempenho.

Os autores identificam dois principais gargalos de otimização que impedem o treinamento totalmente em FP4, distintos do simples erro de quantização:

1. Oscilação de Pesos: Um modo de falha dominante onde pequenas atualizações em pesos mestres de alta precisão fazem com que suas representações quantizadas alternem repetidamente entre bins adjacentes (por exemplo, entre 0 e 0,5). Esse fenômeno interrompe a otimização, particularmente à medida que as taxas de aprendizado decaem, impedindo que o modelo convirja para um mínimo local estável.
2. Características de Outlier: Um pequeno número de canais de ativação exibe magnitudes massivas, criando faixas dinâmicas que o FP4 não consegue representar com precisão. Isso leva a fatores de escala elevados que esmagam a precisão de valores menores dentro do mesmo grupo de quantização.

Métodos existentes, como a receita NVFP4 da NVIDIA, frequentemente dependem de precisão mista (mantendo alguns blocos em BF16) ou não abordam totalmente essas dinâmicas de otimização em um cenário de pré-treinamento do zero.

2. Metodologia: TetraJet-v2

O TetraJet-v2 é um método FQT de 4 bits de ponta a ponta projetado especificamente para NVFP4, aplicando o formato a ativações, pesos e gradientes em todas as camadas lineares. A metodologia aborda os gargalos identificados através de três componentes principais:

A. Quantização Dupla de Blocos Não Viciada para Camadas Lineares NVFP4

Para maximizar a precisão dentro das restrições numéricas do NVFP4 (que usa grupos de 16 elementos com escala de 8 bits), os autores propõem um esquema de quantização dupla de blocos:

• Estrutura: Um bloco externo de 128 elementos é escalado globalmente, enquanto blocos internos de 16 elementos são escalados localmente. Isso evita o cálculo do valor absoluto máximo sobre todo o tensor, melhorando a compatibilidade com hardware e a precisão em comparação com a escala por tensor.
• Estimativa de Gradiente Não Viciada: O design da camada linear garante estimativa de gradiente não viciada ao:
  • Alinhar tensores na passagem de trás (estimando $\hat{X}$ em vez de $X$).
  • Usar Arredondamento Estocástico na passagem de trás para garantir $E[\text{round}(x)] = x$, prevenindo viés nas expectativas de gradiente.
  • Adotar uma forma de grupo de pesos $1 \times 16$ (não quadrada) para maior precisão, exigindo alinhamento específico para a matriz de pesos transposta.

B. OsciReset: Supressão de Oscilação

Para abordar a oscilação de pesos, os autores introduzem o OsciReset, o primeiro algoritmo especificamente projetado para suprimir esse gargalo no pré-treinamento de LLMs.

• Detecção: O método rastreia a trajetória de otimização dos pesos, calculando uma métrica de "Risco de Oscilação" definida como a razão entre a distância percorrida pelo peso quantizado ($dist_Q$) e a distância percorrida pelo peso mestre ($dist_M$). Alto risco indica que um peso está alternando de bins apesar de pequenas atualizações no peso mestre.
• Mecanismo: Quando a taxa de aprendizado é baixa e o risco de oscilação excede um limiar, o peso mestre do elemento oscilante é redefinido para o centro de seu bin de quantização atual.
• Vantagem: Diferentemente de métodos anteriores que congelam pesos ou usam regularização pesada (o que prejudica a otimização global), o OsciReset permite que o peso continue otimizando em uma nova trajetória menos propensa à oscilação, sem degradar o desempenho imediatamente.

C. OutControl: Controle de Outliers

Para gerenciar outliers, o TetraJet-v2 emprega uma estratégia híbrida:

• Transformação Hadamard Aleatória (RHT): Aplicada aos cálculos da passagem de trás ($dX$e$dW$) para distribuir a energia de outliers de forma mais uniforme entre os canais, mitigando a variância de quantização.
• Retenção Estática de Outliers: Reconhecendo que canais de outliers são estruturalmente consistentes (canais específicos sempre têm normas altas), o método identifica os principais $p\%$ (por exemplo, 10%) de canais de outliers durante uma breve janela de calibração. Esses canais são mantidos em precisão superior (por exemplo, FP8) tanto para passagens de frente quanto de trás, enquanto o restante da rede permanece em FP4. Essa seleção estática evita a sobrecarga computacional de ajustes dinâmicos por lote.

3. Contribuições Principais

1. Quantização Dupla de Blocos Não Viciada: Um design inovador de camada linear NVFP4 com estimativa de gradiente não viciada e formas de bloco ótimas ($1 \times 16$) para LLMs, superando receitas NVFP4 existentes.
2. OsciReset: O primeiro algoritmo eficaz para suprimir a oscilação de pesos em LLMs durante o pré-treinamento, resolvendo uma fonte primária de instabilidade no treinamento de baixa precisão.
3. OutControl: Um algoritmo de precisão mista que combina RHT para a passagem de trás com retenção estática de outliers para passagens de frente e de trás, preservando a precisão.
4. Avaliação Abrangente: Experimentos extensivos de pré-treinamento em modelos OLMo-2 (até 370M de parâmetros, 212B de tokens) demonstrando a eficácia da pilha completa.

4. Resultados Experimentais

Os autores avaliaram o TetraJet-v2 em modelos OLMo-2 (70M, 150M e 370M de parâmetros) treinados do zero.

• Redução da Lacuna de Desempenho: O TetraJet-v2-full (incorporando todos os métodos propostos) reduziu a lacuna de desempenho em relação ao treinamento de precisão total (BF16) em uma média de 51,3% entre os tamanhos de modelo.
• Perplexidade (PPL): No modelo de 370M treinado com 212B de tokens, o TetraJet-v2-full alcançou uma PPL de validação de 25,11, comparado a 23,70 para BF16 e 26,20 para a linha de base NVFP4 da NVIDIA.
• Tarefas de Descendente: O método alcançou o desempenho médio mais alto entre todos os métodos FP4 em um conjunto de benchmarks (ARC, BoolQ, MMLU, etc.), com um modelo de 370M atingindo 43,60% de precisão em média.
• Eficiência: Implementações de kernel CUDA alcançaram um aceleração de ponta a ponta de 1,67× sobre FP8 (TransformerEngine) em hardware RTX 5090, verificando ganhos práticos de desempenho apesar da complexidade adicional de supressão de oscilação e controle de outliers.
• Estudos de Ablação:
  • A remoção do OsciReset levou a degradação severa de desempenho e um rebote na perda.
  • O OutControl melhorou a precisão tanto na passagem de frente quanto na de trás.
  • O método mostrou-se robusto a escolhas de hiperparâmetros (por exemplo, limiar de oscilação, tempo de início).

5. Significado e Alegações

O artigo afirma que o TetraJet-v2 representa um passo significativo rumo a tornar o pré-treinamento totalmente de baixa precisão (4 bits) de LLMs viável e competitivo com o treinamento de alta precisão.

• Otimização vs. Representação: Os autores enfatizam que a lacuna de desempenho no treinamento de baixa precisão não se deve apenas ao erro de quantização (representação), mas é criticamente impulsionada por questões de otimização (oscilação e outliers) que métodos anteriores falharam em abordar sistematicamente.
• Generalização: Diferentemente de técnicas anteriores de mitigação de oscilação projetadas para Vision Transformers (que frequentemente dependem de ajuste fino ou ajuste específico de hiperparâmetros), o OsciReset é apresentado como um método simples e eficaz aplicável ao pré-treinamento de LLMs do zero com ajuste mínimo.
• Praticidade: Ao alcançar uma aceleração de 1,67× sobre FP8 enquanto reduz a lacuna de precisão em relação ao BF16 em mais de 50%, o trabalho demonstra que o treinamento totalmente quantizado pode oferecer tanto eficiência computacional quanto qualidade de modelo, potencialmente reduzindo os custos de hardware e energia da pesquisa de modelos fundamentais.

Os autores reconhecem limitações relacionadas à escala de seus experimentos (até 370M de parâmetros) e à atual falta de suporte de hardware mainstream para estratégias de precisão mista como FP6 $\times$ FP4, que eles exploraram como uma direção futura potencial para ganhos adicionais de precisão.