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

Code

• github.com/areal-project/areal
  OfficialIn paper★ 0

Abstract

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 τ^2-bench, AReaL-DTA improves training throughput by up to 8.31 over dense training and up to 1.70 over sparse training. Our code is available at https://github.com/areal-project/AReaL/tree/feat/dta.

Reproductions