🤖 machine learning

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

LongStraw 是一种架构感知型执行栈，它通过在不使用自动求导的情况下实现共享提示词评估以及顺序响应重放来优化内存使用，从而在固定的 GPU 预算下实现了百万级 Token 的强化学习后训练，并成功在 Qwen3.6-27B 和 GLM-5.2 等上下文长度超过 200 万 Token 的模型上展示了执行能力。

原作者： 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

发布于 2026-07-17

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： 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

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正在试图教一个超级聪明的机器人去破解一个宏大的、多步骤的谜题。为了做到这一点，机器人在开始猜测答案之前，需要阅读相当于一整个图书馆规模的线索（即“上下文”）。在人工智能的世界里，这被称为“长文本训练”（long-context training）。通常，当我们训练这些机器人时，我们会喂给它们一本包含大量线索的书，然后让它们写下一个简短的答案。问题在于，机器人的大脑（它的计算机内存）会被书中的细节填得满满当当，以至于在它还没写完答案之前，内存就耗尽了。这就像是在试图一边解谜，一边试图把整座图书馆都捧在手里；最终，你的手臂会因为太累而无法动弹。

为了解决这个问题，科学家们一直在构建越来越大的芯片库（GPU）来容纳所有信息。但如果你没有一台超级计算机怎么办？如果你只有固定数量的芯片，比如标准的 8 个或 32 个呢？这就是名为《LongStraw》的论文所要解决的谜题。它在问：我们能否教一个机器人阅读一个两百万字的篇章，并用仅有的固定数量的计算机芯片写出解决方案，而不需要购买更多芯片？答案不是“是的，我们把芯片做大了”，而是“是的，我们改变了拿书的方式”。

研究人员构建了一个精巧的系统，叫做 LongStraw。把它想象成一位组织极其严密的图书管理员。与其试图在机器人写答案时将整本两百万字的篇章都摊开在桌子上（这会导致桌面崩溃），这位图书管理员会先完整地阅读一遍故事，记住故事的核心“神韵”或“状态”，然后把沉重的书收起来。现在，机器人只需要看着那张写着微小“神韵”的便签就能写出简短的答案。如果机器人需要核对某个特定细节，图书管理员会迅速抽出那一页，展示给机器人看，然后立即将其放回原处。这样一来，即使故事再大，桌面也永远不会变得杂乱无章。

该论文在两种不同类型的机器人大脑上测试了这个想法：一个叫作 Qwen，另一个叫作 GLM。他们尝试让这些机器人阅读一个长度为 2,097,152 个 token（大约是两百万个单词或字符）的提示词（即故事），然后写出响应。他们是在严格的预算下完成的：使用 8 块 H20 GPU 来运行 Qwen 机器人，以及使用 32 块 H20 GPU 来运行 GLM 机器人。

以下是他们的发现：

在“阅读”部分是有效的： 他们成功处理了整个两百万 token 的故事，而没有导致计算机内存崩溃。他们证明了通过聪明地处理何时保留信息以及何时释放信息，你可以将这项庞大的任务放在固定数量的芯片上完成。
节省了空间，但牺牲了时间： 通过把沉重的书收起来，只在需要时才抽出页面，他们节省了大量的内存。例如，当他们将机器人“猜测”的数量（组大小）从 2 增加到 8 时，内存占用几乎没有变化（仅增加了约 0.21 GB），但完成任务所需的时间变长了，因为他们必须一页一页地抽出页面，而不是一次性全部取出。
这是一个“概念验证”，而非完美的终点： 论文非常诚实地说明了它没有做到的事情。虽然该系统成功运行了模拟，并展示了机器人能够阅读长篇故事并写出答案，但它并未完全证明机器人学会了更新大脑的完美方式。连接机器人不同“大脑部分”（梯度）的一些数学运算，并没有在所有芯片之间实现像完美训练过程那样的完全同步。这就像是展示了一台汽车引擎可以在极少的燃料下运转，但尚未证明它能在整个赛道上行驶而不熄火。

研究人员将他们的结果称为“执行收据”（execution receipt）。这张收据证明了：“是的，我们成功地在特定的硬件上运行了这个庞大的任务，且没有崩溃。”他们并没有声称自己赢得了比赛或造出了最快的车；他们只是证明了这辆车确实可以在这条特定的赛道上，用这么多的油开起来。

简而言之，LongStraw 表明你不需要超级计算机来训练处理宏大故事的 AI。你只需要一种更聪明的内存管理方式，用一点点额外的时间来换取大量的空间节省。它为那些预算有限的小型团队和研究人员打开了大门，让他们可以实验这些巨型 AI 模型，证明只要你足够聪明，使用固定的、较小的工具包也能成就伟大的事业。

技术摘要：LongStraw：在固定 GPU 预算下实现超过 2M Token 的长上下文强化学习（RL）

1. 问题陈述

在推理系统支持的上下文长度（接近百万级 token）与强化学习（RL）后训练所使用的上下文长度之间，出现了一个显著的差距。后者通常仍被限制在 256K token 或更低。对于那些需要通过长序列累积观察、工具输出和先前决策的 AI Agent 而言，这一差距至关重要。

将 RL 后训练（特别是组相对策略优化，即 GRPO）扩展到百万级 token 上下的核心挑战在于固定 GPU 预算下的内存问题。与推理不同（推理中前向图可以在状态缓存后丢弃），训练必须为多个响应进行评分并对它们进行反向传播。标准方法需要同时在内存中保留完整的 Prompt 和所有的 Response 图，这会导致二次方级的注意力计算和巨大的激活值存储开需。现有的技术如 FlashAttention、LoRA 或 Ring Attention 虽然解决了特定的瓶颈，但并不能本质上允许在不扩展到更大集群的情况下，让固定的 GPU GRPO 运行能够容纳 2M+ token 的上下文。

2. 方法论：LongStraw 执行栈

LongStraw 是一个架构感知的执行栈，旨在通过权衡额外的重放时间来减少 GPU 显存使用量，从而实现长 Prompt 与短 Response 计算的解耦。其核心设计原则是：仅当后续的 response token 依赖于某个 tensor 时，才跨越 prompt 边界保留该 tensor。

2.1 执行调度

系统针对分组更新（grouped update）分为四个阶段运行：

Prompt 捕获（无梯度）： 对共享的 Prompt（ $P$ 个 token）进行一次性评估，并禁用自动微分。系统仅保存用于未来 token 的架构特定条件状态（例如：循环状态、KV page、潜变量表示），并释放所有瞬态激活值（FFN 中间变量、Attention 暂存空间、MoE 路由缓冲区）。
预步评分（Pre-Step Scoring）： 使用冻结的 Prompt 状态，为所有组内的成员评估旧策略（old-policy）和参考策略（reference-policy）的对数概率。
串行 Response 重放（Serial Response Replay）： 系统不再同时维护 $G$ 个 response 分支的图，而是通过自动微分（autograd）逐个重放 response 分支。它重用只读的捕获 Prompt 状态，计算损失，进行反向传播，并立即释放该 response 图。
梯度累积与优化： 从所有 $G$ 个成员中累积局部梯度，随后每个 worker 执行一次优化器调用。

2.2 特定架构的实现

LongStraw 将该策略适配到了两种不同的模型家族：

Qwen3.6-27B（混合循环/全注意力架构）：
- 状态保留： 在 GPU 上保留紧凑的 Gated DeltaNet (GDN) 循环状态和上下文并行 (CP) 分片的 Key/Value (KV) pages，用于 16 层全注意力层。
- 机制： 使用 NF4 QLoRA。48 层 GDN 提供固定大小的状态，而 16 层全注意力层使用物理页面压缩（physical page compaction）。系统在正向传播中执行全局 LSE/输出合并，但承认存在缺失 K/V adapter 梯度同步的问题。
- 硬件： 在 8 台 H20 GPU 上执行（CP8）。
GLM-5.2（压缩注意力/MoE 架构）：
- 状态保留： 将长 Prompt 状态（MLA 潜变量 pages 和 DSA indexer-key pages）移动到 CPU 内存。在 response 重放期间，每次仅将一个解码器层加载到 GPU。
- 机制： 在 32 个 rank 之间采用“Zigzag”上下文页面布局（CP32）和专家并行（EP32）。它利用全层激活检查点（whole-layer activation checkpointing）来限制 response 图的大小。系统逐层重建 IndexShare 调度方案。
- 硬件： 在 32 台 H20 GPU 上执行（TP1/CP32/EP32）。

3. 核心贡献

共享 Prompt 执行设计： 一种新颖的调度方案，通过存储一次 Prompt 状态（脱离 autograd）并使 response 工作串行化，将活跃的训练图从 $P + \sum R_i$ 降低到 $\max(R_i)$ 。
双架构实现： 为两种结构迥异的模型成功实例化了该设计：一种是稠密混合循环/注意力堆栈（Qwen），另一种是压缩注意力 MoE 堆栈（GLM）。
固定预算执行凭证（Execution Receipts）： 通过实证证据表明，在不增加 GPU 数量的情况下，GRPO 执行路径可以完成超过 2M token（对于 Qwen 最高达 4.25M）的 Prompt，同时明确记录了在分布式梯度组合方面的局限性。

4. 实验结果

4.1 Qwen3.6-27B (8 H20s)

上下文长度： 成功执行了 2,097,152 个 token 的上下文（2,088,960 prompt + 8,192 response）。
显存效率： 将组大小从 $G=2$ 增加到 $G=8$ 时，峰值分配显存仅增加了 0.21 GB (0.213%)，证实了显存规模随最大的 response 分支变化，而非随组大小变化。
扩展能力： 在相同的 8 台 H20 上，压力测试达到了 4.25M 个位置（4,448,256 prompt + 8,192 response），每个 rank 的峰值显存约为 83 GB。
局限性： 分布式更新并非完全一致；K/V adapter 梯度未在 rank 间同步，且 prompt 状态的梯度已被脱离（detached）。

4.2 GLM-5.2 (32 H20s)

上下文长度： 在所有 78 层中验证了 2,097,152 个 token prompt 的端到端执行路径。
显存策略： CPU 驻留的 Prompt 状态（每个 rank 约 5.8 GiB）使得 GPU 在 78 层反向传播期间仍能保持在预算内。
局限性： 被接受的凭证中的 DSA（动态稀疏注意力）算子是局部于每个 CP 分片的（即从 65k 个 token 中选择前 2048 个，而非全局 2M 上下文）。此外，该路径绕过了 Megatron 的 finalize_model_grads，导致 CP 复制的非专家（non-expert）adapter 梯度在优化器步骤前未进行规约（reduce）。

5. 重要性与主张

论文将 LongStraw 定位为一种固定预算的可行性包络线（feasibility envelope），而非“世界纪录”级别的上下文长度。

状态生命周期 vs. 稀疏性： 作者认为，固定预算的容量是由 tensor 的生命周期和物理所有权决定的，而不仅仅是靠稀疏性（MoE）或高效的注意力算子。通过允许 prompt 中间变量在捕获后立即失效，系统避免了全序列 autograd 带来的显存爆炸。
物理所有权： 该工作强调，仅有逻辑分片是不够的；物理内存分配必须与分片策略相匹配（例如，将页面复制到大小合适的 tensor 中），才能有效地释放内存。
适度的声明： 作者明确指出，这些是执行凭证（execution receipts），而非对训练正确性的证明。Prompt 状态是脱离的（缺失 $\partial z_P / \partial \theta$ 项），且分布式梯度组合尚不完整。结果证明了该执行路径能够在硬件约束内运行，从而使研究人员能够在不进行大规模扩展的情况下探索长上下文 RL，但目前尚未解决分布式梯度一致性或策略改进的问题。

总之，LongStraw 证明了通过精细管理状态生命周期和物理所有权，在固定且相对较小的 GPU 集群上实现百万级 token 的 RL 后训练是可行的。这标志着瓶颈已从原始显存容量转向了调度与状态管理。