diff --git a/docs/docs.json b/docs/docs.json index 43af57818fd..bd9f7e30269 100644 --- a/docs/docs.json +++ b/docs/docs.json @@ -259,7 +259,6 @@ { "group": "Recipes", "pages": [ - "examples/fully-async", "examples/geo3k-vlm", "examples/geo3k-vlm/multi-turn", "examples/multi-lora", @@ -275,6 +274,7 @@ "group": "Infra Features", "root": "examples/infra-features", "pages": [ + "examples/infra-features/fully-async", "examples/infra-features/low-precision", "examples/infra-features/p2p-weight-transfer", "examples/infra-features/random-async", diff --git a/docs/examples/fully-async.md b/docs/examples/fully-async.md deleted file mode 100644 index 8770bbe2135..00000000000 --- a/docs/examples/fully-async.md +++ /dev/null @@ -1,54 +0,0 @@ ---- -title: "Fully Asynchronous Rollout Example" -description: "Demonstrates fully asynchronous rollout generation for higher efficiency." -# Generated from examples/fully_async/README.md by scripts/tools/sync_example_docs.py. Edit that README, not this file. ---- -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` (`FullyAsyncRolloutFn`, a class-based rollout function that owns a persistent background worker). It requires `MILES_EXPERIMENTAL_ROLLOUT_REFACTOR=1`. - -## Files -* `run-qwen3-4b-fully_async.sh`: example launch script with Qwen3‑4B. -* `run_qwen3_5_4b_fully_async_eval.py`: Qwen3.5‑4B with async checkpoint eval — `--eval-backend fleet` (dedicated eval fleet) or `--eval-backend external` (fn-launched sglang server). -* `external_eval_fn.py`: reference `CheckpointEvalFn` — launches/attaches an external sglang server and evals snapshots on it. - -## Prerequisite -First set up model & environment following the Qwen3-4B example. - -## Quick Start -```bash -cd miles -bash examples/fully_async/run-qwen3-4b-fully_async.sh -``` -You should see log lines like: -``` -Started fully-async rollout worker -``` - -## How It Works (Very Short) -* First train call: the rollout fn starts a persistent worker task on the shared rollout event loop. -* The worker keeps up to `--rollout-batch-size` groups in flight using `generate_and_rm_group`. -* Completed groups are pushed into a queue; each step drains until it has `--rollout-batch-size` groups. -* Aborted or too-stale groups are recycled back into the data source. - -## Evaluation -Without extra GPUs, eval shares the rollout engines (producer pauses during the blocking -eval). For eval that never pauses training, `run_qwen3_5_4b_fully_async_eval.py` shows both -checkpoint-pinned backends behind the same contract: `--eval-backend fleet` (in-job eval -fleet via `--eval-num-gpus`) and `--eval-backend external` (`--eval-function-path` pointed -at `external_eval_fn.ExternalSglangEvalFn`, which launches or attaches its own sglang -server). See the fully-async docs for the posture trade-offs. - -## Limitations -* Ordering is best effort (sorted at the end by index). - -## Config Differences (3 Key Points) -To enable the fully async pattern there are only three changes compared to a normal run: - -1. Use the async training driver: `train_async.py` (not `train.py`). -2. Enable the class-based rollout API: `MILES_EXPERIMENTAL_ROLLOUT_REFACTOR=1`. -3. Pass `--fully-async`. - -Why is it still "fully" async although `train_async.py` itself schedules rollouts step‑by‑step? - -Because the real generation work is done by a **persistent background worker** owned by `FullyAsyncRolloutFn`. Each call from `train_async.py` only drains already completed samples from the worker's output queue; the worker has been continuously generating since the first call. Thus rollout production (model inference) and training consume happen in parallel with minimal waiting. diff --git a/docs/examples/index.md b/docs/examples/index.md index cbdce2eb5c3..088f516b252 100644 --- a/docs/examples/index.md +++ b/docs/examples/index.md @@ -9,7 +9,6 @@ A few are purely demonstrative, but most are verifiable against a concrete perfo End-to-end training workflows — the place to start. -- **[fully_async](/examples/fully-async)**: Demonstrates fully asynchronous rollout generation for higher efficiency. - **[geo3k_vlm](/examples/geo3k-vlm)**: Training VLMs with FSDP using GRPO on the GEO3K dataset. - **[multi_turn](/examples/geo3k-vlm/multi-turn)**: The same dataset over multiple turns, with the model cropping images through an interactive environment. - **[lora](https://github.com/radixark/miles/tree/main/examples/lora)**: LoRA fine-tuning with the Megatron backend. @@ -25,6 +24,7 @@ End-to-end training workflows — the place to start. Runtime and infrastructure plumbing rather than training recipes — how miles moves data and weights around. +- **[fully_async](/examples/infra-features/fully-async)**: Demonstrates fully asynchronous rollout generation for higher efficiency. - **[low_precision](/examples/infra-features/low-precision)**: Examples of FP8 training and inference, plus INT4 QAT, for improved throughput and stability. - **[p2p_weight_transfer](/examples/infra-features/p2p-weight-transfer)**: Point-to-point weight transfer between training and rollout engines. - **[random_async](/examples/infra-features/random-async)**: Dataset-free stress test of the async rollout ↔ trainer loop. diff --git a/docs/examples/infra-features/fully-async.md b/docs/examples/infra-features/fully-async.md new file mode 100644 index 00000000000..cd2f78e5e30 --- /dev/null +++ b/docs/examples/infra-features/fully-async.md @@ -0,0 +1,28 @@ +--- +title: "Fully Asynchronous Rollout Example" +description: "Demonstrates fully asynchronous rollout generation for higher efficiency." +# Generated from examples/infra_features/fully_async/README.md by scripts/tools/sync_example_docs.py. Edit that README, not this file. +--- +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` (`FullyAsyncRolloutFn`, a class-based rollout function that owns a persistent background worker). It requires `MILES_EXPERIMENTAL_ROLLOUT_REFACTOR=1`. + +## Files +* `run_qwen3_5_4b_fully_async_eval.py`: Qwen3.5‑4B with async checkpoint eval — `--eval-backend fleet` (dedicated eval fleet) or `--eval-backend external` (fn-launched sglang server, `examples.infra_features.fully_async.external_eval_fn.ExternalSglangEvalFn`). +* `run_qwen3_30b_a3b_fully_async.py`: the same pattern on a 30B MoE — `tp=8`, `ep=8`, one 8-GPU rollout engine. +* `external_eval_fn.py`: reference `CheckpointEvalFn` — launches/attaches an external sglang server and evals snapshots on it. + +## Quick Start +Each launcher downloads its own checkpoint and converts it, then submits the job: +```bash +python examples/infra_features/fully_async/run_qwen3_5_4b_fully_async_eval.py +``` +You should see log lines like: +``` +Started fully-async rollout worker +``` + +## At a larger scale +[`examples/experimental/openenv/glm52_tbench2`](https://github.com/radixark/miles/tree/main/examples/experimental/openenv/glm52_tbench2) runs +the same flag on a frontier-scale agentic workload: GLM-5.2 744B-A40B on terminal-bench-2, +16 GB300 nodes split 8 training / 8 inference, one Daytona sandbox per episode. diff --git a/docs/examples/infra-features/random-async.md b/docs/examples/infra-features/random-async.md index b2c01b377d0..1080b338c37 100644 --- a/docs/examples/infra-features/random-async.md +++ b/docs/examples/infra-features/random-async.md @@ -3,7 +3,7 @@ title: "Random fully-async example" description: "Dataset-free stress test of the async rollout ↔ trainer loop." # Generated from examples/infra_features/random_async/README.md by scripts/tools/sync_example_docs.py. Edit that README, not this file. --- -Minimal sibling of `examples/fully_async/`. Exercises the entire async +Minimal sibling of `examples/infra_features/fully_async/`. Exercises the entire async rollout ↔ trainer loop **without any real dataset, real reward model, or meaningful generation** — useful as an agent infrastructure stress test for bigger agentic workloads. diff --git a/docs/models/glm/glm5-2.md b/docs/models/glm/glm5-2.md index 71e1ca25960..c3e3f4a09b9 100644 --- a/docs/models/glm/glm5-2.md +++ b/docs/models/glm/glm5-2.md @@ -146,4 +146,4 @@ The launcher exposes these as flags: - [Low Precision RL](/advanced/low-precision) — opt-in via `--fp8-rollout`. - [Speculative Decoding](/advanced/speculative-decoding) — opt-in via `--enable-mtp`. - [LoRA](/advanced/lora) — via `scripts/run_glm5_2_744b_a40b_lora.py`. -- [Fully Async Rollout](/examples/fully-async) — the terminal-bench-2 agentic example (§4.2) runs fully async. +- [Fully Async Rollout](/examples/infra-features/fully-async) — the terminal-bench-2 agentic example (§4.2) runs fully async. diff --git a/docs/user-guide/fully-async.md b/docs/user-guide/fully-async.md index d96dde15c5d..794760c4978 100644 --- a/docs/user-guide/fully-async.md +++ b/docs/user-guide/fully-async.md @@ -35,9 +35,8 @@ Four launch scripts show the mode end to end, from a single-node smoke test to a | Script | What it covers | |---|---| -| [`run-qwen3-4b-fully_async.sh`](https://github.com/radixark/miles/blob/main/examples/fully_async/run-qwen3-4b-fully_async.sh) | The smallest complete run: Qwen3-4B on one engine per GPU, with `--max-weight-staleness` shown as a commented-out option | -| [`run_qwen3_30b_a3b_fully_async.py`](https://github.com/radixark/miles/blob/main/examples/fully_async/run_qwen3_30b_a3b_fully_async.py) | The same pattern on a 30B MoE, with `tp=8`, `ep=8`, and one 8-GPU rollout engine | -| [`run_qwen3_5_4b_fully_async_eval.py`](https://github.com/radixark/miles/blob/main/examples/fully_async/run_qwen3_5_4b_fully_async_eval.py) | Both checkpoint eval backends behind one flag, `--eval-backend fleet` or `--eval-backend external` | +| [`run_qwen3_30b_a3b_fully_async.py`](https://github.com/radixark/miles/blob/main/examples/infra_features/fully_async/run_qwen3_30b_a3b_fully_async.py) | The same pattern on a 30B MoE, with `tp=8`, `ep=8`, and one 8-GPU rollout engine | +| [`run_qwen3_5_4b_fully_async_eval.py`](https://github.com/radixark/miles/blob/main/examples/infra_features/fully_async/run_qwen3_5_4b_fully_async_eval.py) | Both checkpoint eval backends behind one flag, `--eval-backend fleet` or `--eval-backend external` | | [`run_glm5_2_744b_a40b_daytona.py`](https://github.com/radixark/miles/blob/main/examples/experimental/openenv/glm52_tbench2/run_glm5_2_744b_a40b_daytona.py) | GLM-5.2 744B-A40B on 16 GB300 nodes, split 8 training and 8 inference, with multi-turn terminal-bench-2 episodes in per-task Daytona sandboxes. It runs 128 in-flight trajectories against a 64-sample train batch and evaluates on the shared rollout engines | ### Customizations @@ -247,7 +246,7 @@ them explicitly with `--eval-sglang-*` if the fleet is large enough to want them ### Mode 3: External backend The contract lives in [`miles/rollout/checkpoint_eval.py`](https://github.com/radixark/miles/blob/main/miles/rollout/checkpoint_eval.py), with a -reference implementation in [`examples/fully_async/external_eval_fn.py`](https://github.com/radixark/miles/blob/main/examples/fully_async/external_eval_fn.py). +reference implementation in [`examples/infra_features/fully_async/external_eval_fn.py`](https://github.com/radixark/miles/blob/main/examples/infra_features/fully_async/external_eval_fn.py). Subclass `CheckpointEvalFn` and implement `evaluate_checkpoint(checkpoint_dir, input)`. The trainer hands over a snapshot path per eval point and owns dispatch, logging, and diff --git a/examples/README.md b/examples/README.md index b127b859ba6..c8ef2d7cdbe 100644 --- a/examples/README.md +++ b/examples/README.md @@ -6,7 +6,6 @@ These examples are runnable starting points for your own RL workflow. A few are End-to-end training workflows — the place to start. -- **[fully_async](./fully_async)**: Demonstrates fully asynchronous rollout generation for higher efficiency. - **[geo3k_vlm](./geo3k_vlm)**: Training VLMs with FSDP using GRPO on the GEO3K dataset. - **[multi_turn](./geo3k_vlm/multi_turn)**: The same dataset over multiple turns, with the model cropping images through an interactive environment. - **[lora](./lora)**: LoRA fine-tuning with the Megatron backend. @@ -22,6 +21,7 @@ End-to-end training workflows — the place to start. Runtime and infrastructure plumbing rather than training recipes — how miles moves data and weights around. +- **[fully_async](./infra_features/fully_async)**: Demonstrates fully asynchronous rollout generation for higher efficiency. - **[low_precision](./infra_features/low_precision)**: Examples of FP8 training and inference, plus INT4 QAT, for improved throughput and stability. - **[p2p_weight_transfer](./infra_features/p2p_weight_transfer)**: Point-to-point weight transfer between training and rollout engines. - **[random_async](./infra_features/random_async)**: Dataset-free stress test of the async rollout ↔ trainer loop. diff --git a/examples/fully_async/README.md b/examples/fully_async/README.md deleted file mode 100644 index db61e7a0aa6..00000000000 --- a/examples/fully_async/README.md +++ /dev/null @@ -1,51 +0,0 @@ -# Fully Asynchronous Rollout Example - -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` (`FullyAsyncRolloutFn`, a class-based rollout function that owns a persistent background worker). It requires `MILES_EXPERIMENTAL_ROLLOUT_REFACTOR=1`. - -## Files -* `run-qwen3-4b-fully_async.sh`: example launch script with Qwen3‑4B. -* `run_qwen3_5_4b_fully_async_eval.py`: Qwen3.5‑4B with async checkpoint eval — `--eval-backend fleet` (dedicated eval fleet) or `--eval-backend external` (fn-launched sglang server). -* `external_eval_fn.py`: reference `CheckpointEvalFn` — launches/attaches an external sglang server and evals snapshots on it. - -## Prerequisite -First set up model & environment following the Qwen3-4B example. - -## Quick Start -```bash -cd miles -bash examples/fully_async/run-qwen3-4b-fully_async.sh -``` -You should see log lines like: -``` -Started fully-async rollout worker -``` - -## How It Works (Very Short) -* First train call: the rollout fn starts a persistent worker task on the shared rollout event loop. -* The worker keeps up to `--rollout-batch-size` groups in flight using `generate_and_rm_group`. -* Completed groups are pushed into a queue; each step drains until it has `--rollout-batch-size` groups. -* Aborted or too-stale groups are recycled back into the data source. - -## Evaluation -Without extra GPUs, eval shares the rollout engines (producer pauses during the blocking -eval). For eval that never pauses training, `run_qwen3_5_4b_fully_async_eval.py` shows both -checkpoint-pinned backends behind the same contract: `--eval-backend fleet` (in-job eval -fleet via `--eval-num-gpus`) and `--eval-backend external` (`--eval-function-path` pointed -at `external_eval_fn.ExternalSglangEvalFn`, which launches or attaches its own sglang -server). See the fully-async docs for the posture trade-offs. - -## Limitations -* Ordering is best effort (sorted at the end by index). - -## Config Differences (3 Key Points) -To enable the fully async pattern there are only three changes compared to a normal run: - -1. Use the async training driver: `train_async.py` (not `train.py`). -2. Enable the class-based rollout API: `MILES_EXPERIMENTAL_ROLLOUT_REFACTOR=1`. -3. Pass `--fully-async`. - -Why is it still "fully" async although `train_async.py` itself schedules rollouts step‑by‑step? - -Because the real generation work is done by a **persistent background worker** owned by `FullyAsyncRolloutFn`. Each call from `train_async.py` only drains already completed samples from the worker's output queue; the worker has been continuously generating since the first call. Thus rollout production (model inference) and training consume happen in parallel with minimal waiting. diff --git a/examples/fully_async/run-qwen3-4b-fully_async.sh b/examples/fully_async/run-qwen3-4b-fully_async.sh deleted file mode 100644 index 07dd2b6ee5f..00000000000 --- a/examples/fully_async/run-qwen3-4b-fully_async.sh +++ /dev/null @@ -1,143 +0,0 @@ -#!/bin/bash - -# for rerun the task -pkill -9 sglang -sleep 3 -ray stop --force -pkill -9 ray -pkill -9 python -sleep 3 -pkill -9 ray -pkill -9 python - -set -ex - -# will prevent ray from buffering stdout/stderr -export PYTHONUNBUFFERED=1 - -NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) -if [ "$NVLINK_COUNT" -gt 0 ]; then - HAS_NVLINK=1 -else - HAS_NVLINK=0 -fi -echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" - -SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" -MODEL_ARGS_LINE="$(python3 "${SCRIPT_DIR}/../../miles/utils/external_utils/model_args_utils.py" "qwen3-4B")" || exit 1 -read -ra MODEL_ARGS <<< "${MODEL_ARGS_LINE}" -CKPT_ARGS=( - --hf-checkpoint /root/Qwen3-4B - #--hf-checkpoint /root/Qwen3-4B-FP8 - --ref-load /root/Qwen3-4B_torch_dist - --load /root/Qwen3-4B_miles/ - --save /root/Qwen3-4B_miles/ - --save-interval 20 -) - -PROMPT_SET=/path/to/dapo-math-17k.jsonl - -ROLLOUT_ARGS=( - --fully-async - --prompt-data ${PROMPT_SET} - --input-key prompt - --label-key label - --apply-chat-template - --rollout-shuffle - - --rm-type dapo - --reward-key score - - --num-rollout 3000 - --rollout-batch-size 32 - --n-samples-per-prompt 8 - --rollout-max-response-len 8192 - --rollout-temperature 1 - - --global-batch-size 256 - --balance-data - - # for staleness control - #--max-weight-staleness 2 -) - -PERF_ARGS=( - --tensor-model-parallel-size 2 - --sequence-parallel - --pipeline-model-parallel-size 1 - --context-parallel-size 1 - --expert-model-parallel-size 1 - --expert-tensor-parallel-size 1 - - --recompute-granularity full - --recompute-method uniform - --recompute-num-layers 1 - - # --micro-batch-size 1 - --use-dynamic-batch-size - --max-tokens-per-gpu 9216 -) - -GRPO_ARGS=( - --advantage-estimator grpo - --use-kl-loss - --kl-loss-coef 0.00 - --kl-loss-type low_var_kl - --entropy-coef 0.00 - --eps-clip 0.2 - --eps-clip-high 0.28 - - --use-tis -) - -OPTIMIZER_ARGS=( - --optimizer adam - --lr 1e-6 - --lr-decay-style constant - --weight-decay 0.1 - --adam-beta1 0.9 - --adam-beta2 0.98 -) - -SGLANG_ARGS=( - --rollout-num-gpus-per-engine 1 -) - -MISC_ARGS=( - # default dropout in megatron is 0.1 - --attention-dropout 0.0 - --hidden-dropout 0.0 - # should be good for model performance - --accumulate-allreduce-grads-in-fp32 - --attention-softmax-in-fp32 - # need to comment this when using model with MLA - --attention-backend flash -) - -# launch the master node of ray in container -export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"} -ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats - -RUNTIME_ENV_JSON="{ - \"env_vars\": { - \"PYTHONPATH\": \"/root/Megatron-LM/\", - \"MILES_EXPERIMENTAL_ROLLOUT_REFACTOR\": \"1\", - \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", - \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\" - } -}" - -ray job submit --address="http://127.0.0.1:8265" \ - --runtime-env-json="${RUNTIME_ENV_JSON}" \ - -- python3 train_async.py \ - --actor-num-nodes 1 \ - --actor-num-gpus-per-node 4 \ - --rollout-num-gpus 4 \ - ${MODEL_ARGS[@]} \ - ${CKPT_ARGS[@]} \ - ${ROLLOUT_ARGS[@]} \ - ${OPTIMIZER_ARGS[@]} \ - ${GRPO_ARGS[@]} \ - ${PERF_ARGS[@]} \ - ${SGLANG_ARGS[@]} \ - ${MISC_ARGS[@]} diff --git a/examples/infra_features/fully_async/README.md b/examples/infra_features/fully_async/README.md new file mode 100644 index 00000000000..22370fe2d5c --- /dev/null +++ b/examples/infra_features/fully_async/README.md @@ -0,0 +1,31 @@ +# Fully Asynchronous Rollout Example + + +> **Read the docs:** [Fully Async RL](https://miles.radixark.com/docs/user-guide/fully-async) +> covers the schedule, the data buffer, the three evaluation modes, and every `--fully-async` +> argument. + + +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` (`FullyAsyncRolloutFn`, a class-based rollout function that owns a persistent background worker). It requires `MILES_EXPERIMENTAL_ROLLOUT_REFACTOR=1`. + +## Files +* `run_qwen3_5_4b_fully_async_eval.py`: Qwen3.5‑4B with async checkpoint eval — `--eval-backend fleet` (dedicated eval fleet) or `--eval-backend external` (fn-launched sglang server, `examples.infra_features.fully_async.external_eval_fn.ExternalSglangEvalFn`). +* `run_qwen3_30b_a3b_fully_async.py`: the same pattern on a 30B MoE — `tp=8`, `ep=8`, one 8-GPU rollout engine. +* `external_eval_fn.py`: reference `CheckpointEvalFn` — launches/attaches an external sglang server and evals snapshots on it. + +## Quick Start +Each launcher downloads its own checkpoint and converts it, then submits the job: +```bash +python examples/infra_features/fully_async/run_qwen3_5_4b_fully_async_eval.py +``` +You should see log lines like: +``` +Started fully-async rollout worker +``` + +## At a larger scale +[`examples/experimental/openenv/glm52_tbench2`](../../experimental/openenv/glm52_tbench2) runs +the same flag on a frontier-scale agentic workload: GLM-5.2 744B-A40B on terminal-bench-2, +16 GB300 nodes split 8 training / 8 inference, one Daytona sandbox per episode. diff --git a/examples/fully_async/__init__.py b/examples/infra_features/fully_async/__init__.py similarity index 100% rename from examples/fully_async/__init__.py rename to examples/infra_features/fully_async/__init__.py diff --git a/examples/fully_async/external_eval_fn.py b/examples/infra_features/fully_async/external_eval_fn.py similarity index 100% rename from examples/fully_async/external_eval_fn.py rename to examples/infra_features/fully_async/external_eval_fn.py diff --git a/examples/fully_async/run_qwen3_30b_a3b_fully_async.py b/examples/infra_features/fully_async/run_qwen3_30b_a3b_fully_async.py similarity index 100% rename from examples/fully_async/run_qwen3_30b_a3b_fully_async.py rename to examples/infra_features/fully_async/run_qwen3_30b_a3b_fully_async.py diff --git a/examples/fully_async/run_qwen3_5_4b_fully_async_eval.py b/examples/infra_features/fully_async/run_qwen3_5_4b_fully_async_eval.py similarity index 98% rename from examples/fully_async/run_qwen3_5_4b_fully_async_eval.py rename to examples/infra_features/fully_async/run_qwen3_5_4b_fully_async_eval.py index 43e9b20e92d..38db9f1daa6 100644 --- a/examples/fully_async/run_qwen3_5_4b_fully_async_eval.py +++ b/examples/infra_features/fully_async/run_qwen3_5_4b_fully_async_eval.py @@ -96,7 +96,7 @@ def execute(args: ScriptArgs): if args.eval_backend == "fleet": eval_args += "--eval-num-gpus 1 --eval-num-gpus-per-engine 1 " else: - eval_args += "--eval-function-path examples.fully_async.external_eval_fn.ExternalSglangEvalFn " + eval_args += "--eval-function-path examples.infra_features.fully_async.external_eval_fn.ExternalSglangEvalFn " # The fn launches its own sglang server on the last GPU, outside the Ray split. eval_env = {"MILES_EXTERNAL_EVAL_GPUS": str(args.num_gpus_per_node - 1)} diff --git a/examples/infra_features/random_async/README.md b/examples/infra_features/random_async/README.md index 6ddf6ad0cb5..564200ae8fb 100644 --- a/examples/infra_features/random_async/README.md +++ b/examples/infra_features/random_async/README.md @@ -1,6 +1,6 @@ # Random fully-async example -Minimal sibling of `examples/fully_async/`. Exercises the entire async +Minimal sibling of `examples/infra_features/fully_async/`. Exercises the entire async rollout ↔ trainer loop **without any real dataset, real reward model, or meaningful generation** — useful as an agent infrastructure stress test for bigger agentic workloads. diff --git a/miles/rollout/checkpoint_eval.py b/miles/rollout/checkpoint_eval.py index 97020aa472d..0a2939fa976 100644 --- a/miles/rollout/checkpoint_eval.py +++ b/miles/rollout/checkpoint_eval.py @@ -10,7 +10,7 @@ ``EvalSkip(reason)`` to skip a point with attribution instead of counting as a crash. Requires ``train_async.py`` and a snapshot source (``--eval-hf-dir`` or ``--save-hf``). -``examples/fully_async/external_eval_fn.py`` is a working implementation. +``examples/infra_features/fully_async/external_eval_fn.py`` is a working implementation. """ import abc diff --git a/tests/e2e/megatron/test_qwen3_4b_fully_async_eval.py b/tests/e2e/megatron/test_qwen3_4b_fully_async_eval.py index ac5d32d5d7f..b158c3928a4 100644 --- a/tests/e2e/megatron/test_qwen3_4b_fully_async_eval.py +++ b/tests/e2e/megatron/test_qwen3_4b_fully_async_eval.py @@ -92,7 +92,7 @@ def execute(eval_mode: str): if eval_mode == "fleet": eval_args += "--eval-num-gpus 1 --eval-num-gpus-per-engine 1 " elif eval_mode == "external": - eval_args += "--eval-function-path examples.fully_async.external_eval_fn.ExternalSglangEvalFn " + eval_args += "--eval-function-path examples.infra_features.fully_async.external_eval_fn.ExternalSglangEvalFn " eval_env = {"MILES_EXTERNAL_EVAL_GPUS": str(NUM_GPUS - 1)} perf_args = ( diff --git a/tests/fast/rollout/test_checkpoint_eval.py b/tests/fast/rollout/test_checkpoint_eval.py index 498d0cdf8f2..aaaa0a33c8e 100644 --- a/tests/fast/rollout/test_checkpoint_eval.py +++ b/tests/fast/rollout/test_checkpoint_eval.py @@ -469,14 +469,14 @@ def remote(self, rollout_id): assert len(dispatcher.pending) == 0 -# ---------------- example fn (examples/fully_async/external_eval_fn.py) ---------------- +# ------- example fn (examples/infra_features/fully_async/external_eval_fn.py) ------- @pytest.fixture def external_fn_env(monkeypatch): import importlib - mod = importlib.import_module("examples.fully_async.external_eval_fn") + mod = importlib.import_module("examples.infra_features.fully_async.external_eval_fn") calls = [] server = SimpleNamespace(loaded_version=None) diff --git a/tests/snapshots/launch_scripts/sh/examples/fully_async/run-qwen3-4b-fully_async.sh.txt b/tests/snapshots/launch_scripts/sh/examples/fully_async/run-qwen3-4b-fully_async.sh.txt deleted file mode 100644 index cf90b6930ee..00000000000 --- a/tests/snapshots/launch_scripts/sh/examples/fully_async/run-qwen3-4b-fully_async.sh.txt +++ /dev/null @@ -1,192 +0,0 @@ -# returncode: 0 - -### 0 -"pkill" -"-9" -"sglang" - -### 1 -"sleep" -"3" - -### 2 -"ray" -"stop" -"--force" - -### 3 -"pkill" -"-9" -"ray" - -### 4 -"pkill" -"-9" -"python" - -### 5 -"sleep" -"3" - -### 6 -"pkill" -"-9" -"ray" - -### 7 -"pkill" -"-9" -"python" - -### 8 -"nvidia-smi" -"topo" -"-m" - -### 9 -"python3" -"/examples/fully_async/../../miles/utils/external_utils/model_args_utils.py" -"qwen3-4B" - -### 10 -"ray" -"start" -"--head" -"--node-ip-address" -"127.0.0.1" -"--num-gpus" -"8" -"--disable-usage-stats" - -### 11 -"ray" -"job" -"submit" -"--address=http://127.0.0.1:8265" -"--runtime-env-json={\n \"env_vars\": {\n \"PYTHONPATH\": \"/root/Megatron-LM/\",\n \"MILES_EXPERIMENTAL_ROLLOUT_REFACTOR\": \"1\",\n \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\",\n \"NCCL_NVLS_ENABLE\": \"0\"\n }\n}" -"--" -"python3" -"train_async.py" -"--actor-num-nodes" -"1" -"--actor-num-gpus-per-node" -"4" -"--rollout-num-gpus" -"4" -"--swiglu" -"--num-layers" -"36" -"--hidden-size" -"2560" -"--ffn-hidden-size" -"9728" -"--num-attention-heads" -"32" -"--group-query-attention" -"--num-query-groups" -"8" -"--use-rotary-position-embeddings" -"--disable-bias-linear" -"--normalization" -"RMSNorm" -"--norm-epsilon" -"1e-6" -"--rotary-base" -"1000000" -"--vocab-size" -"151936" -"--kv-channels" -"128" -"--qk-layernorm" -"--hf-checkpoint" -"/root/Qwen3-4B" -"--ref-load" -"/root/Qwen3-4B_torch_dist" -"--load" -"/root/Qwen3-4B_miles/" -"--save" -"/root/Qwen3-4B_miles/" -"--save-interval" -"20" -"--fully-async" -"--prompt-data" -"/path/to/dapo-math-17k.jsonl" -"--input-key" -"prompt" -"--label-key" -"label" -"--apply-chat-template" -"--rollout-shuffle" -"--rm-type" -"dapo" -"--reward-key" -"score" -"--num-rollout" -"3000" -"--rollout-batch-size" -"32" -"--n-samples-per-prompt" -"8" -"--rollout-max-response-len" -"8192" -"--rollout-temperature" -"1" -"--global-batch-size" -"256" -"--balance-data" -"--optimizer" -"adam" -"--lr" -"1e-6" -"--lr-decay-style" -"constant" -"--weight-decay" -"0.1" -"--adam-beta1" -"0.9" -"--adam-beta2" -"0.98" -"--advantage-estimator" -"grpo" -"--use-kl-loss" -"--kl-loss-coef" -"0.00" -"--kl-loss-type" -"low_var_kl" -"--entropy-coef" -"0.00" -"--eps-clip" -"0.2" -"--eps-clip-high" -"0.28" -"--use-tis" -"--tensor-model-parallel-size" -"2" -"--sequence-parallel" -"--pipeline-model-parallel-size" -"1" -"--context-parallel-size" -"1" -"--expert-model-parallel-size" -"1" -"--expert-tensor-parallel-size" -"1" -"--recompute-granularity" -"full" -"--recompute-method" -"uniform" -"--recompute-num-layers" -"1" -"--use-dynamic-batch-size" -"--max-tokens-per-gpu" -"9216" -"--rollout-num-gpus-per-engine" -"1" -"--attention-dropout" -"0.0" -"--hidden-dropout" -"0.0" -"--accumulate-allreduce-grads-in-fp32" -"--attention-softmax-in-fp32" -"--attention-backend" -"flash"