Skip to content

[RFC] PPOTrainer with TransferQueue Integration #5400

Description

@wuxibin89

Author: verl team
Status: Draft
Last updated: 2026-02-25.

Summary

This RFC describes the design of the synchronous PPO trainer with TransferQueue integration implemented in verl/trainer/main_ppo_sync.py. It introduces a new PPO training pipeline that decouples control flow from data flow by using TransferQueue for tensor transfer and a ReplayBuffer for metadata-based sampling, while preserving the single-controller orchestration model of verl. The design aligns with the TransferQueue and Hybrid-Controller architecture described in verl 0.7 release blog.

Motivation

Problem

In the original RayPPOTrainer (see main_ppo.py and ray_trainer.py), the single-controller RLTrainer handles both control and data flow. All experience data (DataProto / padded tensors) are routed through the trainer when dispatching to and collecting from workers. This creates a single-point bottleneck when:

  • Moving large batches of tensors (e.g., long sequences, multi-modal data) between rollout, reward, actor, and critic.
  • Scaling to many workers or large batch sizes; the controller’s memory and bandwidth become the limit.

The issue is especially acute for multimodal training (images, video, audio) and for algorithms that move large tensors per sample (e.g., router replay). Previous use of the Ray object store for large tensors did not provide tensor-oriented optimizations (zero-copy, RDMA, fine-grained column access).

Goal

  • Decouple control from data: The trainer only dispatches instructions and metadata (e.g., batch keys, partition id, field names). Actual tensor data transfer is handled by TransferQueue via reference passing and optimized paths (zero-copy, optional RDMA).
  • Keep single-controller semantics: Orchestration, scheduling, and debugging remain in one place; we preserve the existing dispatch/collect pattern at the metadata level.
  • Support flexible rollout and batching: Enable per-prompt n sampling, multiple outputs per agent loop, and zero-padding/zero-copy representation where applicable.

Design Overview

ppo_trainer_tq.png

verl-core defines four components: Model Engine, Rollout Engine, Checkpoint Engine, and Transfer Queue. The sync PPO trainer with TransferQueue:

  • Uses the Model Engine (actor, critic, reference) unchanged; workers still receive metadata and read/write data via TransferQueue.
  • Uses the Rollout Engine in server mode (e.g., vLLM/SGLang) with an AgentLoop-based client; rollout workers put agent-loop outputs into TransferQueue instead of returning large payloads to the controller.
  • Uses TransferQueue as the data plane: controller and workers use BatchMeta-style metadata (here, KVBatchMeta: keys, partition_id, tags) and TransferQueue’s KV APIs for get/put/clear.
  • Uses the Checkpoint Engine for syncing weights from trainer to rollout replicas (e.g., colocated setup).
# Controller (PPOTrainer) – metadata only
batch = next(dataloader)
self.agent_loop_manager.generate_sequences(batch)   # dispatch prompts
batch = self.replay_buffer.sample(partition_id="train", global_steps=...)
output = self.actor_rollout_wg.compute_log_prob(batch)  # dispatch KVBatchMeta
# ... then tq.kv_batch_get / compute / tq.kv_batch_put on controller

Workers resolve batch keys via TransferQueue (tq.get(batch) / tq.put(output)), so heavy tensors never go through the controller.

High-Level Data Flow

  1. Rollout (generate sequences)

    • Trainer sends a TensorDict of prompts (with global_steps, uid, etc.) to AgentLoopManagerTQ, which forwards chunks to AgentLoopWorkerTQ (Ray) workers.
    • Each worker runs agent loops (e.g., single-turn or tool-use) and, on completion, writes outputs into TransferQueue via tq.async_kv_batch_put (keys like {uid}_{session_id}_{index}), with tags (e.g. global_steps, status, seq_len). No large tensor is returned to the trainer.
    • The trainer only records that “this global_steps batch was sent” (e.g., by adding entries to ReplayBuffer with status: "running").
  2. ReplayBuffer (sample by step)

    • A ReplayBuffer runs a background thread that periodically calls tq.kv_list() and merges per-partition metadata into an in-memory view.
    • The trainer calls replay_buffer.sample(partition_id="train", global_steps=self.global_steps) which blocks until all prompts for that step have status == "success" (or similar), then returns a KVBatchMeta (keys, partition_id, tags) — still no tensor data at the controller.
  3. Training phases (old_log_prob, ref_log_prob, values, advantage, update_actor/critic)

    • For each phase, the trainer:
      • Dispatches KVBatchMeta to the relevant WorkerGroup (e.g., actor_rollout_wg.compute_log_prob(batch)).
      • Workers use tq.kv_batch_get(keys, partition_id, fields) to pull only the needed fields, compute, then tq.kv_batch_put(...) to write results back.
      • On the controller, after the worker call returns, the trainer may call tq.kv_batch_get to read fields for advantage computation or metrics, then tq.kv_batch_put to write advantages/returns, etc.
    • So: control (which keys, which phase) stays on the trainer; data lives in TransferQueue and is accessed by both controller and workers via KV APIs.
  4. Cleanup

    • After the step (and after metrics), the trainer calls tq.kv_clear(keys, partition_id) and replay_buffer.remove(partition_id, keys) so the same keys are not reused incorrectly.

This yields a clear separation: control and metadata on the single controller; tensor storage and transfer in TransferQueue.

Key Components

1. TransferQueue (Data Plane)

  • Initialization: tq.init(config.transfer_queue) is called once in the entry path (e.g., in TaskRunner.run() before creating PPOTrainer). Workers that need TransferQueue (e.g., AgentLoopWorkerTQ) call tq.init() in their constructor (with no config or with worker-appropriate config if needed).
  • APIs used:
    • KV-style: tq.kv_list(), tq.async_kv_put, tq.async_kv_batch_put, tq.kv_batch_get, tq.kv_batch_put, tq.kv_clear.
    • Partitions: e.g. "train" and "val" for training vs validation.
  • Data format: Values are stored as TensorDict-compatible structures; the trainer and workers use TensorDict / list_of_dict_to_tensordict and to_padded_tensor() where needed. TransferQueue is optimized for PyTorch tensors (zero-copy, optional RDMA) and supports multiple storage backends (see TransferQueue doc).

2. ReplayBuffer (Metadata / Readiness)

  • Role: Bridges “which samples are ready for this step?” with the trainer’s need to get a KVBatchMeta for that step without pulling tensor data through the controller.
  • Mechanism:
    • A daemon thread periodically calls tq.kv_list() and updates an in-memory structure: partition_id -> { key -> tags }.
    • add(partition_id, items) records keys with tags (e.g. global_steps, status).
    • sample(partition_id, global_steps=...) blocks (with sleep) until for that partition and step, every key has a terminal status (e.g. "success"); then returns KVBatchMeta(partition_id, keys, tags).
    • remove(partition_id, keys) drops those keys from the in-memory view after the step is done.
  • Design choice: Sampling is by global_steps so that the sync PPO trainer still behaves as “one global step = one prompt batch, all rollouts for that batch must finish before we train.” This keeps on-policy semantics while moving data off the controller.

3. AgentLoopWorkerTQ and AgentLoopManagerTQ

  • AgentLoopWorkerTQ (Ray remote):

    • Extends AgentLoopWorker; in __init__ calls tq.init() and sets up background tasks for agent loops.
    • generate_sequences(batch): spawns per-sample agent loops (e.g., _run_prompt_run_agent_loop), each of which may produce multiple outputs (e.g., rollout.n samples or multi-turn).
    • When an agent loop finishes, _agent_loop_postprocess builds fields (e.g., prompts, responses, reward scores, masks) and calls tq.async_kv_batch_put(keys, fields, tags, partition_id). Keys are like {uid}_{session_id}_{index} so multi-output and per-prompt n are supported.
    • Status tags ("running" / "success" / "failure") and global_steps are written so ReplayBuffer and trainer can reason about readiness.
  • AgentLoopManagerTQ:

    • Holds a ReplayBuffer reference.
    • generate_sequences(prompts): marks prompts as pending in the replay buffer (add(partition_id, { uid: { global_steps, status: "running" } })), then dispatches prompt chunks to AgentLoopWorkerTQ workers via Ray. It does not wait for tensor results; workers push results into TransferQueue.
  • AgentLoopOutput

    • uid: auto generated uuid for each prompt from dataset.
    • session_id: [0, n), spawned n AgentLoop for each prompt, the n can be dynamically adjusted according to task difficulty.
    • index: [0, m), generated m AgentLoopOut from each AgentLoop, e.g context compression, sub-agent, etc.
ppo_trainer_tq.png

So: rollout data never flows back through the trainer; it goes directly into TransferQueue. The trainer only tracks metadata and later samples via ReplayBuffer.

4. PPOTrainer (Controller)

  • Initialization: Creates tokenizer, dataloaders, ReplayBuffer, and (optionally) KL controller. No TransferQueue init here (done at process level in TaskRunner.run()).
  • Worker setup: Same as existing PPO: resource pools, WorkerGroups (actor_rollout, critic, etc.), checkpoint manager, reward loop manager, and AgentLoopManagerTQ with the shared ReplayBuffer.
  • Step (training):
    1. Build batch from dataloader; assign uid and global_steps.
    2. agent_loop_manager.generate_sequences(batch) (fire-and-forget rollout).
    3. replay_buffer.sample(partition_id="train", global_steps=self.global_steps) → get KVBatchMeta.
    4. Optional colocated reward; then balance batch (e.g., seqlen across DP ranks).
    5. Old log prob: dispatch KVBatchMeta to actor_rollout_wg.compute_log_prob(batch); workers get/put via tq; trainer may tq.kv_batch_get / tq.kv_batch_put for fields needed for metrics and next phase.
    6. Ref log prob (if used): same pattern with ref policy.
    7. Values (if critic): critic_wg.infer_batch(batch); workers get/put; trainer reads values from TQ for advantage.
    8. Advantage: trainer tq.kv_batch_get required fields, computes advantages/returns, then tq.kv_batch_put nested advantages/returns (and optionally token_level_rewards, rollout_is_weights).
    9. Update critic then update actor via WorkerGroups; workers again use tq.kv_batch_get / tq.kv_batch_put for minibatch data.
    10. Metrics: trainer tq.kv_batch_get needed fields, then computes and logs metrics.
    11. Cleanup: tq.kv_clear(keys, partition_id) and replay_buffer.remove(partition_id, keys).

Validation follows the same idea: dispatch validation prompts to AgentLoopManagerTQ, sample from ReplayBuffer for partition_id="val", then fetch from TQ for logging and cleanup.

5. WorkerGroup and KVBatchMeta

  • The single-controller RayWorkerGroup and decorators already support KVBatchMeta: when a remote method returns or when results are collected, the framework treats KVBatchMeta as metadata (e.g., concatenating per-worker KVBatchMeta with KVBatchMeta.concat(output)). See verl/single_controller/base/decorator.py.
  • Workers (e.g., ActorRolloutRefWorker / TrainingWorker) that are used with this trainer receive KVBatchMeta and must:
    • Call tq.kv_batch_get(keys, partition_id, fields) to obtain the actual TensorDict for their shard.
    • Run model forward/backward.
    • Call tq.kv_batch_put(keys, partition_id, fields) to write outputs back.
  • So the contract is: “input/output are identified by (keys, partition_id); payload is in TransferQueue.” The controller never sees the full tensor batch for these phases.

Differences from main_ppo.py (RayPPOTrainer)

Aspect main_ppo.py main_ppo_sync.py
Data carrier DataProto (padded tensors through controller) KVBatchMeta + TransferQueue (metadata through controller, tensors in TQ)
Rollout result path Controller collects DataProto from workers Workers put into TransferQueue; controller only has keys/tags
Sampling One batch per step from dataloader, then generate then train ReplayBuffer samples by global_steps (and optionally partition); supports variable n per prompt
Padding / copy Controller often holds padded batches Zero-padding and zero-copy where supported by TQ backend
Agent loop Only return one output per sample Multiple outputs per prompt (e.g. {uid}_{session_id}_{index})

References

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions