OneComp: One-Line Revolution for Generative AI Model Compression

Resumo Técnico: OneComp

1. O Problema

A implementação de modelos de fundação (Foundation Models) para IA generativa enfrenta barreiras críticas de footprint de memória, latência e custo de hardware. Modelos com dezenas a centenas de bilhões de parâmetros excedem a capacidade de memória de hardware comercial, tornando a inferência em precisão total proibitiva para a maioria das organizações.

Embora a Quantização Pós-Treinamento (PTQ) seja uma estratégia chave para mitigar esses gargalos, sua implementação prática é fragmentada. Os praticantes precisam navegar por um ecossistema complexo de algoritmos de quantização, orçamentos de precisão, estratégias de calibração baseadas em dados e regimes de execução dependentes de hardware. Não existe uma "receita" universal; a combinação ótima varia conforme o modelo, a tarefa e o hardware. Isso cria uma lacuna significativa entre o que é teoricamente alcançável na pesquisa e o que é viável na produção.

2. Metodologia: O Framework OneComp

O OneComp é um framework de código aberto que transforma o fluxo de trabalho de compressão em um pipeline reprodutível e adaptativo a recursos. A premissa central é tratar o primeiro checkpoint quantizado como um "pivô" deployável, onde cada estágio subsequente melhora o mesmo modelo, garantindo que a qualidade aumente à medida que mais recursos computacionais são investidos.

O pipeline segue três princípios de design:

1. API Path-to-Plan: Deriva automaticamente o fluxo de trabalho a partir da arquitetura do modelo e do orçamento de GPU, eliminando a necessidade de configuração manual.
2. Motor Adaptativo a Recursos: Fornece melhoria monotônica de qualidade, escalando de PTQ por camada até refinamento global.
3. Arquitetura Extensível: Permite a integração de novos algoritmos sem redesenhar o pipeline.

Etapas do Pipeline OneComp:

• Investigação e Planejamento (AutoBit):
  • Analisa a sensibilidade de cada camada à quantização.
  • Resolve um problema de otimização inteira linear (usando o solver SCIP) para alocar precisão mista heterogênea (diferentes larguras de bits por módulo) dentro de um orçamento global de memória, em vez de usar uma largura de bit uniforme.
• Pré-processamento (Transformações Equivalentes):
  • Aplica transformações para redistribuir a energia de "outliers" (valores extremos) que degradam a quantização.
  • Inclui escalonamento por canal (SmoothQuant), rotações baseadas em Hadamard ou aprendidas (SpinQuant) e balanceamento de pesos (Sinkhorn).
• Estágios de Quantização (Refinamento Progressivo):
  • PTQ por Camada (Layer-wise): O ponto de entrada padrão. Processa uma camada linear de cada vez. O OneComp incorpora técnicas avançadas para mitigar a propagação de erros:
    • QEP (Quantization Error Propagation): Corrige o alvo de quantização de cada camada para compensar erros introduzidos por camadas anteriores quantizadas.
    • LPCD (Submodule-Aware Coordinate Descent): Otimiza conjuntamente submódulos acoplados (ex: projeções Q/K/V ou Gate/Up/Down) para manter a consistência funcional.
    • Formatos: Suporta JointQ (otimização conjunta de escala e inteiros para 3-4 bits) e formatos binários estruturados DBF/MDBF para regimes extremos (1-2 bits).
  • PTQ por Bloco (Block-wise): Refina um bloco inteiro do Transformer de cada vez.
    • Compensação: Ajusta parâmetros restantes em FP16 dentro de um bloco parcialmente quantizado.
    • Refinamento: Otimiza os parâmetros quantizados de um bloco completo usando destilação de MSE e, opcionalmente, refinamento cruzado entre blocos adjacentes.
  • PTQ Global: Otimiza todos os parâmetros do modelo simultaneamente (quando há memória suficiente), minimizando a divergência KL entre o modelo professor (precisão total) e o aluno (quantizado), com alinhamento de estados intermediários.
• Ajuste Fino (Fine-tuning):
  • Opcionalmente, aplica QLoRA para ajustar parâmetros de baixo rank (adapters) sobre o backbone quantizado, recuperando erros residuais ou adaptando-se a domínios específicos.

3. Contribuições Chave

• Pipeline Unificado e Automatizado: Elimina a necessidade de configuração manual de hiperparâmetros complexos, oferecendo uma interface de "uma linha" para gerar modelos prontos para produção.
• Planejador de Precisão Mista (AutoBit): Demonstra que a alocação heterogênea de bits, baseada em estatísticas de ativação, supera significativamente a quantização uniforme, especialmente em orçamentos agressivos.
• Técnicas de Correção de Erro: Introduz e valida QEP e LPCD, que reduzem drasticamente a degradação de desempenho em PTQ por camada, aproximando-a de métodos mais caros sem exigir hardware adicional.
• Suporte a Regimes Extremos: Apresenta MDBF (Multi-Envelope DBF), uma evolução do formato binário estruturado que supera o DBF original ao modelar envelopes de magnitude de maior rank, mantendo precisão significativa em 1-2 bits.
• Melhoria Monotônica: Estabelece um paradigma onde cada estágio adicional (Camada -> Bloco -> Global -> Fine-tuning) melhora consistentemente a qualidade do modelo.

4. Resultados Experimentais

Os experimentos foram realizados em famílias de modelos LLaMA e Qwen, utilizando métricas de Perplexidade (WikiText-2) e Acurácia Zero-Shot (ARC, PIQA, WinoGrande).

• Precisão Mista vs. Uniforme: O AutoBit superou consistentemente a quantização uniforme. No LLaMA-3-8B com 4.16 bits/peso, a versão "activation-aware" atingiu uma perplexidade de 6.85 (próximo ao baseline de 6.14), enquanto a quantização uniforme degradou para 9.53. Em 4.00 bits, a quantização uniforme colapsou, mas o AutoBit manteve desempenho utilizável.
• Correção de Propagação de Erro: A combinação de QEP e Submodule (LPCD) restaurou a qualidade do modelo em configurações de 3 e 4 bits, reduzindo a perplexidade de valores catastróficos (ex: >600) para níveis próximos ao baseline, superando o GPTQ padrão.
• JointQ vs. GPTQ: O JointQ superou o GPTQ padrão em todas as configurações de 3 e 4 bits, especialmente em quantização por canal, demonstrando que a otimização conjunta de escalas e inteiros preserva melhor a estrutura relevante para a tarefa.
• Quantização Extrema (1-2 bits): O formato MDBF superou o DBF em todos os modelos testados (de 0.6B a 8B parâmetros), mantendo perplexidade e acurácia significativas onde a quantização uniforme falharia completamente.
• Escalabilidade de Recursos: A qualidade do modelo aumentou monotonicamente à medida que se avançava do estágio de camada para bloco e global, confirmando que o pipeline permite que os usuários parem em qualquer ponto compatível com seu orçamento de hardware sem perder a qualidade do estágio anterior.

5. Significado e Impacto

O OneComp representa uma mudança de paradigma na compressão de modelos generativos:

• Democratização: Torna técnicas de ponta de pesquisa acessíveis a engenheiros não especialistas, automatizando a complexidade da seleção de algoritmos.
• Eficiência Econômica: Permite que modelos massivos sejam implantados em hardware limitado (ex: uma única GPU consumer), reduzindo custos de infraestrutura e energia.
• Padronização: Oferece um framework aberto e extensível que pode se tornar o padrão da indústria para pipelines de compressão, integrando futuras inovações (como quantização de ativações e cache KV) de forma coesa.

Em suma, o OneComp preenche a lacuna entre a inovação algorítmica e a implantação de produção, transformando a quantização de um processo manual e fragmentado em um fluxo de trabalho automatizado, adaptativo e de alta qualidade.