Add --async-max-concurrent-samples to decouple fully-async generation concurrency from batch size

[yueming-yuan]

Problem

In fully async mode the generation concurrency is hard-wired to the training batch: AsyncRolloutWorker keeps at most rollout_batch_size groups in flight, i.e. exactly one training batch worth of trajectories. There is no way to run the generation side ahead of consumption (e.g. to keep a larger engine fleet or a larger sandbox pool saturated) without also inflating the training batch.

Change

• New --async-max-concurrent-samples (default None): maximum number of concurrently generating trajectories in fully async mode. The worker converts it to in-flight groups (value // n_samples_per_prompt). Unset keeps the legacy bound (rollout_batch_size groups), so existing runs are unaffected.
• Asserted against --sglang-server-concurrency at worker construction so a misconfigured value fails fast instead of silently throttling at the HTTP client.
• Dropped the dead AsyncRolloutWorker(concurrency=...) constructor parameter — it was stored and never read; the real bound is the one above.

Consumption is intentionally untouched: each train step still collects rollout_batch_size groups. The new knob only controls how far generation runs ahead, trading higher engine/sandbox utilization for higher average weight staleness (combine with --max-weight-staleness to bound the tail).

Validation

GLM-5.2 744B fully-async disaggregated run (8 train + 8 inference nodes, agentic tbench2 episodes): --rollout-batch-size 8 --n-samples-per-prompt 8 --async-max-concurrent-samples 128 doubles in-flight trajectories from 64 to 128 (16 groups); unset falls back to the previous behavior.

🤖 Generated with Claude Code

[gemini-code-assist]

Code Review

This pull request introduces a new command-line argument --async-max-concurrent-tasks to decouple generation concurrency from the training batch size in fully async mode, updating AsyncRolloutWorker to apply this limit. The review feedback recommends replacing the assert statement in the constructor with ValueError exceptions to properly validate that the maximum concurrent tasks are positive and do not exceed the server concurrency.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

According to the repository guidelines, constructor and function argument validation should use ValueError instead of assert. Additionally, we should validate that --async-max-concurrent-tasks is a positive integer.

Suggested change

	if args.async_max_concurrent_tasks is not None:
	assert args.async_max_concurrent_tasks <= args.sglang_server_concurrency, (
	f"--async-max-concurrent-tasks ({args.async_max_concurrent_tasks}) must not exceed "
	f"--sglang-server-concurrency ({args.sglang_server_concurrency})"
	)
	if args.async_max_concurrent_tasks is not None:
	if args.async_max_concurrent_tasks <= 0:
	raise ValueError(
	f"--async-max-concurrent-tasks ({args.async_max_concurrent_tasks}) must be positive"
	)
	if args.async_max_concurrent_tasks > args.sglang_server_concurrency:
	raise ValueError(
	f"--async-max-concurrent-tasks ({args.async_max_concurrent_tasks}) must not exceed "
	f"--sglang-server-concurrency ({args.sglang_server_concurrency})"
	)

References

1. Use ValueError instead of assert for validating function or constructor arguments (such as checking for positive or non-negative values).

[Shi-Dong]

I'd suggest replacing this assertion with a warning: i.e. if async_max_concurrent_samples is too large, then we replace it with min(async_max_concurrent_samples, sglang_server_concurrency), and log a warning. Or perhaps it's fine to remove this constraint altogether?

The reason is that we want async_max_concurrent_samples to be large (otherwise it'll just be redundant), but this assertion makes a training run prone to being tripped by an upper bound.

[Shi-Dong]

It'd be great if we could add a guideline here saying that ideally it should be set larger than global batch size.