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

Declaração do Problema

A rápida escalabilidade dos Modelos de Linguagem Grandes (LLMs) intensificou a demanda por algoritmos de treinamento eficientes em hardware. Embora o FP8 tenha surgido como uma alternativa viável de baixa precisão, o treinamento com formatos de 4 bits (como MXFP4) permanece um desafio significativo. Uma dificuldade primária é que os valores extremos (outliers) de ativação e gradiente amplificam os erros de quantização, potencialmente desestabilizando a otimização. Estudos iniciais indicam que a quantização grosseira de outliers de ativação pode dominar a faixa de quantização, degradando severamente a qualidade de ponto flutuante. Apesar dos esforços para mitigar o ruído via precisão mista ou agendamentos em etapas, uma questão fundamental permanece: Por que o treinamento FP4 de pipeline completo frequentemente diverge mesmo quando as ativações forward e os gradientes de ativação parecem estáveis? Além disso, não está claro quais componentes específicos do pipeline de treinamento (propagação forward, gradientes de ativação ou gradientes de peso) são mais responsáveis pela degradação da convergência e quais mecanismos de estabilização são eficazes.

Metodologia

Os autores conduziram um estudo controlado sobre o pré-treinamento do Llama 3.1–8B utilizando o conjunto de dados C4, aderindo à configuração de pré-treinamento LLM do MLPerf. O estudo foi executado em GPUs AMD Instinct MI355X utilizando suporte nativo a tensor cores FP4, distinguindo este trabalho de pesquisas anteriores que dependiam de emulação de software.

O desenho experimental central envolveu uma "escada" progressiva de habilitação do MXFP4:

1. Linha de Base: Treinamento completo em FP8.
2. Quantização por Etapas: O MXFP4 foi progressivamente habilitado em três componentes, mantendo os outros fixos em FP8:
   • Apenas propagação forward (Fprop).
   • Fprop + gradientes de ativação (Dgrad).
   • Fprop + Dgrad + gradientes de peso (Wgrad).
3. Avaliação de Estabilização: Sob o pipeline completo MXFP4 (Fprop + Dgrad + Wgrad), os autores testaram várias intervenções para restaurar a estabilidade:
   • Arredondamento Estocástico: Para introduzir aleatoriedade.
   • Rotações Hadamard Aleatórias: Para descorrelacionar outliers.
   • Rotações Hadamard Determinísticas: Utilizando especificamente matrizes $16 \times 16$ (H16) e $32 \times 32$ (H32).

Métricas: A convergência foi definida como atingir uma perplexidade de validação de 3.3. A métrica primária para eficiência foi o sobrecarga de tokens (o aumento relativo no número de tokens necessários para atingir o alvo em comparação com a linha de base FP8). Métricas secundárias incluíram o rendimento (throughput) por etapa de treinamento e a aceleração de ponta a ponta.

Principais Contribuições e Descobertas

1. Identificação da Fonte de Instabilidade

O estudo isola a quantização do Gradiente de Peso (Wgrad) como o principal motor da degradação da convergência.

• Habilitar o MXFP4 apenas para Fprop incorreu em um sobrecarga de tokens de 8–9%.
• Habilitar Fprop + Dgrad resultou em um sobrecarga semelhante de 10–11%.
• Habilitar o pipeline completo (Fprop + Dgrad + Wgrad) sem estabilização causou um salto na sobrecarga de tokens para 26–27% ou resultou em divergência completa ("Não Convergência").
• Conclusão: A instabilidade não é inerente à quantização de Fprop ou Dgrad, mas é dominada pela quantização de Wgrad.

2. Avaliação de Estratégias de Estabilização

Os autores testaram se adicionar estocasticidade ou rotações estruturadas poderia estabilizar o pipeline completo:

• Arredondamento Estocástico: Falhou em estabilizar o pipeline completo; a execução divergiu.
• Rotações Hadamard Aleatórias: Similarmente falhou em convergir quando aplicado ao pipeline completo.
• Rotações Hadamard Determinísticas: Restaurou com sucesso a otimização estável. Aplicar uma transformação Hadamard determinística (especificamente H16) ao pipeline completo MXFP4 reduziu a sobrecarga de tokens de volta para 8–9%, rastreando de perto a linha de base FP8.
• Insight: Intervenções estocásticas parecem amplificar o erro de quantização efetivo em vez de mitigar os efeitos dos outliers. A instabilidade é impulsionada por erros de micro-escala estruturados ao longo de caminhos de gradiente sensíveis, que rotações determinísticas podem corrigir.

3. Eficiência e Rendimento

Com a questão de estabilidade resolvida via Hadamard Determinístico (H16), os autores demonstraram os ganhos práticos de eficiência do FP4:

• Rendimento: A configuração MXFP4 + H16 alcançou um rendimento de etapa de treinamento 20% maior em comparação com a linha de base FP8 nas GPUs MI355X.
• Aceleração de Ponta a Ponta: Como a sobrecarga de tokens foi minimizada (8–9%), a vantagem de rendimento traduziu-se em uma aceleração de ponta a ponta de 9–10%.
• Insight Chave: As acelerações do FP4 são "limitadas pela estabilidade". Sem controlar a instabilidade do Wgrad, o aumento na exigência de tokens para convergir nega os benefícios brutos de computação de menor precisão.

Significado e Alegações

O artigo afirma fornecer a primeira demonstração de pré-treinamento completo em MXFP4 em hardware da classe MI355X sem emulação de software. Seu significado reside na mudança da compreensão da dinâmica de treinamento FP4:

• Mecanismo de Falha: A instabilidade no pré-treinamento FP4 não é devido à estocasticidade insuficiente, mas sim a erros de micro-escala estruturados em caminhos de gradiente sensíveis (especificamente Wgrad).
• Receita Prática: Uma rotação Hadamard determinística é um estabilizador necessário e eficaz para o treinamento MXFP4 de pipeline completo, permitindo a realização de ganhos de eficiência de hardware.
• Limitação Contextual: Os autores notam modestamente que esta receita de estabilização é dependente do cenário. Embora eficaz para o Llama 3.1–8B no conjunto de dados C4, o comportamento pode variar entre diferentes modelos, conjuntos de dados ou métodos de ajuste fino, e não deve ser assumido como universal sem evidências adicionais.

Em resumo, o trabalho estabelece que substituir FP8 por MXFP4 é viável para pré-treinamento em grande escala, desde que os erros de quantização do Wgrad sejam controlados por intervenções estruturais determinísticas em vez de injeção de ruído estocástico.