MXNorm: Reusing MXFP block scales for efficient tensor normalisation

O artigo apresenta o MXNorm, uma alternativa eficiente ao RMSNorm que reutiliza escalas de blocos MXFP8 para reduzir drasticamente o tamanho das operações de redução, mantendo a precisão do treinamento e oferecendo acelerações práticas em kernels de modelos como o Llama 3.

O essencial

Imagine que você está construindo um arranha-céu gigante (um modelo de Inteligência Artificial) usando tijolos muito pequenos e leves (dados de baixa precisão). O objetivo é construir o prédio o mais rápido possível.

Nos últimos anos, a tecnologia evoluiu tanto que a parte de "empilhar tijolos" (multiplicação de matrizes) ficou extremamente rápida. É como se tivéssemos robôs que colocam tijolos na velocidade da luz.

O Problema: O "Chefe" que demora demais
No entanto, antes de empilhar os tijolos, há um passo crucial: o "Chefe" precisa verificar se a parede está reta e nivelada. Na linguagem de IA, isso se chama Normalização.
O problema é que, enquanto os robôs de empilhar ficaram 80 vezes mais rápidos, o "Chefe" que verifica o nível só ficou cerca de 5 vezes mais rápido. Ele se tornou o gargalo, o ponto lento que segura todo o progresso. Além disso, para fazer esse trabalho de nível, o "Chefe" precisa de ferramentas pesadas e precisas (alta precisão), o que gasta muita energia e tempo.

A Solução: MXNorm (O "Chefe" Inteligente)
Os autores deste paper, da Graphcore, criaram uma solução chamada MXNorm. A ideia é genial e simples:

1. O Truque da Reutilização: Para colocar os tijolos na nova tecnologia (formato MXFP), os robôs já precisam calcular um "fator de escala" para cada bloco de tijolos. É como se eles já medissem o tamanho do bloco antes de empilhá-lo.
2. A Pulo do Gato: Em vez de o "Chefe" (Normalização) fazer uma nova medição completa e demorada para ver se a parede está reta, o MXNorm diz: "Ei, vocês já mediram o tamanho do bloco para o robô, não? Usem essa mesma medida para nivelar a parede também!"

A Analogia do Buffet
Pense em um buffet de comida:

• Método Antigo (RMSNorm): Você pega um prato, coloca a comida, depois vai até a balança, pesa tudo, calcula a média, ajusta o prato e só então serve. É preciso, mas lento.
• Método Novo (MXNorm): O garçom já sabe o tamanho do prato e a quantidade de comida que vai servir (porque ele já organizou a comida em porções padronizadas). Ele usa essa informação prévia para garantir que o prato não fique pesado demais, sem precisar ir à balança novamente.

O que eles descobriram?

• Funciona de verdade: Eles testaram isso em modelos gigantes (como o Llama 3, com bilhões de parâmetros) e descobriram que a qualidade da "parede" (a precisão do modelo) não caiu. O prédio ficou tão sólido quanto antes.
• Velocidade: Como o "Chefe" não precisa mais fazer o trabalho pesado de medir tudo do zero, o processo ficou muito mais rápido. Em alguns casos, a velocidade aumentou em até 2,4 vezes apenas nessa etapa.
• O Segredo do Sucesso: Eles perceberam que usar uma média simples (p=1) às vezes falhava em modelos gigantes, permitindo que "tijolos fora do lugar" (valores extremos) estragassem a parede. Mas, usando uma média mais robusta (p=2, que é como uma média quadrática), eles conseguiram manter a estabilidade perfeita, igual ao método antigo.

Resumo em uma frase:
O MXNorm é uma técnica inteligente que reaproveita medições que já estavam sendo feitas para outra tarefa, eliminando a necessidade de um passo lento e pesado, permitindo que os modelos de IA cresçam mais rápido e consumam menos energia, sem perder qualidade.

É como se você descobrisse que, para saber se a sua casa está nivelada, não precisa usar um nível de bolha novo; basta olhar para a sombra que a casa já projeta no chão e usar isso como referência.

Resumo técnico

1. O Problema

O desempenho de matrizes multiplicadas (matmuls) em cargas de trabalho de aprendizado profundo melhorou drasticamente nos últimos anos, impulsionado por aceleradores de IA que utilizam formatos de ponto flutuante de baixa precisão (como FP8 e FP4). No entanto, outras operações fundamentais, como reduções (necessárias para normalização) e operações elementares, não acompanharam esse ritmo de melhoria.

• Gargalo de Desempenho: Enquanto a velocidade de matmul aumentou em até 80x (de V100 para GB200), as reduções e operações elementares melhoraram apenas em fatores de 5x a 9x, limitadas pela largura de banda de memória e throughput dos núcleos CUDA.
• Ineficiência na Normalização: Camadas de normalização, como o RMSNorm (padrão em modelos como Llama), exigem calcular a Raiz Média Quadrática (RMS) de todo o vetor de ativação antes de quantizá-lo para formatos de baixa precisão (MXFP). Isso cria um gargalo de redução desnecessário, pois a estatística (RMS) é calculada em alta precisão e depois descartada ao converter para o formato quantizado.
• Objetivo: Eliminar a redundância de calcular estatísticas separadas para normalização e quantização, reduzindo o custo computacional e de memória dessas operações.

2. Metodologia: MXNorm

Os autores propõem o MXNorm, uma substituição "drop-in" (plug-and-play) para o RMSNorm que reutiliza as escalas de bloco já calculadas durante a conversão para o formato MXFP8 (Microscaling).

• Conceito Central: Em vez de calcular o RMS completo de uma linha de ativação, o MXNorm estima o inverso do RMS utilizando apenas os máximos absolutos por bloco (block absmaxes) que são necessários para gerar as escalas de quantização MXFP.
• Aproximação Matemática:
  • O método baseia-se no teorema de que a média de potências generalizada dos máximos absolutos dos blocos converge para o RMS da distribuição, multiplicado por uma constante dependente do tamanho do bloco e da distribuição.
  • A fórmula de estimativa do inverso do RMS ($\tilde{\rho}$) é dada por:
    $$\tilde{\rho}_t = \tilde{c}(p, B) \cdot \left( \frac{1}{K} \sum_{k=1}^{K} \tilde{m}_{tk}^p \right)^{-1/p}$$
    Onde $\tilde{m}_{tk}$ é o máximo absoluto do bloco $k$, $K$ é o número de blocos, $B$ é o tamanho do bloco, e $p$ é o expoente da norma (testado com $p=1$ e $p=2$).
• Fusão de Operações: O MXNorm funde a normalização e a quantização em uma única passagem sobre os dados.
  • Calcula-se o máximo absoluto por bloco.
  • Estima-se o RMS global a partir desses máximos.
  • Normaliza-se e quantiza-se simultaneamente, gerando as escalas e valores quantizados diretamente.
• MXNormLinear: Para lidar com o parâmetro de ganho afim ($\gamma$) típico do RMSNorm, o método aplica o ganho diretamente na matriz de pesos da camada linear subsequente, evitando a necessidade de multiplicação elementar em formato quantizado, mantendo a eficiência.

3. Contribuições Chave

1. Redução de Redução: Diminui o tamanho da operação de redução necessária para a normalização em 32x, pois opera apenas sobre os máximos dos blocos em vez de todos os elementos.
2. Fusão de Kernel: Permite fundir a normalização com a conversão para MXFP, eliminando a necessidade de ler/escritura intermediária de tensores em alta precisão.
3. Análise de Estabilidade: Demonstra que a estabilidade do treinamento depende criticamente dos limites (bounds) impostos pela normalização. O uso de $p=2$ (média quadrática) fornece limites mais apertados e estáveis, similares ao RMSNorm, enquanto $p=1$ (média aritmética) pode levar a instabilidades em modelos grandes devido a limites de saída mais amplos.
4. Implementação Prática: Fornece uma implementação compatível com PyTorch (torch.compile) que não exige alterações arquiteturais complexas nos modelos LLM.

4. Resultados

Os autores validaram o MXNorm no pré-treinamento de modelos Llama 3 com tamanhos de 125M, 1B e 8B parâmetros.

• Precisão de Treinamento:
  • Em modelos menores (125M e 1B), o MXNorm ($p=1$ e $p=2$) mostrou sensibilidade à taxa de aprendizado e convergência de perda muito semelhantes ao baseline RMSNorm.
  • Em modelos de 8B parâmetros, o MXNorm com $p=2$ alcançou uma perda de treinamento final (2.126) praticamente idêntica ao RMSNorm (2.132).
  • O MXNorm com $p=1$ falhou em escalar para 8B, apresentando picos de perda (loss spikes) e convergência inferior, devido à estimativa menos precisa do RMS e limites de saída mais largos que exacerbam outliers.
• Desempenho (Speedup):
  • Em hardware comercial (NVIDIA GB200), o MXNorm forneceu um speedup de até 2.4x no kernel de normalização isolado em comparação ao RMSNorm + MXCast.
  • Em camadas completas de transformadores do Llama 3 8B, isso se traduziu em um speedup de 1.3% para MXFP8 e 2.6% para NVFP4. Embora pareça pequeno, em sistemas de larga escala, isso representa ganhos significativos de throughput.
• Análise de Outliers: A análise revelou que picos de perda no MXNorm ($p=1$) estavam correlacionados com a explosão de características de ativação (outliers), sugerindo que a capacidade de limitar valores extremos (fornecida por $p=2$) é crucial para a estabilidade em modelos grandes.

5. Significado e Conclusão

O MXNorm representa um avanço importante na otimização de hardware para LLMs em baixa precisão. Ao reconhecer que a normalização e a quantização compartilham a mesma estatística de escala (máximos de bloco), o método remove um gargalo de computação e memória que se tornaria crítico à medida que os aceleradores de IA se tornam ainda mais rápidos em multiplicações de matrizes.

• Impacto: Permite que modelos grandes sejam treinados e inferidos com formatos de ponto flutuante de 4 ou 8 bits sem sacrificar a estabilidade do treinamento ou a precisão final.
• Futuro: O trabalho sugere que outras operações não aceleradas por matmuls (como posições rotacionais e unidades lineares gateadas) também devem ser reavaliadas para fusão com operações de quantização, seguindo a filosofia de "reutilização de estatísticas" proposta pelo MXNorm.

Em resumo, o MXNorm é uma solução eficiente que alinha a complexidade computacional da normalização com a arquitetura de baixa precisão moderna, garantindo escalabilidade e estabilidade para a próxima geração de modelos de linguagem.