Resumo Técnico: Reduzindo o Gargalo de Memória da GPU com Compressão Sem Perdas para ML

Declaração do Problema

As cargas de trabalho modernas de aprendizado de máquina (ML), incluindo Redes Neurais de Grafos (GNNs), Modelos de Recomendação de Aprendizado Profundo (DLRMs) e Grandes Modelos de Linguagem (LLMs), frequentemente processam conjuntos de dados que variam de centenas de gigabytes a terabytes. Esses volumes de dados excedem a capacidade da memória disponível na GPU, necessitando de transferências de dados sob demanda da memória da CPU, disco ou armazenamento em rede via o barramento de interconexão PCIe.

A principal restrição de desempenho nesses cenários é a largura de banda do PCIe, que é mais de uma ordem de magnitude inferior à largura de banda da memória da GPU. Consequentemente, a execução é frequentemente limitada pela latência de transferência de dados, em vez da capacidade de computação da GPU. Embora técnicas como cache e sobreposição de computação com transferências ofereçam um alívio parcial, elas não podem mitigar totalmente o gargalo, especialmente quando ocorrem falhas de cache (cache misses).

As soluções existentes frequentemente dependem de compressão com perda (ex: quantização) para reduzir o volume de dados. No entanto, métodos com perda introduzem uma degradação de precisão imprevisível que varia conforme o modelo e a carga de trabalho. Em implantações comerciais, mesmo perdas mínimas de precisão (ex: 0,1%) são muitas vezes inaceitáveis devido aos riscos de receita, tornando a compressão com perda difícil de adotar como uma solução de uso geral.

A compressão sem perdas oferece uma alternativa amigável para implantação que preserva a precisão do modelo, mas enfrenta desafios significativos de implementação em pipelines de ML:

1. Competição de Latência: As latências de transferência CPU-GPU são competitivas com as latências de descompressão da GPU; algoritmos de compressão tradicionais frequentemente adicionam um overhead que anula as economias de transferência.
2. Restrições de Memória: A descompressão muitas vezes requer grandes buffers intermediários ou estruturas de metadados complexas que competem com os inputs do modelo pela escassa memória da GPU.
3. Descompasso de Paralelismo: Muitos algoritmos sem perdas padrão (ex: codificação de Huffman, LZ4) dependem de buscas em dicionários ou padrões de acesso à memória irregulares que dificultam o paralelismo de warps da GPU e os acessos de memória coalescidos.

Metodologia: Invariant Bit Packing (IBP)

Os autores propõem o Invariant Bit Packing (IBP), um novo algoritmo de compressão sem perdas projetado especificamente para transferências de tensores de ML. O insight central é que os tensores de ML, mesmo quando numericamente diversos, frequentemente exibem bits invariantes de baixa entropia através de grupos de tensores. Por exemplo, em representações de ponto flutuante, os bits de sinal e expoente frequentemente permanecem constantes em muitos tensores devido à regularização e dinâmica de treinamento.

Algoritmo Central

O IBP opera em três estágios principais:

1. Pré-processamento (Geração de Metadados):

   • O sistema analisa um lote de tensores de entrada para identificar posições de bits que são invariantes (ou seja, possuem o mesmo valor) em uma alta porcentagem ($T$) dos tensores.
   • Ele gera uma Máscara de tamanho fixo (indicando quais posições de bits são invariantes) e um Bitval (armazenando os valores invariantes).
   • Este metadado é armazenado na memória da GPU (tipicamente do tamanho de kilobytes) e é reutilizado para todo o conjunto de dados, evitando o overhead de metadados por tensor.

2. Compressão:

   • Os tensores de entrada são processados em blocos (chunks). Para cada bloco, o algoritmo verifica se os bits nas posições invariantes correspondem ao Bitval armazenado.
   • Se eles coincidirem, os bits invariantes são removidos (eliminados), e os bits restantes são deslocados para preencher a lacuna.
   • Um bit de participação é prefixado a cada bloco para indicar se ele foi comprimido ou se foi deixado sem compressão (para lidar com casos onde a compressão não gera economia de espaço).
   • A compressão é realizada na GPU usando Memória Virtual Unificada de zero-copy, permitindo que a GPU acesse a memória da CPU diretamente sem a necessidade de configuração explícita de DMA iniciada pela CPU para pequenas transferências.

3. Descompressão:

   • A descompressão é otimizada para o paralelismo de warp da GPU. Um warp de 32 threads lê cooperativamente um segmento contíguo de dados comprimidos da memória da CPU para um espaço de trabalho de memória compartilhada (shared memory).
   • As threads utilizam primitivas de warp (comunicação de baixa latência dentro de um warp) para realizar um prefix scan, calculando o deslocamento de bits inicial para o bloco atribuído a cada thread.
   • As threads então reconstroem seus dados de forma independente, reinserindo os bits invariantes do Bitval global nas posções corretas.
   • O processo utiliza acessos de memória assíncronos para sobrepor a busca de dados com a lógica de descompressão, ocultando a latência do PCIe.

Integração do Sistema

Os autores integram o IBP em pipelines de ML de duas formas principais:

• Compressão de Cache: Comprimindo tensores armazenados em caches de software da GPU (ex: em frameworks de GNN) para aumentar a capacidade efetiva do cache.
• Compressão de Transferência PCIe: Comprimindo tensores na memória da CPU antes da transferência. Para lidar com tensores comprimidos de tamanho variável, o sistema utiliza compressão in-place (alocando buffers de tamanho fixo, mas lendo apenas a parte comprimida) e mantém uma máscara de bits na memória da GPU para rastrear quais tensores estão comprimidos.

Principais Contribuições

1. Análise de Invariância de Dados de ML: O artigo identifica e quantifica a presença de bits invariantes em diversas cargas de trabalho de ML (GNN, DLRM, LLM) e tipos de dados (FP32, FP16, BF16), demonstrando que distribuições estruturadas permitem compressão significativa sem perdas sem perda de precisão.
2. Algoritmo Invariant Bit Packing (IBP): Uma técnica de compressão sem perdas otimizada para GPU que remove bits invariantes, armazena-os em metadados globais mínimos e utiliza descompressão paralela de warp para alcançar alto throughput.
3. Integração de Sistema: Implementação do IBP em populares frameworks de ML (DGL, Legion, Colossal-AI, FlexGen, InfiniGen) via APIs de fácil uso (extensões PyTorch e bibliotecas CUDA).
4. Avaliação Abrangente: Benchmarking extensivo contra bibliotecas aceleradas por GPU de última geração (nvCOMP, ndzip-gpu) e sistemas nativos através de conjuntos de dados do mundo real.

Resultados

A avaliação foi conduzida em uma GPU NVIDIA A100 com PCIe Gen 4.

• Desempenho de Compressão: O IBP alcança economia de espaço de 10–12% para conjuntos de dados densos de GNN, 8% para pesos de DLRM e até 27% para pesos de LLM e KV-caches (particularmente com tipos de dados BF16).
• Throughput: O IBP supera consistentemente as bibliotecas de compressão aceleradas por GPU existentes. Para conjuntos de dados esparsos de GNN, ele alcança um aumento de 9.7x no throughput de transferência em comparação com transferências não comprimidas.
• Acelerações de Carga de Trabalho:
  • Treinamento de GNN: Aceleração média de 74% (até 2.7x em casos específicos) ao reduzir o tempo de transferência PCIe e aumentar as taxas de acerto de cache.
  • Inferência de DLRM: Aceleração média de 180% para buscas de embeddings, impulsionada principalmente por transferências eficientes iniciadas pela GPU e redução do overhead da CPU.
  • Inferência de LLM: Aceleração média de 24% para inferência com offloading. Para offloading de pesos (BF16), a latência foi reduzida em 22%; para offloading de KV-cache, a latência foi reduzida em até 27%.
• Overhead: O overhead de descompressão é mínimo; para economias de espaço abaixo de 50% (típicas para dados densos de ML), o throughput permanece acima de 95% da taxa ideal de transferência não comprimida. O tempo de pré-processamento é moderado e ocorre fora do caminho crítico.

Significância e Alegações

O artigo afirma que o IBP fornece uma solução prática e sem perdas para o gargalo de memória da GPU sem os riscos de precisão associados à quantização ou compressão com perda. Ao aproveitar a estrutura inerente dos dados de ML (bits invariantes) e otimizar para a arquitetura da GPU (paralelismo de warp, acesso coalescido, zero-copy), o IBP efetivamente oculta o overhead de descompressão atrás da latência de transferência do PCIe.

Os autores enfatizam que o IBP é um algoritmo de uso geral aplicável a vários tipos de dados (float32, float16, bfloat16, vetores esparsos) e cargas de trabalho (treinamento e inferência). A capacidade de integrar o IBP com mudanças mínimas de código (frequentemente menos de 35 linhas de Python) nos frameworks existentes demonstra seu potencial para ser adotado em implantações comerciais de ML para escalar modelos além dos limites de memória da GPU, mantendo o desempenho e a precisão.