Skip to content

Move fully-async rollout from examples into miles/rollout #1716

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

Merged
yueming-yuan merged 1 commit into from
Jul 31, 2026
Merged

Move fully-async rollout from examples into miles/rollout #1716

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
2 changes: 1 addition & 1 deletion docs/developer/architecture.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -87,7 +87,7 @@ sequenceDiagram
```

This is the sync path. Async (`train_async.py` + `--rollout-function-path
fully_async_rollout.generate_rollout_fully_async`) breaks the request from the trainer
miles.rollout.fully_async_rollout.generate_rollout_fully_async`) breaks the request from the trainer
loop and uses a continuously-running worker.

## Where common changes go
Expand Down
6 changes: 3 additions & 3 deletions docs/examples/fully-async.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,8 +31,8 @@ instead of the sum.
## Files

```text
miles/rollout/fully_async_rollout.py # AsyncRolloutWorker + entry function
examples/fully_async/
├── fully_async_rollout.py # AsyncRolloutWorker + entry function
├── run-qwen3-4b-fully_async.sh # launch script (Qwen3-4B)
└── run_qwen3_30b_a3b_fully_async.py # MoE variant
```
Expand All @@ -59,7 +59,7 @@ Just two flags:
```diff
- python3 train.py ...
+ python3 train_async.py ...
+ --rollout-function-path fully_async_rollout.generate_rollout_fully_async
+ --rollout-function-path miles.rollout.fully_async_rollout.generate_rollout_fully_async
```

Everything else — model args, optimizer, GRPO config — stays the same.
Expand All @@ -68,7 +68,7 @@ Everything else — model args, optimizer, GRPO config — stays the same.

The interesting code is small. Here's the global worker manager:

```python fully_async_rollout.py
```python miles/rollout/fully_async_rollout.py
_global_worker = None
_worker_lock = threading.Lock()

Expand Down
2 changes: 1 addition & 1 deletion docs/user-guide/fully-async.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@ function that owns the background worker:
```diff
- python3 train.py ...
+ python3 train_async.py ...
+ --rollout-function-path fully_async_rollout.generate_rollout_fully_async
+ --rollout-function-path miles.rollout.fully_async_rollout.generate_rollout_fully_async

@guapisolo guapisolo Jul 29, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

I suggest we directly add a args --fully-async to make sure all user know there is a fully async example.

@guapisolo guapisolo Jul 29, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Also we need to add a fully async CI on dapo math, will submit a PR later.

```

Everything else belongs in the same [argument groups](/user-guide/argument-groups) as a
Expand Down
2 changes: 1 addition & 1 deletion docs/user-guide/training-script-walkthrough.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -292,7 +292,7 @@ Enable it with two changes to the launch script:
```diff
- python3 train.py ...
+ python3 train_async.py ...
+ --rollout-function-path fully_async_rollout.generate_rollout_fully_async
+ --rollout-function-path miles.rollout.fully_async_rollout.generate_rollout_fully_async
```

| Mode | Per-iteration latency | Throughput | When to use |
Expand Down
5 changes: 3 additions & 2 deletions examples/fully_async/README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,9 @@

This example shows a simple way to make rollout generation **fully asynchronous**: a single global worker is created once and then keeps running in the background, continuously pulling prompts and launching generation tasks. Training only needs to fetch already finished results. This removes the per‑step wait that happens in the normal synchronous style.

The implementation lives in the core library at `miles/rollout/fully_async_rollout.py` (global async worker + `generate_rollout_fully_async` entry).

## Files
* `fully_async_rollout.py`: global async worker + `generate_rollout_fully_async` entry.
* `run-qwen3-4b-fully_async.sh`: example launch script with Qwen3‑4B.

## Prerequisite
Expand Down Expand Up @@ -37,7 +38,7 @@ To enable the fully async pattern there are only two changes compared to a norma
1. Use the async training driver: `train_async.py` (not `train.py`).
2. Set the rollout function path:
```bash
--rollout-function-path fully_async_rollout.generate_rollout_fully_async
--rollout-function-path miles.rollout.fully_async_rollout.generate_rollout_fully_async
```

Why is it still "fully" async although `train_async.py` itself schedules rollouts step‑by‑step?
Expand Down
4 changes: 2 additions & 2 deletions examples/fully_async/run-qwen3-4b-fully_async.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@ CKPT_ARGS=(
PROMPT_SET=/path/to/dapo-math-17k.jsonl

ROLLOUT_ARGS=(
--rollout-function-path fully_async_rollout.generate_rollout_fully_async
--rollout-function-path miles.rollout.fully_async_rollout.generate_rollout_fully_async
--prompt-data ${PROMPT_SET}
--input-key prompt
--label-key label
Expand Down Expand Up @@ -120,7 +120,7 @@ ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-s

RUNTIME_ENV_JSON="{
\"env_vars\": {
\"PYTHONPATH\": \"/root/Megatron-LM/:${SCRIPT_DIR}\",
\"PYTHONPATH\": \"/root/Megatron-LM/\",
\"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\",
\"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\"
}
Expand Down
7 changes: 2 additions & 5 deletions examples/fully_async/run_qwen3_30b_a3b_fully_async.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -61,7 +61,7 @@ def execute(args: ScriptArgs):
)

rollout_args = (
"--rollout-function-path fully_async_rollout.generate_rollout_fully_async "
"--rollout-function-path miles.rollout.fully_async_rollout.generate_rollout_fully_async "
f"--prompt-data {args.data_dir}/dapo-math-17k/dapo-math-17k.jsonl "
"--input-key prompt "
"--label-key label "
Expand Down Expand Up @@ -150,9 +150,6 @@ def execute(args: ScriptArgs):
f"{args.extra_args} "
)

import os

fully_async_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)))
U.execute_train(
train_args=train_args,
num_gpus_per_node=args.num_gpus_per_node,
Expand All @@ -161,7 +158,7 @@ def execute(args: ScriptArgs):
megatron_path=args.megatron_path,
extra_env_vars={
"FLASHINFER_DISABLE_VERSION_CHECK": "1",
"PYTHONPATH": f"{args.megatron_path}:{fully_async_dir}",
"PYTHONPATH": args.megatron_path,
},
)
Comment on lines 159 to 163

@gemini-code-assist gemini-code-assist Bot Jul 18, 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.

medium

Omitting the repository root (U.repo_base_dir) from PYTHONPATH can lead to a ModuleNotFoundError when attempting to import miles (such as miles.rollout.fully_async_rollout) if the script is executed from a different directory or in environments where the package is not installed globally. Adding U.repo_base_dir to PYTHONPATH ensures consistency with other launch scripts and prevents potential import issues.

Suggested change
extra_env_vars={
"FLASHINFER_DISABLE_VERSION_CHECK": "1",
"PYTHONPATH": f"{args.megatron_path}:{fully_async_dir}",
"PYTHONPATH": args.megatron_path,
},
)
extra_env_vars={
"FLASHINFER_DISABLE_VERSION_CHECK": "1",
"PYTHONPATH": f"{args.megatron_path}:{U.repo_base_dir}",
},
)


Expand Down
2 changes: 1 addition & 1 deletion examples/infra_features/random_async/random_async_rollout.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -267,7 +267,7 @@ def stop_global_worker() -> None:
class AsyncRandomRolloutWorker:
"""Background asyncio loop that fills an output queue with random sample groups.

Mirrors ``examples/fully_async/fully_async_rollout.AsyncRolloutWorker`` but
Mirrors ``miles.rollout.fully_async_rollout.AsyncRolloutWorker`` but
skips the data buffer and reward model entirely.
"""

Expand Down
9 changes: 4 additions & 5 deletions examples/swe-agent/run-glm47-flash-agentic-async.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
"""GLM-4.7-Flash fully-async agentic training with SWE-bench data.

Disaggregated fully-async variant for agentic tasks: training and rollout run
on separate nodes concurrently. Uses train_async.py and the fully_async_rollout
module so that weight updates do not block generation. Agent tasks are dispatched
on separate nodes concurrently. Uses train_async.py and the
miles.rollout.fully_async_rollout module so that weight updates do not block generation. Agent tasks are dispatched
to a Harbor-based agent server.

GLM-4.7-Flash architecture: 47 layers, 20 attention heads, 64 routed experts,
Expand Down Expand Up @@ -35,7 +35,6 @@
import miles.utils.external_utils.command_utils as U

SCRIPT_DIR = Path(__file__).resolve().parent
FULLY_ASYNC_DIR = (Path(__file__).resolve().parent.parent / "fully_async").resolve()

# Cluster-wide GPU-node ceiling for the ckpt-conversion job. Kept below the
# raw node count so ckpt conversion doesn't starve the rest of the cluster.
Expand Down Expand Up @@ -155,7 +154,7 @@ def execute(args: ScriptArgs):
)

rollout_args = (
"--rollout-function-path fully_async_rollout.generate_rollout_fully_async "
"--rollout-function-path miles.rollout.fully_async_rollout.generate_rollout_fully_async "
f"--prompt-data {args.prompt_data} "
"--input-key prompt "
"--metadata-key metadata "
Expand Down Expand Up @@ -346,7 +345,7 @@ def execute(args: ScriptArgs):
miles_root = U.repo_base_dir

extra_env_vars = {
"PYTHONPATH": f"{args.megatron_path}:{SCRIPT_DIR}:{FULLY_ASYNC_DIR}:{miles_root}",
"PYTHONPATH": f"{args.megatron_path}:{SCRIPT_DIR}:{miles_root}",
"MILES_EXPERIMENTAL_ROLLOUT_REFACTOR": "1",
"NCCL_NVLS_ENABLE": os.environ.get("HAS_NVLINK", "0"),
"SGLANG_ENABLE_TP_MEMORY_INBALANCE_CHECK": "true",
Expand Down
0 examples/fully_async/fully_async_rollout.py → miles/rollout/fully_async_rollout.py
View file
Open in desktop
File renamed without changes.
2 changes: 1 addition & 1 deletion train_async.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@
logger = logging.getLogger(__name__)


# The framework supports other asynchronous approaches such as fully async (which is shown in examples/fully_async).
# The framework supports other asynchronous approaches such as fully async (see miles/rollout/fully_async_rollout.py).
async def train(args):
assert not args.colocate, "Colocation is not supported for async training."
validate_async_off_policy_correction(args)
Expand Down
Loading