HARTS: Efficient Agentic Reinforcement Learning for Hybrid-Attention Models over Arbitrary Rollout Trees —— 精读

论文链接：https://arxiv.org/abs/2608.28158

发表时间：2026 年 8 月 28 日（arXiv:2608.28158v1）

发表机构：蚂蚁集团（Ant Group）——纯企业团队，七位作者（Boyuan Meng、Peihua Bao、Hong Liu、Xiaowei Zhu、Chao Wang、Gen Li、Zhenxuan Pan）全部来自蚂蚁集团

领域标签：cs.LG（机器学习）/ 分布式训练系统 / Agentic RL 基础设施

备注：实验基于 AReaL Agentic RL 框架 + Megatron 训练栈 + Ling-3.0-tiny 混合注意力 MoE 模型，单机 8 卡 H20-3e GPU 完成。未见开源代码仓库。

一、论文背景

1.1 Agentic RL 为什么会产生「树」

大模型强化学习训练早已不是「一问一答」的朴素形态。当模型作为 Agent 在环境中工作时——比如挂在一个 Claude Code 风格的脚手架里解 SWE-bench 任务——它会做多轮交互、调用工具、并行探索、失败后重试。一次采样（rollout）因此往往不是一条独立的线性序列，而是一棵不规则的树：深度、分叉位置、扇出都不固定，每条从根到叶的路径是一条训练轨迹，而多条轨迹会共享任务描述、环境状态和早期交互历史。

论文给了一个量化直觉：在其 Claude Code 脚手架跑 SWE-bench 产生的负载中，把树展开成独立轨迹后，非重放 token 行数约是紧凑执行的 5.63 倍。也就是说，传统「逐轨迹独立训练」的方式，把同一段共享前缀在前向和反向里重复算了五遍多——共享前缀冗余是该负载的主要训练成本。

1.2 混合注意力模型：前缀共享的新障碍

如果只是全注意力（full attention）模型，问题相对好办：树结构主要决定每个 query 能看到哪些祖先的 key/value，用一个树掩码或等价的 KV 视图就能表达。这也是现有工作（如 Tree Training）的路线。

但 2025 年以来，主流大模型越来越多采用混合注意力（hybrid attention）架构：把全注意力层与线性注意力层（RetNet、GLA、DeltaNet、Gated DeltaNet、Kimi Delta Attention/KDA 等）交错排列，兼顾长上下文效率与表达能力。线性注意力沿着 token 顺序携带循环状态，其高性能训练内核采用分块并行（chunkwise）形式：在固定长度 B 的块内并行计算，块间传递循环状态。

这带来一个根本冲突：树的分叉位置几乎不可能恰好落在固定块边界上。一个侧分支要复用父路径的状态，必须从最近的块边界状态恢复，并在内核里「重放」不足一块的祖先状态更新，才能保持与逐轨迹训练一致的浮点路径。在打包式分块执行接口下，一次调用内部产生的边界状态，必须等调用返回后才能初始化另一个独立序列——这又引出调用规划问题。于是前缀共享在线性注意力上同时是状态恢复、调用规划、可微性和计算密度四个问题的组合。

1.3 现有系统的缺口

论文对两个最接近的系统做了尖锐剖析：

• Tree Training（arXiv:2511.00413）：其全注意力路径用 DFS 序列化加树掩码可以做到较稠密的计算，但其 GDN/线性注意力路径按树节点粒度路由状态、执行计算——短而不规则的节点段会碎裂成大量小串行调用，且节点边界会改变标准分块训练的数值路径。其论文只在全注意力 Qwen3-30B/32B 上报告实验，没有真实混合注意力训练实验。更结构性的问题是，其严格零冗余划分要求跨微批保留父计算图和边界状态，这与广泛使用的激活重计算（activation recomputation，为省显存丢弃前向激活、反向时重算）直接冲突。
• AReaL-DTA（arXiv:2602.00482）：面向标准全注意力 Transformer，按树节点粒度动态执行，只支持数据并行，不提供稠密执行，也未评估混合注意力模型。

简言之：高效执行真实混合注意力模型 + 多维并行 + 激活重计算 + 任意 rollout 树，这四件事没有系统同时做到。HARTS 补的就是这个缺口。

二、论文定位和关联工作

2.1 研究谱系梳理

（1）rollout 树与前缀冗余方向。 Tree Training 是该方向的代表：用 OR-Tools 做容量约束划分，在全注意力 Qwen3 上展示了树打包的加速潜力；AReaL-DTA 用动态树注意力降低内存占用。二者是 HARTS 最直接的对照物，HARTS 对它们的批评集中在混合注意力支持、并行策略兼容性和激活重计算兼容性三点上。

（2）序列打包与前缀复用方向。 Sequence packing 通过去 padding 提高利用率，但不复用相同前缀的模型计算；推理时的 prefix caching 复用 KV 状态，但不涉及训练反向传播和逐语义位置的 RL 目标。HARTS 做的是训练期的前缀计算复用，需要额外解决可微性。

（3）高性能注意力内核方向。 FlashAttention 与各类分块线性注意力内核优化单次算子调用，但不决定任意树上的分支状态恢复、可微状态转移和调用间依赖。HARTS 的定位是规划层：在这些内核之上规划树的共享执行，并把线性注意力训练接口扩展为可选返回块边界状态、反向传播梯度穿过这些状态。

（4）Agentic RL 系统方向。 AReaL（NeurIPS 2025）是本文的集成基座——异步大规模 RL 系统；GRPO/PPO 提供策略目标。HARTS 的贡献在于让这些既有的 RL 目标和训练栈无感知地吃到前缀共享收益。

2.2 定位对比表

定位结论：HARTS 是首个在真实混合注意力模型上证明任意 rollout 树前缀共享加速的系统，把前缀共享从一个「全注意力上的打包技巧」升级为贯穿批构建、双注意力执行、MoE 计算和目标接口的系统设计。

三、问题定义

3.1 从具体场景到抽象结构

具体问题：如何在不改变 RL 目标语义、不破坏数值一致性、不牺牲多维并行和激活重计算的前提下，让混合注意力模型对树状 rollout 只算一次共享前缀？

核心洞察：论文发现前缀共享的本质是一个**「去重坐标系」问题**。轨迹空间的重复（S 个语义位置）与计算空间的唯一性（C 个紧凑 token）可以通过一个映射 m: S→C 分离——只要把所有计算都搬到紧凑坐标系，把所有语义差异（不同标签、优势、掩码、损失权重）都留在标量目标层面，就能既共享计算又保留逐位置训练目标。

类比：这类似于编译器中的公共子表达式消除（CSE）——不改变程序语义，只消除重复计算；但这里的难点在于注意力状态是有序的、跨块传递的，消除重复计算需要重排执行顺序并保证浮点路径一致。

3.2 形式化定义

• 语义位置 s=(x,i)：轨迹 x 中第 i 个 token 出现，全体集合 S，|S|=S。
• 紧凑 token：承载唯一输出的 token，全体集合 C，|C|=C。紧凑 token 工作量只计唯一的非重放行，不含有界的线性注意力重放、MLA 临时祖先 KV 视图和后端 padding。
• 给定：任意 rollout 树展开出的完整轨迹集合 X、DP 度 D、微批容量 T、混合注意力模型（含分块线性注意力，块长 B）。
• 求：一个训练执行方案（微批划分 + DP 副本指派 + 槽位调度 + 线性注意力调用序列 + 状态交接方式 + 目标恢复方式），最小化端到端训练时间。
• 约束：① 每个语义位置的当前 log 概率、RL 元数据、MoE 损失权重与逐轨迹训练语义一致；② 分块状态划分与浮点路径与常规训练一致；③ 支持激活重计算（前向图不跨微批驻留）；④ 兼容 DP/TP/SP/PP/EP。

3.3 抽象的精妙之处

把「去重」与「语义」正交分解后，词表维度和专家维度的重计算（O(SV) 与按 S 展开的专家计算）被压到 O(CV+S)，而 RL 算法侧完全无感——GRPO 的群组相对优势、旧/参考 log 概率、掩码和归约规则一个字都不用改。这个接口的干净是全文能够成立的关键支点。

四、问题解法

HARTS 的解法按数据流分四个阶段，对应论文 3.2–3.5 节。

4.1 前缀感知的微批规划与调度（规划层一）

类比：装箱问题（bin packing）+ 负载均衡的联合优化，但「物品大小」不是轨迹长度而是去重后的紧凑工作量。

机制：将轨迹按字典序稳定排序后合并为压缩 trie，其唯一节点数即紧凑工作量。相邻轨迹的最长公共前缀（LCP）构成一个可以 RMQ 索引 O(1) 查询的结构——在 LCP 谷底（浅分叉处）切分微批的重复代价最低。候选方案按字典序三元组排序：总紧凑工作量 → 槽位关键路径工作量 → 最大副本累积工作量。

搜索策略分档：

之后做自适应精炼（边界微调 + 稀疏修复，移动/交换少量轨迹），最后贪心完成 DP 副本指派：微批按工作量排序，每 D 个组成一个槽，槽内重批分给累积工作量少的副本。规划保证精确覆盖、容量可行、K mod D = 0。

4.2 线性注意力执行的最少调用规划（规划层二）

这是全文最「算法味」的部分。

执行模型：把 rollout 树收缩为段树（segment tree，段 = 无内部分叉的最长连续 token 路径）。在打包式分块调用接口下，一次调用内产生的块边界状态只有等调用返回后才能初始化另一个独立输入序列——因此每个分叉处最多一个孩子可以作为直接延续留在同一输入序列里，其余孩子都要开新序列。

类比：这像多级流水线的调度——上一道工序的产出（边界状态）只有下线后才能喂给另一条并行产线。

两个状态就绪情形与递推：定义 T_pre(v)（入口锚点状态在当前调用前可用时的最少轮数）和 T_in(v)（v 作为直接延续进入当前调用时的最少轮数）。叶子段二者均为 1；内点按公式 (13)(14) 递推——直接延续子树不等待，侧分支延迟一轮。整棵树的最少调用数即 T_pre(root)。算法用维护孩子 T_pre 最大值/次大值的技巧做到线性时间 O(|V|+|E|)，恢复阶段自顶向下分配轮数预算并选择直接延续孩子。

重放的定位：分叉落在块内时，侧分支从最近的块边界状态恢复并重放 [a(p), p) 区间。重放长度由块大小和分叉位置唯一决定，不是优化变量；强制总重放量 R = Σ(|ch(v)|−1)(e_v mod B) 与直接延续的选择无关。规划器优化的是串行调用深度，重放量固定。

数值保真：重放只在 KDA 内核里更新循环状态，不产生也不写回语义输出；祖先的投影、因果卷积、MLP/MoE 计算与最终输出都只物化一次。这保证了与逐轨迹训练完全一致的 chunkwise 状态划分。

4.3 混合注意力层的共享执行（运行时层）

一份紧凑物理布局贯通所有层：输入序列内 token 连续、同轮序列直接拼成一个变长 KDA 调用、重放源和请求边界状态可直接寻址。

MLA（全注意力）：投影和 RoPE 每个紧凑 token 只算一次（RoPE 用原轨迹位置 ID）；分支的注意力视图采用非对称打包 Q/KV——|Q|=l（当前分支 token），|KV|=h+l（祖先+当前），用 Transformer Engine 的右下角因果对齐保证 query i 恰好看到 j≤h+i 的条目（即完整祖先前缀+分支内因果三角）。输出用索引 scatter 写回唯一紧凑位置。祖先 KV 视图只是注意力核内的临时展开，论文坦诚没有消除它。

KDA（线性注意力）：每轮恰好一次打包调用；根序列从零状态出发，其余序列从早前轮返回的边界状态出发；接口扩展为按需物化请求的块边界状态集合 {S[c]}（公式 16），且反向时初始状态梯度会自动累积回产生该边界的早期计算。因果卷积的上下文用真实祖先的前 w−1 个投影激活临时前缀构造，段标识符防止卷积窗口跨越序列边界。

MoE：路由器直接接收 C 个紧凑隐状态；确定性 top-k 不丢 token 路由下，每个紧凑 token 只做一次路由、EP 通信、专家计算与合并。

激活重计算兼容：所有边界状态在一层的前向内产生和消费，不跨微批驻留父图。反向时按固定计划重算同样的紧凑布局、轮次、重放区间和状态连接，无需重新规划。

4.4 不展开地恢复 RL 与 MoE 语义（目标接口层）

紧凑到语义的 log 概率：每个紧凑 logit 行只算一次词表级的 max 和 sum-exp，再按原始批序读取各语义位置的目标 logit 得到 λ_s。词表维工作从 O(SV) 降到 O(CV)+O(S)，从不构造展开的 S×V logits。反向时按链式法则把所有映射到同一紧凑行的语义位置梯度累加回去（公式 19）——因为是独立形成的梯度求和而非乘以复用计数，不同标签/优势/掩码的贡献被完整保留。策略目标 L_policy=F(λ, ξ) 的接口不变，GRPO/PPO 无感。

语义多重性恢复 MoE 损失权重：定义 q_c = 映射到紧凑 token c 的语义位置数（Σq_c = S）。路由器 z-loss 以 q_c 为 token 权重，负载均衡辅助损失以 q_c 加权路由概率质量与 top-k 指派计数，专家偏置更新用同样的加权指派计数。多重性不改变路由决策和专家计算，只恢复训练权重。

五、评估指标与实验证据

5.1 实验设置

• 模型：Ling-3.0-tiny——24 层混合注意力 MoE，6 组 ×（3 KDA + 1 MLA），128 专家 top-8，词表 157,184，BF16；8×H20-3e GPU。
• 负载：Claude Code 脚手架解 SWE-bench 任务产生的 Agentic RL rollout 树，微批容量 T=65,536，全部开启全层激活重计算。
• 对照：逐轨迹展开执行的基线（同配置）。

5.2 指标体系与主结果

压缩没有一比一转化为加速的原因论文交代得很清楚：KDA 分叉重放、紧凑运行时物化、边界状态的物化与梯度流量、打包调用启动开销、TP/PP/EP 通信同步、MLA 展开的祖先 KV 视图都不随 C 等比例缩小。

规划质量（vs Tree Training 的 OR-Tools 方案，十个 rollout 树步）：紧凑 token 从 80,213,867 降到 77,581,958（−3.28%），逐步压缩比 1.007–1.079×，规划时间加速 7.24–9.73×；槽内 P95 不平衡从 34.64% 降到 2.14%，当前步总量最大不平衡从 5.68% 降到 1.11%。这说明联合规划不仅省算力，还显著改善了同步关键路径。

5.3 数值保真（三层证据）

1. 重放消融（固定 checkpoint 与输入，对比四条执行路径）：无重放方案（直接传父末态、不恢复 chunk 相位）在交接深度 4–5 处的对称 KL 误差超出基线重跑噪声 32.4–36.1%，平均 |ΔNLL| 超出 17.0–18.7%；HARTS Replay 在所有深度都贴近基线自重跑变化。这直接证明「最近块边界恢复+有界重放」是数值对齐的必要设计。
2. 全模型 logit 相似度：五种并行配置下 HARTS 与基线的逐 token 全词表 logit 余弦相似度均大于 0.9997（无论开不开激活重计算），与基线自重跑的 0.99972–0.99979 几乎无差；标量损失绝对差小于 1e-8。
3. 在线 RL 有效性：τ3-Bench 在线训练前 120 步，HARTS 与基线的奖励曲线趋势与波动范围相似。

5.4 实验设计的说服力

三层数值证据形成了从算子（重放）、到模型（全 logits）、到学习动态（在线奖励）的完整证据链；而加速实验覆盖 DP/TP/SP/PP/EP 的多种组合，证明收益不是特定并行配置下的巧合。唯一可议之处是在线训练只对比了 120 步的奖励趋势而非最终收敛性能——但作为系统论文，这个程度的 RL 有效性验证已经高于同类工作。

六、效果优势的根源解释

对比对象：逐轨迹展开基线。它为何长期有效？因为它简单——每个轨迹独立执行，状态划分和浮点路径天然与推理一致，任何并行策略都能直接套用。它的根本局限在于：计算的组织单位（轨迹）与计算的实际内容（共享前缀）不匹配，5.63 倍的冗余是这个不匹配的直接代价。

HARTS 的根本性改变不是「做了前缀共享」（Tree Training 已做），而是把前缀共享从「全注意力层的打包问题」重构为「贯穿整个模型与目标接口的坐标系切换」，由此打通了四条因果链：

因果链一：工作量度量改变 → 规划质量提升。 OR-Tools 方案先按 FFD 分配数据副本、再在各副本内独立划分，无法联合优化槽内同步关键路径；HARTS 以紧凑 token 数为工作量、以（总量, 槽位关键路径, 副本累积）字典序为目标联合规划，才有了 P95 不平衡 34.6%→2.1% 的改善。同步槽以最慢副本为节拍，这个改善直接转化为实际时间收益。

因果链二：重放设计 → 数值保真。 线性注意力的 chunkwise 状态是有相位的：直接传父末态等于跳过了 chunk 边界对齐，状态内容与逐轨迹训练不再一致，误差沿分支深度累积（消融中深度 4–5 处误差超噪声 32%+）。HARTS 的「最近边界锚点 + 不足一块重放」把误差钉死在重跑噪声水平，这解释了 logit 余弦 >0.9997 与损失差 <1e-8——不是近似得巧，而是结构上保持了相同的浮点路径。

因果链三：紧凑布局贯通 → 压缩高效转化为加速。 如果只有注意力层共享而 MLP/MoE 仍按轨迹展开，词表维与专家维的 O(SV) 计算会吞掉大部分收益。HARTS 的单一物理布局让 MLA（非对称 Q/KV 视图）、KDA（按轮打包调用）、MoE（紧凑路由+多重性权重）、LM 头（紧凑 logits+映射交叉熵）全部在 C 行坐标上运行，5.62× 的行压缩才能兑现 4.81× 的实测加速。

因果链四：最少调用规划 → 计算密度。 Tree Training 的节点粒度执行把工作碎裂成大量小串行内核，GPU 利用率低；HARTS 的 T_pre/T_in 递推保证在打包分块执行模型下串行调用数最少（附录案例显示实际树 2–3 次调用即可完成），每轮打包成稠密变长调用，这是加速得以「兑现」的执行效率前提。

反事实检验：去掉重放（No-replay 消融）→ 深分支数值误差显著累积；换回 OR-Tools 规划 → 紧凑工作多 3.28% 且负载严重失衡；若展开恢复 logits（朴素方案）→ S×V 的 logits 与梯度流量会毁掉共享收益。三个反事实分别对应三项关键设计的必要性。

七、必要知识反推

假设让一个完全没有背景的人重做这项工作，他最少需要掌握什么？

7.1 领域知识层

• Agentic RL 的 rollout 形态：理解 Agent 采样为何是树、前缀共享率有多高（5.63× 这个数字本身就是动机成立的前提——需要先实测负载才有立项依据）。
• 混合注意力模型的结构：MLA 的 KV 机制、KDA 的 chunkwise 状态递推（公式 1）、因果卷积的局部上下文需求。不理解这些，就无法意识到「树掩码方案在线性注意力上失效」这个核心矛盾。
• 激活重计算的内存-计算交换：理解「丢弃前向激活、反向重算」为何与「跨微批驻留父图」结构性冲突——这是批评 Tree Training 实用性的理论支点。

7.2 方法论知识层

• 组合优化：装箱/划分问题的 DP 求解、LCP/RMQ 结构、贪心与局部搜索（合并/分裂演化、best-fit/max-shared 放置）。
• 树上的动态规划：T_pre/T_in 双状态递推 + 最大/次大值维护的线性化技巧、自顶向下预算恢复——这是算法理论工具箱的标准功。
• 自动微分机理：链式法则下的梯度累积（共享 logit 行的多语义位置梯度求和、边界状态的初始状态梯度回流）决定了「可微交接」能否实现。
• RL 训练目标的形式接口：能把 GRPO/PPO 抽象为 L_policy=F(λ, ξ)，才能发现「只要 λ 序与语义不变目标就无感」这个接口级洞察。

7.3 工程知识层

• Megatron 式多维并行：DP 槽位同步语义、TP/SP 的行划分、EP 的 all-to-all 通信同步、PP 的微批次序一致性——规划器必须在这些约束内工作。
• GPU 内核执行效率：调用启动开销、变长打包调用 vs 碎裂小调用的利用率差异（这是「最少调用」目标的现实动因）。
• 系统验证方法学：如何设计数值保真实验（对称 KL、平均 |ΔNLL|、logit 余弦、与自重跑噪声对标）——「与重跑噪声比较」是个值得学习的基线选择。

7.4 知识融合的关键节点

真正的化学反应发生在两处：一是把「前缀共享」重新表述为语义位置到紧凑行的映射 m:S→C，这让 RL 语义恢复、MoE 多重性、log 概率压缩统一成一个坐标问题；二是发现线性注意力分叉的**「重放长度固定、就绪时机可优」**的分解——把一个看似纠缠的调度问题拆成「不可变的数值约束」加「可优化的调用排序」，最少调用规划才有了干净的成立条件。前者是系统视角与算法视角的融合，后者是对浮点路径语义的深刻理解与调度理论的融合。

八、论文中可以提取的通用性灵感

灵感一：语义与计算的坐标系分离。 核心思想：把「重复的计算」搬到去重坐标系，把「不同的语义」留在标量层，两侧各自最优。 论文证据：语义位置 S 与紧凑 token C 的映射；O(SV)→O(CV+S) 的词表计算压缩；多重性 q_c 恢复 MoE 损失权重而不展开专家计算。 推广场景：① 多查询共享受检索文档的 RAG 服务（文档编码一次、逐查询恢复上下文）；② 数据库视图物化（公共子查询物化+谓词下推）；③ 编译器 CSE 与值编号；④ 视频编码中的帧间预测（共享参考帧+残差按需恢复）。

灵感二：不可变数值约束与可优化调度的正交分解。 核心思想：把问题里「由物理/数值性质决定、不可动」的部分固定下来，只对剩余自由度做优化，既保证正确性又收窄搜索空间。 论文证据：重放区间由块大小和分叉位置唯一决定（非优化变量），规划器只优化调用顺序；强制重放总量 R 与直接延续选择无关。 推广场景：① 课程编排中先修关系固定、只排时间表；② 硬件综合中时序约束固定、优化布局布线；③ 科学工作流中数据依赖固定、优化任务放置。

灵感三：最少串行深度的树执行模型。 核心思想：处理树状依赖的产出可被并发消费时，每个分叉选一条「直接延续」骨干、其余分支推迟一轮，递推可证明达到最少串行轮数。 论文证据：T_pre/T_in 递推及线性时间恢复；附录三个真实树的案例（2–3 次调用完成）。 推广场景：① 版本控制的部分重放与分支合并策略；② 依赖图上的波前并行调度；③ 增量计算框架中的构建依赖并行化。

灵感四：保真必须显式付费，不能免费搭车。 核心思想：任何改变计算顺序/复用中间结果的优化，都应配备一个「与噪声基线对标」的数值验证，必要时用有界重算换取路径一致。 论文证据：No-replay 在深分支 SymKL 超噪声 32.4–36.1% vs Replay 贴合噪声；logit 余弦与「基线自重跑」而非理想值 1.0 比较。 推广场景：① 算子融合与重排的数值回归测试；② 混合精度训练的损失轨迹比对；③ 分布式训练的确定性复现工程。

灵感五：规模放大的摊销逻辑。 核心思想：当「单位收益」随规模增长而「固定开销」增长更慢时，优化收益会随规模自发扩大，这类系统设计应优先投向大模型场景。 论文证据：论文明确论证共享一个前缀 token 省下的模型计算随层数/隐维/专家计算增长，而规划、布局物化、内核启动等控制成本增长更慢。 推广场景：① 推理批调度器的复杂策略在大流量下才划算；② 分布式缓存的元数据开销摊销；③ 编译器激进优化只在长运行程序上回本。

附录：术语速查

• rollout tree / rollout 树：一次任务采样产生的不规则树，每条根到叶路径为一条训练轨迹。
• 语义位置 / 紧凑 token：前者是轨迹-序-偏移的坐标（可重复），后者是去重后的唯一计算行。
• chunkwise 状态传播：线性注意力训练内核按固定块推进循环状态的内部过程。
• 重放（replay）：侧分支从最近块边界状态恢复、重算不足一块祖先状态更新以对齐数值路径的固定操作。
• 直接延续（direct continuation）：分叉处留在当前输入序列中的那个孩子；其余孩子为侧分支。
• 语义多重性 q_c：映射到同一紧凑 token 的语义位置数，用于恢复 MoE 损失的 token 权重。
• F/B/Grad 与 Core 加速：前者不含规划延迟的完整训练路径加速，后者把规划时间计入。