From 685ec9995656d491e9705196b7095511d714dcc9 Mon Sep 17 00:00:00 2001 From: Shi Dong Date: Sun, 26 Jul 2026 19:17:57 -0700 Subject: [PATCH 1/2] docs: fix factual errors in fully-async example page --- docs/examples/fully-async.md | 79 ++++++++++++++++++++++-------------- 1 file changed, 48 insertions(+), 31 deletions(-) diff --git a/docs/examples/fully-async.md b/docs/examples/fully-async.md index b706ece316c..541a333e3ee 100644 --- a/docs/examples/fully-async.md +++ b/docs/examples/fully-async.md @@ -49,7 +49,9 @@ You should see: ```text Creating new global async worker... Continuous async rollout worker started -[trainer] iter 1/3000 | drained 32 samples (queued: 18) +Starting async rollout generation for 32 groups +... +Rollout completed in 41.23s! Global worker queue size: 3 ``` ## What changes vs. the default recipe @@ -77,8 +79,7 @@ def get_global_worker(args, data_buffer): with _worker_lock: if _global_worker is None or not _global_worker.worker_thread.is_alive(): print("Creating new global async worker...") - _global_worker = AsyncRolloutWorker(args, data_buffer, - concurrency=args.sglang_server_concurrency) + _global_worker = AsyncRolloutWorker(args, data_buffer) _global_worker.start() return _global_worker ``` @@ -91,32 +92,42 @@ Key points: * **`atexit` hook.** Worker is torn down when the process exits — no orphaned generation tasks. -The worker itself keeps `--rollout-batch-size` tasks in flight using -`generate_and_rm_group`: +The worker loop (`continuous_worker_loop`, condensed) keeps up to +`--rollout-batch-size` groups in flight using `generate_and_rm_group` — or +`--async-max-concurrent-samples ÷ n_samples_per_prompt` groups when that cap is set: ```python -async def _producer(self): - while not self._stop: - if len(self._inflight) < self.target_inflight: - self._inflight.add(asyncio.create_task(self._launch_one())) - done, self._inflight = await asyncio.wait( - self._inflight, return_when=asyncio.FIRST_COMPLETED - ) - for task in done: - self._output_queue.put(task.result()) +while self.running: + # reap finished tasks + active_tasks -= {task for task in active_tasks if task.done()} + + # top up the in-flight set with fresh groups from the data buffer + while len(active_tasks) < max_concurrent_tasks and self.running: + for group in self.data_buffer.get_samples(1): + task = asyncio.create_task(generate_and_rm_group( + self.args, group, + sampling_params=self.state.sampling_params.copy(), + evaluation=False, + )) + task.add_done_callback(...) # puts the result on self.output_queue + active_tasks.add(task) + + await asyncio.sleep(1) ``` And the trainer-side entry simply drains: ```python -async def generate_rollout_fully_async(args, rollout_id, *, evaluation=False): - worker = get_global_worker(args, data_buffer) - samples = [] - while len(samples) < args.global_batch_size: - samples.append(worker._output_queue.get(timeout=600)) - return RolloutFnTrainOutput(samples=samples) +def generate_rollout_fully_async(args, rollout_id, data_buffer, evaluation=False): + if evaluation: + raise ValueError("Evaluation mode not supported in simple async rollout") + return run(generate_rollout_async(args, rollout_id, data_buffer)) ``` +`generate_rollout_async` collects completed groups from the worker's output queue +until it has `--rollout-batch-size` of them, recycles any aborted groups back to the +data buffer, and returns the batch sorted by prompt index. + ## What's happening underneath ```mermaid @@ -148,7 +159,8 @@ populated, the trainer never blocks on generation. | Knob | Effect | |---|---| -| `--rollout-batch-size` | Worker target in-flight count | +| `--rollout-batch-size` | Worker target in-flight group count | +| `--async-max-concurrent-samples` | Hard cap on in-flight samples (overrides the batch-size default) | | `--sglang-server-concurrency` | Per-engine concurrency cap | | `--num-steps-per-rollout` | Increase to consume more per drain (off-policy) | @@ -161,14 +173,16 @@ because there's nothing waiting to be consumed. ## What to watch +The example reports its state through stdout rather than dedicated metrics: + ```text -async/queue_depth stable (50-200 typical) -async/producer_throughput_qps consistent -async/consumer_drain_seconds < producer cycle time +Rollout completed in 41.23s! Global worker queue size: 3 +Warning: No progress for 30.0s. Queue size: 0, Collected: 12/32 ``` -If `consumer_drain_seconds > producer_cycle_time`, your trainer is starving the queue — -check GPU utilization. +A queue size that stays above zero after each drain means generation is keeping up +with training. The no-progress warning means the trainer is waiting on generation. +In wandb, compare `perf/rollout_time` against `perf/actor_train_time` as usual. ## Limitations @@ -184,9 +198,11 @@ check GPU utilization. ### Async on a 30 B MoE -`run_qwen3_30b_a3b_fully_async.py` shows the same pattern with `tp=4 ep=8` and -`--sglang-enable-ep-moe`. The only practical difference is increasing -`--rollout-batch-size` to 64+ to keep the larger engine pool fed. +`run_qwen3_30b_a3b_fully_async.py` shows the same pattern on a 30B MoE: +`--tensor-model-parallel-size 8` and `--expert-model-parallel-size 8` on the training +side, a single 8-GPU SGLang engine (`--rollout-num-gpus-per-engine 8`), and +`--use-tis`. It also demonstrates the two weight-sync combinations for async runs: +`--pause-generation-mode in_place` with `broadcast`, or `retract` with `p2p`. ### Async + R3 @@ -201,5 +217,6 @@ it uses `generate_and_rm_group` under the hood. ### Async + partial rollout -If you also use `--partial-rollout`, half-finished trajectories are saved to disk and -resumed — useful when the worker is killed mid-flight. +If you also use `--partial-rollout`, unfinished trajectories are recycled back to the +in-memory data buffer and resume generating in a later batch instead of being thrown +away — useful when weight updates abort in-flight generation. From bd23693c3415d0b78ee91083604441edd0f054a3 Mon Sep 17 00:00:00 2001 From: Shi Dong Date: Mon, 27 Jul 2026 09:47:39 -0700 Subject: [PATCH 2/2] docs: correct weight-sync combo count and partial-rollout resumption details --- docs/examples/fully-async.md | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/docs/examples/fully-async.md b/docs/examples/fully-async.md index 541a333e3ee..4faf80cf09e 100644 --- a/docs/examples/fully-async.md +++ b/docs/examples/fully-async.md @@ -201,8 +201,9 @@ In wandb, compare `perf/rollout_time` against `perf/actor_train_time` as usual. `run_qwen3_30b_a3b_fully_async.py` shows the same pattern on a 30B MoE: `--tensor-model-parallel-size 8` and `--expert-model-parallel-size 8` on the training side, a single 8-GPU SGLang engine (`--rollout-num-gpus-per-engine 8`), and -`--use-tis`. It also demonstrates the two weight-sync combinations for async runs: -`--pause-generation-mode in_place` with `broadcast`, or `retract` with `p2p`. +`--use-tis`. It also demonstrates the weight-sync combinations for async runs — +`--pause-generation-mode` (`in_place`/`retract`) × `--update-weight-transfer-mode` +(`broadcast`/`p2p`) — every pairing except `in_place` + `p2p`, which the script rejects. ### Async + R3 @@ -217,6 +218,8 @@ it uses `generate_and_rm_group` under the hood. ### Async + partial rollout -If you also use `--partial-rollout`, unfinished trajectories are recycled back to the -in-memory data buffer and resume generating in a later batch instead of being thrown -away — useful when weight updates abort in-flight generation. +When weight updates abort in-flight generation, this example already recycles the +aborted groups back to the in-memory data buffer — but it calls `reset_for_retry()` +first, so they regenerate from the prompt. The stock `--partial-rollout` flag, which +instead resumes from the partial response, applies to the built-in rollout function; +to get true resumption here, recycle the aborted groups without resetting them.