Resumo Técnico: Desmistificando a Instabilidade Numérica na Inferência de LLMs com HEAL

Declaração do Problema

À medida que os Grandes Modelos de Linguagem (LLMs) são implantados em domínios de missão crítica, como finanças, medicina e direito, a reprodutibilidade do output tornou-se um requisito rigoroso de sistema. Embora os profissionais utilizem tipicamente o decodificação greedy para eliminar a estocasticidade algorítmica, implantações empíricas usando precisões de 16 bits (ex: BF16) ainda exibem uma divergência catastrófica de output entre hardwares heterogêneos (ex: diferentes gerações de GPUs NVIDIA ou aceleradores AMD).

As estratégias de mitigação existentes enfrentam um dilema inaceitável:

1. Determinismo Estrito Bit a Bit: Forçar ordens de redução de ponto flutuante determinísticas ou usar kernels especializados restringe a implantação a caminhos de execução estreitos, desabilita recursos modernos de hardware e sacrifica o throughput.
2. Pipelines Globais de Alta Precisão: Realizar todos os cálculos em FP32 elimina o ruído de downcasting, mas acarreta penalidades sistêmicas proibitivas. Servir inteiramente em FP32 pode aumentar o Tempo por Token de Saída (TPOT) em até 13,1× e dobra a pegada de memória do cache de Chave-Valor (KV), tornando isso economicamente impraticável para o serviço em larga escala.

O desafio central é alcançar uma reprodutibilidade robusta e pronta para implantação sem sacrificar a eficiência dos aceleradores de hardware de 16 bits.

Metodologia: HEAL (Hybrid Error ALleviation)

Os autores propõem o HEAL, uma intervenção direcionada que aproxima a reprodutibilidade de nível FP32 enquanto adere às restrições de hardware. A metodologia baseia-se em uma investigação sistemática das causas raízes da não-reprodutibilidade, movendo-se de uma falha macroscópica para origens microscópicas.

1. Análise de Causa Raiz

Através de perfilamento em nível SASS e decomposição de erro, o artigo identifica que a não-reprodutibilidade não é impulsionada pela aritmética interna do kernel (que opera predominantemente em FP32 via Tensor Cores), mas sim por erros de truncamento introduzidos durante o downcasting nas fronteiras dos kernels.

• Origem Microscópica: A instabilidade provém do truncamento obrigatório de registradores de alta precisão para formatos de armazenamento de menor precisão nas fronteiras dos kernels (ex: escrever resultados de acumuladores FP32 de volta para a memória BF16).
• Acúmulo: Essas perdas microscópicas de precisão fracionária acumulam-se exponencialmente através das camadas da rede. Uma vez que o erro acumulado ultrapassa um limiar específico, ele altera as log-probabilidades finais o suficiente para alterar a seleção do token top-1, causando o "result flipping" (inversão de resultado).
• Insight Chave: A reprodutibilidade robusta não requer paridade matemática estrita com um ground truth em FP32; ela apenas requer limitar o ruído numérico abaixo do limiar que desencadeia a inversão de resultado.

2. A Solução HEAL

O HEAL aborda os dois gargalos primários de um pipeline global de FP32 (Memória para Atenção, Computação para GEMM) através de dois mecanismos direcionados:

A. Mitigando o Truncamento de Atenção via Quantização INT16

• Problema: O upcasting do cache de Chave-Valor (KV) para FP32 dobra o uso de memória, criando um gargalo de memória.
• Solução: O HEAL reconhece que formatos de ponto flutuante padrão subutilizam sua largura de bits para tensores Q, K e V. Ele aplica quantização INT16 para armazenar esses tensores.
• Mecanismo: Durante a execução, o tensor INT16 é algebricamente desempacotado em dois tensores FP16 ($X_{hi}$ e $X_{lo}$). O tensor inferior é deslocado por um fator de $2^{10}$, permitindo que o par $(X_{hi} + X_{lo}/2^{10})$ emule 20 bits de precisão fracionária contínua.
• Benefício: Isso preserva a pegada de memória estrita de 16 bits do cache KV enquanto reduz significativamente o erro de truncamento, efetivamente prevenindo a divergência sem expandir o uso de memória.

B. Mitigando o Truncamento de GEMM via Compensação de Erro Algébrica

• Problema: O upcasting de inputs de GEMM para FP32 força o hardware a ignorar os Tensor Cores de alto rendimento de 16 bits, causando uma severa regressão de desempenho.
• Solução: O HEAL sintetiza multiplicações de matrizes de alta precisão usando uma estratégia de compensação de erro algébrica executada inteiramente em Tensor Cores de 16 bits.
• Mecanismo: As ativações de entrada são decompostas em um valor primário de 16 bits e um termo de erro residual. O sistema realiza duas multiplicações de matriz paralelas usando inputs de 16 bits padrão e soma os resultados. Isso captura o erro de truncamento enquanto mantém o rendimento do hardware nativo de 16 bits.
• Tratamento de Overflow: Para lidar com a estreita faixa dinâmica do FP16, o HEAL aplica redimensionamento direcionado a canais discrepantes (outliers) raros (que constituem <1% dos canais fixos) para evitar o overflow catastrófico sem adicionar overhead significativo.

Principais Contribuições

1. Diagnóstico de Instabilidade: O artigo demonstra que a não-reprodutibilidade intrínseca é fundamentalmente impulsionada por erros de truncamento de fronteira, em vez de não-associatividade aritmética interna. Estabelece que a reprodutibilidade robusta requer limitar este ruído abaixo de um limiar específico de inversão de resultado.
2. Framework HEAL: Introduz uma solução de duas frentes que evita penalidades de hardware ao salvaguardar a Atenção via quantização INT16 e sintetizar GEMM de alta precisão via compensação de erro algébrica.
3. MCR-Bench: Introduz o Mission-Critical Reproducibility Benchmark, um novo conjunto de avaliação voltado para a reprodutibilidade em domínios de alto risco (medicina, direito, finanças) onde a consistência do resultado é primordial.
4. Trade-off Eficiência-Reprodutibilidade: Demonstra que o HEAL alcança reprodutibilidade de nível FP32 para tarefas de downstream em variados tipos de hardware, reduzindo o overhead de desempenho em até 7,1× comparado a pipelines completos de FP32.

Resultados

Avaliado em cinco modelos (Qwen3-8B/14B/32B, Llama-3.1-8B, DeepSeek-R1-Distill-Qwen-14B) e três tipos de hardware (NVIDIA A100, H100, AMD MI300X):

• Reprodutibilidade: O HEAL-Full (o pipeline híbrido completo) alcança uma taxa de inversão de resposta (answer flip rate) e um MSE de Hidden Final quase idênticos ao baseline de FP32 (ex: 0,0% de taxa de inversão para Llama-8B e DeepSeek-14B, comparado a 11,7% e 12,8% para o BF16 nativo).
• Overhead de Desempenho:
  • Baseline FP32: Aumenta o TPOT de 4,2× a 13,1×.
  • HEAL-Full: Incorre em apenas 40%–109% de overhead de TPOT em relação ao baseline nativo de BF16.
  • Memória: O HEAL mantém a pegada de memória de 16 bits para o cache KV, evitando a expansão de 2× de memória exigida pelo upcasting total de FP32.
• Agnosticismo de Hardware: O erro numérico acumulado permanece consistente entre diferentes tipos de hardware sob configurações de precisão fixa, confirmando que o HEAL neutraliza efetivamente a divergência específica de hardware.

Significância

O artigo afirma que a suposição de que um custo de "tudo ou nada" de pipelines de alta precisão total é necessário para uma reprodutibilidade robusta é incorreta. Ao identificar que a não-reprodutibilidade é uma sensibilidade localizada a características discrepantes específicas (truncamento de fronteira) em vez de um problema de ruído generalizado, o HEAL restaura o comportamento determinístico através de hardwares heterogêneos. Crucialmente, o método alcança essa estabilidade ao atualizar menos de 1% do computo total (via compensação algébrica) e otimizar formatos de armazenamento, oferecendo um caminho prático para implantar LLMs em ambientes de missão crítica onde saídas numéricas exatas são exigidas sem o ônus econômico da inferência total em FP32.