HyperAIHyperAI

Command Palette

Search for a command to run...

LongStraw:固定GPU预算下超越200万Token的长上下文强化学习

摘要

推理时支持的上下文长度与强化学习后训练中实际使用的上下文长度之间的差距日益扩大。推理系统正接近百万级Token上下文,而后训练工作负载通常仍停留在256K Token或以下,并依赖部署时的长度泛化。这一差距对AI智能体尤为重要,因为其观察、工具输出、文档和先前决策会在长轨迹中不断累积。与推理不同,训练必须对以相同历史为条件的多个响应进行评分和反向传播。二次注意力计算和长生命周期的反向状态使GPU内存成为扩展训练上下文的主要瓶颈。我们提出LongStraw,一个面向架构感知的执行栈,用于在固定GPU预算下进行百万级Token的强化学习后训练,并以分组相对策略优化(GRPO)实例化。LongStraw在不启用自动微分的情况下一次性评估共享提示,仅保留后续Token所需的模型特定状态,并在自动微分下逐一重放短响应分支。这将实时训练图从完整的提示和响应序列缩减为单个响应分支,以额外的重放时间换取更低的GPU内存使用。我们在两个显著不同的模型系列上实现了LongStraw:混合循环与全注意力的Qwen3.6-27B,以及压缩注意力混合专家模型GLM-5.2。在8块H20 GPU上,LongStraw完成了Qwen在210万位置上的分组评分和响应反向传播,组大小分别为2和8。将组大小从2增加到8仅增加0.21 GB的峰值分配内存。另一项压力测试将执行范围扩展至446万位置。在32块H20 GPU上,我们验证了GLM-5.2所有78层在210万Token提示下的端到端LongStraw执行路径。这些结果表明,状态生命周期和物理所有权是强化学习后训练实际上下文限制的关键决定因素。当前实验确立了执行能力而非完整的训练正确性,因为捕获的提示状态是分离的,且部分分布式前向和梯度组合路径尚未完成。通过减少对更大GPU集群的依赖,LongStraw降低了长上下文训练的硬件门槛,使更多研究人员和小型团队能够在有限的加速器预算下探索这一方向。

一句话总结

复旦大学 MindLab 提出 LongStraw,一种架构感知的执行栈,它在不启用自动微分的情况下对共享 prompt 进行一次评估,仅保留模型特定的状态,并逐个重放 response 分支,从而在固定 GPU 预算下,支持对 Qwen3.6-27B 和 GLM-5.2 进行基于 GRPO 的强化学习后训练,上下文长度可达 446 万 token,降低了长上下文 AI agent 训练的硬件门槛。

核心贡献

  • LongStraw 是一种执行栈,通过将 prompt 状态从反向传播中分离,并在自动微分下串行重放短的 response 分支,将共享 prompt 的评估与 response 的反向传播解耦,从而将全序列训练图转化为单 response 图,以减少 GPU 内存占用。
  • 在 Qwen3.6-27B(8 块 H20 GPU)和 GLM-5.2(32 块 H20 GPU)上实现,LongStraw 完成了对 210 万 token prompt 的分组评分和反向传播,可扩展到 446 万个位置,且从分组大小 2 到 8 内存仅增加 0.21 GB,并在 GLM-5.2 的全部 78 层上进行端到端反向传播。
  • 实验表明,RL 后训练的实际上下文限制由张量状态生命周期和物理所有权决定,而 LongStraw 所展示的执行能力降低了在有限加速器预算下进行长上下文训练的硬件门槛。

引言

AI agent 越来越多地执行包含推理、动作和观察的长时间多步轨迹,这些轨迹都承载在 prompt 上下文中。使用分组相对策略优化(GRPO)对这些 agent 进行后训练,需要为共享同一长 prompt 的多个 response 打分,但 prompt 图、response 图、缓存的注意力状态以及分布式通信的内存占用总和,会迅速超出固定的 GPU 内存预算。现有的内存高效注意力和参数高效微调方法可以降低个体成本,而像 Ring Attention 这样的扩展系统通过增加更多加速器来扩展序列长度;但两者都无法在严格固定的设备数量下,单独实现 GRPO 更新的拟合。

作者提出了 LongStraw,一种共享 prompt 的执行设计,将长 prompt 与短 response 的计算解耦。prompt 在不启用自动微分的情况下评估一次,存储其架构特定的状态,然后每个 response 在自动微分下依次重放,并在优化器步骤前累积梯度。这种串行重放避免了同时保持 prompt 图和所有 response 图。该方法在两种不同的模型系列上实现:一个结合了循环层和全注意力层的密集混合 Qwen 栈,以及一个 GLM 压缩注意力/混合专家(MoE)栈,证明了在固定 GPU 预算下长上下文 GRPO 的可行性主要是一个张量生命周期和所有权问题,而非原始上下文长度。

方法

作者通过更新图而非特定的注意力核来定义系统问题。分组更新包括五个逻辑上不同的阶段:prompt 捕获、步骤前评分、优势构建、策略重放和优化器事务。

参考框架图:

在传统的全序列自动微分中,活跃图跨越 prompt 和 response token。作者转而捕获一个只读的 prompt 状态,在未改变的步骤前参数下计算旧得分和参考得分,并一次重放一个 response 图。这种串行重放将活跃的策略自动微分激活限制在一个 response 内,尽管分组索引输入和冻结分数仍会随着所提供的分组而增长。裁剪的策略项遵循 PPO,而分组相对优势和成员归一化则遵循 GRPO 目标。

两个模型在两个独立的轴上存在差异:前馈轴(密集层与 MoE)和 token 混合轴(GDN/全注意力与 MLA/DSA)。

如下图所示:

Qwen3.6-27B 结合了 48 个循环 GDN 层和 16 个全注意力层,并采用密集 FFN。GLM-5.2 结合了 21 个索引计算层和 57 个 IndexShare 层,其中包含三个密集 FFN 和 75 个 MoE FFN。这些架构选择决定了参数驻留、token 路由、激活缓冲区的形状,以及保留的 prompt 状态和 response 时的集合通信操作。

LongStraw 是一个执行栈,只有当后续的 response token 依赖某个张量时,才在 prompt 边界上保留该张量。运行时执行一个事务:在当前策略下以禁用自动微分的方式运行 prompt,保存模型特定的 prompt 状态并释放临时张量。完成的 prompt 状态被视为只读。对于每个分组成员,在当前策略下重建短的 response 路径(启用自动微分),重用只读的 prompt 状态,反向传播成员损失,并立即释放该成员的图。所有反向传播完成后,保留累积的本地梯度,并在每个 worker 上调用一次优化器。

对于 Qwen,紧凑的 GDN 和 KV 状态保留在 GPU 上。对于 GLM,上下文并行(CP)局部的 MLA 和索引器键页保留在 CPU 上。在 response 重放期间,GLM 将一个层所需的页暂存,执行短 response 层,然后释放暂存副本。对短 response 图应用全层检查点来限制激活生命周期,重新计算注意力投影、稀疏选择和 MoE 路径,而不保留完整的层图。

GLM 的最终分组运行是预算受限演进过程的终点,它保持 32 块 H20 的分配不变,并一次转移一个限制资源。

参考框架图:

该演进过程隔离了依赖类别,包括全图内存、前缀状态容量、可微分重放、跨层架构状态、设备放置和分组排序。第一阶段定位了全图瓶颈,表明长 prompt 无法保留在可微分图中。第二阶段通过在不启用自动微分的情况下捕获 MLA 隐变量和 DSA 索引器键页,扩展了前缀状态容量。第三阶段建立了基于 CPU 驻留页的单层可微分重放。第四阶段实现了全层架构闭合,修复了 IndexShare 生命周期和激活生命周期问题。第五阶段确立了并行所有权和 CPU 页。最后,全新的分组运行执行两个串行反向传播,并在每个 worker 上调用一次分布式优化器,尽管 CP 局部的 DSA 和缺少 CP 梯度最终化仍然是限制。

实验

实验评估了密集 Qwen 和 MoE GLM 模型在固定预算下执行长上下文 GRPO 重放,证实了在高达 210 万和 425 万上下文位置时,所有 worker 上的评分、成形反向传播和优化器调用均已完成。这些记录确立了执行能力和前向保真度,但未达到分布式更新一致性和完整梯度对等,因为复制适配器的关键跨 rank 梯度归约要么缺失,要么被绕过。发现表明,将 prompt 捕获与可微分重放分离,并使用分页注意力,能够在有限的 GPU 内存内适应极端上下文,但这并不构成一个连贯的分布式训练步骤,也不构成模型质量声明。

Prompt 边界状态被分为持久张量(在无梯度前缀捕获中保持不变)和捕获后丢弃的临时工作。持久组件,如 Qwen 全注意力 KV 页和 GLM MLA 隐变量页,在各 rank 上紧凑存储,并随 prompt 长度呈亚线性伸缩,而临时激活在 response 重放期间按块重新计算,以将峰值内存保持在固定预算内。这种分离通过用 CPU 卸载和重新计算换取 GPU 驻留,实现了无需增加设备的长上下文执行。Qwen 全注意力 KV 页在 8 个 rank 上以紧凑的 GPU 页存储,每个 rank 按 O(P/8) 伸缩,并在全局 LSE/输出合并期间被每个 response 查询读取。GLM 持久状态被卸载到 CPU,MLA 隐变量页和 DSA 索引器键页每个 rank 按 O(P/32) 伸缩,并在 response 重放期间逐层恢复,同时在全层检查点下重新计算临时注意力和 MoE 工作。

上下文并行和专家并行分别分布 prompt 状态和路由专家,但单独任何一种布局都不能产生同步的分布式更新。在所检查的配置中,缺少适配器或非专家权重的梯度归约,导致每个 rank 执行自己的优化器步骤,因此全局更新不一致。实验展示了在 8 个 GPU 上 425 万上下文和 32 个 GPU 上 210 万上下文的训练可行性包络,不提供收敛或学习保证。CP8 运行全局 response 注意力前向传播,但在反向传播中仅归约 dQ,dK/dV 和适配器梯度未被归约;每个 rank 调用 AdamW,因此未建立一致的全局适配器更新。EP32 在 32 个 rank 上分派路由专家,但 DSA 评分和稀疏注意力保持对每个 CP 分片局部,且自定义代码路径跳过了 finalize_model_grads,导致 CP 复制非专家适配器梯度未被归约。在两种布局中,prompt 被分离,response 目标对每个上下文分片是局部的,因此分布式优化器步骤不会产生全局一致的模型更新。Qwen 包络在 8 块 H20 GPU 上达到 4,456,448 个位置,而 GLM 包络在 32 块 H20 GPU 上达到 2,097,152 个位置;两者都是固定预算可行性包络,而非收敛训练或上下文长度记录的证据。

两份经审计的固定预算执行记录确认了在接近 200 万上下文位置时的完整训练步骤(前向传播、反向传播和优化器调用):Qwen 在 8 块 H20 上,使用 209 万上下文 prompt 和 8,192 个 response 输入;GLM 在 32 块 H20 上,使用 210 万上下文 prompt 和仅 3 个 response 输入。两者仅记录 worker 本地事件,未经验证的分布式一致性,且 GLM 的全程峰值内存未被捕获。相同的 8 块 H20 预算后来将 Qwen 扩展到完整的 425 万上下文记录,表明 200 万是一个已实现的工作点,而非容量上限。Qwen 在 8 块 H20 上的 G=8 记录在 6,785 秒内完成八个串行成员和一次本地 AdamW 调用,每个 rank 峰值内存为 97.711 GB。GLM 在 32 块 H20 上的记录在 2,975 秒内完成(加上模型创建),但缺少全程峰值内存测量;捕获窗口诊断显示每个 rank 为 112.6–145.1 GB。所有记录都使用分离的 prompt 状态和本地优化器步骤,因此分布式参数更新不保证连贯。

GLM 重跑协议固定了 MinT Runtime、Megatron-LM、Megatron-Bridge、verl 和 GLM-5.2 模型快照的确切来源标识,均声明为干净检出或带有两个已知补丁。这些标识是前瞻性的,约束未来的重跑于可审计的软件栈,但不会追溯性地验证先前的 200 万执行记录。结果进一步表明,在用于 GLM 200 万记录的相同 8 块 H20 预算上,Qwen 达到了 425 万上下文,这强调了该工作是加速器受限的可行性包络,而非上下文长度记录。声明的源代码栈包括 MinT Runtime、Megatron-LM、Megatron-Bridge、verl 和 GLM-5.2 的确切提交,带有干净检出和两个文档化的补丁。这些固定的标识是前瞻性的:它们将未来的重跑绑定到可审计的软件栈,但不会追溯地证明历史 200 万记录。GLM-5.2 模型快照的原生最大长度为 1,048,576 个位置,而 200 万 prompt 是该限制的两倍;在相同的 8 块 H20 预算上,Qwen 达到 425 万,证实了固定预算包络的框架。

GLM 最终运行跟踪清单记录了 128 个跟踪文件,分为旧阶段和策略阶段,其中旧阶段仅捕获前向层结束事件,而策略阶段在所有 rank 上捕获全部 78 层的前向和反向事件。每个策略 rank 都对全部 78 层进行了检查点,总计 4,992 个前向和 4,992 个反向事件,尽管跟踪只是执行记录,不能保证数值正确性。策略阶段在所有 rank 上发出 4,992 个前向和 4,992 个反向层结束事件,每个 rank 的 78 层覆盖率为 100%。旧阶段每个 rank 仅包含 78 个前向事件,无反向事件,不提供策略反向覆盖。跟踪作为执行级别的记录,确认了每个 worker 上的终端路径完成,但不验证梯度正确性或跨 rank 归约。

实验评估了一种节省内存的设计,该设计将 prompt 状态分为持久、紧凑存储的张量(卸载或重新计算)和临时、重新计算的激活,从而在 8–32 块 H20 GPU 上实现对高达 425 万 token 上下文的完整训练步骤。探索了上下文并行和专家并行两种布局,但缺少梯度归约意味着每个 rank 执行自己的优化器步骤,因此没有产生全局连贯的模型更新,且收敛性无法保证。经审计的执行记录确认了在 200 万和 425 万上下文大小下的完整前向、反向和优化器调用,尽管这些是固定预算可行性包络,而非有效长上下文训练的证据。GLM 重跑协议固定了源代码标识,以实现前瞻性可重现性,但不会追溯验证早期运行的准确性。


用 AI 构建 AI

从创意到上线——通过免费 AI 协同编码、开箱即用的环境和最优惠的 GPU 价格,加速您的 AI 开发。

AI 协同编码
开箱即用的 GPU
最优定价

HyperAI Newsletters

订阅我们的最新资讯
我们会在北京时间 每周一的上午九点 向您的邮箱投递本周内的最新更新
邮件发送服务由 MailChimp 提供