Hugging Face's Lewis Tunstall says tree packing has landed in AsyncGRPO in TRL. It forwards each unique token of forked multi-turn rollouts once instead of re-forwarding shared history, and reached the same reward with 26x less trainer compute on a real agentic RL run on MiMo Envs.
The post explains why shared prefixes across rollouts waste compute and quantifies the trainer-side gain from packing them once, which is useful for anyone running long agentic RL training.
Async RL works best with large batch sizes, but then the fwd/bwd pass becomes a major bottleneck, especially with agentic training where rollouts are >100k tokens long 😢
A neat solution is to use a tree data structure for shared prefixes across trajectories. In our experiments, this speeds up the fwd/bwd by >26x 🔥
Goated systems-level improvement to TRL by Amine!
🌲 Tree packing has landed in AsyncGRPO in TRL! ⚡ 26× less trainer compute to reach the same reward on agentic RL. GRPO starts from the same prompt G times, and forked multi-turn rollouts share most of their history. Sequence packing forwards those shared tokens once per sample. Tree packing forwards each unique token once. The first version used FlexAttention, but the shipped one is a single FA3 varlen call. Each token retains its own target, advantage, and previous log-prob. Extensively tested against sequence packing on dense, sliding-window, softcap, MoE, and VLM models. When it can't pack anything, it costs ~1%. On a REAL agentic RL run on MiMo Envs, tree packing reached the same reward with 26× less trainer compute 👇
Source: Lewis Tunstall · x.comPublished
