AREAL-DTA: Dynamic Tree Attention for Efficient Reinforcement Learning of Large Language Models

TL;DR

AREAL-DTA introduces a dynamic tree attention method that significantly improves the efficiency of reinforcement learning training for large language models by exploiting prefix sharing and dynamic traversal, leading to up to 8.31 times higher throughput.

Contribution

It presents a novel DFS-based execution strategy and distributed batching mechanism for scalable, efficient RL training of large language models using prefix sharing.

Findings

Achieves up to 8.31× higher training throughput.

Efficiently exploits prefix sharing in RL training.

Scales across multiple GPUs with dynamic prefix tree processing.

Abstract

Reinforcement learning (RL) based post-training for large language models (LLMs) is computationally expensive, as it generates many rollout sequences that could frequently share long token prefixes. Existing RL frameworks usually process these sequences independently, repeatedly recomputing identical prefixes during forward and backward passes during policy model training, leading to substantial inefficiencies in computation and memory usage. Although prefix sharing naturally induces a tree structure over rollouts, prior tree-attention-based solutions rely on fully materialized attention masks and scale poorly in RL settings. In this paper, we introduce AREAL-DTA to efficiently exploit prefix sharing in RL training. AREAL-DTA employs a depth-first-search (DFS)-based execution strategy that dynamically traverses the rollout prefix tree during both forward and backward computation, materializing only a single root-to-leaf path at a time. To further improve scalability, AREAL-DTA incorporates a load-balanced distributed batching mechanism that dynamically constructs and processes prefix trees across multiple GPUs. Across the popular RL post-training workload, AREAL-DTA achieves up to $8.31\times$ in $\tau^2$-bench higher training throughput.