PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Boyuan Meng, Peihua Bao, Hong Liu, Xiaowei Zhu, Chao Wang, Gen Li, Zhenxuan Pan

arXiv:2608.28158v1cs.LGcs.DC

TL;DR

Agentic RL rollout trees contain shared histories that conventional training recomputes, and existing systems lack dense, differentiable execution for hybrid-attention models with activation recomputation. HARTS jointly plans compact tree execution and chunkwise state recovery, achieving 4.81–4.87× F/B/Grad speedup while maintaining close numerical fidelity and similar early reward trends.

  • Problem

    Agentic RL produces irregular rollout trees with shared prefixes, while existing approaches do not jointly support efficient, differentiable hybrid-attention execution compatible with activation recomputation.

  • Method

    HARTS compresses shared prefixes, jointly plans microbatches and distributed schedules, and uses minimum-call chunkwise state recovery with bounded replay and differentiable handoffs.

  • Results

    4.81–4.87× F/B/Grad speedup is achieved across multiple parallel configurations with activation recomputation on SWE-bench-derived Agentic RL rollout trees.

  • Takeaways & Limitations

    HARTS demonstrates arbitrary-rollout-tree prefix-sharing speedups on a real hybrid-attention model while preserving per-token log-probabilities and similar early reward trends.

  • Takeaways & Limitations

    Residual imbalance, bounded replay, state materialization, launch overhead, communication, synchronization, and MLA ancestor-KV expansion reduce elapsed-time savings.

Abstract

from arXiv · show

Agentic reinforcement learning (RL) often produces irregular rollout trees with shared histories. Training root-to-leaf trajectories independently recomputes these shared prefixes. Existing systems primarily target full-attention models and lack dense, differentiable hybrid-attention execution compatible with activation recomputation. We present HARTS (Hybrid-Attention RL over Tree Structures). HARTS jointly plans microbatches, data-parallel (DP) replica assignments, and microbatch-slot schedules using non-replay compact-token work after prefix compression. For chunkwise linear attention, a linear-time algorithm coordinates chunk-boundary state recovery and replay and produces the minimum number of sequential linear-attention calls under our packed execution model. HARTS preserves the chunkwise state partitioning of trajectory-wise training: it does not repeat projections, MLP/MoE computation, or final outputs, and performs only bounded state replay for numerical alignment. Per round, HARTS batches all branches into one packed call, propagates gradients through differentiable state handoffs, supports activation recomputation, and restores per-token log-probabilities. For deterministic, no-token-drop top-$k$ MoE routing, semantic multiplicities restore MoE-objective token weights and load statistics. Existing RL objectives retain their interface. To our knowledge, HARTS is the first system to demonstrate arbitrary-rollout-tree prefix-sharing speedups on a real hybrid-attention model. On an Agentic RL workload generated from SWE-bench tasks, HARTS achieves $4.81$--$4.87\times$ forward/backward/gradient speedup with activation recomputation across multiple parallel configurations. Its numerical differences are comparable to baseline self-rerun variation, and its reward trend is similar to the baseline over the first 120 steps of $τ^3$-Bench training.

1 Introduction

Agentic RL rollout trees contain shared prefixes that conventional trajectory-wise training recomputes, while existing systems do not jointly address efficient hybrid-attention execution. HARTS plans and executes shared tree computation densely, achieving substantial speedups with activation recomputation and comparable numerical behavior.

  • Motivation: Agentic RL workloads can produce irregular rollout trees whose shared prefixes make trajectory-wise expansion approximately 5.63× larger than compact execution.The workload is generated from Claude Code scaffold solving SWE-bench tasks.
  • Motivation: Hybrid-attention tree execution requires ancestor visibility for full attention plus state recovery, replay, call planning, and differentiability for chunkwise linear attention.Arbitrary forks rarely coincide with fixed chunk boundaries, creating additional execution constraints.
  • Evaluation: 4.81–4.87× F/B/Grad speedup is achieved across multiple parallel configurations with activation recomputation on a hybrid-attention MoE workload.The experiments use Ling-3.0-tiny with MLA and KDA layers on SWE-bench-derived rollout trees.
  • Evaluation: Mean token-wise logit cosine similarity exceeds 0.9997 and the first 120 steps of τ3-Bench training show reward trends similar to baseline.The numerical differences are reported as close to baseline self-rerun variation.
  • HARTS: HARTS jointly plans prefix-aware microbatches, DP-replica assignments, and microbatch-slot schedules using non-replay compact-token work.The planner balances total, critical-path, and replica work after shared-prefix compression.
  • HARTS: Its packed linear-attention execution uses bounded chunk-state replay, preserves differentiable state handoffs, and avoids repeating projections, MLP/MoE computation, and final outputs.Under the packed execution model, a linear-time algorithm selects a schedule with the minimum number of sequential linear-attention calls.

2 Background and Related Work

Agentic RL rollouts form irregular trees with shared histories, requiring compact representations that preserve semantic training targets. Hybrid-attention training additionally requires chunkwise recurrent-state handling, while prior systems and kernels leave important execution and parallelism gaps.

  • Rollout Trees in Agentic RL: A single Agentic RL task can generate an irregular rollout tree in which every root-to-leaf path is a training trajectory.Explorations may share issue descriptions, repository state, and early interactions before diverging.
  • Rollout Trees in Agentic RL: A semantic position identifies a token occurrence in a trajectory, while compact tokens represent unique shared-prefix computation executed once.Compact rows store the hidden state and later logits for these unique outputs.
  • Rollout Trees in Agentic RL: Shared physical rows do not merge training targets: semantic positions can retain different labels, advantages, masks, log-probabilities, and weights.Prefix sharing removes duplicate computation while preserving per-occurrence objectives.
  • Hybrid-Attention Training: Hybrid-attention models interleave full- and linear-attention layers, with linear attention carrying recurrent state across fixed-size chunks.The formulation permits either dense or MoE feed-forward networks.
  • Hybrid-Attention Training: When a fork lies inside a chunk, HARTS resumes from the preceding boundary and replays the shorter ancestor interval with branch tokens in one call.This preserves the state partitioning and floating-point path of conventional trajectory-wise training while passing boundary states differentiably.
  • Related Work: Prior tree-training approaches lack demonstrated real hybrid-attention experiments, dense interleaved-layer execution, or broad parallelism support.Tree Training uses node-granularity linear-attention execution, while AReaL-DTA targets full-attention Transformers and supports only data parallelism.
  • Related Work: Sequence packing, inference-time prefix caching, and optimized attention kernels do not jointly solve training backward, semantic RL targets, or branch-state recovery.These techniques optimize narrower parts of the pipeline rather than tree-structured hybrid-attention training.
  • Distributed Execution: Prefix reuse and distributed load must be planned together because slot time follows the heaviest microbatch and communication also affects elapsed time.Microbatch membership, DP assignments, and slot order are therefore coupled planning decisions.

3 Design: HARTS

HARTS plans prefix-sharing execution for arbitrary rollout trees by jointly optimizing compact-token microbatches, DP assignments, and schedules, then preserves RL and MoE semantics without expanding duplicate computation.

  • Overview: HARTS accepts complete rollout-tree trajectories and maintains mappings from semantic positions to compact rows representing unique shared-prefix computation.Shared prefixes reuse hidden states and logits while retaining trajectory-specific targets and metadata.
  • Prefix-Aware Planning: The planner evaluates microbatch candidates using total compact work, slot-critical work, and maximum cumulative replica work.These objectives prioritize prefix compression, synchronization-critical work, and per-replica balance, respectively.
  • Prefix-Aware Planning: Compact trie work equals the number of unique token nodes and supports capacity checks, candidate comparison, and contiguous partition construction.Splitting at a boundary increases work by the longest common prefix length of the adjacent trajectories.
  • Prefix-Aware Planning: The natural partition recursively splits over adjacent-LCP valleys until compact-token work fits capacity, while the planner also considers DP-compatible candidate ranges.The first capacity-feasible partition is called the natural partition P0, with K0 microbatches.
  • Prefix-Aware Planning: Candidate generation combines exact contiguous-partition search with heuristic greedy plans using multiple processing orders and placement policies.Exact search optimizes minimum work over its restricted candidate space; greedy plans add small-K alternatives and align microbatch counts to DP.
  • Semantic Preservation: HARTS restores original-order log-probabilities and semantic gradients from compact rows, reducing vocabulary-sized work from O(SV) to O(CV).Forward and backward preserve distinct labels, advantages, masks, and weights without constructing expanded S×V logits or gradients.

4 Implementation

HARTS implements prefix-sharing execution across hybrid-attention and MoE layers using a persistent compact token ordering, planned packed calls, state handoffs, and activation recomputation.

  • Planning and parallelism: The planner assigns microbatches, DP replicas, and execution slots, while the training stack executes these fixed plans without retraversing rollout trees.Planner outputs include per-round linear-attention sequences, replay intervals, and state sources.
  • Compact execution: HARTS fixes one compact physical ordering across MLA, KDA, MoE, and the language-model head, writing semantic outputs back to compact rows.Temporary ancestor-KV, replay, context, or padded views are derived within layers without changing the persistent ordering.
  • MLA execution: MLA projects and applies RoPE once per compact token, then temporarily expands only the ancestor key/value views required for branch attention.Separate packed Q and KV boundaries with bottom-right causal alignment preserve ancestor visibility and return outputs to unique compact positions.
  • KDA execution: KDA issues one packed call per execution round, reuses requested boundary states, and propagates initial-state gradients through differentiable state handoffs.When forks fall inside chunks, replay restores KDA state while avoiding repeated projection and convolution work; convolution histories are isolated across packed sequences.
  • MoE execution: Deterministic no-token-drop top-k MoE routes each compact row once, while semantic multiplicities restore objective weights, routing statistics, and expert-load counts.The compact layout avoids expanding shared prefixes according to multiplicity.
  • Recomputation and RL interface: Activation recomputation re-executes the fixed compact layout, linear-attention calls, replay, and state handoffs without retaining activations or computation graphs across microbatches.HARTS also restores current per-token log-probabilities in original token order for GRPO and compatible objectives such as PPO.

5 Evaluation

Across rollout-tree batches and parallel configurations, HARTS reduces compact work, improves scheduling balance, and delivers substantial training speedups while preserving numerical fidelity and reward trends.

  • Prefix-Aware Microbatch Planning and Schedule Construction: HARTS reduces compact tokens from 80,213,867 to 77,581,958 across ten steps, a 3.28% reduction relative to OR-Tools.Per-step compact ratios range from 1.007421× to 1.079273×.
  • Prefix-Aware Microbatch Planning and Schedule Construction: 7.236×–9.732× planning-time speedups show that HARTS plans faster than the OR-Tools baseline across the evaluated rollout-tree steps.The reported ratios are OR-Tools/HARTS planning-time ratios.
  • Prefix-Aware Microbatch Planning and Schedule Construction: HARTS improves both within-slot and current-step-total DP-replica work balance while using fewer scheduled slots.Both typical and high-percentile within-slot imbalance, as well as cumulative-step imbalance, decrease.
  • Efficiency across Parallel Configurations: 4.81–4.87× F/B/Grad speedup and 4.39–4.63× Core speedup accompany 5.62–5.63× non-replay compact-row compression across three parallel configurations.Core speedup includes HARTS planning time, while neither metric includes the optimizer update.
  • Training Fidelity and RL Effectiveness: HARTS maintains numerical fidelity with MeanCos above 0.9997 and absolute scalar-loss differences below 10^-8 across the evaluated parallel configurations.MeanCos remains close to baseline self-rerun variation, including with activation recomputation.
  • Training Fidelity and RL Effectiveness: Over 120 τ^3-Bench steps, HARTS and baseline reward curves show similar overall trends and variation.The result is reported for the complete online Agentic RL pipeline.

6 Conclusion

HARTS enables prefix-aware hybrid-attention training over arbitrary rollout trees while preserving numerical alignment and activation recomputation. Across parallel configurations, it delivers substantial speedups, with remaining gains limited by replay, imbalance, communication, and launch costs.

  • HARTS assigns prefix-aware microbatches to DP replicas and slots, while a linear-time planner coordinates chunk-boundary state reuse and replay.The system uses one compact layout across model layers and differentiable state handoffs compatible with activation recomputation.
  • 5.62–5.63× non-replay compact-row compression yields 4.81–4.87× F/B/Grad speedup and 4.39–4.63× Core speedup.Experiments use an Agentic RL workload generated from SWE-bench tasks, with full-layer activation recomputation and several parallel configurations.
  • Numerical differences remain close to baseline self-rerun variation, with mean token-wise full-vocabulary logit cosine above 0.9997.The first 120 steps of online τ^3-Bench training also show a reward trend similar to the baseline.
  • Residual imbalance, bounded KDA replay, state materialization, launch overhead, communication, synchronization, and MLA KV expansion separate compression from measured speedup.Synchronized slots run at the pace of their slowest DP replica, reducing elapsed-time savings when replicas are imbalanced.
  • HARTS speedups are expected to increase with model scale as shared-prefix computation grows faster than planning, layout, and launch overheads.The paper frames this as an expectation that larger models may approach the ideal suggested by non-replay compact-row compression.

A Linear-Attention Execution Planning Cases

The execution-planning cases illustrate how HARTS packs rollout-tree segments into a small number of chunkwise linear-attention calls. Replay labels and packed-member assignments expose how branches become ready across successive calls.

  • Figures 6–8 apply Algorithm 1 with B=64, report mandatory replay total R, and show packed members with call assignments.The cases use captions to summarize call counts and replay organization.
  • Case I: Case I uses three calls with R=373; circles denote segment token counts, solid edges direct continuations, and dashed edges independent members with replay lengths.Replay lengths are shown under chunk_size=64.
  • Case I: Case I’s first call packs five members with replay 26, while calls 2–3 contain thirteen members plus one final member and 347 replay tokens.The cited call listings provide the member chains and replay values for each packed call.
  • Case II: Case II uses two calls with R=813, unfolding the three child subtrees of shared root R2 vertically for readability.Replay labels indicate how many branches become ready together for the second variable-length call.
  • Case III: Case III uses two calls with R=780, combining the long G continuation chain with three replay-started members before parallelizing side leaves.The second call contains 22 singletons with replay 769.
Loading 2608.28158v1…