AREAL-DTA：用于大语言模型高效强化学习的动态树注意力机制
AREAL-DTA: Dynamic Tree Attention for Efficient Reinforcement Learning of Large Language Models
AI总结：
针对RL后训练中rollout序列共享前缀导致计算冗余的问题，提出基于深度优先搜索的动态树注意力机制，结合负载均衡分布式批处理，实现高效前缀共享，训练吞吐量提升最高8.31倍。
摘要：
基于强化学习（RL）的大语言模型（LLM）后训练计算成本高昂，因为其生成大量rollout序列，这些序列经常共享长token前缀。现有的RL框架通常在策略训练期间独立处理这些序列，即在策略梯度计算的前向和后向传播中重复计算相同的前缀，导致计算资源和内存使用的严重低效。尽管前缀共享自然地在rollout上形成树结构，但打包的树掩码方法在RL设置中扩展性差。在本文中，我们介绍AReaL-DTA，它高效地利用了RL训练中的前缀共享。AReaL-DTA采用基于深度优先搜索（DFS）的执行策略，在前向和后向计算期间动态遍历rollout前缀树，每次只具体化一条从根到叶的路径。为了进一步提高可扩展性，AReaL-DTA结合了负载均衡的分布式批处理机制，跨多个GPU动态构建和处理前缀树。在τ²-bench上，AReaL-DTA相比密集训练提高了高达8.31倍的训练吞吐量，相比稀疏训练提高了高达1.70倍。我们的代码可在该https URL获取。
英文摘要：
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\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.
机构、总结、中文摘要均由AI分析而来，可能存在错误，请以论文正文为准。
