🤖 machine learning

LongStraw: Long-Context RL Beyond 2M Tokens under a Fixed GPU Budget

O LongStraw é uma pilha de execução consciente da arquitetura que permite o aprendizado por reforço de pós-treinamento de um milhão de tokens sob um orçamento fixo de GPU ao otimizar o uso de memória através da avaliação de prompt compartilhada sem autograd e replicação sequencial de resposta, demonstrando com sucesso capacidade de execução em modelos como Qwen3.6-27B e GLM-5.2 com comprimentos de contexto excedendo 2 milhões de tokens.

Autores originais: Changhai Zhou, Kieran Liu, Yuhua Zhou, Qian Qiao, Jun Gao, Harry Zhang, Irvine Lu, Nolan Ho, Lucian Li, Andrew Lei, Cleon Cheng, Steven Chiang, Yihang Zeng, Di Zhang, Rio Yang, Kaijie Chen, Andrew Che

Publicado 2026-07-17

📖 5 min de leitura🧠 Leitura aprofundada

CC BY 4.0

Autores originais: Changhai Zhou, Kieran Liu, Yuhua Zhou, Qian Qiao, Jun Gao, Harry Zhang, Irvine Lu, Nolan Ho, Lucian Li, Andrew Lei, Cleon Cheng, Steven Chiang, Yihang Zeng, Di Zhang, Rio Yang, Kaijie Chen, Andrew Chen, Pony Ma, Weizhong Zhang, Cheng Jin

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 robô superinteligente a resolver um mistério massivo de várias etapas. Para fazer isso, o robô precisa ler uma biblioteca inteira de pistas (o "contexto") antes mesmo de começar a adivinhar a solução. No mundo da Inteligência Artificial, isso é chamado de "treinamento de contexto longo". Normalmente, quando treinamos esses robôs, alimentamos-nos com um livro enorme de pistas e depois pedimos que escrevam uma resposta curta. O problema é que o cérebro do robô (sua memória de computador) fica tão cheio de detalhes do livro que ele fica sem espaço antes mesmo de conseguir terminar de escrever a resposta. É como tentar resolver um mistério enquanto segura a biblioteca inteira em suas mãos; eventualmente, seus braços ficam cansados demais para se mover.

Para resolver isso, cientistas têm construído bibliotecas cada vez maiores de chips de computador (GPUs) para conter toda a informação. Mas e se você não tiver um supercomputador? E se você tiver apenas um número fixo e pequeno de chips, como um conjunto padrão de oito ou trinta e dois? Este é o enigma que o artigo "LongStraw" aborda. Ele pergunta: Podemos ensinar um robô a ler uma história de dois milhões de palavras e escrever uma solução usando apenas um conjunto pequeno e fixo de chips de computador, sem comprar mais? A resposta não é "sim, nós tornamos os chips maiores", mas sim "sim, nós mudamos a forma como seguramos o livro".

Os pesquisadores construíram um sistema inteligente chamado LongStraw. Pense nisso como um bibliotecário que é incrivelmente organizado. Em vez de tentar manter toda a história de dois milhões de palavras aberta sobre a mesa enquanto o robô escreve sua resposta (o que derrubaria a mesa), o bibliotecário lê a história inteira uma vez, memoriza apenas a "vibe" ou o "estado" essencial da história e, então, guarda o livro pesado. Agora, o robô só precisa escrever sua resposta curta olhando para essa pequena nota de "vibe". Se o robô precisar verificar um detalhe específico, o bibliotecário puxa rapidamente apenas aquela página, mostra ao robô e, imediatamente, a coloca de volta. Dessa forma, a mesa nunca fica desordenada, mesmo que a história seja enorme.

O artigo testa essa ideia em dois tipos diferentes de cérebros de robô: um chamado Qwen e outro chamado GLM. Eles tentaram fazer esses robôs lerem um prompt (a história) que tinha 2.097.152 tokens de comprimento (o que é aproximadamente dois milhões de palavras ou caracteres) e depois escreverem uma resposta. Eles fizeram isso com um orçamento rigoroso: oito GPUs H20 para o robô Qwen e 32 GPUs H20 para o robô GLM.

Aqui está o que eles descobriram:

Funciona para a parte da "leitura": Eles conseguiram processar com sucesso toda a história de dois milhões de tokens sem que a memória do computador travasse. Eles provaram que é possível encaixar essa tarefa massiva em um número fixo de chips sendo inteligente sobre quando segurar a informação e quando deixá-la ir.
Economiza espaço, mas leva tempo: Ao guardar o livro pesado e apenas puxar as páginas quando necessário, eles economizaram uma enorme quantidade de memória. Por exemplo, quando aumentaram o número de "palpites" do robô (o tamanho do grupo) de 2 para 8, o uso de memória quase não mudou (apenas cerca de 0,21 GB a mais), mas levou mais tempo para terminar o trabalho porque eles tiveram que puxar as páginas uma por uma em vez de todas de uma vez.
É uma "prova de conceito", não um término perfeito: O artigo é muito honesto sobre o que ele não fez. Embora o sistema tenha conseguido executar a simulação e mostrado que o robô podia ler a história longa e escrever a resposta, ele não provou totalmente que o robô aprendeu a maneira perfeita de atualizar seu cérebro. Parte da matemática que conecta as diferentes "partes do cérebro" do robô (gradientes) não foi totalmente sincronizada entre todos os chips da maneira que uma sessão de treinamento perfeita exigiria. É como mostrar que um motor de carro pode funcionar com uma pequena quantidade de combustível, mas ainda não provar que ele pode percorrer uma pista de corrida inteira sem engasgar.

Os pesquisadores chamam o resultado deles de um "recibo de execução". É um recibo que diz: "Sim, nós executamos com sucesso esta tarefa massiva neste hardware específico sem travar". Eles não alegaram ter vencido uma corrida ou construído o carro mais rápido; eles apenas provaram que o carro pode dirigir nesta pista específica com esta quantidade específica de combustível.

Em resumo, o LongStraw mostra que você não precisa de um supercomputador para treinar IA em histórias massivas. Você só precisa de uma maneira mais inteligente de gerenciar sua memória, trocando um pouco de tempo extra por muita economia de espaço. Isso abre as portas para equipes menores e pesquisadores com orçamentos limitados experimentarem esses modelos gigantes de IA, provando que você pode fazer grandes coisas com um kit de ferramentas pequeno e fixo, se for inteligente o suficiente.

Resumo Técnico: LongStraw: RL de Longo Contexto Além de 2M de Tokens sob um Orçamento Fixo de GPU

1. Definição do Problema

Um hiato significativo surgiu entre os comprimentos de contexto suportados por sistemas de inferência (aproximando-se de contextos de milhões de tokens) e aqueles utilizados durante o pós-treinamento de Aprendizado por Reforço (RL), que frequentemente permanecem limitados a 256K tokens ou menos. Esse hiato é particularmente crítico para agentes de IA cujas trajetórias acumulam observações, saídas de ferramentas e decisões prévias ao longo de sequências longas.

O desafio central em estender o pós-treinamento de RL (especificamente o Group Relative Policy Optimization, ou GRPO) para contextos de milhões de tokens sob um orçamento fixo de GPU é a memória. Diferente da inferência, onde o grafo de forward pode ser descartado após o cache de estado, o treinamento deve pontuar múltiplas respostas e realizar o backpropagation através delas. Abordagens padrão exigem a retenção do prompt completo e de todos os grafos de resposta simultaneamente na memória, levando a uma computação de atenção quadrática e ao armazenamento massivo de ativações. Técnicas existentes como FlashAttention, LoRA ou Ring Attention abordam gargalos específicos, mas não permitem inerentemente que uma execução de GRPO com orçamento fixo de GPU comporte um contexto de 2M+ de tokens sem escalar para clusters maiores.

2. Metodologia: Stack de Execução LongStraw

O LongStraw é uma stack de execução consciente da arquitetura, projetada para desacoplar o prompt longo da computação da resposta curta, trocando tempo de replay adicional por redução no uso de memória da GPU. O princípio central de design é reter um tensor através da fronteira do prompt apenas quando um token de resposta posterior dependa dele.

2.1 O Cronograma de Execução

O sistema opera em quatro fases distintas para uma atualização agrupada:

Captura de Prompt (No-Grad): O prompt compartilhado ( $P$ tokens) é avaliado uma única vez com diferenciação automática desabilitada. O sistema salva apenas o estado condicional específico da arquitetura necessário para tokens futuros (ex: estados recorrentes, páginas KV, representações latentes) e libera todas as ativações transitórias (intermediários de FFN, scratch de atenção, buffers de roteamento de MoE).
Pontuação de Pre-Step: As probabilidades logarítmicas da política antiga e da política de referência são avaliadas para todos os membros do grupo usando o estado do prompt congelado.
Replay de Resposta Serial: Em vez de manter $G$ grafos de resposta simultaneamente, o sistema reaplica um ramo de resposta por vez sob autograd. Ele reutiliza o estado de prompt capturado (somente leitura), computa a perda, realiza o backpropagation e libera imediatamente o grafo da resposta.
Acúmulo de Gradiente & Otimização: Os gradientes locais de todos os $G$ membros são acumulados, seguidos por uma única chamada de otimizador por worker.

2.2 Implementações Específicas de Arquitetura

O LongStraw adapta esta estratégia para duas famílias de modelos distintas:

Qwen3.6-27B (Híbrido Recorrente/Full-Attention):
- Retenção de Estado: Mantém estados recorrentes compactos de Gated DeltaNet (GDN) e páginas de Chave/Valor (KV) particionadas por paralelismo de contexto (CP) para as 16 camadas de full-attention na GPU.
- Mecanismo: Utiliza NF4 QLoRA. As 48 camadas GDN fornecem um estado de tamanho fixo, enquanto as 16 camadas de full-attention utilizam compactação de página física. O sistema realiza um merge global de LSE/output para o forward pass, mas reconhece a falta de sincronização de gradiente de adaptador K/V.
- Hardware: Executado em 8 H20s (CP8).
GLM-5.2 (Compressed Attention/MoE):
- Retenção de Estado: Move o estado do prompt longo (páginas latentes de MLA e páginas de indexador-chave de DSA) para a memória da CPU. Apenas uma camada de decoder é posicionada na GPU por vez durante o replay da resposta.
- Mecanismo: Utiliza um layout de página de contexto "Zigzag" através de 32 ranks (CP32) e Paralelismo de Especialista (EP32). Emprega checkpointing de ativação de camada inteira para limitar o grafo da resposta. O sistema reconstrói cronogramas de IndexShare camada por camada.
- Hardware: Executado em 32 H20s (TP1/CP32/EP32).

3. Principais Contribuições

Design de Execução de Prompt Compartilhado: Um cronograma inovador que armazena o estado do prompt uma única vez (desconectado do autograd) e serializa o trabalho da resposta, reduzindo o grafo de treinamento vivo de $P + \sum R_i$ para $\max(R_i)$ .
Implementação de Dupla Arquitetura: Instanciação bem-sucedida deste design para dois modelos estruturalmente diferentes: uma stack densa híbrida recorrente/atenção (Qwen) e uma stack MoE de atenção comprimida (GLM).
Recibos de Execução de Orçamento Fixo: Evidência empírica demonstrando que caminhos de execução de GRPO podem completar para prompts que excedem 2M de tokens (e até 4.25M para Qwen) sem aumentar a contagem de GPUs, documentando explicitamente as limitações em composição de gradiente distribuído.

4. Resultados Experimentais

4.1 Qwen3.6-27B (8 H20s)

Comprimento de Contexto: Executou com sucesso um contexto de 2.097.152 tokens (2.088.960 prompt + 8.192 resposta).
Eficiência de Memória: Aumentar o tamanho do grupo de $G=2$ para $G=8$ aumentou a memória alocada de pico em apenas 0,21 GB (0,213%), confirmando que a memória escala com o maior ramo de resposta, não com o tamanho do grupo.
Capacidade Estendida: Um teste de estresse atingiu 4,25M de posições (4.448.256 prompt + 8.192 resposta) nos mesmos 8 H20s, atingindo o pico de ~83 GB por rank.
Limitação: A atualização distribuída não é totalmente coerente; gradientes de adaptador K/V não são sincronizados entre os ranks, e o gradiente do estado do prompt é desconectado.

4.2 GLM-5.2 (32 H20s)

Comprimento de Contexto: Validou o caminho de execução end-to-end para um prompt de 2.097.152 tokens através de todas as 78 camadas.
Estratégia de Memória: O estado do prompt residente na CPU (aprox. 5,8 GiB por rank) permitiu que a GPU permanecesse dentro do orçamento durante o backward pass de 78 camadas.
Limitação: O operador DSA (Dynamic Sparse Attention) no recibo aceito é local a cada shard de CP (selecionando top-2048 de 65k tokens em vez do contexto global de 2M). Além disso, o caminho ignora o finalize_model_grads do Megatron, deixando gradientes de adaptador não-especialistas replicados em CP sem redução antes do passo do otimizador.

5. Significância e Alegações

O artigo posiciona o LongStraw não como um "recorde mundial" de comprimento de contexto, mas como um envelope de viabilidade de orçamento fixo.

Tempo de Vida do Estado vs. Esparsidade: Os autores argumentam que a capacidade de orçamento fixo é determinada pelo tempo de vida do tensor e pela propriedade física, não meramente pela esparsidade (MoE) ou kernels de atenção eficientes em memória. Ao permitir que os intermediários do prompt "morram" imediatamente após a captura, o sistema evita a explosão de memória do autograd de sequência completa.
Propriedade Física: O trabalho enfatiza que o particionamento lógico é insuficiente; a alocação de memória física deve corresponder à estratégia de particionamento (ex: copiar páginas para tensores do tamanho correto) para liberar a memória de forma eficaz.
Alegações Modestas: Os autores declaram explicitamente que estes são recibos de execução, não provas de correção de treinamento. O estado do prompt é desconectado (faltando o termo $\partial z_P / \partial \theta$ ), e a composição de gradiente distribuído é incompleta para ambos os modelos. Os resultados estabelecem que o caminho de execução cabe dentro das restrições de hardware, permitindo que pesquisadores explorem RL de contexto longo sem escala massiva, mas não alegam ter resolvido a paridade de gradiente distribuído ou os aspectos de melhoria de política ainda.

Em resumo, o LongStraw demonstra que, com o gerenciamento cuidadoso do tempo de vida do estado e da propriedade física, o pós-treinamento de RL de milhões de tokens é alcançável em clusters de GPU relativamente pequenos e de orçamento fixo, deslocando o gargalo da capacidade bruta de memória para o agendamento e gerenciamento de estado.