💬 NLP

Out of the Memory Barrier: A Highly Memory Efficient Training System for LLMs with Million-Token Contexts

O artigo apresenta o OOMB, um sistema de treinamento altamente eficiente em memória que permite o treinamento de modelos de linguagem de grande escala com contextos de milhões de tokens em uma única GPU, utilizando um framework recorrente em blocos com recomputação de ativações e otimizações avançadas de gerenciamento de cache KV para manter a sobrecarga de memória constante.

Autores originais: Wenhao Li, Daohai Yu, Gen Luo, Yuxin Zhang, Fei Chao, Rongrong Ji, Yifan Wu, Jiaxin Liu, Ziyang Gong, Zimu Liao

Publicado 2026-03-03

📖 5 min de leitura🧠 Leitura aprofundada

CC BY 4.0

Autores originais: Wenhao Li, Daohai Yu, Gen Luo, Yuxin Zhang, Fei Chao, Rongrong Ji, Yifan Wu, Jiaxin Liu, Ziyang Gong, Zimu Liao

Artigo original sob licença CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Esta é uma explicação gerada por IA do artigo abaixo. Não foi escrita nem endossada pelos autores. Para precisão técnica, consulte o artigo original. Ler aviso legal completo

Imagine que você está tentando ensinar um aluno superinteligente (um Modelo de Linguagem Grande, ou LLM) a ler e entender um livro gigantesco, com milhões de páginas. O problema não é que o aluno demore para ler; o problema é que ele precisa segurar todas as páginas na memória ao mesmo tempo para não esquecer o que leu antes.

Se o livro tiver 256.000 páginas, a "mesa de trabalho" (a memória do computador) fica tão cheia que o aluno desmaia antes mesmo de terminar o primeiro capítulo. É assim que funciona o treinamento de IA hoje: a memória explode quando o contexto (o tamanho do texto) aumenta.

O paper que você enviou apresenta uma solução brilhante chamada OOMB ("Fora da Barreira de Memória"). Vamos explicar como isso funciona usando analogias do dia a dia.

1. O Problema: A Mesa de Trabalho Cheia

Normalmente, para treinar uma IA com textos longos, o computador tenta guardar tudo o que foi processado na memória RAM da placa de vídeo (GPU). É como se você tentasse ler um livro de 1 milhão de páginas segurando todas as páginas abertas na sua mão ao mesmo tempo. É impossível.

2. A Solução OOMB: O Método do "Caderno de Resumos"

A equipe criou um sistema inteligente que muda a forma como a IA "lê" e "estuda". Eles usam três truques principais:

Truque 1: Ler em "Blocos" e Jogar Fora (Recomputação)

Em vez de segurar todo o livro, o OOMB divide o texto em pequenos pedaços (blocos).

Como funciona: A IA lê um bloco, tira as notas principais (o que chamamos de "KV Cache"), e joga o bloco de texto fora imediatamente.
O Pulo do Gato: Quando chega a hora de corrigir os erros (o "retrocesso" ou backpropagation), a IA não precisa do texto original. Ela recria o texto na hora apenas para aquele momento, como se estivesse relembrando uma história que acabou de ouvir.
Analogia: É como se você lesse um parágrafo, anotasse a ideia principal no caderno, e jogasse o papel no lixo. Se precisar revisar, você reescreve o parágrafo mentalmente a partir da anotação. Isso mantém sua mesa de trabalho sempre vazia, não importa o tamanho do livro.

Truque 2: O Armário de Arquivos (Gerenciamento de Memória Paginada)

Agora que a IA não guarda o texto, ela precisa guardar as "anotações" (o KV Cache) de tudo o que já leu. Como essas anotações crescem sem parar, o sistema usa um armário de arquivos inteligente.

Como funciona: Em vez de tentar empurrar todas as anotações em uma única pilha gigante (o que quebraria a mesa), o sistema guarda as anotações em "pastas" (páginas) organizadas.
Analogia: Imagine que você tem um armário infinito. Em vez de tentar colocar todos os documentos em uma única gaveta enorme, você usa pastas numeradas. Quando uma pasta enche, você coloca outra ao lado, sem bagunçar o resto. Isso evita que o sistema fique lento procurando espaço livre.

Truque 3: O Estagiário Rápido (Offloading Assíncrono)

Mesmo com as pastas, às vezes o armário fica muito cheio para caber na mesa de trabalho (memória da GPU).

Como funciona: O sistema tem um "estagiário" (a memória do processador/CPU) que fica ao lado. Enquanto a IA está focada em calcular a próxima parte do texto, o estagiário pega as pastas mais antigas e as leva para o armário ao lado (CPU) sem que a IA precise parar.
Analogia: É como um cozinheiro (GPU) que está cortando legumes. Enquanto ele corta, um ajudante (CPU) pega os legumes que já foram cortados e os leva para a geladeira, para que o cozinheiro nunca fique sem espaço na bancada. O ajudante trabalha tão rápido que o cozinheiro nem percebe que ele saiu.

3. O Truque Extra: Ler Apenas o Importante (Atenção Esparsa)

Para textos extremamente longos (milhões de tokens), o sistema usa um filtro.

Como funciona: Em vez de tentar conectar cada palavra com todas as outras palavras do livro (o que é muito trabalhoso), a IA aprende a focar apenas nas páginas mais relevantes para a palavra atual.
Analogia: Se você está escrevendo um capítulo sobre "cafés", você não precisa reler o capítulo inteiro sobre "espaço sideral" do livro. Você só busca as páginas que falam de café. Isso economiza muita energia e tempo.

O Resultado Final?

Graças a essa combinação de técnicas, o OOMB consegue fazer algo que antes parecia impossível:

Antes: Para treinar um modelo com um contexto de 4 milhões de palavras, você precisaria de um supercomputador gigante com dezenas de placas de vídeo caras.
Agora: Com o OOMB, você consegue fazer isso em uma única placa de vídeo (uma H200), que cabe em um rack de servidor comum.

Resumo em uma frase

O OOMB é como transformar um caminhão de mudanças (que precisa de um galpão gigante para guardar tudo) em um caminhão de entrega expressa que faz várias viagens rápidas, descarta o que não precisa e usa um ajudante para guardar as caixas antigas, permitindo que você transporte uma cidade inteira com um único veículo.

Isso democratiza a pesquisa em Inteligência Artificial, permitindo que universidades e empresas menores treinem modelos capazes de ler livros inteiros sem precisar de milhões de dólares em hardware.

1. O Problema: Barreira de Memória em Contextos Longos

O treinamento de Modelos de Linguagem de Grande Escala (LLMs) com contextos extensos (milhões de tokens) enfrenta um obstáculo crítico: o custo proibitivo de memória de GPU, e não o tempo de treinamento.

Causa Raiz: O consumo de memória das ativações escala linearmente com o comprimento da sequência. Em métodos tradicionais, armazenar todas as ativações para a retropropagação esgota rapidamente a memória da GPU.
Gargalo Atual: Mesmo com técnicas como ZeRO3 e paralelismo de tensores, a memória necessária para o KV Cache (Key-Value Cache) e as ativações torna o treinamento de contextos superiores a 32K-256K inviável em uma única GPU, exigindo clusters massivos.
Limitação de Métodos Existentes: Métodos de extensão de contexto sem treinamento (training-free) falham em raciocínio de longo alcance, e o ajuste fino (fine-tuning) dedicado exige recursos computacionais que muitas instituições não possuem.

2. Metodologia: O Sistema OOMB

Os autores propõem o OOMB (Out Of the Memory Barrier), um sistema de treinamento altamente eficiente que combina uma abordagem de treinamento recursivo por blocos (chunk-recurrent) com otimizações de gerenciamento de memória.

A. Treinamento Recursivo por Blocos com Recálculo de Ativações

Estratégia: A sequência de entrada é dividida em blocos (chunks). Durante o forward pass, as ativações de cada bloco são computadas e imediatamente descartadas.
Recomputação: Para o backward pass, as ativações são recalculadas "on-the-fly" (sob demanda).
Resultado: Isso reduz a complexidade de memória das ativações de $O(N)$ para $O(1)$ (constante), independentemente do comprimento total do contexto. O gargalo de memória é deslocado exclusivamente para o gerenciamento do KV Cache.

B. Gerenciamento de KV Cache Paginado (Paged KV Cache)

Inspirado no vLLM (para inferência), mas adaptado para treinamento com suporte a retropropagação.
Memória Paginada: O KV Cache e seus gradientes são armazenados em páginas não contíguas na memória. Isso elimina a fragmentação de memória e operações caras de realocação/cópia ao adicionar novos pares chave-valor.
Kernels Personalizados: Otimizações em nível de kernel (Triton/CUDA) realizam a acumulação de gradientes in-place, contornando o sistema de autograd do PyTorch para evitar armazenar o KV Cache como uma ativação temporária.

C. Sparse Attention em Nível de Página

Para reduzir a complexidade computacional e de comunicação, o sistema implementa atenção esparsa baseada em páginas.
Mecanismo: Utiliza uma estratégia de recuperação Top-K para identificar as páginas de chaves mais relevantes para cada página de consultas, aproximando a atenção densa.
Benefício: Reduz drasticamente a quantidade de dados transferidos e processados, mantendo a precisão do modelo.

D. Offloading Assíncrono de KV Cache para CPU

Como o KV Cache ainda cresce linearmente com o contexto, o OOMB transfere parte dele para a memória da CPU.
Assincronia: O sistema utiliza streams CUDA dedicados e memória pinned para transferir dados entre CPU e GPU enquanto a computação ocorre, mascarando a latência de transferência.
Estratégia Adaptativa:
- Para atenção densa/local: Pré-busca (pre-fetching) do próximo bloco durante o cálculo do bloco atual.
- Para atenção esparsa: A transferência ocorre em paralelo com os cálculos de projeção de chaves e valores, já que os índices das páginas necessárias são conhecidos após a projeção da consulta.

3. Principais Contribuições

Redução de Memória de Ativação para O(1): O primeiro sistema a manter um footprint de memória de ativação constante durante o treinamento de contextos de milhões de tokens.
Gerenciamento de Memória Paginado para Treinamento: Adaptação bem-sucedida de técnicas de paginação para o KV Cache e seus gradientes durante o treinamento, eliminando fragmentação.
Otimização de Offloading Assíncrono: Um mecanismo que permite treinar contextos massivos em uma única GPU, escondendo a latência de transferência de dados.
Integração Sinérgica: A combinação de chunk-recurrent, paginação, atenção esparsa e offloading cria um sistema onde o ganho de eficiência é maior que a soma das partes.

4. Resultados Experimentais

Os experimentos foram realizados no modelo Qwen2.5-7B em uma única GPU H200.

Eficiência de Memória:
- O overhead de memória de ponta aumenta apenas 10 MB para cada 10.000 tokens adicionais de contexto.
- É possível treinar o Qwen2.5-7B com um contexto de 4 milhões de tokens em uma única GPU H200.
- Em comparação, métodos de paralelismo de contexto (como Ring Attention) exigiriam um cluster grande para tarefas similares.
Desempenho e Latência:
- O sistema mantém uma taxa de transferência (throughput) competitiva.
- O offloading assíncrono adiciona menos de 5% de overhead de latência em cenários de atenção densa.
- Com atenção esparsa, o sistema alcança acelerações de até 13,5x em contextos de 256K tokens comparado a baselines sem otimização.
Validação de Precisão:
- O uso de atenção esparsa não degrada significativamente a perda (loss) do modelo.
- O modelo convergiu com sucesso para contextos de 256K e 1 milhão de tokens, mantendo curvas de perda semelhantes às de atenção densa.

5. Significado e Impacto

O trabalho representa um avanço substancial na eficiência de recursos para o treinamento de LLMs de longo contexto.

Democratização: Permite que pesquisadores e instituições com recursos limitados (apenas uma GPU de ponta) realizem experimentos de treinamento de longo contexto que antes exigiam clusters de centenas de GPUs.
Sustentabilidade: Ao reduzir drasticamente a necessidade de hardware, o método contribui para a redução da pegada de carbono e do consumo energético associado ao treinamento de IA.
Viabilidade Prática: Demonstra que o treinamento de modelos com contextos de milhões de tokens é viável sem comprometer a qualidade do modelo, superando a barreira de memória que limitava a comunidade a contextos de dezenas de milhares de tokens em ambientes de GPU única.

Em resumo, o OOMB remove a barreira de memória que impedia o treinamento escalável de LLMs em contextos ultra-longos, oferecendo uma solução técnica robusta que combina reutilização de memória, paginação e computação assíncrona.