Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
yueming-yuan wants to merge 2 commits into
base: main
Choose a base branch
from
+22 −4
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
834d710
Add --async-max-concurrent-tasks: decouple fully-async generation con…
yueming-yuan Jul 15, 2026
ca751f8
fully_async: drop dead AsyncRolloutWorker.concurrency; assert concurr…
yueming-yuan Jul 15, 2026
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
15 changes: 11 additions & 4 deletions examples/fully_async/fully_async_rollout.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,7 @@ def get_global_worker(args, data_buffer: DataSource):
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 @@ -80,10 +80,14 @@ class AsyncRolloutWorker:
Supports continuous running, independent of rollout function lifecycle
"""

def __init__(self, args, data_buffer: DataSource, concurrency=10):
def __init__(self, args, data_buffer: DataSource):
if args.async_max_concurrent_samples is not None:
assert args.async_max_concurrent_samples <= args.sglang_server_concurrency, (
f"--async-max-concurrent-samples ({args.async_max_concurrent_samples}) must not exceed "
f"--sglang-server-concurrency ({args.sglang_server_concurrency})"
)
Comment on lines +84 to +88

@Shi-Dong Shi-Dong Jul 15, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

self.args = args
self.data_buffer = data_buffer # Directly save data_buffer reference
self.concurrency = concurrency
self.running = True
self.output_queue = queue.Queue(maxsize=1000) # Continuous output queue
self.worker_thread = None
Expand All @@ -94,7 +98,10 @@ async def continuous_worker_loop(self):
print("Continuous async rollout worker started")

active_tasks = set()
max_concurrent_tasks = self.args.rollout_batch_size
if self.args.async_max_concurrent_samples is not None:
max_concurrent_tasks = max(1, self.args.async_max_concurrent_samples // self.args.n_samples_per_prompt)
else:
max_concurrent_tasks = self.args.rollout_batch_size
group_id_counter = 0

while self.running:
Expand Down
11 changes: 11 additions & 0 deletions miles/utils/arguments.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -490,6 +490,17 @@ def add_rollout_arguments(parser):
"async mode. None (default) disables staleness filtering."
),
)
parser.add_argument(
"--async-max-concurrent-samples",
type=int,
default=None,
help=(
"Maximum number of concurrently generating trajectories in fully async mode, "
"decoupling generation concurrency from the training batch size. None (default) "
"keeps the legacy bound of one training batch worth of trajectories "
"(rollout_batch_size groups, i.e. rollout_batch_size * n_samples_per_prompt)."
Comment on lines +498 to +501

@Shi-Dong Shi-Dong Jul 15, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

),
)
parser.add_argument(
"--custom-generate-function-path",
type=str,
Expand Down
Loading