Resumo Técnico: PagedWeight

Declaração do Problema

Modelos de Linguagem de Grande Escala (LLMs) de Mistura de Especialistas (MoE) oferecem alta eficiência e precisão ao ativar apenas um subconjunto de especialistas por token. No entanto, em cenários de serviço, os modelos MoE enfrentam uma tensão de memória crítica: a GPU deve armazenar tanto os pesos massivos do modelo quanto o crescente cache Chave-Valor (KV). Embora técnicas como o PagedAttention gerenciem a fragmentação do cache KV, elas não abordam a significativa pegada de memória dos pesos MoE, que podem ocupar mais de 60% da memória da GPU.

Os métodos de quantização existentes são amplamente estáticos, fixando a precisão antes do início da inferência. Embora existam alguns métodos de adaptação de precisão em tempo de execução, eles não utilizam mudanças de precisão para realocar dinamicamente a memória da GPU entre os pesos e o cache KV. Consequentemente, há uma falta de sistemas que possam liberar memória dinamicamente de pesos de especialistas menos críticos conforme o cache KV se expande, sem sacrificar a precisão da tarefa ou interromper a inferência.

Metodologia: PagedWeight

Os autores propõem o PagedWeight, um novo sistema de gerenciamento para o serviço de MoE LLM que quantiza dinamicamente os pesos do modelo em tempo de execução para equilibrar a precisão dos pesos dos especialistas com os tamanhos do cache KV. O sistema trata os pesos dos especialistas de forma semelhante aos blocos de cache KV paginados, permitindo o descarregamento (offloading) e recarregamento dinâmicos.

Componentes Principais

1. Representação de Peso Paginado (Paged Weight Representation):

   • O PagedWeight opera na granularidade de blocos lineares por especialista (identificados por camada, especialista e tipo de módulo: gate_up ou down).
   • Ele utiliza a tecnologia Any-Precision LLM (APL), representando os pesos em um formato de plano de bits sobreposto com Tabelas de Consulta (LUTs) associadas.
   • Uma Tabela de Páginas de Peso (Weight Page Table) rastreia a largura de bits desejada ($d_i$) e a largura de bits comprometida ($q_i$) para cada bloco linear. As páginas podem existir nos estados residente na GPU, residente na CPU ou de transferência.

2. Planejador de Tempo de Execução Consciente de Qualidade (Quality-Aware Runtime Planner):
   O planejador determina quais pesos quantizar (reduzir a largura de bits) com base em três fatores para minimizar a perda de precisão:

   • Sensibilidade Global Offline: Utiliza pontuações ponderadas pelo Hessiano para estimar a sensibilidade intrínseca de cada bloco linear à quantização.
   • Estatísticas de Roteamento Online: Rastreia a "massa de roteamento" (frequência e peso de seleção) dos especialistas. Especialistas frequentemente roteados ("quentes") são protegidos com pisos de largura de bits mais altos e multiplicadores de dano.
   • Resíduos de Prompt: Um termo de correção específico do prompt previsto via cabeças de regressão linear. Isso ajusta a estimativa de sensibilidade global com base na sequência de entrada atual, reconhecendo que o impacto da quantização varia conforme o prompt.

   O planejador calcula uma pontuação de "dano" para potenciais reduções de largura de bits. Quando a pressão do cache KV aciona um alvo de memória, o planejador seleciona avidamente as reduções de menor dano (maior qualidade por byte economizado) para atingir o alvo.

3. Execução Assíncrona e Otimização de Kernel:

   • Movimentação de Página Assíncrona: Para ocultar a latência, o sistema descarrega as páginas de peso para a RAM da CPU e as recarrega assincronamente. Os compromissos de largura de bits são atualizados apenas em limites seguros (por exemplo, após um plano ser totalmente executado), garantindo que a inferência não seja interrompida.
   • Kernel de Precisão Mista Fundido (Fused Mixed-Precision Kernel): Um kernel CUDA customizado lê diretamente os planos de bits e as LUTs de Any-Precision, suportando diferentes larguras de bits para cada bloco linear dentro de uma única operação fundida para reduzir o overhead.

Principais Contribuições

1. Design do Sistema: A proposta do PagedWeight, um sistema de peso paginado para pesos MoE de Any-Precision online que gerencia larguras de bits comprometidas/desejadas, estados de página e consumo de memória dinamicamente.
2. Algoritmo de Planejamento: Um planejador de tempo de execução consciente da qualidade que sintetiza sensibilidade offline, estatísticas de roteamento online e resíduos de prompt para selecionar estratégias de redução de peso de baixo dano sob pressão de memória.
3. Estratégia de Execução: Implementação de movimentação de página assíncrona para ocultar a latência de descarregamento/recarregamento com perda mínima de throughput, juntamente com um kernel MoE de precisão mista fundido.
4. Validação Empírica: Avaliação abrangente em três modelos MoE (Qwen1.5-MoE-A2.7B, Mixtral-8×7B, Gemma-4-26B-A4B) demonstrando desempenho superior sobre as linhas de base existentes.

Resultados

Os autores avaliaram o PagedWeight contra FP16, Any-Precision LLM (APL), DP-LLM e MxMoE em tarefas de modelagem de linguagem, raciocínio e contexto longo.

• Compromisso Qualidade-Memória: O PagedWeight alcança precisão equivalente ao FP16 com até 72,0% de economia de memória da GPU e uma melhoria de throughput de 1,94× em comparação com as linhas de base.
• Qualidade em Orçamento Fixo: Em orçamentos de memória similares, o PagedWeight melhora a qualidade da tarefa sobre métodos de quantização em até 39,3%, com uma perda de throughput de no máximo 4,1%.
• Desempenho de Contexto Longo: Sob orçamentos de memória apertados (ex: 10 GB), o PagedWeight mantém a qualidade de contexto longo (Passage Retrieval, NarrativeQA) comparável ao FP16, enquanto as linhas de base de quantização estática (como APL) sofrem degradação significativa.
• Throughput: O PagedWeight iguala o throughput das linhas de base de quantização uniforme, com a maior queda observada sendo de 4,1% no batch size 4.
• Ablação: Remover componentes como estatísticas de roteamento, resíduos de prompt ou movimentação de página degrada consistentemente o desempenho, confirmando a necessidade do design completo do sistema.

Significância

O artigo afirma que o PagedWeight navega efetivamente pelo complexo tradeoff entre a precisão da tarefa do modelo, consumo de memória e throughput/latência. Ao tratar os pesos dos especialistas como unidades paginadas e gerenciáveis, em vez de ativos estáticos, o PagedWeight desbloqueia um espaço de tradeoff inexplorado entre a precisão da tarefa e a alocação de memória da GPU entre os pesos e o cache KV. O sistema demonstra que a quantização dinâmica e consciente da qualidade é uma estratégia viável e superior para o serviço de MoE em comparação com a quantização estática ou métodos que não se adaptam à pressão de memória em tempo de execução.