🤖 machine learning

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

LongStraw es una pila de ejecución consciente de la arquitectura que permite el aprendizaje por refuerzo de un millón de tokens en el post-entrenamiento bajo un presupuesto fijo de GPU mediante la optimización del uso de memoria a través de la evaluación de prompts compartidos sin autograd y la reproducción secuencial de respuestas, demostrando con éxito la capacidad de ejecución en modelos como Qwen3.6-27B y GLM-5.2 con longitudes de contexto que superan los 2 millones de tokens.

Autores originales: 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 lectura🧠 Análisis profundo

CC BY 4.0

Autores originales: 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

Artículo original bajo licencia CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Esta es una explicación generada por IA del artículo a continuación. No ha sido escrita ni avalada por los autores. Para mayor precisión técnica, consulte el artículo original. Leer descargo de responsabilidad completo

Imagina que estás intentando enseñarle a un robot superinteligente cómo resolver un misterio masivo de múltiples pasos. Para hacer esto, el robot necesita leer una biblioteca entera de pistas (el "contexto") antes de que siquiera pueda empezar a adivinar la solución. En el mundo de la Inteligencia Artificial, esto se llama "entrenamiento de contexto largo". Normalmente, cuando entrenamos a estos robots, les damos un libro enorme de pistas y luego les pedimos que escriban una respuesta corta. El problema es que el cerebro del robot (su memoria de computadora) se llena tanto con los detalles del libro que se queda sin espacio antes de que pueda siquiera terminar de escribir la respuesta. Es como intentar resolver un misterio mientras sostienes toda la biblioteca en tus manos; eventualmente, tus brazos se cansarán demasiado para moverse.

Para solucionar esto, los científicos han estado construyendo bibliotecas cada vez más grandes de chips de computadora (GPUs) para contener toda la información. Pero, ¿qué pasa si no tienes una supercomputadora? ¿Qué pasa si solo tienes un número pequeño y fijo de chips, como un conjunto estándar de ocho o treinta y dos? Este es el rompecabezas que aborda el artículo "LongStraw". Pregunta: ¿Podemos enseñar a un robot a leer una historia de dos millones de palabras y escribir una solución usando solo un conjunto pequeño y fijo de chips de computadora, sin comprar más? La respuesta no es "sí, hicimos los chips más grandes", sino más bien "sí, cambiamos la forma en que sostenemos el libro".

Los investigadores construyeron un sistema ingenioso llamado LongStraw. Piensa en esto como un bibliotecario que es increíblemente organizado. En lugar de intentar mantener todo el libro de dos millones de palabras abierto sobre el escritorio mientras el robot escribe su respuesta (lo que colapsaría el escritorio), el bibliotecario lee toda la historia una vez, memoriza solo la "vibra" o el "estado" esencial de la historia, y luego guarda el pesado libro. Ahora, el robot solo tiene que escribir su respuesta corta mirando esa pequeña nota de la "vibra". Si el robot necesita verificar un detalle específico, el bibliotecario saca rápidamente solo esa página, se la muestra al robot y luego la vuelve a guardar inmediatamente. De esta manera, el escritorio nunca se desordena, aunque la historia sea enorme.

El artículo pone a prueba esta idea en dos tipos diferentes de cerebros de robot: uno llamado Qwen y otro llamado GLM. Intentaron hacer que estos robots leyeran un prompt (la historia) que era 2,097,152 tokens de largo (lo que es aproximadamente dos millones de palabras o caracteres) y luego escribieran una respuesta. Hicieron esto con un presupuesto estricto: ocho GPUs H20 para el robot Qwen y 32 GPUs H20 para el robot GLM.

Esto es lo que encontraron:

Funciona para la parte de "lectura": Lograron procesar con éxito la historia completa de dos millones de tokens sin que la memoria de la computadora colapsara. Demostraron que puedes realizar esta tarea masiva en un número fijo de chips siendo inteligente sobre cuándo retener la información y cuándo dejarla ir.
Ahorra espacio, pero toma tiempo: Al guardar el pesado libro y solo sacar las páginas cuando es necesario, ahorraron una gran cantidad de memoria. Por ejemplo, cuando aumentaron el número de "conjeturas" del robot (el tamaño del grupo) de 2 a 8, el uso de memoria apenas cambió (solo unos 0.21 GB más), pero tomó más tiempo terminar el trabajo porque tuvieron que sacar las páginas una por una en lugar de todas a la vez.
Es un "prueba de concepto", no un final perfecto: El artículo es muy honesto sobre lo que no hizo. Aunque el sistema ejecutó con éxito la simulación y mostró que el robot podía leer la historia larga y escribir la respuesta, no demostró plenamente que el robot aprendiera la forma perfecta de actualizar su cerebro. Parte de la matemática que conecta las diferentes "partes del cerebro" del robot (gradientes) no estaba totalmente sincronizada entre todos los chips de la manera en que lo requeriría una sesión de entrenamiento perfecta. Es como mostrar que el motor de un coche puede funcionar con una pequeña cantidad de combustible, pero no demostrar aún que puede recorrer toda una pista de carreras sin detenerse.

Los investigadores llaman su resultado un "recibo de ejecución". Es un recibo que dice: "Sí, ejecutamos con éxito esta tarea masiva en este hardware específico sin colapsar". No pretendían haber ganado una carrera o haber construido el coche más rápido; solo demostraron que el coche puede conducir en esta pista específica con esta cantidad específica de gasolina.

En resumen, LongStraw muestra que no necesitas una supercomputadora para entrenar IA en historias masivas. Solo necesitas una forma más inteligente de gestionar tu memoria, intercambiando un poco de tiempo extra por mucho espacio ahorrado. Abre la puerta para que equipos más pequeños y de investigación con presupuestos limitados experimenten con estos modelos de IA gigantes, demostiendo que puedes hacer grandes cosas con un kit de herramientas pequeño y fijo si eres lo suficientemente ingenioso.

Resumen Técnico: LongStraw: RL de Contexto Largo Más Allá de los 2M de Tokens bajo un Presupuesto Fijo de GPU

1. Planteamiento del Problema

Ha surgido una brecha significativa entre las longitudes de contexto soportadas por los sistemas de inferencia (que se acercan a los contextos de millones de tokens) y aquellas utilizadas durante el Aprendizaje por Refuerzo (RL) en la etapa de post-entrenamiento, que a menudo permanecen limitadas a 256K tokens o menos. Esta brecha es particularmente crítica para los agentes de IA cuyas trayectorias acumulan observaciones, salidas de herramientas y decisiones previas a lo largo de secuencias largas.

El desafío central para extender el post-entrenamiento de RL (específicamente la Optimización de Política Relativa de Grupo, o GRPO) a contextos de millones de tokens bajo un presupuesto fijo de GPU es la memoria. A diferencia de la inferencia, donde el grafo de propagación hacia adelante puede descartarse tras el almacenamiento en caché del estado, el entrenamiento debe puntuar múltiples respuestas y retropropagar a través de ellas. Los enfoques estándar requieren retener el prompt completo y todos los grafos de respuesta simultáneamente en memoria, lo que conduce a una computación de atención cuadrática y un almacenamiento masivo de activaciones. Las técnicas existentes como FlashAttention, LoRA o Ring Attention abordan cuellos de botella específicos, pero no permiten inherentemente que una ejecución de GRPO con presupuesto de GPU fijo soporte un contexto de más de 2M de tokens sin escalar hacia clústeres más grandes.

2. Metodología: Stack de Ejecución LongStraw

LongStraw es un stack de ejecución consciente de la arquitectura diseñado para desacoplar el prompt largo de la computación de la respuesta corta, intercambiando tiempo de ejecución adicional por una reducción en el uso de memoria de la GPU. El principio de diseño central es retener un tensor a través de la frontera del prompt solo cuando un token de respuesta posterior dependa de él.

2.1 El Cronograma de Ejecución

El sistema opera en cuatro fases distintas para una actualización agrupada:

Captura de Prompt (Sin Gradiente): El prompt compartido ( $P$ tokens) se evalúa una sola vez con la diferenciación automática desactivada. El sistema guarda solo el estado condicional específico de la arquitectura requerido para futuros tokens (por ejemplo, estados recurrentes, páginas KV, representaciones latentes) y libera todas las activaciones transitorias (intermediarios FFN, scratch de atención, buffers de enrutamiento MoE).
Puntuación de Pre-Paso (Pre-Step Scoring): Las probabilidades logarítmicas de la política antigua y de la política de referencia se evalúan para todos los miembros del grupo utilizando el estado del prompt congelado.
Replay de Respuesta Serial: En lugar de mantener $G$ grafos de respuesta simultáneamente, el sistema reproduce un solo brazo de respuesta a la vez bajo autograd. Reutiliza el estado del prompt capturado de solo lectura, computa la pérdida, realiza la retropropagación y libera inmediatamente el grafo de la respuesta.
Acumulación de Gradientes y Optimización: Los gradientes locales de todos los $G$ miembros se acumulan, seguido de una única llamada al optimizador por trabajador.

2.2 Implementaciones Específicas de Arquitectura

LongStraw adapta esta estrategia a dos familias de modelos distintos:

Qwen3.6-27B (Híbrido Recurrente/Atención Completa):
- Retención de Estado: Mantiene estados recurrentes compactos de Gated DeltaNet (GDN) y páginas de Clave/Valor (KV) compartimentadas por paralelismo de contexto (CP) para las 16 capas de atención completa en la GPU.
- Mecanismo: Utiliza NF4 QLoRA. Las 48 capas GDN proporcionan un estado de tamaño fijo, mientras que las 16 capas de atención completa utilizan la compactación de páginas físicas. El sistema realiza un merge global de LSE/salida para el paso hacia adelante (forward pass), pero reconoce la falta de sincronización de gradientes del adaptador K/V.
- Hardware: Ejecutado en 8 H20s (CP8).
GLM-5.2 (Atención Comprimida/MoE):
- Retención de Estado: Traslada el estado del prompt largo (páginas latentes de MLA y páginas de índice-key de DSA) a la memoria CPU. Solo se escala un nivel de decodificador a la GPU a la vez durante el replay de la respuesta.
- Mecanismo: Utiliza un diseño de disposición de páginas de contexto "Zigzag" a través de 32 rangos (CP32) y Paralelismo de Expertos (EP32). Emplea el checkpointing de activación de capa completa para limitar el grafo de la respuesta. El sistema reconstruye los esquemas de IndexShare capa por capa.
- Hardware: Ejecutado en 32 H20s (TP1/CP32/EP32).

3. Contribuciones Clave

Diseño de Ejecución de Prompt Compartido: Un cronograma novedoso que almacena el estado del prompt una vez (desconectado de autograd) y serializa el trabajo de la respuesta, reduciendo el grafo de entrenamiento vivo de $P + \sum R_i$ a $\max(R_i)$ .
Implementación de Doble Arquitectura: Instanciación exitosa de este diseño para dos modelos estructuralmente diferentes: un stack denso híbrido recurrente/atención (Qwen) y un stack MoE de atención comprimida (GLM).
Recibos de Ejecución de Presupuesto Fijo: Evidencia empírica que demuestra que las rutas de ejecución de GRPO pueden completarse para prompts que superan los 2M de tokens (y hasta 4.25M para Qwen) sin aumentar el número de GPUs, documentando explícitamente las limitaciones en la composición de gradientes distribuidos.

4. Resultados Experimentales

4.1 Qwen3.6-27B (8 H20s)

Longitud de Contexto: Ejecutó con éxito un contexto de 2,097,152 tokens (2,088,960 de prompt + 8,192 de respuesta).
Eficiencia de Memoria: Incrementar el tamaño del grupo de $G=2$ a $G=8$ aumentó la memoria asignada pico en solo 0.21 GB (0.213%), confirmando que la memoria escala con el brazo de respuesta más grande, no con el tamaño del grupo.
Capacidad Extendida: Una prueba de estrés alcanzó las 4.25M de posiciones (4,448,256 de prompt + 8,192 de respuesta) en los mismos 8 H20s, alcanzando un pico de ~83 GB por rango.
Limitación: La actualización distribuida no es totalmente coherente; los gradientes del adaptador K/V no se sincronizan entre rangos, y el gradiente del estado del prompt está desvinculado.

4.2 GLM-5.2 (32 H20s)

Longitud de Contexto: Validó la ruta de ejecución de extremo a extremo para un prompt de 2,097,152 tokens a través de las 78 capas.
Estrategia de Memoria: El estado del prompt residente en CPU (aprox. 5.8 GiB por rango) permitió que la GPU se mantuviera dentro del presupuesto durante el backward pass de 78 capas.
Limitación: El operador DSA (Atención Dispersa Dinámica) en el recibo aceptado es local a cada partición CP (seleccionando top-2048 de 65k tokens en lugar del contexto global de 2M). Además, la ruta omite finalize_model_grads de Megatron, dejando los gradientes de los adaptadores replicados por CP sin reducir antes del paso del optimizador.

5. Significado y Reivindicaciones

El artículo posiciona a LongStraw no como un "récord mundial" de longitud de contexto, sino como un sobre de viabilidad de presupuesto fijo.

Vida Útil del Estado vs. Dispersión: Los autores argumentan que la capacidad de presupuesto fijo está determinada por la vida útil del tensor y la propiedad física, no simplemente por la dispersión (MoE) o los kernels de atención eficientes en memoria. Al permitir que los intermediarios del prompt mueran inmediatamente después de su captura, el sistema evita la explosión de memoria del autograd de secuencia completa.
Propiedad Física: El trabajo enfatiza que la partición lógica es insuficiente; la asignación de memoria física debe coincidir con la estrategia de partición (por ejemplo, copiando páginas a tensores del tamaño adecuado) para liberar la memoria de manera efectiva.
Reivindicaciones Modestas: Los autores declaran explícitamente que estos son recibos de ejecución, no pruebas de corrección de entrenamiento. El estado del prompt está desvinculado (falta el término $\partial z_P / \partial \theta$ ), y la composición de gradientes distribuidos es incompleta para ambos modelos. Los resultados establecen que la ruta de ejecución cabe dentro de las restricciones de hardware, permitiendo a los investigadores explorar el RL de contexto largo sin un escalamiento masivo, pero no afirman haber resuelto aún la paridad de gradientes distribuidos o los aspectos de mejora de la política.

En resumen, LongStraw demuestra que, con una gestión cuidadosa de la vida útil del estado y la propiedad física, el post-entrenamiento de RL de millones de tokens es alcanzable en clústeres de GPU relativamente pequeños y de presupuesto fijo, desplazando el cuello de botella de la capacidad bruta de memoria hacia la programación y la gestión de estados.