Resumo Técnico: MosaicQuant

Declaração do Problema

Embora a quantização de 4 bits ofereça um potencial significativo para reduzir a pegada de memória e acelerar a inferência de Grandes Modelos de Linguagem (LLMs), ela enfrenta uma limitação fundamental: a largura de bits restrita não consegue representar fielmente tanto os valores comuns densos (inliers) quanto os valores raros de grande magnitude (outliers). Essa discrepância leva a uma degradação substancial da precisão.

As soluções existentes normalmente empregam quantização de precisão mista, mantendo os outliers em alta precisão (ex: FP16 ou INT8) enquanto quantizam o restante para 4 bits. No entanto, essa abordagem quebra a uniformidade da execução de baixa bitagem. A necessidade de operações de fallback de alta precisão introduz overhead adicional de movimentação de dados, conversão de formato e conversão de precisão, o que acaba por minar o ganho de velocidade prático prometido pela inferência de baixa bitagem.

Metodologia

Os autores propõem o MosaicQuant, um paradigma de quantização de 4 bits unificado baseado no princípio da desagregação de inlier–outlier. Em vez de preservar os outliers em alta precisão, o MosaicQuant decompõe a matriz de pesos inteiramente em componentes de 4 bits: uma base densa e um resíduo esparso.

1. Desagregação de Inlier–Outlier

O método começa migrando os outliers de ativação para a matriz de pesos usando transformações equivalentes (ex: SmoothQuant ou QuaRot). A matriz de pesos migrada $\hat{W}$ é então decomposta como:
$$\hat{W} \approx Q(\hat{W}) + R$$
Onde:

• Base Densa de 4 bits ($Q(\hat{W})$): Captura fielmente os inliers comuns.
• Resíduo Esparso de 4 bits ($R$): Compensa os erros de quantização introduzidos pelos outliers.

Crucialmente, ambos os componentes operam em 4 bits, preservando a uniformidade do pipeline de execução de baixa bitagem.

2. Alocação de Resíduo por Bloco

O artigo identifica que os erros de quantização não são distribuídos uniformemente entre os canais, mas exibem uma concentração a nível de bloco refinada. Portanto, o MosaicQuant aloca o orçamento de resíduo esparso ao nível do bloco, em vez de ao nível do canal.

Para selecionar os blocos mais críticos sem avaliações de saída completas e dispendiosas, os autores introduzem um $\Delta$-score guiado por Hessiano. Esta métrica aproxima a redução marginal na distorção de saída para cada bloco, ponderando o erro de resíduo com a sensibilidade de ativação (estimativa de Hessiano diagonal) dos canais de entrada. O algoritmo seleciona os top-$K$ blocos com os maiores $\Delta$-scores para receber a compensação de resíduo esparso de 4 bits. Experimentos mostram que um orçamento esparso de apenas 10% dos blocos é suficiente para recuperar mais de 50% da distorção de saída.

3. ZipperEngine: Execução Unificada

Uma representação unificada sozinha é insuficiente se o resíduo esparso for executado via um kernel separado, o que recarregaria as ativações e fundiria saídas parciais, adicionando latência significativa. Para resolver isso, os autores introduzem o ZipperEngine, um motor de execução especializado que:

• Funde Kernels: Integra o cálculo do resíduo esparso diretamente no kernel GEMM de 4 bits denso.
• Pipeline Sobreposto: Utiliza um pipeline tileado e sobreposto onde quantização, computação densa, compensação esparsa e operações de redução-soma são agendadas concorrentemente entre Tensor Cores, CUDA Cores e HBM.
• Resultado: Isso elimina o overhead de despacho de kernel separado e o tráfego de memória redundante, realizando a eficiência de hardware de um pipeline de inferência de 4 bits totalmente unificado.

Principais Contribuições

1. Paradigma MosaicQuant: Uma abordagem de quantização de 4 bits que realiza a desagregação de inlier–outlier por bloco, mitigando a perda de precisão enquanto mantém um padrão de execução de baixa bitagem totalmente funcional.
2. Seleção Guiada por Hessiano: Uma estratégia de seleção de blocos eficiente e consciente das ativações usando um escore $\Delta$-Hessiano para identificar blocos de resíduo críticos, substituindo a avaliação exata de distorção de saída dispendiosa.
3. ZipperEngine: Um motor de execução de 4 bits denso–esparso que funde o cálculo do resíduo esparso no kernel GEMM denso através de um pipeline sobreposto, permitindo ganhos de velocidade práticos mesmo com blocos de resíduo adicionais.

Resultados Experimentais

Os autores avaliaram o MosaicQuant nos modelos LLaMA3 (3B, 8B) e Qwen3 (4B, 8B, 14B, 32B).

• Precisão: O MosaicQuant preserva uma precisão próxima à de FP16. Sob quantização W4A4, ele supera significativamente outras linhas de base de baixa bitagem (RTN, SmoothQuant, QuaRot, SpinQuant, ResQ, Atom). Por exemplo, no Qwen3-32B, o MosaicQuant atinge uma perplexidade de 8.4, comparado a 9.2 para o ResQ e 9.3 para o Atom. Ele também mantém um desempenho robusto onde outros métodos (como RTN e SmoothQuant) frequentemente colapsam em modelos Qwen3.
• Throughput: Em uma NVIDIA RTX 4090, o MosaicQuant alcança até 1.24× de aceleração sobre a linha de base W16A16. Ele entrega consistentemente o maior throughput em tamanhos de lote (batch sizes) de 4 a 32 em comparação ao TensorRT-LLM e outras linhas de base de quantização.
• Overhead: Os ramos esparsos de 4 bits incorrem em overhead mínimo, com o uso de memória aumentando apenas 8.5–9.5% e o overhead de latência permanecendo abaixo de 6.2% devido ao pipeline fundido.

Significância

O artigo afirma que o MosaicQuant resolve o compromisso entre precisão e eficiência de hardware na quantização de 4 bits para LLMs. Ao eliminar a necessidade de operações de fallback de alta precisão e unificar tanto a representação quanto a execução em um único pipeline de baixa bitagem coerente, o MosaicQuant permite o serviço eficiente de LLMs sem sacrificar a regularidade da execução uniforme de 4 bits. Esta abordagem demonstra que a compensação esparsa de baixa precisão pode recuperar efetivamente os erros induzidos por outliers, mantendo as vantagens de throughput da inferência densa de baixa bitagem.