diff --git a/docs/developer/architecture.md b/docs/developer/architecture.md index 972fcb637cc..0ee3bb1bb3c 100644 --- a/docs/developer/architecture.md +++ b/docs/developer/architecture.md @@ -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 diff --git a/docs/examples/fully-async.md b/docs/examples/fully-async.md index b706ece316c..ce3bb6a90be 100644 --- a/docs/examples/fully-async.md +++ b/docs/examples/fully-async.md @@ -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 ``` @@ -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. @@ -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() diff --git a/docs/user-guide/fully-async.md b/docs/user-guide/fully-async.md index 3ac4dd4aee4..df7ce648411 100644 --- a/docs/user-guide/fully-async.md +++ b/docs/user-guide/fully-async.md @@ -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 ``` Everything else belongs in the same [argument groups](/user-guide/argument-groups) as a diff --git a/docs/user-guide/training-script-walkthrough.md b/docs/user-guide/training-script-walkthrough.md index 6176a59c568..1489128b2a1 100644 --- a/docs/user-guide/training-script-walkthrough.md +++ b/docs/user-guide/training-script-walkthrough.md @@ -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 | diff --git a/examples/fully_async/README.md b/examples/fully_async/README.md index 53f08b3ca48..92231766d81 100644 --- a/examples/fully_async/README.md +++ b/examples/fully_async/README.md @@ -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 @@ -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? diff --git a/examples/fully_async/run-qwen3-4b-fully_async.sh b/examples/fully_async/run-qwen3-4b-fully_async.sh index bfd12696bfb..44cbf361446 100644 --- a/examples/fully_async/run-qwen3-4b-fully_async.sh +++ b/examples/fully_async/run-qwen3-4b-fully_async.sh @@ -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 @@ -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}\" } diff --git a/examples/fully_async/run_qwen3_30b_a3b_fully_async.py b/examples/fully_async/run_qwen3_30b_a3b_fully_async.py index 0a2dbca9249..eb4bec395e4 100644 --- a/examples/fully_async/run_qwen3_30b_a3b_fully_async.py +++ b/examples/fully_async/run_qwen3_30b_a3b_fully_async.py @@ -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 " @@ -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, @@ -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, }, ) diff --git a/examples/infra_features/random_async/random_async_rollout.py b/examples/infra_features/random_async/random_async_rollout.py index 4958e0a1209..fcf4e82407c 100644 --- a/examples/infra_features/random_async/random_async_rollout.py +++ b/examples/infra_features/random_async/random_async_rollout.py @@ -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. """ diff --git a/examples/swe-agent/run-glm47-flash-agentic-async.py b/examples/swe-agent/run-glm47-flash-agentic-async.py index 0710b6bbecf..6d35da459b5 100644 --- a/examples/swe-agent/run-glm47-flash-agentic-async.py +++ b/examples/swe-agent/run-glm47-flash-agentic-async.py @@ -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, @@ -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. @@ -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 " @@ -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", diff --git a/examples/fully_async/fully_async_rollout.py b/miles/rollout/fully_async_rollout.py similarity index 100% rename from examples/fully_async/fully_async_rollout.py rename to miles/rollout/fully_async_rollout.py diff --git a/train_async.py b/train_async.py index 62a542d93db..dc827f08d3f 100644 --- a/train_async.py +++ b/train_async.py @@ -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)