HARTS: A Prefix-Sharing Training System for Agentic RL with Hybrid-Attention Models
AI 摘要
arXiv:2608.28158v2 Announce Type: replace Abstract: Agentic reinforcement learning (RL) generates irregular rollout trees, but trajectory-by-trajectory training recomputes their shared interaction histories. Existing training-time prefix-reuse systems do not jointly address state recovery, dense execution, and distributed training for hybrid-attention models. We present HARTS (Hybrid-Attention RL over Tree Structures), to our knowledge the first system to demonstrate training acceleration from arbitrary rollout-tree prefix sharing on real hybrid-attention models. A two-level planner constructs compression-oriented microbatches, then jointly refines their membership, cross-microbatch sharing, placement, and execution order. Within each sharing window, gradients from dependent microbatches accumulate before one backward pass through the shared prefix computation. Differentiable boundary-state handoff and bounded replay preserve the per-trajectory chunkwise numerical path, while a linear-time planner attains the global minimum number of sequential linear-attention calls under our packed chunkwise-call model. The runtime uses a common compact layout across attention and dense or mixture-of-experts (MoE) layers, with support for distributed parallelism and activation recomputation. Recovering log-probabilities for each original trajectory position preserves RL objectives; for deterministic top-k routing without token dropping, occurrence counts restore MoE objective weights and expert-load statistics. On SWE-bench, HARTS delivers 4.77-10.45x the update throughput of trajectory-expanded training across two sampling schemes on 7.9B- and 124B-parameter hybrid-attention MoE models using 8-64 GPUs. On the 7.9B model, full-model logit differences are comparable to baseline self-rerun variation. The 124B model shows stable online learning trends through step 150 on SWE-bench and $\tau^3$-Bench.