Skip to content

feature: add experimental async policy trainer #487

Description

@VLa1111

Motivation

AReno currently provides GRPO and GSPO as stable policy-only algorithms through PolicyOnlyTrainer, whose rollout -> reward -> advantage -> train loop runs synchronously. In practice this leaves the training engine idle waiting on rollout: long-output RLVR (large max_new_tokens) makes train wait for rollout; large n_samples lets rollout dominate step time; agentic RL has unstable tool/reward latency that forces the Python runtime and GPU to wait on each other; and on single-node multi-GPU the train and rollout engines sit on different devices but a synchronous loop cannot form a local pipeline.

AReno already has most of the machinery for this — train_devices, rollout_devices, rollout_tp_size, policy_sync_bucket_mb, a separate rollout engine, and direct GPU policy sync — so the work is to reorganize the execution relationship, not rewrite the training system. Async batches must also carry a rollout policy version so the train loop can bound staleness, and group-relative advantages must keep their prompt-group boundaries, so the queue unit is a full materialized batch rather than individual completions.

This is related to, but deliberately narrower than, rewriting the stable trainer. PR #480 lands the first self-contained piece — the areno.experimental.async_policy primitives — and this issue tracks the overall feature so the series has a home.

Proposed feature

Add async-grpo and async-gspo under areno.experimental with the following scope:

  • add areno.experimental.async_policy with AsyncTrainBatch, AsyncTrainBatchQueue, and staleness helpers (started in feat: add async policy trainer primitives #480);
  • register async-grpo / async-gspo through AlgorithmSpec with experimental=True, leaving stable grpo / gspo unchanged;
  • implement AsyncPolicyTrainer reusing PolicyOnlyTrainer._run_prompt_rollout and _materialize_train_batch, so rollout logprobs, rewards, group-relative advantages, and loss inputs keep existing semantics;
  • run a background RolloutWorker that loads a prompt batch, records the rollout policy version, runs rollout -> reward -> advantage -> TrainSequence, and enqueues an AsyncTrainBatch;
  • run a TrainLoop that dequeues a batch, checks staleness, trains or drops it, increments train_policy_version, and triggers policy sync on a configured interval;
  • expose bounded, additive configuration: async_queue_size (2), async_rollout_workers (1), async_prefetch_batches (1), async_max_staleness (1), async_sync_interval_steps (1), async_queue_timeout_s (300), async_drop_stale_batches (true);
  • propagate worker failures through an error queue and stop cleanly on KeyboardInterrupt, reward exception, worker crash, and max_steps.

The intended data flow is:

prompt source
  -> RolloutWorker: rollout -> reward -> group-relative advantages -> TrainSequence
  -> bounded AsyncTrainBatchQueue
  -> TrainLoop: staleness check -> backend.train -> train_policy_version++ -> policy sync

Example usage would look like:

  areno train \
    --algo async-grpo \
    --ckpt Qwen/Qwen3-0.6B \
    --dataset-path gsm8k:main \
    --dataset-loader-fn examples/math/dataset_loader.py \
    --reward-fn-path examples/math/math_verify_reward.py \
    --train-devices 0,1 \
    --tp-size 2 \
    --rollout-devices 2,3 \
    --rollout-tp-size 1 \
    --batch-size 2 \
    --n-samples 8 \
    --max-new-tokens 512 \
    --async-queue-size 2 \
    --async-max-staleness 1

The remaining work lands as incremental PRs: experimental registration, the async trainer baseline (worker + train loop + fake-backend tests), separate-rollout-engine policy sync, agentic compatibility, and a recipe/benchmark/documentation PR.

Validation plan

The implementation would include deterministic CPU coverage for:

  • bounded-queue semantics: full/empty blocking, timeout, close, error propagation;
  • staleness filtering: a fresh batch trains, a stale batch is dropped or errored;
  • tokens / logprobs / mask / advantages alignment across materialized sequences;
  • algorithm registration: async-grpo uses GRPO loss, async-gspo uses GSPO loss, both experimental=True, stable grpo / gspo unchanged;
  • worker lifecycle and cleanup across all exit paths;
  • a full loop against fake rollout and fake train backends.

A bounded end-to-end run would additionally verify finite losses and gradients, parameter updates, queue/staleness/sync metrics, and no leaked background threads. This proposal does not claim a throughput target; the benchmark will document where async helps and where sync/queue overhead degrades it.

Questions for maintainers

  1. Is an opt-in areno.experimental trainer acceptable when the stable GRPO/GSPO path is left unchanged?
  2. Should staleness be checked on the queue get side (train loop), the put side (worker), or both?
  3. For P0, is it acceptable to validate scheduling semantics on a colocated / fake backend and defer
    separate-rollout-engine policy sync to a follow-up PR?

Alternatives considered

  • Making the default PolicyOnlyTrainer asynchronous would silently change stable GRPO/GSPO behavior and complicate debugging; keeping async opt-in under areno.experimental avoids this.
  • Queueing individual completions instead of full materialized batches would break group-relative advantage boundaries and token/logprob alignment.
  • Dropping the rollout policy version and treating every batch as fresh would ignore on-policy drift; staleness control is required.
  • Implementing a separate multi-node rollout server is out of scope for AReno's single-node, self-contained positioning.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions