Efficient, VRAM-Constrained xLM Inference on Clients

1. Declaração do Problema

O artigo aborda o desafio crítico de executar modelos de IA de grande escala e alta precisão (coletivamente denominados xLMs, incluindo LLMs Densos, LLMs de Mistura de Especialistas (MoE) e Modelos Visão-Linguagem (VLMs)) em sistemas do lado do cliente (desktops e laptops) com Memória de Vídeo (VRAM) limitada.

• A Restrição: As GPUs de clientes frequentemente possuem significativamente menos VRAM do que a memória total exigida por modelos modernos de alta precisão (por exemplo, um modelo de 77GB em uma GPU de 16GB).
• A Limitação das Soluções Existentes:
  • Quantização/Poda: Essas técnicas reduzem o tamanho do modelo, mas sacrificam a precisão, o que é inaceitável para aplicações de alta fidelidade.
  • Descarregamento Estático: O agendamento híbrido existente (por exemplo, em llama.cpp) frequentemente depende de configuração manual (por exemplo, especificar o número de camadas da GPU) e falha em adaptar-se dinamicamente a condições variáveis do sistema (threads da CPU, largura de banda PCIe) ou fases de inferência (preenchimento vs. decodificação).
  • Falta de Universalidade: Nenhum produto existente único lida eficientemente com modos interativos, processamento em lote, VLMs de alta resolução e arquiteturas densas/MoE simultaneamente sob orçamentos estritos de VRAM.

2. Metodologia

Os autores propõem uma técnica de agendamento híbrido novel, guiada por perfis de benchmark, chamada Fragmentação em Pipeline (Pipelined Sharding), combinada com um conjunto de otimizações específicas para visão chamadas VLMOpt.

A. Fragmentação em Pipeline (Inovação Central)

Esta é uma técnica de agendamento totalmente automática que quebra o grafo de inferência no nível de subcamada (por exemplo, dividindo blocos de Atenção e FFN) em vez de apenas no nível de camada. Opera em três fases:

1. Fase de Instalação (Perfilamento):

   • Um conjunto de benchmarks executa no hardware específico do cliente para perfilar o desempenho do kernel (FLOPS) para vários tamanhos de tensores, níveis de quantização e contagens de threads da CPU.
   • Crucialmente, perfila sob carga concorrente (simulando tráfego PCIe e contenção do controlador de memória) para gerar estimativas de tempo precisas para execução híbrida CPU-GPU.
   • Isso cria um pequeno banco de dados de perfis (~170KB) usado para estimativa de custos.

2. Fase de Planejamento (Geração de Agenda):

   • O sistema define um conjunto de "Níveis de Token" (por exemplo, 1, 4, 16, 64 tokens por iteração) representando diferentes cargas de inferência.
   • Para cada nível, gera três planos de agendamento distintos baseados no orçamento atual de VRAM e condições do sistema:
     • Apenas GPU: Todas as subcamadas não fixadas (unpinned) executam na GPU; os pesos são transmitidos via PCIe. Ideal para cargas de computação altas (lotes grandes/contexts) onde a latência PCIe é ocultada pela computação.
     • Estático: Subcamadas de alta prioridade (Atenção, KV Cache) são fixadas na VRAM; as de baixa prioridade (FFN) executam na CPU. Minimiza o tráfego PCIe, mas depende da computação da CPU.
     • Dinâmico: Uma abordagem híbrida onde os pesos para subcamadas da GPU são transmitidos da memória do sistema (sysRAM) enquanto a CPU executa subcamadas anteriores. Isso sobrepõe transferências PCIe com computação.
   • Um estimador de tempo baseado em Roofline usa o banco de dados de perfis para prever a latência de cada plano e seleciona o ótimo para o nível de token específico.

3. Fase de Inferência (Execução):

   • Durante a inferência, o sistema determina a contagem de "novo token" para o lote atual.
   • Consulta a agenda ótima pré-calculada para aquele nível de token e executa-a.
   • Isso permite que o sistema se adapte perfeitamente entre o modo interativo (Tamanho do Lote = 1, poucos tokens) e o modo em lote (Tamanho do Lote > 1, muitos tokens).

B. VLMOpt (Otimização Visão-Linguagem)

Para suportar VLMs de alta resolução (especificamente o modelo Cosmos-Reason1), os autores implementaram três otimizações dentro do llama.cpp:

1. Descarregamento de Tensor de Visão para CPU: Pesos do modelo de visão são fixados na memória da CPU e transmitidos para a GPU apenas quando necessário, reduzindo a pegada de VRAM.
2. Flash Attention para Visão: Implementa Flash Attention no codificador de visão para reduzir a complexidade de memória $O(N^2)$ das pontuações de atenção para imagens de alta resolução.
3. Evitação de Sobreposição de VRAM: Reestrutura a execução para garantir que a codificação de visão seja concluída e libere VRAM antes que a decodificação de linguagem comece, garantindo que o uso máximo de VRAM seja o máximo das duas partes, em vez da sua soma.

3. Contribuições Principais

1. Algoritmo de Fragmentação em Pipeline: Um agendador híbrido automático e novel que se adapta às condições do sistema (threads da CPU, largura de banda PCIe) e estados de inferência (contexto vs. decodificação) sem ajuste manual.
2. Implementação VLMOpt: Um conjunto de otimizações que permite inferência de VLM de alta resolução em hardware de cliente com requisitos drasticamente reduzidos de VRAM.
3. Implementação Open-Source em C++: A solução é implementada em llama.cpp (um framework focado em C++), tornando-a acessível a desenvolvedores de jogos e engenheiros de IA física.
4. Avaliação Abrangente: Testes rigorosos em LLMs Densos, LLMs MoE e VLMs em hardware de cliente diversificado (de laptops a desktops de alto desempenho).

4. Resultados

A avaliação demonstra melhorias significativas de desempenho sobre baselines agressivos (llama.cpp e vLLM ajustados manualmente):

• Modo Interativo (Tamanho do Lote = 1):
  • TTFT (Tempo até o Primeiro Token): Melhorado em até 6,7×.
  • TPS (Tokens por Segundo): Melhorado em até 30× (por exemplo, para o modelo Qwen235B de 77GB).
  • Viabilidade: Permite que modelos massivos (por exemplo, 77GB) rodem interativamente em GPUs com tão pouco quanto 2GB de VRAM (atingindo >5 TPS, o limiar para velocidade de leitura humana).
• Modo em Lote (Tamanho do Lote > 1):
  • Melhorias de throughput de até 8,2× comparado às baselines.
  • O agendador desloca efetivamente o trabalho para a GPU à medida que o tamanho do lote aumenta, evitando gargalos da CPU.
• Desempenho VLM (Cosmos-Reason1):
  • Redução de VRAM: Reduziu a demanda de VRAM em 10× (de ~20GB para ~2GB), permitindo inferência de alta resolução (1440p) em sistemas de cliente onde era anteriormente impossível.
  • Latência: Alcançou até 1,78× de aceleração na Latência de Ponta a Ponta (E2EL) comparado às baselines em orçamentos de VRAM equivalentes.
• Co-Execução: Em cenários executando um jogo (Cyberpunk 2077) e um LLM concorrentemente, a solução permite que o jogo retenha mais ativos na VRAM, melhorando tanto os FPS do jogo quanto o TPS do LLM ao evitar contenção de recursos.

5. Significado

Este trabalho preenche a lacuna entre os altos requisitos de precisão da IA moderna e as restrições de hardware dos dispositivos de cliente de consumo.

• Inferência sem Perdas: Prova que modelos de alta precisão podem rodar em hardware de cliente sem sacrificar a fidelidade do modelo através de quantização ou poda.
• Democratização: Ao permitir VLMs de alta resolução e grandes modelos MoE em laptops e desktops padrão, capacita desenvolvedores a construir "IA Física" avançada e assistentes de IA em jogos sem depender de infraestrutura de nuvem cara.
• Adaptabilidade: A abordagem baseada em níveis de token e perfis fornece uma estrutura robusta para a evolução futura de hardware, garantindo que o desempenho de inferência escale com os recursos disponíveis (threads da CPU, largura de banda PCIe) em vez de estar estaticamente vinculado a uma configuração fixa.

Em resumo, o artigo apresenta uma solução pronta para produção, automatizada, que torna a próxima geração de IA do lado do cliente viável, eficiente e sem perdas.