diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index 8cb815c1f2..0cbd06e02b 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -201,7 +201,7 @@ jobs: strategy: fail-fast: false matrix: - info: [{"num_gpus": 8, "test_file": "test_quick_start_glm4_9B.py"}, {"num_gpus": 8, "test_file": "test_glm4.7_30B_A3B_pd_mooncake.py"}, {"num_gpus": 8, "test_file": "test_qwen3_30B_A3B.py", "use_deepep": "1", "use_fp8_rollout": "1"}, {"num_gpus": 8, "test_file": "test_qwen3.6_35B_A3B_pd_mooncake.py", "use_deepep": "1"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_30B_A3B_r3.py", "use_deepep": "1", "use_fp8_rollout": "1"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_30B_A3B_r3.py"}, {"num_gpus": 8, "test_file": "test_qwen3_4B_ppo.py"}, {"num_gpus": 8, "test_file": "test_qwen3_4B_ppo_disaggregate.py"}, {"num_gpus": 8, "test_file": "test_qwen3_4B_ppo_train_critic_only.py"}, {"num_gpus": 8, "test_file": "test_moonlight_16B_A3B.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_moonlight_16B_A3B_r3.py"}, {"num_gpus": 8, "test_file": "test_mimo_7B_mtp_only_grad.py"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_debug_rollout_then_train.py"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_opd_sglang.py"}] + info: [{"num_gpus": 8, "test_file": "test_quick_start_glm4_9B.py"}, {"num_gpus": 8, "test_file": "test_glm4.7_30B_A3B_pd_mooncake.py"}, {"num_gpus": 8, "test_file": "test_qwen3_30B_A3B.py", "use_deepep": "1", "use_fp8_rollout": "1"}, {"num_gpus": 8, "test_file": "test_qwen3.6_35B_A3B_pd_mooncake.py", "use_deepep": "1"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_30B_A3B_r3.py", "use_deepep": "1", "use_fp8_rollout": "1"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_30B_A3B_r3.py"}, {"num_gpus": 8, "test_file": "test_qwen3_4B_ppo.py"}, {"num_gpus": 8, "test_file": "test_qwen3_4B_ppo_disaggregate.py"}, {"num_gpus": 8, "test_file": "test_qwen3_4B_ppo_train_critic_only.py"}, {"num_gpus": 8, "test_file": "test_moonlight_16B_A3B.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_moonlight_16B_A3B_r3.py"}, {"num_gpus": 8, "test_file": "test_mimo_7B_mtp_only_grad.py"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_debug_rollout_then_train.py"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_opd_sglang.py"}, {"num_gpus": 4, "test_file": "test_qwen2.5_0.5B_fully_async_short.py"}] defaults: run: working-directory: ${{ github.workspace }} diff --git a/.github/workflows/pr-test.yml.j2 b/.github/workflows/pr-test.yml.j2 index 50fe472273..ad3f4c2970 100644 --- a/.github/workflows/pr-test.yml.j2 +++ b/.github/workflows/pr-test.yml.j2 @@ -33,6 +33,7 @@ {'test_file': 'test_mimo_7B_mtp_only_grad.py', 'num_gpus': 8}, {'test_file': 'test_qwen2.5_0.5B_debug_rollout_then_train.py', 'num_gpus': 8}, {'test_file': 'test_qwen2.5_0.5B_opd_sglang.py', 'num_gpus': 8}, + {'test_file': 'test_qwen2.5_0.5B_fully_async_short.py', 'num_gpus': 4}, ], }, 'e2e-test-precision': { diff --git a/examples/fully_async/README.md b/examples/fully_async/README.md index 36a36cb7d9..975ee96b91 100644 --- a/examples/fully_async/README.md +++ b/examples/fully_async/README.md @@ -1,45 +1,81 @@ -# Fully Asynchronous Rollout Example +# Fully-Async 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. +End-to-end demo of slime's fully-async rollout path. A background asyncio +worker keeps a fixed pool of in-flight generations across rollout boundaries, +so the next training step doesn't wait for the slowest in-flight sample. +The worker itself lives in `slime.rollout.fully_async_rollout`; this +directory is just the launch script + CI test. ## 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 -First set up model & environment following the Qwen3-4B example. +* `run-qwen2.5-0.5B-fully_async.sh` — single-node, 4-GPU, three-rollout demo + with Qwen2.5-0.5B-Instruct on dapo-math-17k. Fast enough to be the CI + smoke test for the fully-async path. + +The same script doubles as `tests/test_qwen2.5_0.5B_fully_async_short.py` in +CI. + +## Prerequisites + +``` +/root/models/Qwen2.5-0.5B-Instruct/ # HF checkpoint +/root/models/Qwen2.5-0.5B-Instruct_torch_dist/ # tools/convert_hf_to_torch_dist.py +/root/datasets/dapo-math-17k/dapo-math-17k.jsonl +``` + +## Run -## Quick Start ```bash cd slime -bash examples/fully_async/run-qwen3-4b-fully_async.sh +bash examples/fully_async/run-qwen2.5-0.5B-fully_async.sh ``` -You should see log lines like: + +You should see: + ``` -Creating new global async worker... -Continuous async rollout worker started +fully-async rollout 0: target=8 queue_warm=0 +fully-async rollout 0: done in ...s, queue_left=... ``` -## How It Works (Very Short) -* First call: create `AsyncRolloutWorker` (thread + asyncio loop). -* Loop keeps up to `--rollout-batch-size` tasks in flight using `generate_and_rm_group`. -* Completed groups are pushed into a queue; caller drains until it has enough samples. -* Worker is stopped automatically at process exit. - -## Limitations -* No evaluation mode. -* Ordering is best effort (sorted at the end by index). -* Minimal error handling. +## How To Plug Your Own Generate Into This -## Config Differences (2 Key Points) -To enable the fully async pattern there are only two changes compared to a normal run: +Two pieces flip the standard pipeline into fully-async: -1. Use the async training driver: `train_async.py` (not `train.py`). +1. Use the async training driver: `python3 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 slime.rollout.fully_async_rollout.generate_rollout_fully_async + ``` + +For custom per-sample logic, use slime's standard plug-in points — they +work unchanged under fully-async: + +``` +--custom-generate-function-path your.module.generate # (args, sample, sampling_params) -> Sample | list[Sample] +--custom-rm-path your.module.reward # (args, sample | list[Sample]) -> float | list[float] +``` -Why is it still "fully" async although `train_async.py` itself schedules rollouts step‑by‑step? +See `examples/swe_codex/` for a non-trivial example that plugs in a +multi-turn agent (Claude Code in a Docker-Proxy sandbox) this way. + +## Worker Internals (Very Short) + +* First call: create a process-wide `AsyncRolloutWorker` (thread + asyncio + loop). The worker is shared across all subsequent `generate_rollout` + calls so its queue stays warm. +* Loop keeps up to `args.sglang_server_concurrency` tasks in flight using + `generate_and_rm_group`. +* Completed groups land on an output queue; each `generate_rollout` call + drains until it has `rollout_batch_size` groups and returns them sorted + by `sample.index`. +* Groups containing an `ABORTED` sample are pushed back into + `data_buffer.add_samples` instead of being shipped to training. +* Worker is stopped automatically at process exit via `atexit`. + +## Limitations -Because the real generation work is done by a **persistent background worker** created in `generate_rollout_fully_async`. 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. +* No evaluation mode (would conflict with the continuous-running model). +* Ordering across rollouts is best-effort — within a rollout, groups are + sorted by index before being handed to training. +* TODO: partial-rollout-style resume for `ABORTED` trajectories is not + yet wired; for now the trajectory is re-queued and starts over. diff --git a/examples/fully_async/fully_async_rollout.py b/examples/fully_async/fully_async_rollout.py deleted file mode 100644 index 7208365c18..0000000000 --- a/examples/fully_async/fully_async_rollout.py +++ /dev/null @@ -1,263 +0,0 @@ -import asyncio -import atexit -import queue -import threading -import time - -# Import core functions from sglang_rollout directly to avoid code duplication -from slime.rollout.sglang_rollout import GenerateState, generate_and_rm_group -from slime.utils.async_utils import run -from slime.utils.types import Sample - -# Global worker manager -_global_worker = None -_worker_lock = threading.Lock() - - -def get_global_worker(args, data_buffer): - """Get or create global worker""" - global _global_worker - 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.start() - return _global_worker - - -def stop_global_worker(): - """Stop global worker""" - global _global_worker - with _worker_lock: - if _global_worker is not None: - _global_worker.stop() - _global_worker = None - - -class AsyncRolloutWorker: - """ - Simplified asynchronous rollout worker, using threads instead of processes - Supports continuous running, independent of rollout function lifecycle - """ - - def __init__(self, args, data_buffer, concurrency=10): - 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 - self.state = GenerateState(args) - - async def continuous_worker_loop(self): - """Continuous work loop - constantly get data from data_buffer and process""" - print("Continuous async rollout worker started") - - active_tasks = set() - max_concurrent_tasks = self.args.rollout_batch_size - group_id_counter = 0 - - while self.running: - try: - # Clean up completed tasks - if active_tasks: - done_tasks = {task for task in active_tasks if task.done()} - for task in done_tasks: - try: - task.result() # Results are already handled in callbacks - except Exception as e: - print(f"Task failed with exception: {e}") - active_tasks -= done_tasks - - # If active task count hasn't reached limit, try to get new data and start tasks - while len(active_tasks) < max_concurrent_tasks and self.running: - samples = self.data_buffer.get_samples(1) - - for group in samples: - group_id = group_id_counter - group_id_counter += 1 - - # Create new async task - task = asyncio.create_task( - generate_and_rm_group( - self.args, - group, - sampling_params=self.state.sampling_params.copy(), - evaluation=False, - ) - ) - - # Add completion callback - def make_callback(gid): - def task_done_callback(done_task): - result = done_task.result() - self.output_queue.put((gid, result)) - - return task_done_callback - - task.add_done_callback(make_callback(group_id)) - active_tasks.add(task) - break - - # Brief sleep to avoid busy waiting - await asyncio.sleep(1) - - except Exception as e: - print(f"Error in continuous worker loop: {e}") - await asyncio.sleep(1) - - if active_tasks: - print(f"Waiting for {len(active_tasks)} continuous tasks to complete...") - await asyncio.wait(active_tasks) - - print("Continuous async rollout worker stopped") - - def worker_thread_func(self): - """Worker function running in independent thread""" - asyncio.run(self.continuous_worker_loop()) - - def start(self): - """Start continuous work mode""" - if self.worker_thread is None or not self.worker_thread.is_alive(): - self.worker_thread = threading.Thread(target=self.worker_thread_func, daemon=True) - self.worker_thread.start() - print("Started continuous async worker thread") - - def stop(self): - """Stop worker thread""" - self.running = False - if self.worker_thread and self.worker_thread.is_alive(): - self.worker_thread.join(timeout=5) - print("Stopped async worker thread") - - def get_completed_groups(self) -> list[tuple]: - """Get completed sample groups""" - completed = [] - while True: - try: - result = self.output_queue.get_nowait() - completed.append(result) - except queue.Empty: - break - return completed - - def get_queue_size(self) -> int: - """Get current output queue size""" - return self.output_queue.qsize() - - -async def generate_rollout_async(args, rollout_id: int, data_buffer) -> list[list[Sample]]: - """ - Simplified asynchronous rollout generation - using global continuous worker - """ - assert args.rollout_global_dataset - - # Get global worker, which will run continuously - worker = get_global_worker(args, data_buffer) - - # Simplified: directly use rollout_batch_size as target - target_data_size = args.rollout_batch_size - - data = [] - completed_groups = {} - do_print = True - - print(f"Starting async rollout generation for {target_data_size} groups") - print(f"Global worker queue size: {worker.get_queue_size()}") - - # Main loop: collect results from global worker's output queue - start_time = time.time() - last_progress_time = start_time - no_progress_timeout = 30.0 # Warn if no progress for 30 seconds - - while len(data) < target_data_size: - # Collect completed results - completed = worker.get_completed_groups() - - made_progress = False - for group_id, group in completed: - completed_groups[group_id] = group - made_progress = True - - if made_progress: - last_progress_time = time.time() - - # Process completed groups in order (try to maintain order, but not strict requirement) - processed_any = False - - # Process all available completed groups - available_ids = list(completed_groups.keys()) - for group_id in available_ids: - if len(data) >= target_data_size: - break - - group = completed_groups.pop(group_id) - - # If any sample in the group was aborted, return the whole group to the data buffer - # and do not forward it to the training engine. - try: - any_aborted = any([sample.status == Sample.Status.ABORTED for sample in group]) - except Exception: - any_aborted = False - - if any_aborted: - try: - # add back to buffer so it can be retried or handled by buffer policy - data_buffer.add_samples([group]) - print(f"Returned aborted group {group_id} to data buffer", flush=True) - except Exception as e: - print(f"Failed to return aborted group {group_id} to buffer: {e}", flush=True) - # don't count as processed for training - continue - - if do_print: - print( - f"First rollout sample: {[group[0].prompt + group[0].response]}, " - f"label: {group[0].label}, reward: {group[0].reward}", - flush=True, - ) - do_print = False - - # Simplified: directly add samples, no filters used - data.append(group) - processed_any = True - - # Check progress - current_time = time.time() - if current_time - last_progress_time > no_progress_timeout: - print( - f"Warning: No progress for {no_progress_timeout}s. " - f"Queue size: {worker.get_queue_size()}, " - f"Collected: {len(data)}/{target_data_size}" - ) - last_progress_time = current_time - - # If no results were processed, brief sleep to avoid busy waiting - if not processed_any: - await asyncio.sleep(0.01) - - duration = time.time() - start_time - print(f"Rollout completed in {duration:.2f}s! Global worker queue size: {worker.get_queue_size()}") - - if data: - print( - f"Finish rollout: {[data[-1][0].prompt + data[-1][0].response]}, " - f"label: {data[-1][0].label}, reward: {data[-1][0].reward}", - flush=True, - ) - - data = sorted(data, key=lambda group: group[0].index) - return data - - -def generate_rollout_fully_async(args, rollout_id, data_buffer, evaluation=False): - if evaluation: - raise ValueError("Evaluation mode not supported in simple async rollout") - - completed_samples = run(generate_rollout_async(args, rollout_id, data_buffer)) - return completed_samples - - -# Register exit cleanup function - -atexit.register(stop_global_worker) diff --git a/examples/fully_async/run-qwen2.5-0.5B-fully_async.sh b/examples/fully_async/run-qwen2.5-0.5B-fully_async.sh new file mode 100755 index 0000000000..f843f28501 --- /dev/null +++ b/examples/fully_async/run-qwen2.5-0.5B-fully_async.sh @@ -0,0 +1,135 @@ +#!/bin/bash +# Tiny end-to-end fully-async GRPO example using Qwen2.5-0.5B-Instruct on the +# dapo-math-17k dataset. Designed to run on a single 4-GPU node in a few +# minutes — the same script the CI uses for ``test_qwen2.5_0.5B_fully_async_short``. +# +# Prerequisites: +# /root/models/Qwen2.5-0.5B-Instruct/ (HF checkpoint) +# /root/models/Qwen2.5-0.5B-Instruct_torch_dist/ (from tools/convert_hf_to_torch_dist.py) +# /root/datasets/dapo-math-17k/dapo-math-17k.jsonl + +# clean any leftover ray/sglang +pkill -9 sglang 2>/dev/null || true +sleep 3 +ray stop --force 2>/dev/null || true +pkill -9 ray python 2>/dev/null || true +sleep 3 + +set -ex + +export PYTHONBUFFERED=16 + +NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) +HAS_NVLINK=$([ "$NVLINK_COUNT" -gt 0 ] && echo 1 || echo 0) +echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +source "${SCRIPT_DIR}/../../scripts/models/qwen2.5-0.5B.sh" + +MODEL_DIR=${MODEL_DIR:-/root/models/Qwen2.5-0.5B-Instruct} +DATA_PATH=${DATA_PATH:-/root/datasets/dapo-math-17k/dapo-math-17k.jsonl} + +CKPT_ARGS=( + --hf-checkpoint "${MODEL_DIR}" + --ref-load "${MODEL_DIR}_torch_dist" + --save /tmp/slime_fully_async_demo/ + --save-interval 9999 +) + +ROLLOUT_ARGS=( + # ↓↓↓ This is the only knob you need to flip to go fully-async ↓↓↓ + --rollout-function-path slime.rollout.fully_async_rollout.generate_rollout_fully_async + + --prompt-data "${DATA_PATH}" + --input-key prompt + --label-key label + --apply-chat-template + --rollout-shuffle + + --rm-type deepscaler + + --num-rollout 3 + --rollout-batch-size 8 + --n-samples-per-prompt 4 + --rollout-max-response-len 1024 + --rollout-temperature 1 + + --global-batch-size 32 + --balance-data +) + +PERF_ARGS=( + --tensor-model-parallel-size 1 + --sequence-parallel + --pipeline-model-parallel-size 1 + --context-parallel-size 1 + --expert-model-parallel-size 1 + --expert-tensor-parallel-size 1 + + --use-dynamic-batch-size + --max-tokens-per-gpu 4096 +) + +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 +) + +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 + --sglang-mem-fraction-static 0.55 +) + +MISC_ARGS=( + --attention-dropout 0.0 + --hidden-dropout 0.0 + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + --attention-backend flash +) + +# launch the master node of ray in container +NUM_GPUS=${NUM_GPUS:-4} +export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"} +ray start --head --node-ip-address "${MASTER_ADDR}" --num-gpus "${NUM_GPUS}" --disable-usage-stats + +RUNTIME_ENV_JSON="{ + \"env_vars\": { + \"PYTHONPATH\": \"/root/Megatron-LM/:${SCRIPT_DIR}\", + \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", + \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\" + } +}" + +# fully-async splits actor / rollout onto disjoint GPUs (no colocation). +ACTOR_GPUS=${ACTOR_GPUS:-1} +ROLLOUT_GPUS=${ROLLOUT_GPUS:-$((NUM_GPUS - ACTOR_GPUS))} + +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 "${ACTOR_GPUS}" \ + --rollout-num-gpus "${ROLLOUT_GPUS}" \ + ${MODEL_ARGS[@]} \ + "${CKPT_ARGS[@]}" \ + "${ROLLOUT_ARGS[@]}" \ + "${OPTIMIZER_ARGS[@]}" \ + "${GRPO_ARGS[@]}" \ + "${PERF_ARGS[@]}" \ + "${SGLANG_ARGS[@]}" \ + "${MISC_ARGS[@]}" 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 778f58b2ad..0000000000 --- a/examples/fully_async/run-qwen3-4b-fully_async.sh +++ /dev/null @@ -1,139 +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 PYTHONBUFFERED=16 - -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)" -source "${SCRIPT_DIR}/../../scripts/models/qwen3-4B.sh" - -CKPT_ARGS=( - --hf-checkpoint /root/Qwen3-4B - #--hf-checkpoint /root/Qwen3-4B-FP8 - --ref-load /root/Qwen3-4B_torch_dist - --load /root/Qwen3-4B_slime/ - --save /root/Qwen3-4B_slime/ - --save-interval 20 -) - -PROMPT_SET=/path/to/dapo-math-17k.jsonl - -ROLLOUT_ARGS=( - --rollout-function-path fully_async_rollout.generate_rollout_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 -) - -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/:${SCRIPT_DIR}\", - \"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/slime/rollout/fully_async_rollout.py b/slime/rollout/fully_async_rollout.py new file mode 100644 index 0000000000..a54f4083aa --- /dev/null +++ b/slime/rollout/fully_async_rollout.py @@ -0,0 +1,256 @@ +"""Fully-async rollout for slime. + +Decouples ``max_concurrent_tasks`` from ``rollout_batch_size``: a background +asyncio worker keeps a fixed pool of in-flight trajectories across rollout +boundaries, so the next training step doesn't have to wait for the slowest +in-flight sample to finish. + +Use with ``--rollout-function-path slime.rollout.fully_async_rollout.generate_rollout_fully_async``. +Plug in per-sample logic via ``--custom-generate-function-path`` and +per-sample reward via ``--custom-rm-path`` — the worker calls slime's stock +:func:`generate_and_rm_group` which dispatches to those. + +Concurrency is sourced from ``args.sglang_server_concurrency`` and scaled by +the number of sglang engines (``rollout_num_gpus // rollout_num_gpus_per_engine``) +to match the per-sample semaphore cap in :mod:`slime.rollout.sglang_rollout`. + +The worker is intentionally oblivious to slime's higher-level pause / +weight-update signalling (e.g. ``GenerateState.aborted``). Each in-flight +generation short-circuits on those signals on its own and surfaces +:data:`Sample.Status.ABORTED`; the only piece the worker owns is +**redirecting ABORTED groups back to ``data_buffer``** instead of shipping +them to training, so the next rollout (with refreshed weights) can pick +them up. +""" + +from __future__ import annotations + +import asyncio +import atexit +import logging +import queue +import threading +import time + +from slime.rollout.sglang_rollout import GenerateState, generate_and_rm_group +from slime.utils.async_utils import run +from slime.utils.types import Sample + +__all__ = [ + "AsyncRolloutWorker", + "generate_rollout_fully_async", +] + +logger = logging.getLogger("slime.rollout.fully_async") + + +# Global worker, shared across rollout calls so the queue stays warm. +_global_worker: AsyncRolloutWorker | None = None +_worker_lock = threading.Lock() + + +def _get_global_worker(args, data_buffer) -> AsyncRolloutWorker: + global _global_worker + with _worker_lock: + if _global_worker is None or not _global_worker.worker_thread.is_alive(): + logger.info("starting fully-async rollout worker") + num_engines = max(1, args.rollout_num_gpus // args.rollout_num_gpus_per_engine) + _global_worker = AsyncRolloutWorker( + args, data_buffer, concurrency=args.sglang_server_concurrency * num_engines + ) + _global_worker.start() + return _global_worker + + +def _stop_global_worker() -> None: + global _global_worker + with _worker_lock: + if _global_worker is not None: + _global_worker.stop() + _global_worker = None + + +atexit.register(_stop_global_worker) + + +class AsyncRolloutWorker: + """Background thread + asyncio loop that continuously consumes groups + from ``data_buffer`` and runs :func:`generate_and_rm_group` on each.""" + + def __init__(self, args, data_buffer, concurrency: int = 10): + self.args = args + self.data_buffer = data_buffer + self.concurrency = concurrency + self.running = True + self.output_queue: queue.Queue[tuple[int, list[Sample]]] = queue.Queue(maxsize=1000) + self.worker_thread: threading.Thread | None = None + self.state = GenerateState(args) + + # -- public -------------------------------------------------------------- + + def start(self) -> None: + if self.worker_thread is None or not self.worker_thread.is_alive(): + self.worker_thread = threading.Thread(target=self._thread_main, name="fully-async-rollout", daemon=True) + self.worker_thread.start() + + def stop(self) -> None: + self.running = False + if self.worker_thread and self.worker_thread.is_alive(): + self.worker_thread.join(timeout=5) + + def get_completed_groups(self) -> list[tuple[int, list[Sample]]]: + completed: list[tuple[int, list[Sample]]] = [] + while True: + try: + completed.append(self.output_queue.get_nowait()) + except queue.Empty: + break + return completed + + def queue_size(self) -> int: + return self.output_queue.qsize() + + # -- internals ----------------------------------------------------------- + + def _thread_main(self) -> None: + asyncio.run(self._loop()) + + async def _loop(self) -> None: + active_tasks: set[asyncio.Task] = set() + max_concurrent = self.concurrency + gid_counter = 0 + + while self.running: + try: + # Reap done tasks + if active_tasks: + done = {t for t in active_tasks if t.done()} + for t in done: + try: + t.result() # results already handled in callback + except Exception as e: # noqa: BLE001 + logger.warning("fully-async task crashed: %r", e) + active_tasks -= done + + # Top up. + while len(active_tasks) < max_concurrent and self.running: + groups = self.data_buffer.get_samples(1) + if not groups: + break + for group in groups: + gid = gid_counter + gid_counter += 1 + task = asyncio.create_task( + generate_and_rm_group( + self.args, + group, + sampling_params=self.state.sampling_params.copy(), + evaluation=False, + ) + ) + task.add_done_callback(self._make_done_cb(gid)) + active_tasks.add(task) + + await asyncio.sleep(1) + except Exception as e: # noqa: BLE001 + logger.exception("fully-async loop iteration error: %s", e) + await asyncio.sleep(1) + + if active_tasks: + logger.info( + "fully-async: waiting for %d in-flight tasks to drain", + len(active_tasks), + ) + try: + await asyncio.wait(active_tasks, timeout=30) + except Exception: # noqa: BLE001 + pass + + def _make_done_cb(self, gid: int): + def _cb(done_task: asyncio.Task) -> None: + try: + result = done_task.result() + except Exception: # noqa: BLE001 + logger.exception("fully-async: process task raised") + return + if not isinstance(result, list): + logger.warning( + "fully-async: generate_and_rm_group returned %r, expected list[Sample]; dropping", + type(result).__name__, + ) + return + # Aborted group → requeue, don't ship to training. + if any(getattr(s, "status", None) == Sample.Status.ABORTED for s in result): + try: + self.data_buffer.add_samples([result]) + except Exception: # noqa: BLE001 + logger.exception("fully-async: failed to requeue aborted group") + return + self.output_queue.put((gid, result)) + + return _cb + + +async def _generate_rollout_async(args, rollout_id: int, data_buffer) -> list[list[Sample]]: + assert args.rollout_global_dataset + worker = _get_global_worker(args, data_buffer) + + target = args.rollout_batch_size + logger.info( + "fully-async rollout %d: target=%d queue_warm=%d", + rollout_id, + target, + worker.queue_size(), + ) + + collected: dict[int, list[Sample]] = {} + started = time.time() + last_log = started + LOG_EVERY = 30.0 + + while len(collected) < target: + # Pull whatever's done. + drained = 0 + for gid, group in worker.get_completed_groups(): + collected[gid] = group + drained += 1 + + if not drained: + await asyncio.sleep(0.05) + + now = time.time() + if now - last_log > LOG_EVERY: + logger.info( + "fully-async rollout %d: collected %d/%d, queue=%d, elapsed=%.1fs", + rollout_id, + len(collected), + target, + worker.queue_size(), + now - started, + ) + last_log = now + + # Order by sample.index for determinism (slime convention). + def _key(group: list[Sample]) -> int: + for s in group: + idx = getattr(s, "index", None) + if idx is not None: + return int(idx) + return 0 + + out = sorted(collected.values(), key=_key)[:target] + logger.info( + "fully-async rollout %d: done in %.1fs, queue_left=%d", + rollout_id, + time.time() - started, + worker.queue_size(), + ) + return out + + +def generate_rollout_fully_async(args, rollout_id, data_buffer, evaluation: bool = False): + """Slime ``--rollout-function-path`` entrypoint.""" + + if evaluation: + raise ValueError("fully-async rollout doesn't support evaluation mode") + return run(_generate_rollout_async(args, rollout_id, data_buffer)) diff --git a/slime/rollout/sglang_rollout.py b/slime/rollout/sglang_rollout.py index 44c03858ff..c7f86b98ca 100644 --- a/slime/rollout/sglang_rollout.py +++ b/slime/rollout/sglang_rollout.py @@ -309,7 +309,14 @@ async def generate_and_rm( ) async def generate_and_rm_group( args: Namespace, group: list[Sample], sampling_params: dict[str, Any], evaluation: bool = False -) -> list[Sample]: +) -> list[Sample] | list[list[Sample]]: + # ``generate_and_rm`` may return either a ``Sample`` or a ``list[Sample]`` + # depending on whether the ``--custom-generate-function-path`` callable + # emits one trainable sample or several (e.g. multi-turn agent rollouts + # that fan out into multiple prefix-chained samples). The asyncio.gather + # below preserves whichever shape each task produced, so the group is + # ``list[Sample]`` for plain rollouts and ``list[list[Sample]]`` for + # the fan-out case. state = GenerateState(args) if state.aborted: diff --git a/tests/test_qwen2.5_0.5B_fully_async_short.py b/tests/test_qwen2.5_0.5B_fully_async_short.py new file mode 100644 index 0000000000..8c5bcd231b --- /dev/null +++ b/tests/test_qwen2.5_0.5B_fully_async_short.py @@ -0,0 +1,128 @@ +"""CI smoke test for the fully-async rollout path. + +Mirrors ``test_qwen2.5_0.5B_async_short`` (Qwen2.5-0.5B + dapo-math-17k + +3 rollouts of GRPO) but flips the rollout function over to +``slime.rollout.fully_async_rollout.generate_rollout_fully_async`` so the +fully-async worker path gets exercised end-to-end. + +Kept intentionally minimal so it runs in the same time budget as the +existing 0.5B short tests. +""" + +import os + +import slime.utils.external_utils.command_utils as U + +TIGHT_DEVICE_MEMORY = U.get_bool_env_var("SLIME_TEST_TIGHT_DEVICE_MEMORY", "1") + +MODEL_NAME = "Qwen2.5-0.5B-Instruct" +MODEL_TYPE = "qwen2.5-0.5B" +NUM_GPUS = 4 + + +def prepare(): + U.exec_command("mkdir -p /root/models /root/datasets") + U.exec_command(f"hf download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}") + U.hf_download_dataset("zhuzilin/dapo-math-17k") + + +def execute(): + ckpt_args = f"--hf-checkpoint /root/models/{MODEL_NAME}/ " f"--ref-load /root/models/{MODEL_NAME}/ " + + rollout_args = ( + # The only line that differs from test_qwen2.5_0.5B_async_short.py: + # use the public fully-async rollout function. + "--rollout-function-path slime.rollout.fully_async_rollout.generate_rollout_fully_async " + "--prompt-data /root/datasets/dapo-math-17k/dapo-math-17k.jsonl " + "--input-key prompt " + "--label-key label " + "--apply-chat-template " + "--rollout-shuffle " + "--rm-type deepscaler " + "--num-rollout 3 " + "--rollout-batch-size 8 " + "--n-samples-per-prompt 4 " + "--rollout-max-response-len 8192 " + "--rollout-temperature 0.8 " + "--global-batch-size 32 " + "--balance-data " + ) + + perf_args = ( + "--tensor-model-parallel-size 1 " + "--sequence-parallel " + "--pipeline-model-parallel-size 1 " + "--context-parallel-size 1 " + "--expert-model-parallel-size 1 " + "--expert-tensor-parallel-size 1 " + "--use-dynamic-batch-size " + "--max-tokens-per-gpu 4096 " + ) + + 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 " + ) + + 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 " + f"--sglang-mem-fraction-static {0.55 if TIGHT_DEVICE_MEMORY else 0.65} " + "--sglang-cuda-graph-max-bs 32 " + "--sglang-enable-metrics " + ) + + ci_args = "--ci-test " + + misc_args = ( + "--attention-dropout 0.0 " + "--hidden-dropout 0.0 " + "--accumulate-allreduce-grads-in-fp32 " + "--attention-softmax-in-fp32 " + "--attention-backend flash " + "--actor-num-nodes 1 " + "--actor-num-gpus-per-node 1 " + "--rollout-num-gpus 3 " + "--megatron-to-hf-mode bridge " + ) + + train_args = ( + f"{ckpt_args} " + f"{rollout_args} " + f"{optimizer_args} " + f"{grpo_args} " + f"{U.get_default_wandb_args(__file__)} " + f"{perf_args} " + f"{sglang_args} " + f"{ci_args} " + f"{misc_args} " + ) + + U.execute_train( + train_args=train_args, + num_gpus_per_node=NUM_GPUS, + megatron_model_type=MODEL_TYPE, + train_script="train_async.py", + ) + + +if __name__ == "__main__": + prepare() + os.environ.pop("http_proxy", None) + os.environ.pop("https_proxy", None) + os.environ.pop("HTTP_PROXY", None) + os.environ.pop("HTTPS_PROXY", None) + execute()