Skip to content
Closed
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
82 changes: 51 additions & 31 deletions docs/examples/fully-async.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
```
Expand All @@ -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
Expand Down Expand Up @@ -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) |

Expand All @@ -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

Expand All @@ -184,9 +198,12 @@ 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 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

Expand All @@ -201,5 +218,8 @@ 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.
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.
Loading