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

Reinforcement learning (RL)-based post-training for large language models (LLMs) is computationally expensive, as it generates many rollout sequences that frequently share long token prefixes. Existing RL frameworks usually process these sequences independently during policy training, i.e., repeatedly recomputing identical prefixes in both the forward and backward passes of policy gradient computation, leading to substantial inefficiencies in computation resources and memory usage. Although prefix sharing naturally induces a tree structure over rollouts, packed tree-mask approaches scale poorly in RL settings. In this paper, we introduce AReaL-DTA, which efficiently exploits 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. On $\tau^2$-bench, AReaL-DTA improves training throughput by up to $8.31\times$ over dense training and up to $1.70\times$ over sparse training. Our code is available at https://github.com/areal-project/AReaL/tree/feat/dta.

Training large language models to improve through trial and feedback is expensive because the model often explores many different trajectories with shared prefixes, such as the same prompt, instructions, or conversation history. Existing training systems usually process these trajectories separately, so they repeatedly redo the same work many times. This paper presents AReaL-DTA, a training system that notices shared prefixes across many trajectories and reuses that shared work instead of recomputing it. The system processes these related trajectories in a dfs order, keeping only the currently needed path in memory and releasing unused parts as soon as possible. It also balances the work across multiple GPUs so that large training jobs can run efficiently. On the $\tau^2$-bench benchmark, AReaL-DTA makes the model-training part up to $8.31\times$ faster than standard dense training and up to $1.70\times$ faster than a sparse tree-based baseline. For the full reinforcement learning pipeline, it improves end-to-end training speed by up to $2.28\times$, while also reducing peak GPU memory use.