refactor(async): budget staleness in rollout batches; version in optimizer updates

examples/ar/qwen3_drpo_4b_base_dapo_sglang_async.yaml

@@ -23,17 +23,12 @@ num_rollouts: 800
 
 # Train slab share (rollout gets the rest). 0.5 -> equal split.
 train_fraction: 0.5
-# Async knobs (one unified loop): max_inflight = concurrent generations (overlap
-# depth; 1 ≈ one-step pipeline). buffer_max_staleness = weight-syncs a buffered
-# group may cross: 0 = on-policy (ratio≈1), >0 = off-policy continuous buffer.
+# Async control. max_staleness counts whole rollout batches the behavior policy
+# may trail the train policy by; admission depth is max_staleness + 1.
+# max_staleness=0 keeps the behavior/train versions aligned at batch admission;
+# max_inflight remains a resource cap.
 max_inflight: 1
-buffer_max_staleness: 0
-
-# =1 is the on-policy reward-parity baseline vs the colocate ARTrainer (a sync
-# sits between every gen/train pair -> NO overlap). Set >1 to overlap
-# generate(N+1) with train(N): data is then off-policy by <= interval weight
-# versions, absorbed by DRPO's rollout-anchored ratio (old_logp_source: rollout).
-weight_sync_interval: 1
+max_staleness: 0
 
 eval_interval: 10
 adv_normalization_scope: group

examples/ar/qwen3_grpo_4b_base_dapo_sglang_async.yaml

@@ -26,19 +26,15 @@ num_rollouts: 800
 # train-bound, lower for generation-bound workloads (mind DIVISIBILITY above).
 train_fraction: 0.5
 
-# Async knobs (one unified loop): max_inflight = concurrent generations (overlap
-# depth; 1 ≈ one-step pipeline). buffer_max_staleness = weight-syncs a buffered
-# group may cross: 0 = on-policy (ratio≈1), >0 = off-policy continuous buffer.
+# Async control. max_staleness counts whole rollout batches the behavior policy
+# may trail the train policy by; admission depth is max_staleness + 1.
+# max_staleness=0 means every batch starts training against the same train
+# version that generated it (rollout/replay numerical differences may still make
+# ratio only approximately 1) — note that with num_updates_per_batch=4 below,
+# updates 2-4 of a batch are still 1-3 updates off-policy by design, which the
+# frozen pi_old anchor and clip_range cover. max_inflight remains a resource cap.
 max_inflight: 1
-buffer_max_staleness: 0
-
-# Push the full dense weights into SGLang every N rollouts. =1 is the on-policy
-# reward-parity baseline vs the colocate ARTrainer (a sync sits between every
-# gen/train pair, so there is NO overlap). Set >1 to actually overlap
-# generate(N+1) with train(N): the data is then off-policy by <= interval
-# weight versions, absorbed by the rollout-anchored PPO ratio (old_logp =
-# segment.log_probs), exactly the slime train_async trade-off.
-weight_sync_interval: 1
+max_staleness: 0
 
 # AIME eval every 10 rollouts + a pre-train baseline (rollout -1).
 eval_interval: 10

examples/diffusion/bagel/bagel_vllmomni_async.yaml

@@ -43,29 +43,23 @@ workers_per_device: 1
 layout: separate
 train_fraction: 0.5            # 4 train GPUs + 4 rollout GPUs on an 8-GPU pool
 
-# LoRA weight-sync cadence (loop concern, read by the trainer). >1 is required
-# for async overlap: interval=1 drains every step (every rollout is a cold sync
-# boundary, no generation can overlap a train step). interval=4 gives 3 of every
-# 4 rollouts a generation overlapped with training. old_logp_source=rollout keeps
-# the emitted π_old as a valid importance-sampling anchor; the ratio may move as
-# the resident rollout policy and current train policy diverge.
-weight_sync_interval: 4
-
-# ---- async knobs (AsyncDiffusionTrainer) ----
+# ---- async control ----
 # max_inflight: concurrent generations. MUST be 1: the trajectory-segment
 # cross-slab transfer (NCCL send) runs on the rollout worker; a second in-flight
 # generation co-tenanting that worker blocks the send behind it (~150s/rollout).
 # With max_inflight=1 the trainer polls (reaps) before topping up launches, so
 # each generation transfers in the idle window before the next launch, then
-# that next generation overlaps the train step. Real overlap needs
-# weight_sync_interval>1 (interval=1 drains every step).
-# buffer_max_staleness: how many weight syncs a buffered group may cross.
-#   0 = never crosses a regular rollout-weight sync (~174s/rollout on BAGEL 4+4).
-#   2 = throughput-optimal continuous buffer (~148s/rollout, matches vllm
-#       colocate); sync-boundary cold rollouts disappear. The emitted π_old
-#       remains a valid anchor, while the measured ratio may reflect policy lag.
+# that next generation overlaps the train step. Real overlap needs admission
+# depth > 1, i.e. max_staleness >= 1.
+#
+# max_staleness counts whole rollout batches, so admission depth is
+# max_staleness + 1 = 12 regardless of num_updates_per_batch. Note this is also
+# the weight-sync period: 11 batches (22 optimizer updates) elapse between
+# publications, which is a long way past AReaL's recommended eta <= 8 — and
+# FlowGRPO's rollout-anchored ratio is their "naive PPO" arm, with no decoupled
+# objective to absorb the drift. Lower this first if the ratio starts to swing.
 max_inflight: 1
-buffer_max_staleness: 2
+max_staleness: 11
 
 logging:
   report_to_wandb: false        # flip to true to enable wandb (rank-0/driver only)
@@ -220,7 +214,7 @@ sampling:
 # LoRA weight sync → vLLM-Omni rollout. Separate slabs ⇒ RemoteLoraWeightSync:
 # rank 0 ships the freshly-trained LoRA adapter to each cross-slab rollout Worker
 # by Ray RPC (no NCCL rendezvous, no name_remap — pushes the adapter directly).
-# The train loop drives cadence via weight_sync_interval.
+# The train loop publishes only when the optimizer-update lag budget binds.
 sync:
   _target_: unirl.distributed.weight_sync.lora.RemoteLoraWeightSync
   verify: true                         # checksum read-back asserts the synced LoRA landed; catches a wrong prefix

unirl/rollout/README.md

@@ -64,14 +64,14 @@ wrong objective.
   plus a `sync:` block; *colocate* — a dedicated engine sharing GPUs with train,
   plus offload/onload and `sync:`.
 - **Driver-side async engines** (`engine/asynchronous.py`, the driver-side half next
-  to `engine/synchronous.py`'s worker-side sync contracts). Both engines expose the
-  same consumer verbs the async trainers program against: `poll` / `drain_freshest` /
-  `pop_evicted` / `quiesce` + engine-owned `weight_version`. `AsyncBatchRolloutEngine`
-  (batch granularity; non-blocking `Handle.launch_nowait` generations, stamps
-  versions at launch, used by `AsyncARTrainer`/`AsyncDiffusionTrainer`) and
-  `AsyncAgenticRolloutEngine` (trajectory granularity over the agentic rank-0
+  to `engine/synchronous.py`'s worker-side sync contracts). `AsyncBatchRolloutEngine`
+  (batch granularity; non-blocking `Handle.launch_nowait` generations, stamps the
+  synced train version at launch, and exposes completion-order FIFO train batches
+  to `AsyncARTrainer`/`AsyncDiffusionTrainer`) and `AsyncAgenticRolloutEngine`
+  (trajectory granularity over the agentic rank-0
   coordinator; normalizes the `[0]` unwraps, assembles n-sibling GRPO groups,
-  stamps versions at completion, used by the partial/async agentic trainers).
+  stamps sync-generation versions at completion, and retains the
+  `drain_freshest`/`pop_evicted` surface used by partial/async agentic trainers).
 
 **Extending it:** a new single-turn engine adds `engine/<name>/config.py` (a
 `BaseEngineConfig` whose `make_engine(**deps)` lazily imports and builds it) and
@@ -94,7 +94,9 @@ implements its weight-receive method and a matching `sync:` handler in
   engine also can't live on a `layout: separate` slab — `_build_rollout` raises.
 - **Quiesce before weight sync / eval / checkpoint on the batch async path** —
   `AsyncBatchRolloutEngine.quiesce()` drains every in-flight generation; a
-  weight + KV update corrupts one mid-flight. The agentic quiesce is a
+  weight + KV update corrupts one mid-flight. A batch weight publication also
+  requires the completed FIFO to be empty; hard-boundary admission guarantees
+  this rather than silently discarding data. The agentic quiesce is a
   turn-boundary `abort` + final poll, folded into
   `AsyncAgenticRolloutEngine.quiesce()`; its `sync_weights()` rejects a live
   drive, then pairs the weight push with the version bump and logs the sync.