From 67764c486e7dab5cd25d3d7adea821faf8b7d804 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Tue, 9 Jun 2026 02:04:44 +0000 Subject: [PATCH 1/5] translation: mechanical slime->vime / sglang->vllm cleanup missed in #107 Low-risk follow-up to the #107 sync: finish the mechanical translation and trim vime-invented prose back to slime's brevity. No runtime-behavior changes. - tests, vime/ray/rollout.py, vime/rollout/_fanout_test_helpers.py, vime/utils/dp_schedule.py, scripts/run-minimax-m2.sh, .github/workflows/pr-test.yml: finish slime->vime / sglang->vllm / SLIME_TEST_*->VIME_TEST_* symbol translation left over from #107. - docs/{en,zh}/developer_guide/profiling.md, examples/qwen3-30B-A3B.md, examples/qwen3-4B.md: trim vime-invented sections to mirror slime; keep only genuine vLLM engine differences (profiler needs --vllm-profiler-config, EPLB config, cudagraph-capture-sizes). - docs/{en,zh}/developer_guide/debug.md: replace the (vime-expanded) Ray Distributed Debugger section with a one-line pointer to verl's Ray Debugging Tutorial. - examples/{fully_async,geo3k_vlm,geo3k_vlm_multi_turn}, scripts/run-qwen2.5- 0.5B-reproducibility.sh, vime/utils/external_utils/command_utils.py: drop vime-invented comments/prose. Keep the vLLM-0.22.0 enforce-eager note (Qwen3-VL logprob parity, vllm#43617) -- a real engine requirement. docs/{en,zh}/advanced/speculative-decoding.md intentionally left unchanged. Co-Authored-By: Claude Opus 4.8 (1M context) --- .github/workflows/pr-test.yml | 8 +- docs/en/developer_guide/debug.md | 43 +-- docs/en/developer_guide/profiling.md | 321 ++--------------- docs/en/examples/qwen3-30B-A3B.md | 79 +---- docs/en/examples/qwen3-4B.md | 29 +- docs/zh/developer_guide/debug.md | 43 +-- docs/zh/developer_guide/profiling.md | 323 ++---------------- docs/zh/examples/qwen3-30B-A3B.md | 80 +---- docs/zh/examples/qwen3-4B.md | 29 +- examples/fully_async/README.md | 3 +- examples/geo3k_vlm/README.md | 8 +- examples/geo3k_vlm_multi_turn/env_geo3k.py | 2 +- scripts/run-minimax-m2.sh | 4 +- scripts/run-qwen2.5-0.5B-reproducibility.sh | 3 - tests/test_agent_adapters.py | 46 +-- tests/test_agent_sdk_adapters.py | 14 +- tests/test_cp_utils.py | 2 +- tests/test_loss_cp_invariance.py | 22 +- tests/test_metric_report.py | 2 +- tests/test_metric_report_dist.py | 4 +- tests/test_qwen2.5_0.5B_fanout_short.py | 8 +- ...test_qwen3_4B_streaming_partial_rollout.py | 2 +- vime/ray/rollout.py | 2 +- vime/rollout/_fanout_test_helpers.py | 2 +- vime/utils/dp_schedule.py | 2 +- vime/utils/external_utils/command_utils.py | 10 +- 26 files changed, 169 insertions(+), 922 deletions(-) diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index 0868048c8..0f34222ee 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -556,10 +556,10 @@ jobs: env: GITHUB_COMMIT_NAME: ${{ github.sha }}_${{ github.event.pull_request.number || 'non-pr' }} WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }} - SLIME_TEST_ENABLE_INFINITE_RUN: ${{ (github.event_name == 'workflow_dispatch' && github.event.inputs.infinite_run) || 'false' }} - SLIME_TEST_USE_DEEPEP: ${{ matrix.info.use_deepep || '0' }} - SLIME_TEST_USE_FP8_ROLLOUT: ${{ matrix.info.use_fp8_rollout || '0' }} - SLIME_TEST_ENABLE_EVAL: ${{ matrix.info.enable_eval || '1' }} + VIME_TEST_ENABLE_INFINITE_RUN: ${{ (github.event_name == 'workflow_dispatch' && github.event.inputs.infinite_run) || 'false' }} + VIME_TEST_USE_DEEPEP: ${{ matrix.info.use_deepep || '0' }} + VIME_TEST_USE_FP8_ROLLOUT: ${{ matrix.info.use_fp8_rollout || '0' }} + VIME_TEST_ENABLE_EVAL: ${{ matrix.info.enable_eval || '1' }} steps: - name: Checkout repository diff --git a/docs/en/developer_guide/debug.md b/docs/en/developer_guide/debug.md index 212affd16..193d9a3db 100644 --- a/docs/en/developer_guide/debug.md +++ b/docs/en/developer_guide/debug.md @@ -70,45 +70,4 @@ When running large scale RL, we will occationally meet the IMA in vLLM, there ar ## Step-by-Step Debugging with Ray Distributed Debugger -Ray provides a [distributed debugger](https://docs.ray.io/en/latest/ray-observability/ray-distributed-debugger.html) based on debugpy that lets you set breakpoints in the driver process and step through code interactively. - -1. Install debugpy: - - ```bash - pip install debugpy==1.8.0 - ``` - -2. Enable `RAY_DEBUG_POSTMORTEM` in your launch script: - - ```bash - export RAY_DEBUG_POSTMORTEM=1 - - RUNTIME_ENV_JSON="{ - \"env_vars\": { - ... - \"RAY_DEBUG_POSTMORTEM\": \"${RAY_DEBUG_POSTMORTEM:-0}\" - } - }" - - ray job submit --address="http://127.0.0.1:8265" \ - --runtime-env-json="${RUNTIME_ENV_JSON}" \ - -- python3 train.py [args...] - ``` - -3. Add `ray.init()` before `breakpoint()` in `train.py`: - - ```python - if __name__ == "__main__": - ray.init() - breakpoint() - args = parse_args() - train(args) - ``` - - `ray.init()` is required because the distributed debugger depends on `core_worker`, which is only available after Ray initialization. Without it, `breakpoint()` raises `AttributeError: 'Worker' object has no attribute 'core_worker'`. - -4. Connect via VS Code: - - Install the [Ray Distributed Debugger](https://marketplace.visualstudio.com/items?itemName=ray-project.ray-distributed-debugger) extension in VS Code. Run your launch script to submit the job. Once the job hits `breakpoint()`, open the Ray Dashboard panel in VS Code and click the active breakpoint to attach the debugger. You can then step through code, inspect variables, and set additional breakpoints directly in the editor. - -> **Note**: Remove `ray.init()` and `breakpoint()` after debugging. An explicit `ray.init()` without arguments may cause issues in multi-node training where Ray injects specific namespace and runtime environment configurations via `ray job submit`. +See verl's [Ray Debugging Tutorial](https://verl.readthedocs.io/en/latest/start/ray_debug_tutorial.html). diff --git a/docs/en/developer_guide/profiling.md b/docs/en/developer_guide/profiling.md index 4e7df22e6..51bf9e014 100644 --- a/docs/en/developer_guide/profiling.md +++ b/docs/en/developer_guide/profiling.md @@ -1,24 +1,12 @@ # Profiling -In vime, you can profile the **rollout (vLLM inference)** path in detail using vLLM's profiling HTTP API. Profiling targets the vLLM engine side, not the Megatron training side. +In vime, we can perform detailed performance analysis of the rollout process using the profiling interface provided by vLLM. -Typical flow: +## 1. Sleeping the Rollout Process -- Start train (`sleep_rollout` + `vllm-profiler-config`) -- Wait until vLLM engines and the router are ready -- Read router/worker addresses from logs -- `start_profile` -- Send a few inference requests -- (Optional) `stop_profile`; or traces flush automatically when `max_iterations` is reached -- Inspect trace files under `torch_profiler_dir` +For more flexible stress testing and profiling, it is often useful to make the vime rollout process enter a waiting state after initialization, instead of starting generation immediately. - - -## 1. Put Rollout into a Wait State (`sleep_rollout`) - -For flexible stress testing and profiling, rollout usually waits after initialization instead of generating immediately. - -Replace `rollout_function_path` in `train.py` startup args—no code changes required: +You can achieve this by replacing the `rollout_function_path` in your startup arguments without modifying the source code: ```bash python train.py \ @@ -26,305 +14,60 @@ python train.py \ ... (other arguments) ``` -This puts the rollout process in an infinite wait loop so you can send HTTP requests or run stress tools manually. - -## 2. Enable the vLLM Profiler (at train startup) +This function will make the rollout process enter an infinite wait loop, allowing you to manually send requests or run stress testing tools. -vLLM registers `/start_profile` and `/stop_profile` only when started with `--profiler-config`. In vime, pass **`--vllm-profiler-config`** through to the `vllm serve` subprocess. +## 2. Enabling the vLLM Profiler -### 2.1 Pass the full config as JSON +vLLM only registers the `/start_profile` and `/stop_profile` endpoints when started with a profiler config. In vime, pass it through to the `vllm serve` subprocess with `--vllm-profiler-config`: ```bash --vllm-profiler-config '{"profiler":"torch","torch_profiler_dir":"/root/logs/vllm_profile","max_iterations":3,"ignore_frontend":true}' ``` -Common JSON fields: - -| Field | Description | -|------|------| -| `profiler` | `"torch"` or `"cuda"` | -| `torch_profiler_dir` | Trace output directory (absolute path) | -| `max_iterations` | Worker auto-stops and flushes after more than N steps (condition is `> N`) | -| `ignore_frontend` | Recommended `true`: profile workers only, lower frontend overhead | - -**Avoid RPC timeout on `stop_profile`:** vLLM APIServer talks to EngineCore/workers over internal RPC. Manually calling `stop_profile` to flush traces can take minutes, while the default `VLLM_RPC_TIMEOUT` is only **10 seconds** (10000 ms), which can interrupt flush or leave traces incomplete. For profiling, set **30 minutes** (1800000 ms). - -Set this variable **before starting train and launching vLLM**, in the Ray worker environment (a local shell `export` may not reach the Ray job). Pass it via `runtime-env-json` on `ray job submit`, for example: - -```bash -export VLLM_RPC_TIMEOUT="${VLLM_RPC_TIMEOUT:-1800000}" - -RUNTIME_ENV_JSON="{ - \"env_vars\": { - \"PYTHONPATH\": \"/root/vime:/root/Megatron-LM\", - \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", - \"VLLM_RPC_TIMEOUT\": \"${VLLM_RPC_TIMEOUT}\" - } -}" - -ray job submit --address=\"http://127.0.0.1:8265\" \ - --runtime-env-json=\"${RUNTIME_ENV_JSON}\" \ - -- python3 train.py \ - ... \ - --vllm-profiler-config '{\"profiler\":\"torch\",\"torch_profiler_dir\":\"/root/logs/vllm_profile\",...}' -``` - - -### 2.2 Verify it took effect - -After train starts, confirm all three in logs (missing any means the profiler is not enabled correctly): - -1. **Args parsed**: `vllm_profiler_config ... profiler='torch'` (and `torch_profiler_dir` path). -2. **Forwarded to vLLM subprocess**: `Launching vLLM server: ... --profiler-config {"profiler":"torch",...}`. -3. **HTTP routes registered**: vLLM startup route list includes `/start_profile` and `/stop_profile` (otherwise `POST /start_profile` returns 404). - -## 3. Get Router and Worker Addresses - -vLLM engines (workers) register on the vllm-router. Example startup log: - -```text -Router launched at 127.0.0.1:3521, Prometheus port: 4153 -Ports for engine 0: {'host': '127.0.0.1', 'port': 15000, ...} -Starting vLLM server on http://127.0.0.1:15000 -``` - -**Note: the router port may change on every job** (random in 3000–4000 by default). Do not reuse the previous port. Verify with curl: - -```bash -curl http://127.0.0.1:3521/workers -``` - -Returns each worker's `url` and `is_healthy`. +**Key Fields:** +* `profiler`: `"torch"` or `"cuda"`. +* `torch_profiler_dir`: Trace output directory (absolute path). +* `max_iterations`: Worker auto-stops and flushes after this many steps. +* `ignore_frontend`: Recommended `true`; profile workers only. -## 4. Use `tools/profile_rollout.py` +## 3. Obtaining vLLM Engine List -The script reads the router's `/workers` list and calls `/start_profile` or `/stop_profile` on every worker. +vLLM engines (workers) are registered with the router. You can retrieve the list of all active engines by accessing the `/workers` endpoint of the router. -### Start Profiling - -```bash -cd /root/vime -python tools/profile_rollout.py \ - --router-url http://127.0.0.1:3521 \ - --action start +The router address is typically printed in the startup logs: ``` - -### Stop Profiling (optional) - -If `--vllm-profiler-config` sets `max_iterations`, the worker **auto-stops and flushes** after enough steps. In practice, traces often appear under `torch_profiler_dir` right after inference—you **do not** need to call `stop_profile` manually. Use this only to end collection early: - -```bash -python tools/profile_rollout.py \ - --router-url http://127.0.0.1:3521 \ - --action stop +Router launched at 127.0.0.1:3000 ``` -## 5. Send Inference Requests - -While `sleep_rollout` is waiting: - -1. `profile_rollout.py --action start` -2. Send a few completion requests to the router or **directly to a worker** (2–4 is enough; traces get large) -3. (Optional) `profile_rollout.py --action stop`; or wait for `max_iterations` to auto-flush -4. Inspect traces under `torch_profiler_dir` - -Example request (`model` is the HF checkpoint path): - +You can use `curl` to view the workers: ```bash -curl -X POST http://127.0.0.1:15000/v1/completions \ - -H "Content-Type: application/json" \ - -d '{"model":"/root/models/Qwen3-4B","prompt":"Hello","max_tokens":32}' +curl http://127.0.0.1:3000/workers ``` +## 4. Using Automated Profiling Tool -## 6. View Traces - -### Perfetto - -1. Open [https://ui.perfetto.dev/](https://ui.perfetto.dev/) -2. **Open trace file**, pick `*.trace.json.gz` -3. Inspect GPU kernels, CPU ops, and the timeline - -### Chrome Tracing +To simplify profiling across multiple engines simultaneously, we provide an automated script: `tools/profile_rollout.py`. -Open `chrome://tracing` in the browser and **Load** a trace file. +### Starting Profiling -### Analysis Tool +By default, this tool starts profiling on all workers: ```bash -cd /root/vime -python tools/analyze_profile.py --profile-dir /root/logs/vllm_profile --all-ranks +python tools/profile_rollout.py --router-url http://127.0.0.1:3000 --action start ``` +### Stopping Profiling Manually -## 7. Troubleshooting - -| Symptom | Fix | -|------|------| -| `POST /start_profile` 404 | Pass `--vllm-profiler-config` as JSON; restart the job | -| Start OK but empty output dir | Confirm curl hits a worker and returns 200; increase `max_iterations` or send more requests | -| Router 503 | Confirm the current job's router port; connect directly to a worker | -| Slow or timed-out stop | Increase `VLLM_RPC_TIMEOUT`; reduce request count | - -## 8. Full Runnable Example - -The script below assumes a **container** environment, vime at `/root/vime`, models and data under `/root/models` and `/root/data`. Two parts: - -1. **`launch_train_for_profiling`**: start train with the profiler (`sleep_rollout`, minimal single-GPU colocate example—adjust GPU layout for your machine). -2. **`run_profiling_session`**: run profiling from **another terminal** after train is ready. - -Save as `/root/vime/run_profiling_demo.sh` and run: +If you set `max_iterations`, the worker auto-stops and flushes. To stop early: ```bash -#!/usr/bin/env bash -# -# Full vime rollout profiling example -# Usage: -# bash /root/vime/run_profiling_demo.sh launch # terminal 1: start train -# bash /root/vime/run_profiling_demo.sh profile # terminal 2: capture traces after train is ready -# -set -euo pipefail - -VIME_ROOT="${VIME_ROOT:-/root/vime}" -HF_CKPT="${HF_CKPT:-/root/models/Qwen3-4B}" -REF_LOAD="${REF_LOAD:-/root/models/Qwen3-4B_torch_dist}" -PROMPT_DATA="${PROMPT_DATA:-/root/data/gsm8k/train.parquet}" -LOG_ROOT="${LOG_ROOT:-/root/logs/vime_profiling}" -PROFILE_DIR="${PROFILE_DIR:-/root/logs/vllm_profile}" -TRAIN_LOG="${LOG_ROOT}/train_profiling.log" -ROUTER_HOST="${ROUTER_HOST:-127.0.0.1}" - -mkdir -p "${LOG_ROOT}" "${PROFILE_DIR}" - -VLLM_PROFILER_CONFIG_JSON="$(printf \ - '{"profiler":"torch","torch_profiler_dir":"%s","max_iterations":3,"ignore_frontend":true}' \ - "${PROFILE_DIR}")" - -launch_train_for_profiling() { - cd "${VIME_ROOT}" - - # Clean up old Ray / vLLM processes (comment out if not needed) - ray stop --force || true - pkill -9 -f "vllm serve" || true - sleep 2 - - ray start --head --node-ip-address 127.0.0.1 --num-gpus 2 --disable-usage-stats - - source "${VIME_ROOT}/scripts/models/qwen3-4B.sh" - - export VLLM_RPC_TIMEOUT="${VLLM_RPC_TIMEOUT:-1800000}" - - RUNTIME_ENV_JSON="{ - \"env_vars\": { - \"PYTHONPATH\": \"${VIME_ROOT}:/root/Megatron-LM\", - \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", - \"VLLM_RPC_TIMEOUT\": \"${VLLM_RPC_TIMEOUT}\" - } - }" - - echo "=== Launching train; log: ${TRAIN_LOG} ===" - echo "=== After engines are up, run: bash $0 profile ===" - - ray job submit --address="http://127.0.0.1:8265" \ - --runtime-env-json="${RUNTIME_ENV_JSON}" \ - -- python3 train.py \ - --train-backend megatron \ - --colocate \ - --actor-num-nodes 1 \ - --actor-num-gpus-per-node 1 \ - --rollout-num-gpus 1 \ - --rollout-num-gpus-per-engine 1 \ - --rollout-backend vllm \ - --rollout-function-path vime.rollout.sleep_rollout.sleep \ - --hf-checkpoint "${HF_CKPT}" \ - --ref-load "${REF_LOAD}" \ - --prompt-data "${PROMPT_DATA}" \ - --input-key question \ - --label-key label \ - --apply-chat-template \ - --rm-type deepscaler \ - --num-rollout 1 \ - --rollout-batch-size 4 \ - --n-samples-per-prompt 1 \ - --rollout-max-response-len 512 \ - --global-batch-size 4 \ - --vllm-gpu-memory-utilization 0.7 \ - --vllm-profiler-config "${VLLM_PROFILER_CONFIG_JSON}" \ - ${MODEL_ARGS[@]} \ - 2>&1 | tee "${TRAIN_LOG}" -} - -discover_router_url() { - local line port - line="$(grep -E 'Router launched at' "${TRAIN_LOG}" | tail -1 || true)" - if [[ -z "${line}" ]]; then - echo "ERROR: Router not found in ${TRAIN_LOG}. Is train still starting?" >&2 - exit 1 - fi - # Router launched at 127.0.0.1:3521, Prometheus port: ... - port="$(echo "${line}" | sed -n 's/.*Router launched at [^:]*:\([0-9]*\).*/\1/p')" - echo "http://${ROUTER_HOST}:${port}" -} - -discover_worker_url() { - local router_url="$1" - python3 - <<'PY' "${router_url}" -import json, sys, urllib.request -router = sys.argv[1] -with urllib.request.urlopen(f"{router}/workers", timeout=10) as r: - workers = json.load(r).get("workers", []) -if not workers: - raise SystemExit("No workers registered") -print(workers[0]["url"]) -PY -} - -run_profiling_session() { - cd "${VIME_ROOT}" - - local router_url worker_url model="${HF_CKPT}" - router_url="$(discover_router_url)" - worker_url="$(discover_worker_url "${router_url}")" - - echo "=== ROUTER=${router_url} WORKER=${worker_url} PROFILE_DIR=${PROFILE_DIR} ===" - - echo "=== 1/3 start_profile (all workers via router) ===" - python tools/profile_rollout.py --router-url "${router_url}" --action start - - echo "=== 2/3 send completions (direct to worker; 3 requests) ===" - for i in 1 2 3; do - curl -sS -X POST "${worker_url}/v1/completions" \ - -H "Content-Type: application/json" \ - -d "{\"model\":\"${model}\",\"prompt\":\"Hello ${i}\",\"max_tokens\":32}" \ - | head -c 400 - echo - done - - echo "=== 3/3 list trace files (max_iterations=3 auto-stop; add --action stop if needed) ===" - sleep 2 - find "${PROFILE_DIR}" -type f \( -name '*.json*' -o -name 'profiler_out_*' \) | sort - echo "Open *.trace.json.gz in https://ui.perfetto.dev/ or run:" - echo " python tools/analyze_profile.py --profile-dir ${PROFILE_DIR} --all-ranks" -} - -case "${1:-}" in - launch) launch_train_for_profiling ;; - profile) run_profiling_session ;; - *) - echo "Usage: $0 {launch|profile}" >&2 - exit 1 - ;; -esac +python tools/profile_rollout.py --router-url http://127.0.0.1:3000 --action stop ``` -**Steps:** - -```bash -# Terminal 1: start train (wait until logs show Router launched at ...) -bash /root/vime/run_profiling_demo.sh launch - -# Terminal 2: capture traces -bash /root/vime/run_profiling_demo.sh profile -``` +## 5. Running Stress Tests -Adjust paths at the top of the script (`/root/models/...`, `/root/data/...`) and GPU layout (`actor-num-gpus-per-node`, `rollout-num-gpus`, etc.) as needed. +While the Rollout process is in a waiting state via `sleep_rollout`, you can: +1. Start profiling using `tools/profile_rollout.py`. +2. Use stress testing tools to send requests to the router or directly to the engines. +3. Wait for profiling to complete (if `max_iterations` was set) or stop it manually. +4. Collect the `.json` trace files from the `torch_profiler_dir` and view them using `chrome://tracing` in Chrome or [Perfetto](https://ui.perfetto.dev/). diff --git a/docs/en/examples/qwen3-30B-A3B.md b/docs/en/examples/qwen3-30B-A3B.md index bca56238d..9fd7ebd2b 100644 --- a/docs/en/examples/qwen3-30B-A3B.md +++ b/docs/en/examples/qwen3-30B-A3B.md @@ -74,82 +74,15 @@ Here, we will briefly introduce the MoE-related parts in the [run-qwen3-30B-A3B. ### Multi-Node Support -The following uses **two machines with 8 GPUs each (16 GPUs total)** as the starting example; scripts and parameters scale to **N nodes**. Key differences from single-node: +For a multi-node environment, the following modifications are necessary: -- Place weights, checkpoints, and data on storage visible to every node (e.g. NFS). -- Set `MASTER_ADDR` to the head **LAN IP** (not `127.0.0.1`). -- Omit CPU Adam (multi-node uses a distributed optimizer; do not use `--optimizer-cpu-offload`). -- `global-batch-size` must equal `rollout-batch-size × n-samples-per-prompt`. + - Place the training model and data on a path accessible by all nodes. + - Set the `MASTER_ADDR` to an address that is accessible by all nodes. + - Remove configurations related to CPU Adam. This is because a distributed optimizer is used, which significantly reduces the optimizer's video memory (VRAM) usage in a multi-node setup. -#### Topology +In addition, you can make the following changes: -| Component | Dual-node defaults | -|-----------|-------------------| -| Cluster | `ACTOR_NUM_NODES=2`, `ACTOR_NUM_GPUS_PER_NODE=8` | -| Megatron training | TP=8, EP=8, CP=2 (experts sharded across nodes) | -| vLLM rollout | Cross-node TP=16 (`rollout-num-gpus-per-engine = nodes × GPUs per node`) | -| Scheduling | Ray cluster + `--colocate` mode | - -Convert checkpoints with Megatron parallelism matching training (dual-node: TP=8, EP=8). Checkpoint EP must match `--expert-model-parallel-size`, or `load_checkpoint` may hang or resharding may be extremely slow. - -#### Start the Ray Cluster - -Start Ray **outside** the training script on each node. Join all workers first; verify `ray status` reports the expected GPU count, then submit training from the head. Dual-node example: - -```bash -# === Head node === -export MASTER_ADDR= -ray start --head --node-ip-address="${MASTER_ADDR}" --num-gpus 8 --disable-usage-stats \ - --dashboard-host=0.0.0.0 --dashboard-port=8265 - -# === Each worker node === -export MASTER_ADDR= -ray start --address="${MASTER_ADDR}:6379" --node-ip-address= --num-gpus 8 -``` - -See [Quick Start — Multi-node training](../get_started/quick_start.md#multi-node-training-for-large-scale-moe-models) for more details. - -#### Run Training - -After the Ray cluster is ready, on the **head node** set multi-node env vars and run the **same script as single-node** (`ACTOR_NUM_NODES>1` skips Ray startup and applies multi-node defaults): - -```bash -export MASTER_ADDR= -export ACTOR_NUM_NODES=2 -export ACTOR_NUM_GPUS_PER_NODE=8 -cd /root/vime -bash scripts/run-qwen3-30B-A3B.sh -``` - -2-step smoke test: - -```bash -NUM_ROLLOUT=2 ENABLE_R3=0 bash scripts/run-qwen3-30B-A3B.sh -``` - -To scale to N nodes (e.g. 4×8), join all workers to Ray, set `ACTOR_NUM_NODES=4` on the head, and tune `MEGATRON_TP` / `MEGATRON_EP` / `MEGATRON_CP` / `ROLLOUT_NUM_GPUS_PER_ENGINE` for total GPU count. - -#### Key Multi-Node Parameters - -| Variable | Dual-node default | Description | -|----------|-------------------|-------------| -| `ACTOR_NUM_NODES` | 2 (default 1 for single-node) | Total nodes including head; script skips Ray startup when >1 | -| `ACTOR_NUM_GPUS_PER_NODE` | 8 | GPUs per node | -| `MEGATRON_TP` / `MEGATRON_EP` / `MEGATRON_CP` | 8 / 8 / 2 | Megatron parallelism | -| `ROLLOUT_NUM_GPUS_PER_ENGINE` | total GPUs | vLLM engine GPU count | -| `ENABLE_R3` | 1 | set to 0 to disable R3 | - -Default batch: `rollout-batch-size=4`, `n-samples-per-prompt=2`, `global-batch-size=8`; vLLM uses `--vllm-moe-backend triton`. - -#### Multi-Node Troubleshooting - -- **Worker cannot join Ray / NCCL failures**: check `MASTER_ADDR`, container `/etc/hosts` (hostname must not map to `127.0.0.1`), `NCCL_SOCKET_IFNAME` / `GLOO_SOCKET_IFNAME`. -- **`Not enough samples X for global_batch_size Y`**: keep `global-batch-size` equal to `rollout-batch-size × n-samples-per-prompt`. -- **GPU memory full but no processes**: restart the container or run `ray stop --force` to clear stale vLLM contexts. - -#### EPLB - -When the total number of GPUs is not a multiple or divisor of the total number of experts, enable vLLM's EPLB (Expert Parallelism Load Balancer) and configure redundant experts via `--vllm-eplb-config`. For example, in a 24-GPU scenario: + - When the total number of GPUs is not a multiple or divisor of the total number of experts, you can enable vLLM's EPLB (Expert Parallelism Load Balancer) and configure redundant experts via `--vllm-eplb-config` to add redundant experts. For example, in a 24-GPU scenario, you can configure it as follows: ```bash VLLM_ARGS=( diff --git a/docs/en/examples/qwen3-4B.md b/docs/en/examples/qwen3-4B.md index ddb6151b2..3068ea221 100644 --- a/docs/en/examples/qwen3-4B.md +++ b/docs/en/examples/qwen3-4B.md @@ -193,7 +193,7 @@ OPTIMIZER_ARGS=( #### VLLM\_ARGS -Parameters for vLLM inference. vime uses vLLM as the rollout backend by default (`rollout.py` launches `VLLMEngine`; the default rollout function is `vime.rollout.vllm_rollout.generate_rollout`), so no extra backend flag is needed. `--rollout-num-gpus-per-engine` corresponds to each vLLM engine's `tensor_parallel_size`. Other vLLM parameters are passed to vime with a `--vllm-` prefix (for example, `--vllm-max-model-len`). +These are the parameters required by vLLM. Here, `--rollout-num-gpus-per-engine` basically corresponds to vLLM's `tensor_parallel_size`. Other vLLM parameters are passed to vime by adding the `--vllm-` prefix. ```bash VLLM_ARGS=( @@ -202,9 +202,7 @@ VLLM_ARGS=( ) ``` -When rollout concurrency is high, tune the vLLM scheduler via the `--vllm-` prefix—for example, `--vllm-max-num-seqs` and `--vllm-max-num-batched-tokens`. Add `--vllm-enforce-eager` for debugging or to work around CUDA graph limits. - -⚠️ vime uses the vLLM router to schedule multiple vLLM servers. With co-located training and inference (`--colocate`), weights are synchronized via CUDA IPC; with decoupled training and inference, the trainer synchronizes weights with vLLM engines over NCCL. +⚠️ vime uses the vLLM router to schedule multiple vLLM servers. `dp_size` is not supported when DP attention is disabled. ### Dynamic Sampling @@ -278,22 +276,21 @@ ray job submit ... \ ... ``` -In this case, 2 GPUs will be allocated for training, and 6 GPUs will be allocated for inference. Like `--actor-num-gpus-per-node`, `--rollout-num-gpus` is a **Ray resource argument** passed to `train.py`: the framework uses it to build the placement group and assign the first bundles to training actors and the remaining bundles to rollout engines (see `vime/ray/placement_group.py`). **Under co-located mode (`--colocate`), this argument is ignored** and is set automatically to `actor_num_gpus_per_node * actor_num_nodes`. Do not put `--rollout-num-gpus` in `VLLM_ARGS`. +In this case, 2 GPUs will be allocated for training, and 6 GPUs will be allocated for inference. -For decoupled training and inference, `VLLM_ARGS` only needs inference-backend settings, for example: +⚠️ If the concurrency on each vLLM server is too high, it may exceed vLLM's default CUDA graph concurrency limit, which will affect inference speed. You can adjust this in the following two ways: -```bash -VLLM_ARGS=( - --rollout-num-gpus-per-engine 2 - --vllm-gpu-memory-utilization 0.9 - --vllm-max-num-seqs 256 - --vllm-max-num-batched-tokens 8192 -) -``` +1. Use `--vllm-server-concurrency` to limit the maximum number of concurrent requests sent to a single vLLM server. For example: + + ```bash + --vllm-server-concurrency 160 + ``` -Add `--vllm-enforce-eager` when debugging or to work around CUDA graph limits. +2. Use `--vllm-cudagraph-capture-sizes` to increase the number of CUDA graphs initialized by vLLM. For example: -⚠️ When using co-located training and inference, Megatron will always occupy some GPU memory. Reduce vLLM's memory footprint with `--vllm-gpu-memory-utilization`, and reserve headroom for training with `--train-memory-margin-bytes`. + ```bash + --vllm-cudagraph-capture-sizes 1 2 4 8 $(seq 16 8 256) + ``` ### Asynchronous Training diff --git a/docs/zh/developer_guide/debug.md b/docs/zh/developer_guide/debug.md index c41a3a95c..21148b038 100644 --- a/docs/zh/developer_guide/debug.md +++ b/docs/zh/developer_guide/debug.md @@ -68,45 +68,4 @@ vime 支持将训练部分和推理部分分开进行调试，从而实现： ## 使用 Ray Distributed Debugger 单步调试 -Ray 提供了基于 debugpy 的[分布式调试器](https://docs.ray.io/en/latest/ray-observability/ray-distributed-debugger.html)，可以在 driver 进程中设置断点并单步执行代码。 - -1. 安装 debugpy： - - ```bash - pip install debugpy==1.8.0 - ``` - -2. 在启动脚本中启用 `RAY_DEBUG_POSTMORTEM`： - - ```bash - export RAY_DEBUG_POSTMORTEM=1 - - RUNTIME_ENV_JSON="{ - \"env_vars\": { - ... - \"RAY_DEBUG_POSTMORTEM\": \"${RAY_DEBUG_POSTMORTEM:-0}\" - } - }" - - ray job submit --address="http://127.0.0.1:8265" \ - --runtime-env-json="${RUNTIME_ENV_JSON}" \ - -- python3 train.py [args...] - ``` - -3. 在 `train.py` 中的 `breakpoint()` 前添加 `ray.init()`： - - ```python - if __name__ == "__main__": - ray.init() - breakpoint() - args = parse_args() - train(args) - ``` - - 必须先调用 `ray.init()`，因为分布式调试器依赖 `core_worker`，只有 Ray 初始化后才可用。否则 `breakpoint()` 会报错 `AttributeError: 'Worker' object has no attribute 'core_worker'`。 - -4. 通过 VS Code 连接调试器： - - 在 VS Code 中安装 [Ray Distributed Debugger](https://marketplace.visualstudio.com/items?itemName=ray-project.ray-distributed-debugger) 扩展。运行启动脚本提交 job 后，当 job 执行到 `breakpoint()` 暂停后，在 VS Code 的 Ray Dashboard 面板中点击活跃的断点即可 attach 调试器，之后可以直接在编辑器中单步执行、查看变量、设置新断点。 - -> **注意**：调试完成后务必移除 `ray.init()` 和 `breakpoint()`。不带参数的 `ray.init()` 在多节点训练中可能导致问题，因为 `ray job submit` 会注入特定的 namespace 和 runtime environment 配置。 +请参考 verl 的 [Ray Debugging Tutorial](https://verl.readthedocs.io/en/latest/start/ray_debug_tutorial.html)。 diff --git a/docs/zh/developer_guide/profiling.md b/docs/zh/developer_guide/profiling.md index 2717c2218..c74fbff27 100644 --- a/docs/zh/developer_guide/profiling.md +++ b/docs/zh/developer_guide/profiling.md @@ -1,24 +1,12 @@ -# 性能分析（Profiling） +# 性能分析 (Profiling) -在vime中，我们可以通过vLLM提供的profiling接口对**rollout（vLLM推理）**过程做详细的性能分析。Profiling针对vLLM engine侧，不是Megatron训练侧。 +在 vime 中，我们可以通过 vLLM 提供的 profiling 接口对 rollout 过程进行详细的性能分析。 -典型流程： +## 1. 使 Rollout 进程进入等待状态 (Sleep Rollout) -- 启动train（sleep_rollout + vllm-profiler-config） -- 等待vLLM engine与router就绪 -- 从日志确认router/worker地址 -- start_profile -- 发送少量推理请求 --（可选）stop_profile；或达到max_iterations后自动落盘 -- 在torch_profiler_dir查看trace文件 +为了更自由地进行压力测试和性能分析，我们通常需要让 vime 的 rollout 进程在初始化完成后进入等待状态，而不是立即开始生成。 - - -## 1. 使Rollout进入等待状态（sleep_rollout） - -为了更灵活地压测和profiling，通常让rollout在初始化完成后进入等待，而不是立即开始生成。 - -在 `train.py` 启动参数中替换 `rollout_function_path` 即可，无需改代码： +你可以通过在启动参数中替换 `rollout_function_path` 来实现，而无需修改代码： ```bash python train.py \ @@ -26,305 +14,60 @@ python train.py \ ... (其他参数) ``` -该函数会让rollout进程进入无限循环等待，便于手动发HTTP请求或运行压测工具。 - -## 2. 启用vLLM Profiler（启动train时配置） +该函数会让 rollout 进程进入无限循环等待状态，方便你手动发送请求或运行压测工具。 -vLLM只有在启动时配置了`--profiler-config`，才会注册`/start_profile`与`/stop_profile`路由。在vime中通过**`--vllm-profiler-config`**转发给`vllm serve`子进程。 +## 2. 启用 vLLM Profiler -### 2.1 使用JSON整包传参 +vLLM 只有在启动时配置了 profiler config 才会注册 `/start_profile` 与 `/stop_profile` 接口。在 vime 中通过 `--vllm-profiler-config` 转发给 `vllm serve` 子进程： ```bash --vllm-profiler-config '{"profiler":"torch","torch_profiler_dir":"/root/logs/vllm_profile","max_iterations":3,"ignore_frontend":true}' ``` -常用JSON字段： - -| 字段 | 说明 | -|------|------| -| `profiler` | `"torch"` 或 `"cuda"` | -| `torch_profiler_dir` | trace输出目录（绝对路径） | -| `max_iterations` | worker记录超过N步后自动stop并落盘（条件为`> N`） | -| `ignore_frontend` | 建议`true`，仅profile worker，降低前端开销 | - -**防止`stop_profile`时RPC超时：** vLLM APIServer与EngineCore/worker之间通过内部RPC通信。手动调用`stop_profile`触发trace落盘可能耗时数分钟，而默认`VLLM_RPC_TIMEOUT`仅**10秒**（10000 ms），容易导致flush中断或trace不完整。Profiling时建议设为**30分钟**（1800000 ms）。 - -该变量须在**启动train、拉起vLLM之前**传入Ray worker环境（仅在本机shell `export`不一定会进入Ray job）。在`ray job submit`的`runtime-env-json`中写入，例如： - -```bash -export VLLM_RPC_TIMEOUT="${VLLM_RPC_TIMEOUT:-1800000}" - -RUNTIME_ENV_JSON="{ - \"env_vars\": { - \"PYTHONPATH\": \"/root/vime:/root/Megatron-LM\", - \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", - \"VLLM_RPC_TIMEOUT\": \"${VLLM_RPC_TIMEOUT}\" - } -}" - -ray job submit --address=\"http://127.0.0.1:8265\" \ - --runtime-env-json=\"${RUNTIME_ENV_JSON}\" \ - -- python3 train.py \ - ... \ - --vllm-profiler-config '{\"profiler\":\"torch\",\"torch_profiler_dir\":\"/root/logs/vllm_profile\",...}' -``` - - -### 2.2 验证是否生效 - -启动train后，在日志中确认以下三点（缺任一项说明profiler未正确启用）： - -1. **参数已解析**：出现`vllm_profiler_config ... profiler='torch'`（及`torch_profiler_dir`路径）。 -2. **已转发给vLLM子进程**：出现`Launching vLLM server: ... --profiler-config {"profiler":"torch",...}`。 -3. **HTTP路由已注册**：vLLM启动时的路由列表中包含`/start_profile`与`/stop_profile`（否则`POST /start_profile`会返回404）。 - -## 3. 获取Router与Worker地址 - -vLLM engine（workers）注册在vllm-router上。启动日志示例： - -```text -Router launched at 127.0.0.1:3521, Prometheus port: 4153 -Ports for engine 0: {'host': '127.0.0.1', 'port': 15000, ...} -Starting vLLM server on http://127.0.0.1:15000 -``` - -**注意：router端口每次job可能变化**（默认在3000–4000随机），不要沿用上次端口。可用curl验证： - -```bash -curl http://127.0.0.1:3521/workers -``` - -返回每个worker的`url`与`is_healthy`。 +**常用字段说明：** +* `profiler`: `"torch"` 或 `"cuda"`。 +* `torch_profiler_dir`: trace 输出目录（绝对路径）。 +* `max_iterations`: worker 记录该步数后自动 stop 并落盘。 +* `ignore_frontend`: 建议 `true`，仅 profile worker。 -## 4. 使用`tools/profile_rollout.py` +## 3. 获取 vLLM 引擎列表 -脚本通过router的`/workers`列表，对所有worker调用`/start_profile`或`/stop_profile`。 +vLLM 引擎（workers）注册在 router 上。你可以通过访问 router 的 `/workers` 接口来获取所有活跃引擎的列表。 -### 启动Profiling - -```bash -cd /root/vime -python tools/profile_rollout.py \ - --router-url http://127.0.0.1:3521 \ - --action start +通常 router 地址会在启动日志中打印： ``` - -### 停止Profiling（可选） - -若在`--vllm-profiler-config`中设置了`max_iterations`，worker在记录足够步数后会**自动stop并落盘**，实践中发完推理后常可直接在`torch_profiler_dir`看到trace，**不必**再手动`stop_profile`。需要提前结束采集时再执行： - -```bash -python tools/profile_rollout.py \ - --router-url http://127.0.0.1:3521 \ - --action stop +Router launched at 127.0.0.1:3000 ``` -## 5. 发送推理请求 - -在sleep_rollout等待期间，执行步骤如下： - -1. `profile_rollout.py --action start` -2. 向router或**直连worker**发送少量completion请求（2～4条即可，trace会很大） -3. （可选）`profile_rollout.py --action stop`；或等待`max_iterations`触发自动落盘 -4. 在`torch_profiler_dir`查看trace - -请求示例（`model`使用HF checkpoint路径）： - +你可以使用 `curl` 查看 workers： ```bash -curl -X POST http://127.0.0.1:15000/v1/completions \ - -H "Content-Type: application/json" \ - -d '{"model":"/root/models/Qwen3-4B","prompt":"Hello","max_tokens":32}' +curl http://127.0.0.1:3000/workers ``` +## 4. 使用自动化 Profiling 工具 -## 6. 查看Trace - -### Perfetto - -1. 打开 [https://ui.perfetto.dev/](https://ui.perfetto.dev/) -2. **Open trace file**，选择`*.trace.json.gz` -3. 查看GPU kernel、CPU算子与时间线 - -### Chrome Tracing +为了简化对多个引擎同时进行 profiling 的操作，我们提供了一个自动化脚本 `tools/profile_rollout.py`。 -浏览器访问`chrome://tracing`，Load加载trace文件。 +### 启动 Profiling -### 分析工具 +默认情况下，该工具会对所有 worker 启动 profiling： ```bash -cd /root/vime -python tools/analyze_profile.py --profile-dir /root/logs/vllm_profile --all-ranks +python tools/profile_rollout.py --router-url http://127.0.0.1:3000 --action start ``` +### 手动停止 Profiling -## 7. 常见问题 - -| 现象 | 处理 | -|------|------| -| `POST /start_profile` 404 | 用JSON传`--vllm-profiler-config`；重启job | -| start成功但目录为空 | 确认curl打到worker且返回200；适当增大`max_iterations`或补发推理 | -| router 503 | 确认当前job的router端口；改直连worker | -| stop很慢或超时 | 增大`VLLM_RPC_TIMEOUT`；减少请求条数 | - -## 8. 完整可运行示例 - -以下脚本假设在**容器内**、vime仓库位于`/root/vime`，模型与数据在`/root/models`、`/root/data`。分两段： - -1. **`launch_train_for_profiling`**：启动带profiler的train（sleep_rollout，单卡colocate最小示例，可按机器改GPU数）。 -2. **`run_profiling_session`**：train就绪后，在**另一个终端**执行profiling。 - -将脚本保存为 `/root/vime/run_profiling_demo.sh` 后执行。 +如果你设置了 `max_iterations`，worker 会自动 stop 并落盘。想要提前停止： ```bash -#!/usr/bin/env bash -# -# vime rollout profiling 完整示例 -# 用法: -# bash /root/vime/run_profiling_demo.sh launch # 终端1：启动train -# bash /root/vime/run_profiling_demo.sh profile # 终端2：train就绪后抓trace -# -set -euo pipefail - -VIME_ROOT="${VIME_ROOT:-/root/vime}" -HF_CKPT="${HF_CKPT:-/root/models/Qwen3-4B}" -REF_LOAD="${REF_LOAD:-/root/models/Qwen3-4B_torch_dist}" -PROMPT_DATA="${PROMPT_DATA:-/root/data/gsm8k/train.parquet}" -LOG_ROOT="${LOG_ROOT:-/root/logs/vime_profiling}" -PROFILE_DIR="${PROFILE_DIR:-/root/logs/vllm_profile}" -TRAIN_LOG="${LOG_ROOT}/train_profiling.log" -ROUTER_HOST="${ROUTER_HOST:-127.0.0.1}" - -mkdir -p "${LOG_ROOT}" "${PROFILE_DIR}" - -VLLM_PROFILER_CONFIG_JSON="$(printf \ - '{"profiler":"torch","torch_profiler_dir":"%s","max_iterations":3,"ignore_frontend":true}' \ - "${PROFILE_DIR}")" - -launch_train_for_profiling() { - cd "${VIME_ROOT}" - - # 清理旧 Ray / vLLM 进程（按需注释） - ray stop --force || true - pkill -9 -f "vllm serve" || true - sleep 2 - - ray start --head --node-ip-address 127.0.0.1 --num-gpus 2 --disable-usage-stats - - source "${VIME_ROOT}/scripts/models/qwen3-4B.sh" - - export VLLM_RPC_TIMEOUT="${VLLM_RPC_TIMEOUT:-1800000}" - - RUNTIME_ENV_JSON="{ - \"env_vars\": { - \"PYTHONPATH\": \"${VIME_ROOT}:/root/Megatron-LM\", - \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", - \"VLLM_RPC_TIMEOUT\": \"${VLLM_RPC_TIMEOUT}\" - } - }" - - echo "=== Launching train; log: ${TRAIN_LOG} ===" - echo "=== After engines are up, run: bash $0 profile ===" - - ray job submit --address="http://127.0.0.1:8265" \ - --runtime-env-json="${RUNTIME_ENV_JSON}" \ - -- python3 train.py \ - --train-backend megatron \ - --colocate \ - --actor-num-nodes 1 \ - --actor-num-gpus-per-node 1 \ - --rollout-num-gpus 1 \ - --rollout-num-gpus-per-engine 1 \ - --rollout-backend vllm \ - --rollout-function-path vime.rollout.sleep_rollout.sleep \ - --hf-checkpoint "${HF_CKPT}" \ - --ref-load "${REF_LOAD}" \ - --prompt-data "${PROMPT_DATA}" \ - --input-key question \ - --label-key label \ - --apply-chat-template \ - --rm-type deepscaler \ - --num-rollout 1 \ - --rollout-batch-size 4 \ - --n-samples-per-prompt 1 \ - --rollout-max-response-len 512 \ - --global-batch-size 4 \ - --vllm-gpu-memory-utilization 0.7 \ - --vllm-profiler-config "${VLLM_PROFILER_CONFIG_JSON}" \ - ${MODEL_ARGS[@]} \ - 2>&1 | tee "${TRAIN_LOG}" -} - -discover_router_url() { - local line port - line="$(grep -E 'Router launched at' "${TRAIN_LOG}" | tail -1 || true)" - if [[ -z "${line}" ]]; then - echo "ERROR: Router not found in ${TRAIN_LOG}. Is train still starting?" >&2 - exit 1 - fi - # Router launched at 127.0.0.1:3521, Prometheus port: ... - port="$(echo "${line}" | sed -n 's/.*Router launched at [^:]*:\([0-9]*\).*/\1/p')" - echo "http://${ROUTER_HOST}:${port}" -} - -discover_worker_url() { - local router_url="$1" - python3 - <<'PY' "${router_url}" -import json, sys, urllib.request -router = sys.argv[1] -with urllib.request.urlopen(f"{router}/workers", timeout=10) as r: - workers = json.load(r).get("workers", []) -if not workers: - raise SystemExit("No workers registered") -print(workers[0]["url"]) -PY -} - -run_profiling_session() { - cd "${VIME_ROOT}" - - local router_url worker_url model="${HF_CKPT}" - router_url="$(discover_router_url)" - worker_url="$(discover_worker_url "${router_url}")" - - echo "=== ROUTER=${router_url} WORKER=${worker_url} PROFILE_DIR=${PROFILE_DIR} ===" - - echo "=== 1/3 start_profile (all workers via router) ===" - python tools/profile_rollout.py --router-url "${router_url}" --action start - - echo "=== 2/3 send completions (direct to worker; 3 requests) ===" - for i in 1 2 3; do - curl -sS -X POST "${worker_url}/v1/completions" \ - -H "Content-Type: application/json" \ - -d "{\"model\":\"${model}\",\"prompt\":\"Hello ${i}\",\"max_tokens\":32}" \ - | head -c 400 - echo - done - - echo "=== 3/3 list trace files (max_iterations=3 auto-stop; add --action stop if needed) ===" - sleep 2 - find "${PROFILE_DIR}" -type f \( -name '*.json*' -o -name 'profiler_out_*' \) | sort - echo "Open *.trace.json.gz in https://ui.perfetto.dev/ or run:" - echo " python tools/analyze_profile.py --profile-dir ${PROFILE_DIR} --all-ranks" -} - -case "${1:-}" in - launch) launch_train_for_profiling ;; - profile) run_profiling_session ;; - *) - echo "Usage: $0 {launch|profile}" >&2 - exit 1 - ;; -esac +python tools/profile_rollout.py --router-url http://127.0.0.1:3000 --action stop ``` -**操作步骤：** - -```bash -# 终端1：启动train（等待vLLM与router就绪，日志出现Router launched at ...） -bash /root/vime/run_profiling_demo.sh launch - -# 终端2：抓trace -bash /root/vime/run_profiling_demo.sh profile -``` +## 5. 进行压力测试 -按需修改脚本顶部的`/root/models/...`、`/root/data/...`与GPU布局（`actor-num-gpus-per-node`、`rollout-num-gpus`等）。 +在 Rollout 进程通过 `sleep_rollout` 处于等待状态时，你可以： +1. 使用 `tools/profile_rollout.py` 启动 profiling。 +2. 使用压测工具向 router 或直接向引擎发送请求。 +3. 等待 profiling 完成（如果设置了 `max_iterations`）或手动停止。 +4. 在 `torch_profiler_dir` 中获取 `.json` trace 文件，并使用 `chrome://tracing` 或 [Perfetto](https://ui.perfetto.dev/) 查看。 diff --git a/docs/zh/examples/qwen3-30B-A3B.md b/docs/zh/examples/qwen3-30B-A3B.md index 039b3920b..525478417 100644 --- a/docs/zh/examples/qwen3-30B-A3B.md +++ b/docs/zh/examples/qwen3-30B-A3B.md @@ -73,82 +73,14 @@ bash scripts/run-qwen3-30B-A3B.sh ### 多机支持 -以下以 **2台机器、每台8卡（共16GPU）** 为入门示例；脚本与参数可扩展到 **N节点**。多机与单机的主要差异： +对于多机环境，需要进行如下的几点修改： +- 将训练模型，数据放在所有机器都可以访问到的路径上； +- 设置各台机器都可以访问到的 `MASTER_ADDR` 之外； +- 去掉 CPU adam 相关的配置，因为使用了 distributed optimizer，所以多机环境下 optimizer 的显存占比会明显下降。 -- 训练模型、数据放在所有节点均可访问的路径（如 NFS）； -- `MASTER_ADDR` 设为 head 节点的 **局域网 IP**（非 `127.0.0.1`）； -- 去掉 CPU Adam（多机使用 distributed optimizer，无需 `--optimizer-cpu-offload`）； -- `global-batch-size` 必须等于 `rollout-batch-size × n-samples-per-prompt`。 +除此之外，还可以进行如下的修改： -#### 拓扑概览 - -| 组件 | 双机默认配置 | -|------|------| -| 集群 | `ACTOR_NUM_NODES=2`，`ACTOR_NUM_GPUS_PER_NODE=8` | -| Megatron训练 | TP=8, EP=8, CP=2（expert 分片跨节点） | -| vLLM Rollout | 跨节点 TP=16（`rollout-num-gpus-per-engine = 节点数 × 每节点GPU`） | -| 调度 | Ray 集群 + `--colocate` 共卡模式 | - -转换 checkpoint 时建议使用与训练一致的 Megatron 并行度（双机示例 TP=8, EP=8）。checkpoint 的 EP 需与 `--expert-model-parallel-size` 一致，否则 `load_checkpoint` 可能极慢或卡住。 - -#### 启动 Ray 集群 - -Ray 集群需在各节点上 **单独启动**，不在训练脚本内。先在所有 worker 节点加入集群，确认 `ray status` 显示预期 GPU 总数后，再在 head 提交训练。示例（双机）： - -```bash -# === Head 节点 === -export MASTER_ADDR= -ray start --head --node-ip-address="${MASTER_ADDR}" --num-gpus 8 --disable-usage-stats \ - --dashboard-host=0.0.0.0 --dashboard-port=8265 - -# === 各 Worker 节点 === -export MASTER_ADDR= -ray start --address="${MASTER_ADDR}:6379" --node-ip-address=<本机_局域网_IP> --num-gpus 8 -``` - -更多说明见 [快速开始 — 多机训练](../get_started/quick_start.md#multi-node-training-for-large-scale-moe-models)。 - -#### 执行训练 - -Ray 集群就绪后，在 **head 节点** 设置多机环境变量并运行 **与单机相同的脚本**（`ACTOR_NUM_NODES>1` 时脚本不会启动 Ray，并使用多机默认参数）： - -```bash -export MASTER_ADDR= -export ACTOR_NUM_NODES=2 -export ACTOR_NUM_GPUS_PER_NODE=8 -cd /root/vime -bash scripts/run-qwen3-30B-A3B.sh -``` - -2 step 冒烟示例： - -```bash -NUM_ROLLOUT=2 ENABLE_R3=0 bash scripts/run-qwen3-30B-A3B.sh -``` - -扩展到 N 节点（例如 4×8）时，在各 worker 加入 Ray 后，于 head 设置 `ACTOR_NUM_NODES=4` 并按总卡数调整 `MEGATRON_TP` / `MEGATRON_EP` / `MEGATRON_CP` / `ROLLOUT_NUM_GPUS_PER_ENGINE`。 - -#### 多机关键参数 - -| 变量 | 双机默认 | 说明 | -|------|----------|------| -| `ACTOR_NUM_NODES` | 2（单机默认为 1） | 训练节点总数（含 head）；>1 时脚本不启动 Ray | -| `ACTOR_NUM_GPUS_PER_NODE` | 8 | 每节点 GPU 数 | -| `MEGATRON_TP` / `MEGATRON_EP` / `MEGATRON_CP` | 8 / 8 / 2 | Megatron 并行 | -| `ROLLOUT_NUM_GPUS_PER_ENGINE` | 总 GPU 数 | vLLM engine 占用卡数 | -| `ENABLE_R3` | 1 | 设为 0 可关闭 R3 路径 | - -脚本默认 batch：`rollout-batch-size=4`，`n-samples-per-prompt=2`，`global-batch-size=8`；vLLM 使用 `--vllm-moe-backend triton`。 - -#### 多机常见问题 - -- **Worker 无法加入 Ray / NCCL 失败**：检查 `MASTER_ADDR`、容器 `/etc/hosts`（hostname 勿指向 `127.0.0.1`）、`NCCL_SOCKET_IFNAME` / `GLOO_SOCKET_IFNAME`。 -- **`Not enough samples X for global_batch_size Y`**：同步调整 `global-batch-size` 与 `rollout-batch-size × n-samples-per-prompt`。 -- **GPU 显存占满但无进程**：重启容器或 `ray stop --force` 清理残留 vLLM 上下文。 - -#### EPLB - -当总卡数并不能被 expert 总数整除时，可以开启 vLLM 的 EPLB（Expert Parallelism Load Balancer），通过 `--vllm-eplb-config` 配置冗余 expert。例如对于 24 卡的场景： +- 当总卡数并不能被 expert 总数乘除时，可以开启 vLLM 的 EPLB（Expert Parallelism Load Balancer），通过 `--vllm-eplb-config` 来增加冗余的 expert，例如对于 24 卡的场景，可以配置： ```bash VLLM_ARGS=( diff --git a/docs/zh/examples/qwen3-4B.md b/docs/zh/examples/qwen3-4B.md index e0fac9059..85720a35a 100644 --- a/docs/zh/examples/qwen3-4B.md +++ b/docs/zh/examples/qwen3-4B.md @@ -193,7 +193,7 @@ OPTIMIZER_ARGS=( #### VLLM_ARGS -vLLM 推理所需的参数。vime 默认使用 vLLM 作为 rollout 后端（`rollout.py` 启动 `VLLMEngine`，默认 rollout 函数为 `vime.rollout.vllm_rollout.generate_rollout`），无需额外指定 backend。`--rollout-num-gpus-per-engine` 对应每个 vLLM engine 的 `tensor_parallel_size`；除此之外的 vLLM 参数均通过添加 `--vllm-` 前缀传给 vime（例如 `--vllm-max-model-len`）。 +vLLM 所需的参数，这里 `--rollout-num-gpus-per-engine` 基本对应 vLLM 的 `tensor_parallel_size`，除此之外的 vLLM 参数均通过添加 `--vllm-` 的前缀来传给 vime。 ```bash VLLM_ARGS=( @@ -202,9 +202,7 @@ VLLM_ARGS=( ) ``` -rollout 并发较高时，还可以通过 `--vllm-` 前缀调节 vLLM scheduler，例如 `--vllm-max-num-seqs`、`--vllm-max-num-batched-tokens`；调试或规避 CUDA graph 相关限制时可加 `--vllm-enforce-eager`。 - -⚠️ vime 会用 vLLM router 来调度多个 vLLM server。训推一体（`--colocate`）时，训练与推理权重经 CUDA IPC 同步；训推分离时，训练侧经 NCCL 与 vLLM engine 同步权重。 +⚠️ vime 会用 vLLM router 来调度多个 vLLM server，在不开启 dp attention 的情况下不支持 `dp_size`。 ### dynamic sampling @@ -278,22 +276,21 @@ ray job submit ... \ ... ``` -此时，就会分配 2 张卡给训练，6 张卡给推理。`--rollout-num-gpus` 与 `--actor-num-gpus-per-node` 一样，是传给 `train.py` 的 **Ray 资源参数**：框架据此创建 placement group，并把前若干 bundle 分给训练 actor、后续 bundle 分给 rollout engine（见 `vime/ray/placement_group.py`）。**共卡模式（`--colocate`）下该参数会被忽略**，并自动设为 `actor_num_gpus_per_node * actor_num_nodes`。请勿把 `--rollout-num-gpus` 写在 `VLLM_ARGS` 中。 +此时，就会分配 2 张卡给训练，6 张卡给推理。 -训推分离时，`VLLM_ARGS` 仅需配置推理后端相关参数，例如： +⚠️ 在进行训推分离的时候，每个 vLLM server 上的并发度太大，超过了 vLLM 默认的 cuda graph 的并发度，影响推理速度。可以用以下 2 种方式进行调整： -```bash -VLLM_ARGS=( - --rollout-num-gpus-per-engine 2 - --vllm-gpu-memory-utilization 0.9 - --vllm-max-num-seqs 256 - --vllm-max-num-batched-tokens 8192 -) -``` +1. 通过 `--vllm-server-concurrency` 限制发给一个 vLLM server 的最大并发量，例如： + + ```bash + --vllm-server-concurrency 160 + ``` -如需调试或规避 CUDA graph 相关限制，可额外加上 `--vllm-enforce-eager`。 +2. 使用 `--vllm-cudagraph-capture-sizes` 增大 vLLM 初始化的 cuda graph 数量，例如： -⚠️ 在训推一体的训练时，megatron 始终会占据一些显存，需要通过 `--vllm-gpu-memory-utilization` 来降低 vLLM 占据的显存比例，并配合 `--train-memory-margin-bytes` 为训练侧预留空间。 + ```bash + --vllm-cudagraph-capture-sizes 1 2 4 8 $(seq 16 8 256) + ``` ### 异步训练 diff --git a/examples/fully_async/README.md b/examples/fully_async/README.md index cd71455a4..7fb271f6f 100644 --- a/examples/fully_async/README.md +++ b/examples/fully_async/README.md @@ -3,8 +3,7 @@ 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. ## Files -The fully-async worker has been **promoted from this example into the core package** — it now lives in -`vime/rollout/fully_async_rollout.py`. This directory keeps only the launch script: +The worker itself lives in `vime.rollout.fully_async_rollout`; this directory is just the launch script: * `run-qwen3-4b-fully_async.sh`: example launch script with Qwen3‑4B. ## Prerequisite diff --git a/examples/geo3k_vlm/README.md b/examples/geo3k_vlm/README.md index df88e7a46..0d82487de 100644 --- a/examples/geo3k_vlm/README.md +++ b/examples/geo3k_vlm/README.md @@ -115,10 +115,4 @@ Our initial geo3k-specific verifier produced "format scores" (**0 and 0.9**) ins We fixed this by switching to the default math RM with clean **binary 0/1 rewards**. If you encounter similar precision issues with non-binary rewards, you can change the reward tensor dtype from `torch.float` to `torch.float16` in `vime/ray/rollout.py` (`_post_process_rewards` method) to truncate precision artifacts. ## B200 -On Blackwell (SM100), vllm automatically dispatches the ViT encoder to -FlashAttention 4 (or FA2 fallback) — no manual override is needed -([vllm/v1/attention/backends/fa_utils.py:81](https://github.com/vllm-project/vllm/blob/main/vllm/v1/attention/backends/fa_utils.py#L81)). -If you hit a kernel issue on a specific model, you can force SDPA with -`--vllm-mm-encoder-attn-backend TORCH_SDPA`. The HF-side -`--attn-implementation flash_attention_2` flag is still relevant when the -model is loaded via Hugging Face Transformers. +On Blackwell (SM100), vLLM automatically selects FlashAttention 4 for the ViT encoder, so no manual override is needed. diff --git a/examples/geo3k_vlm_multi_turn/env_geo3k.py b/examples/geo3k_vlm_multi_turn/env_geo3k.py index 634e8593f..972f0d286 100644 --- a/examples/geo3k_vlm_multi_turn/env_geo3k.py +++ b/examples/geo3k_vlm_multi_turn/env_geo3k.py @@ -32,7 +32,7 @@ class Geo3kEnv(BaseInteractionEnv): an `answer` argument. We run the math reward checker against the ground truth and return feedback for wrong answers. The episode ends when the answer is correct, when max_turns is reached, or when the response has no valid tool call; in the - last case we still try to score a boxed/text answer, matching SkyRL's env. + last case we still try to score a boxed/text answer. """ def __init__(self, *, ground_truth: str | None = None, max_turns: int | None = None): diff --git a/scripts/run-minimax-m2.sh b/scripts/run-minimax-m2.sh index 44937f01c..95380db7c 100644 --- a/scripts/run-minimax-m2.sh +++ b/scripts/run-minimax-m2.sh @@ -34,8 +34,8 @@ BASE_DIR=${BASE_DIR:-"/root"} CKPT_ARGS=( --hf-checkpoint ${BASE_DIR}/MiniMax-M2.5 --ref-load ${BASE_DIR}/MiniMax-M2.5_torch_dist - --load ${BASE_DIR}/MiniMax-M2.5_slime/ - --save ${BASE_DIR}/MiniMax-M2.5_slime/ + --load ${BASE_DIR}/MiniMax-M2.5_vime/ + --save ${BASE_DIR}/MiniMax-M2.5_vime/ --save-interval 20 --megatron-to-hf-mode raw --model-name minimax_m2 diff --git a/scripts/run-qwen2.5-0.5B-reproducibility.sh b/scripts/run-qwen2.5-0.5B-reproducibility.sh index fa4c4a6d8..72dad7608 100644 --- a/scripts/run-qwen2.5-0.5B-reproducibility.sh +++ b/scripts/run-qwen2.5-0.5B-reproducibility.sh @@ -13,8 +13,6 @@ set -ex # will prevent ray from buffering stdout/stderr export PYTHONUNBUFFERED=1 -# Bitwise reproduction depends on a fixed parallel/reduction layout, so the GPU -# count is pinned here (matching the upstream recipe) rather than auto-detected. NUM_GPUS=8 SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" @@ -116,7 +114,6 @@ export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"} ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus ${NUM_GPUS} --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 # Build the runtime environment JSON with proper variable substitution. -# The NCCL_ALGO / NVTE / CUBLAS settings below are required for bitwise determinism. RUNTIME_ENV_JSON="{ \"env_vars\": { \"PYTHONPATH\": \"${VIME_ROOT}:/root/Megatron-LM/\", diff --git a/tests/test_agent_adapters.py b/tests/test_agent_adapters.py index 94e055f67..c234b4c32 100644 --- a/tests/test_agent_adapters.py +++ b/tests/test_agent_adapters.py @@ -109,16 +109,16 @@ def flush() -> None: def test_session_id_comes_from_protocol_fields_not_custom_header(): assert ( openai._request_session_id( - FakeRequest({"X-Slime-Session-Id": "custom"}), + FakeRequest({"X-Vime-Session-Id": "custom"}), {"metadata": {"session_id": "meta-session"}, "user": "body-user"}, ) == "meta-session" ) assert ( - openai._request_session_id(FakeRequest({"X-Slime-Session-Id": "custom"}), {"user": "body-user"}) == "body-user" + openai._request_session_id(FakeRequest({"X-Vime-Session-Id": "custom"}), {"user": "body-user"}) == "body-user" ) assert ( - anthropic._request_session_id(FakeRequest({"X-Slime-Session-Id": "custom", "X-Api-Key": "anthropic-key"})) + anthropic._request_session_id(FakeRequest({"X-Vime-Session-Id": "custom", "X-Api-Key": "anthropic-key"})) == "anthropic-key" ) assert ( @@ -138,7 +138,7 @@ def test_anthropic_translation_keeps_tool_results_and_tool_schema(): "content": [ {"type": "thinking", "thinking": "plan"}, {"type": "text", "text": "ok"}, - {"type": "tool_use", "name": "lookup", "input": {"q": "slime"}}, + {"type": "tool_use", "name": "lookup", "input": {"q": "vime"}}, ], }, {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "u1", "content": "result"}]}, @@ -156,7 +156,7 @@ def test_anthropic_translation_keeps_tool_results_and_tool_schema(): "role": "assistant", "content": "ok", "reasoning_content": "plan", - "tool_calls": [{"function": {"name": "lookup", "arguments": {"q": "slime"}}}], + "tool_calls": [{"function": {"name": "lookup", "arguments": {"q": "vime"}}}], }, {"role": "tool", "content": "result"}, ] @@ -185,7 +185,7 @@ def test_openai_translation_and_responses_input_shapes(): { "id": "call_1", "type": "function", - "function": {"name": "lookup", "arguments": {"q": "slime"}}, + "function": {"name": "lookup", "arguments": {"q": "vime"}}, } ], }, @@ -210,7 +210,7 @@ def test_openai_translation_and_responses_input_shapes(): { "id": "call_1", "type": "function", - "function": {"name": "lookup", "arguments": '{"q": "slime"}'}, + "function": {"name": "lookup", "arguments": '{"q": "vime"}'}, } ], }, @@ -313,7 +313,7 @@ async def fake_generate(prompt_ids, session, body, app, **kwargs): async def run_case(): monkeypatch.setattr(openai, "_generate", fake_generate) - raw = "use it slime" + raw = "use it vime" tokenizer = ToyTokenizer({(451,): raw}) adapter = openai.OpenAIAdapter(tokenizer=tokenizer, vllm_url="http://unused") adapter.open_session("sid-chat-tool-stream", sampling_defaults={"max_new_tokens": 8}) @@ -350,7 +350,7 @@ async def run_case(): assert any(c["choices"][0]["delta"] == {"content": "use it"} for c in chunks) assert tool_delta["tool_calls"][0]["index"] == 0 assert tool_delta["tool_calls"][0]["function"]["name"] == "lookup" - assert tool_delta["tool_calls"][0]["function"]["arguments"] == '{"query": "slime"}' + assert tool_delta["tool_calls"][0]["function"]["arguments"] == '{"query": "vime"}' assert chunks[-1]["choices"][0]["finish_reason"] == "tool_calls" assert segments[0].response_ids == [451] @@ -364,7 +364,7 @@ async def fake_generate(prompt_ids, session, body, app, **kwargs): async def run_case(): monkeypatch.setattr(openai, "_generate", fake_generate) - raw = "look slime" + raw = "look vime" tokenizer = ToyTokenizer({(301,): raw}) adapter = openai.OpenAIAdapter(tokenizer=tokenizer, vllm_url="http://unused") adapter.open_session("sid-responses", sampling_defaults={"max_new_tokens": 8}) @@ -399,7 +399,7 @@ async def run_case(): assert output_types == ["message", "function_call"] assert data["output"][0]["content"][0]["text"] == "look" assert function_call["name"] == "lookup" - assert function_call["arguments"] == '{"query": "slime"}' + assert function_call["arguments"] == '{"query": "vime"}' assert segments[0].response_ids == [301] asyncio.run(run_case()) @@ -414,7 +414,7 @@ async def fake_generate(prompt_ids, session, body, app, **kwargs): async def run_case(): monkeypatch.setattr(openai, "_generate", fake_generate) - raw = "slime" + raw = "vime" tokenizer = ToyTokenizer({(551,): raw}) adapter = openai.OpenAIAdapter(tokenizer=tokenizer, vllm_url="http://unused") adapter.open_session("sid-responses-tool-stream", sampling_defaults={"max_new_tokens": 8}) @@ -450,7 +450,7 @@ async def run_case(): assert created["type"] == "response.created" assert [item["type"] for item in created["response"]["output"]] == ["function_call"] assert completed_call["name"] == "lookup" - assert completed_call["arguments"] == '{"query": "slime"}' + assert completed_call["arguments"] == '{"query": "vime"}' assert segments[0].response_ids == [551] asyncio.run(run_case()) @@ -616,7 +616,7 @@ async def run_case(): upstream_server = TestServer(upstream_app) await upstream_server.start_server() - tool_raw = "slime" + tool_raw = "vime" tokenizer = ScriptedTokenizer( prompts=[ [10, 11], @@ -637,7 +637,7 @@ async def run_case(): headers={"Authorization": "Bearer sid-openai-token"}, json={ "model": "actor", - "input": "find slime", + "input": "find vime", "max_output_tokens": 5, "tools": [ { @@ -657,12 +657,12 @@ async def run_case(): json={ "model": "actor", "input": [ - {"role": "user", "content": "find slime"}, + {"role": "user", "content": "find vime"}, function_call, { "type": "function_call_output", "call_id": function_call["call_id"], - "output": "found slime", + "output": "found vime", }, ], "max_output_tokens": 7, @@ -684,7 +684,7 @@ async def run_case(): assert first.status == 200 assert second.status == 200 assert function_call["name"] == "lookup" - assert function_call["arguments"] == '{"query": "slime"}' + assert function_call["arguments"] == '{"query": "vime"}' assert second_data["output"][0]["content"][0]["text"] == "done" assert [req["token_ids"] for req in upstream.requests] == [[10, 11], [10, 11, 20, 21, 30, 31]] assert upstream.routing_keys == ["sid-openai-token", "sid-openai-token"] @@ -712,7 +712,7 @@ async def run_case(): upstream_server = TestServer(upstream_app) await upstream_server.start_server() - tool_raw = "slime" + tool_raw = "vime" tokenizer = ScriptedTokenizer( prompts=[ [110, 111], @@ -737,7 +737,7 @@ async def run_case(): json={ "model": "actor", "max_tokens": 5, - "messages": [{"role": "user", "content": [{"type": "text", "text": "find slime"}]}], + "messages": [{"role": "user", "content": [{"type": "text", "text": "find vime"}]}], "tools": [ { "name": "lookup", @@ -756,7 +756,7 @@ async def run_case(): "model": "actor", "max_tokens": 7, "messages": [ - {"role": "user", "content": [{"type": "text", "text": "find slime"}]}, + {"role": "user", "content": [{"type": "text", "text": "find vime"}]}, {"role": "assistant", "content": first_data["content"]}, { "role": "user", @@ -764,7 +764,7 @@ async def run_case(): { "type": "tool_result", "tool_use_id": tool_use["id"], - "content": "found slime", + "content": "found vime", } ], }, @@ -786,7 +786,7 @@ async def run_case(): assert first.status == 200 assert second.status == 200 assert tool_use["name"] == "lookup" - assert tool_use["input"] == {"query": "slime"} + assert tool_use["input"] == {"query": "vime"} assert second_data["content"] == [{"type": "text", "text": "anthropic done"}] assert [req["token_ids"] for req in upstream.requests] == [[110, 111], [110, 111, 120, 121, 130]] assert upstream.routing_keys == ["sid-anthropic-token", "sid-anthropic-token"] diff --git a/tests/test_agent_sdk_adapters.py b/tests/test_agent_sdk_adapters.py index a4ad3d4ce..d947cd67a 100644 --- a/tests/test_agent_sdk_adapters.py +++ b/tests/test_agent_sdk_adapters.py @@ -52,7 +52,7 @@ async def fake_generate(prompt_ids, session, body, app, **kwargs): monkeypatch.setattr(openai, "_generate", fake_generate) tokenizer = SDKTokenizer( [ - "slime", + "vime", "final after tool", ] ) @@ -82,7 +82,7 @@ def lookup(query: str) -> str: model_settings=agents.ModelSettings(max_tokens=4), ) try: - result = await agents.Runner.run(agent, "find slime") + result = await agents.Runner.run(agent, "find vime") finally: await client.close() await oai.close() @@ -94,16 +94,16 @@ def lookup(query: str) -> str: assert calls[0]["body"]["tools"][0]["name"] == "lookup" assert calls[1]["body"]["input"][-1] == { "call_id": calls[1]["body"]["input"][-2]["call_id"], - "output": "found slime", + "output": "found vime", "type": "function_call_output", } assert tokenizer.rendered[0][0] == [ {"role": "system", "content": "Use lookup."}, - {"role": "user", "content": "find slime"}, + {"role": "user", "content": "find vime"}, ] assert tokenizer.rendered[1][0][-1] == { "role": "tool", - "content": "found slime", + "content": "found vime", "tool_call_id": calls[1]["body"]["input"][-2]["call_id"], } assert segments[0].metadata["segment_kind"] == "final" @@ -179,7 +179,7 @@ async def fake_generate(prompt_ids, session, body, app, **kwargs): monkeypatch.setattr(openai, "_generate", fake_generate) tokenizer = SDKTokenizer( - ["streamed via sdk slime"] + ["streamed via sdk vime"] ) adapter = openai.OpenAIAdapter(tokenizer=tokenizer, vllm_url="http://unused") client = TestClient(TestServer(adapter.app)) @@ -233,7 +233,7 @@ async def fake_generate(prompt_ids, session, body, app, **kwargs): segments = await adapter.finish_session("sdk-openai-chat-stream") assert "".join(text_parts) == "streamed via sdk" assert tool_names == ["lookup"] - assert tool_arguments == ['{"query": "slime"}'] + assert tool_arguments == ['{"query": "vime"}'] assert finish_reasons == ["tool_calls"] assert usages[-1].prompt_tokens == 2 assert usages[-1].completion_tokens == 1 diff --git a/tests/test_cp_utils.py b/tests/test_cp_utils.py index d37e870a7..3dfaefcad 100644 --- a/tests/test_cp_utils.py +++ b/tests/test_cp_utils.py @@ -15,7 +15,7 @@ from __future__ import annotations -# Import the helpers BEFORE the slime imports so the megatron stub lands +# Import the helpers BEFORE the vime imports so the megatron stub lands # in sys.modules first. pytest's prepend importmode puts this file's # directory (``tests/``) on sys.path, which is what makes the bare-name # import work without an ``__init__.py``. diff --git a/tests/test_loss_cp_invariance.py b/tests/test_loss_cp_invariance.py index 998ba1d6c..9f83e0ae1 100644 --- a/tests/test_loss_cp_invariance.py +++ b/tests/test_loss_cp_invariance.py @@ -7,7 +7,7 @@ Why this matters ---------------- -Slime's loss prescaling + Megatron's per-mb scaling + DDP's grad +Vime's loss prescaling + Megatron's per-mb scaling + DDP's grad averaging compose into one big formula. Any time we touch any one of those three layers the numbers should land in the same place. Until this test existed we only had end-to-end report-formula checks @@ -18,12 +18,12 @@ -------------------------- We reproduce, for each spawned rank, the exact sequence Megatron applies when a 3-tuple ``(loss, num_tokens, log)`` comes back from the loss -function with ``calculate_per_token_loss=False`` — slime's per-rollout- +function with ``calculate_per_token_loss=False`` — vime's per-rollout- mean path: 1. Loss function pre-scales:: loss *= num_microbatches / step_global_batch_size * (dp * cp) - See ``slime/backends/megatron_utils/loss.py:1209-1215``. + See ``vime/backends/megatron_utils/loss.py:1209-1215``. 2. Megatron divides by ``clamp(num_tokens, 1)`` then by ``num_microbatches``:: output_tensor /= torch.clamp(num_tokens, min=1) # num_tokens=1 → no-op @@ -45,7 +45,7 @@ What this test does NOT exercise: the actual Megatron model classes, the real DDP buffer code, fused optimizers, mixed-precision. We use a plain ``nn.Linear`` with manual all-reduce-average to simulate steps 1-4 above. -The contract here is on *our* scaling math (steps 1 + 4 are slime's; +The contract here is on *our* scaling math (steps 1 + 4 are vime's; step 2 is what Megatron does to our 3-tuple). If Megatron later changes step 2 — e.g. drops the ``/= num_microbatches`` — this test won't catch it, but the real GPU integration suite (``test_qwen2.5_0.5B_short.py``) @@ -54,7 +54,7 @@ from __future__ import annotations -# Megatron stub must land in sys.modules first; the slime imports inside +# Megatron stub must land in sys.modules first; the vime imports inside # the worker pick it up via this same module. pytest's prepend importmode # puts ``tests/`` on sys.path so the bare-name import works without an # ``__init__.py``; mp.spawn children inherit the parent's sys.path. @@ -85,8 +85,8 @@ def _grad_norm_worker( """One spawned rank. Builds a tiny ``nn.Linear`` model (deterministic init via ``seed``), - runs slime's per-rollout-mean loss reducer with the rank's share of - the four-rollout fixture, applies the slime-side prescaling, then + runs vime's per-rollout-mean loss reducer with the rank's share of + the four-rollout fixture, applies the vime-side prescaling, then Megatron's per-mb scaling, then ``.backward()``, then a manual all-reduce-average across the dp-with-cp group (mirroring DDP's ``average_in_collective=False`` path with @@ -151,7 +151,7 @@ def _grad_norm_worker( reducer = get_sum_of_sample_mean(my_tl, my_rl, my_masks, my_denoms) loss = reducer(output) - # === Step 1: slime's per-rollout-mean prescaling ====================== + # === Step 1: vime's per-rollout-mean prescaling ====================== # loss.py:1209-1215. ``mpu.get_data_parallel_world_size(with_context_parallel=True)`` # is the dp-with-cp world size, which is ``world_size`` in this setup. loss = loss * num_microbatches / step_global_batch_size * world_size @@ -160,10 +160,10 @@ def _grad_norm_worker( # schedules.py:258-264 — for the 3-tuple, not-per-token-loss path: # output_tensor /= torch.clamp(num_tokens, min=1) # output_tensor /= num_microbatches - # slime passes num_tokens=1 in this path (loss.py:1221), so the + # vime passes num_tokens=1 in this path (loss.py:1221), so the # first divide is a no-op; we keep it explicit to mirror the # source faithfully. - num_tokens_for_scaling = torch.tensor(1.0) # slime's placeholder + num_tokens_for_scaling = torch.tensor(1.0) # vime's placeholder loss = loss / torch.clamp(num_tokens_for_scaling, min=1.0) loss = loss / num_microbatches @@ -218,7 +218,7 @@ def _run_grad_norm_worker(dp_size: int, cp_size: int, tmp_path) -> float: # - (1, 4) deeper CP-only # - (4, 1) deeper DP-only # The full 3*3 matrix lives in test_metric_report_dist.py — here we just -# want enough coverage to catch a sign/factor regression in the slime +# want enough coverage to catch a sign/factor regression in the vime # prescaling math. _PARALLELISM_CASES = [(1, 1), (2, 1), (1, 2), (2, 2), (1, 4), (4, 1)] diff --git a/tests/test_metric_report.py b/tests/test_metric_report.py index 98d7d69c6..8f3f9d753 100644 --- a/tests/test_metric_report.py +++ b/tests/test_metric_report.py @@ -17,7 +17,7 @@ from __future__ import annotations -# Import the helpers BEFORE the slime imports so the megatron stub lands +# Import the helpers BEFORE the vime imports so the megatron stub lands # in sys.modules first. pytest's prepend importmode puts this file's # directory (``tests/``) on sys.path, which is what makes the bare-name # import work without an ``__init__.py``. diff --git a/tests/test_metric_report_dist.py b/tests/test_metric_report_dist.py index 7535027aa..cd1681832 100644 --- a/tests/test_metric_report_dist.py +++ b/tests/test_metric_report_dist.py @@ -23,7 +23,7 @@ from __future__ import annotations # IMPORTANT: import the helpers (and the megatron stub it installs) BEFORE -# any slime import. Spawned workers re-import this module from scratch, so +# any vime import. Spawned workers re-import this module from scratch, so # the same ordering must hold there — see ``stub_megatron_in_worker`` # for the worker-side details. pytest's prepend importmode puts # ``tests/`` on sys.path so the bare-name import works without an @@ -295,7 +295,7 @@ def test_rollout_log_real_distributed_multi_key(dp_size, cp_size, tmp_path): # Keep an explicit reference to silence "unused import" complaints while # documenting that importing the helpers module is load-bearing (it -# installs the megatron stub before slime is touched). +# installs the megatron stub before vime is touched). _ = _cp_dist_helpers diff --git a/tests/test_qwen2.5_0.5B_fanout_short.py b/tests/test_qwen2.5_0.5B_fanout_short.py index 01c5ddbb1..f902467e6 100644 --- a/tests/test_qwen2.5_0.5B_fanout_short.py +++ b/tests/test_qwen2.5_0.5B_fanout_short.py @@ -44,7 +44,7 @@ import vime.utils.external_utils.command_utils as U -TIGHT_DEVICE_MEMORY = U.get_bool_env_var("SLIME_TEST_TIGHT_DEVICE_MEMORY", "1") +TIGHT_DEVICE_MEMORY = U.get_bool_env_var("VIME_TEST_TIGHT_DEVICE_MEMORY", "1") MODEL_NAME = "Qwen2.5-0.5B-Instruct" MODEL_TYPE = "qwen2.5-0.5B" @@ -54,8 +54,8 @@ # through to the Ray-submitted job via an env var so all worker # processes write to the same path. FANOUT_COUNTER_FILE = os.environ.get( - "SLIME_FANOUT_TEST_COUNTER_FILE", - os.path.join(tempfile.gettempdir(), "slime_fanout_test_counter.log"), + "VIME_FANOUT_TEST_COUNTER_FILE", + os.path.join(tempfile.gettempdir(), "vime_fanout_test_counter.log"), ) @@ -190,7 +190,7 @@ def execute(): megatron_model_type=MODEL_TYPE, # Make the counter path visible inside the Ray-submitted job # (helper picks it up via os.environ). - extra_env_vars={"SLIME_FANOUT_TEST_COUNTER_FILE": FANOUT_COUNTER_FILE}, + extra_env_vars={"VIME_FANOUT_TEST_COUNTER_FILE": FANOUT_COUNTER_FILE}, ) # Post-train assertion: compact_generate must have been called exactly diff --git a/tests/test_qwen3_4B_streaming_partial_rollout.py b/tests/test_qwen3_4B_streaming_partial_rollout.py index 898e4c2cc..065f57222 100644 --- a/tests/test_qwen3_4B_streaming_partial_rollout.py +++ b/tests/test_qwen3_4B_streaming_partial_rollout.py @@ -20,7 +20,7 @@ import vime.utils.external_utils.command_utils as U -TIGHT_HOST_MEMORY = U.get_bool_env_var("SLIME_TEST_TIGHT_HOST_MEMORY", "1") +TIGHT_HOST_MEMORY = U.get_bool_env_var("VIME_TEST_TIGHT_HOST_MEMORY", "1") MODEL_NAME = "Qwen3-4B" MODEL_TYPE = "qwen3-4B" diff --git a/vime/ray/rollout.py b/vime/ray/rollout.py index 572f95514..7c1056f90 100644 --- a/vime/ray/rollout.py +++ b/vime/ray/rollout.py @@ -817,7 +817,7 @@ def _validate_rollout_id_annotated(node, depth=0): when a compact / subagent pattern is detected. "Compact" = the rollout function wraps multiple training samples from one - rollout execution into a ``list[Sample]``. In slime's convention the + rollout execution into a ``list[Sample]``. In vime's convention the default rollout shape is ``list[list[Sample]]`` (depth-2: prompt × rollout) so its leaf ``list[Sample]`` lands at depth 1 and we skip validation, preserving backward compatibility. A compact rollout adds a third level: diff --git a/vime/rollout/_fanout_test_helpers.py b/vime/rollout/_fanout_test_helpers.py index 3590ca472..9080ff281 100644 --- a/vime/rollout/_fanout_test_helpers.py +++ b/vime/rollout/_fanout_test_helpers.py @@ -35,7 +35,7 @@ # Each invocation appends one line. The test file reads this after train # completes to assert the framework actually drove the custom path for # every prompt (no silent bypass / no double-submission). -COUNTER_FILE_ENV = "SLIME_FANOUT_TEST_COUNTER_FILE" +COUNTER_FILE_ENV = "VIME_FANOUT_TEST_COUNTER_FILE" async def compact_generate(args, sample, sampling_params): diff --git a/vime/utils/dp_schedule.py b/vime/utils/dp_schedule.py index e30621ca5..283f8631b 100644 --- a/vime/utils/dp_schedule.py +++ b/vime/utils/dp_schedule.py @@ -2,7 +2,7 @@ Pure-Python logic that decides, for one rollout's worth of sample lengths, how to group samples into micro-batches and which DP rank owns each mbs. -Lives outside the ray/sglang-importing modules so it can be unit-tested +Lives outside the ray/vllm-importing modules so it can be unit-tested under CPU-only CI. The scheduling philosophy is **pack first, distribute second**: diff --git a/vime/utils/external_utils/command_utils.py b/vime/utils/external_utils/command_utils.py index ae4e63981..6bab7d9b4 100644 --- a/vime/utils/external_utils/command_utils.py +++ b/vime/utils/external_utils/command_utils.py @@ -108,14 +108,8 @@ def execute_train( master_addr = os.environ.get("MASTER_ADDR", "127.0.0.1") exec_command( - # vLLM renames its VRAM-holding subprocesses via set_process_title() - # (VLLM::EngineCore, VLLM::Worker_TP*, vllm::router), so their cmdline no - # longer contains "vllm serve". Matching only the launcher would leave the - # engine/worker children holding GPU memory and leak it into the next run. - # Match both the launcher and the renamed children; the [v]/[M] bracket - # trick keeps this pattern from matching pkill's own cmdline. This targets - # exactly the vLLM tree, so the old indiscriminate `pkill -9 python` - # (dangerous on colocate/shared nodes) is no longer needed. + # vLLM renames its subprocesses (VLLM::EngineCore / Worker_TP*), so match + # the renamed children too; the [v]/[M] brackets avoid matching pkill itself. "pkill -9 -f '[v]llm serve|VLL[M]::'; " "sleep 3; " f"{'' if external_ray else 'ray stop --force; '}" From deca66d880b8f8ea84e56abb68983476cb76369d Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Tue, 9 Jun 2026 02:25:24 +0000 Subject: [PATCH 2/5] docs: move doc-site trims to their own PR (#201) Reverts docs/** here; the doc-site trims now live in the dedicated docs-only PR #201 so this PR stays scoped to mechanical symbol translation + example/ script/command_utils prose cleanup. Co-Authored-By: Claude Opus 4.8 (1M context) --- docs/en/developer_guide/debug.md | 43 +++- docs/en/developer_guide/profiling.md | 321 +++++++++++++++++++++++--- docs/en/examples/qwen3-30B-A3B.md | 79 ++++++- docs/en/examples/qwen3-4B.md | 29 +-- docs/zh/developer_guide/debug.md | 43 +++- docs/zh/developer_guide/profiling.md | 323 ++++++++++++++++++++++++--- docs/zh/examples/qwen3-30B-A3B.md | 80 ++++++- docs/zh/examples/qwen3-4B.md | 29 +-- 8 files changed, 842 insertions(+), 105 deletions(-) diff --git a/docs/en/developer_guide/debug.md b/docs/en/developer_guide/debug.md index 193d9a3db..212affd16 100644 --- a/docs/en/developer_guide/debug.md +++ b/docs/en/developer_guide/debug.md @@ -70,4 +70,45 @@ When running large scale RL, we will occationally meet the IMA in vLLM, there ar ## Step-by-Step Debugging with Ray Distributed Debugger -See verl's [Ray Debugging Tutorial](https://verl.readthedocs.io/en/latest/start/ray_debug_tutorial.html). +Ray provides a [distributed debugger](https://docs.ray.io/en/latest/ray-observability/ray-distributed-debugger.html) based on debugpy that lets you set breakpoints in the driver process and step through code interactively. + +1. Install debugpy: + + ```bash + pip install debugpy==1.8.0 + ``` + +2. Enable `RAY_DEBUG_POSTMORTEM` in your launch script: + + ```bash + export RAY_DEBUG_POSTMORTEM=1 + + RUNTIME_ENV_JSON="{ + \"env_vars\": { + ... + \"RAY_DEBUG_POSTMORTEM\": \"${RAY_DEBUG_POSTMORTEM:-0}\" + } + }" + + ray job submit --address="http://127.0.0.1:8265" \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + -- python3 train.py [args...] + ``` + +3. Add `ray.init()` before `breakpoint()` in `train.py`: + + ```python + if __name__ == "__main__": + ray.init() + breakpoint() + args = parse_args() + train(args) + ``` + + `ray.init()` is required because the distributed debugger depends on `core_worker`, which is only available after Ray initialization. Without it, `breakpoint()` raises `AttributeError: 'Worker' object has no attribute 'core_worker'`. + +4. Connect via VS Code: + + Install the [Ray Distributed Debugger](https://marketplace.visualstudio.com/items?itemName=ray-project.ray-distributed-debugger) extension in VS Code. Run your launch script to submit the job. Once the job hits `breakpoint()`, open the Ray Dashboard panel in VS Code and click the active breakpoint to attach the debugger. You can then step through code, inspect variables, and set additional breakpoints directly in the editor. + +> **Note**: Remove `ray.init()` and `breakpoint()` after debugging. An explicit `ray.init()` without arguments may cause issues in multi-node training where Ray injects specific namespace and runtime environment configurations via `ray job submit`. diff --git a/docs/en/developer_guide/profiling.md b/docs/en/developer_guide/profiling.md index 51bf9e014..4e7df22e6 100644 --- a/docs/en/developer_guide/profiling.md +++ b/docs/en/developer_guide/profiling.md @@ -1,12 +1,24 @@ # Profiling -In vime, we can perform detailed performance analysis of the rollout process using the profiling interface provided by vLLM. +In vime, you can profile the **rollout (vLLM inference)** path in detail using vLLM's profiling HTTP API. Profiling targets the vLLM engine side, not the Megatron training side. -## 1. Sleeping the Rollout Process +Typical flow: -For more flexible stress testing and profiling, it is often useful to make the vime rollout process enter a waiting state after initialization, instead of starting generation immediately. +- Start train (`sleep_rollout` + `vllm-profiler-config`) +- Wait until vLLM engines and the router are ready +- Read router/worker addresses from logs +- `start_profile` +- Send a few inference requests +- (Optional) `stop_profile`; or traces flush automatically when `max_iterations` is reached +- Inspect trace files under `torch_profiler_dir` -You can achieve this by replacing the `rollout_function_path` in your startup arguments without modifying the source code: + + +## 1. Put Rollout into a Wait State (`sleep_rollout`) + +For flexible stress testing and profiling, rollout usually waits after initialization instead of generating immediately. + +Replace `rollout_function_path` in `train.py` startup args—no code changes required: ```bash python train.py \ @@ -14,60 +26,305 @@ python train.py \ ... (other arguments) ``` -This function will make the rollout process enter an infinite wait loop, allowing you to manually send requests or run stress testing tools. +This puts the rollout process in an infinite wait loop so you can send HTTP requests or run stress tools manually. + +## 2. Enable the vLLM Profiler (at train startup) -## 2. Enabling the vLLM Profiler +vLLM registers `/start_profile` and `/stop_profile` only when started with `--profiler-config`. In vime, pass **`--vllm-profiler-config`** through to the `vllm serve` subprocess. -vLLM only registers the `/start_profile` and `/stop_profile` endpoints when started with a profiler config. In vime, pass it through to the `vllm serve` subprocess with `--vllm-profiler-config`: +### 2.1 Pass the full config as JSON ```bash --vllm-profiler-config '{"profiler":"torch","torch_profiler_dir":"/root/logs/vllm_profile","max_iterations":3,"ignore_frontend":true}' ``` -**Key Fields:** -* `profiler`: `"torch"` or `"cuda"`. -* `torch_profiler_dir`: Trace output directory (absolute path). -* `max_iterations`: Worker auto-stops and flushes after this many steps. -* `ignore_frontend`: Recommended `true`; profile workers only. +Common JSON fields: + +| Field | Description | +|------|------| +| `profiler` | `"torch"` or `"cuda"` | +| `torch_profiler_dir` | Trace output directory (absolute path) | +| `max_iterations` | Worker auto-stops and flushes after more than N steps (condition is `> N`) | +| `ignore_frontend` | Recommended `true`: profile workers only, lower frontend overhead | + +**Avoid RPC timeout on `stop_profile`:** vLLM APIServer talks to EngineCore/workers over internal RPC. Manually calling `stop_profile` to flush traces can take minutes, while the default `VLLM_RPC_TIMEOUT` is only **10 seconds** (10000 ms), which can interrupt flush or leave traces incomplete. For profiling, set **30 minutes** (1800000 ms). + +Set this variable **before starting train and launching vLLM**, in the Ray worker environment (a local shell `export` may not reach the Ray job). Pass it via `runtime-env-json` on `ray job submit`, for example: + +```bash +export VLLM_RPC_TIMEOUT="${VLLM_RPC_TIMEOUT:-1800000}" + +RUNTIME_ENV_JSON="{ + \"env_vars\": { + \"PYTHONPATH\": \"/root/vime:/root/Megatron-LM\", + \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", + \"VLLM_RPC_TIMEOUT\": \"${VLLM_RPC_TIMEOUT}\" + } +}" + +ray job submit --address=\"http://127.0.0.1:8265\" \ + --runtime-env-json=\"${RUNTIME_ENV_JSON}\" \ + -- python3 train.py \ + ... \ + --vllm-profiler-config '{\"profiler\":\"torch\",\"torch_profiler_dir\":\"/root/logs/vllm_profile\",...}' +``` + + +### 2.2 Verify it took effect + +After train starts, confirm all three in logs (missing any means the profiler is not enabled correctly): + +1. **Args parsed**: `vllm_profiler_config ... profiler='torch'` (and `torch_profiler_dir` path). +2. **Forwarded to vLLM subprocess**: `Launching vLLM server: ... --profiler-config {"profiler":"torch",...}`. +3. **HTTP routes registered**: vLLM startup route list includes `/start_profile` and `/stop_profile` (otherwise `POST /start_profile` returns 404). + +## 3. Get Router and Worker Addresses + +vLLM engines (workers) register on the vllm-router. Example startup log: + +```text +Router launched at 127.0.0.1:3521, Prometheus port: 4153 +Ports for engine 0: {'host': '127.0.0.1', 'port': 15000, ...} +Starting vLLM server on http://127.0.0.1:15000 +``` + +**Note: the router port may change on every job** (random in 3000–4000 by default). Do not reuse the previous port. Verify with curl: + +```bash +curl http://127.0.0.1:3521/workers +``` + +Returns each worker's `url` and `is_healthy`. -## 3. Obtaining vLLM Engine List +## 4. Use `tools/profile_rollout.py` -vLLM engines (workers) are registered with the router. You can retrieve the list of all active engines by accessing the `/workers` endpoint of the router. +The script reads the router's `/workers` list and calls `/start_profile` or `/stop_profile` on every worker. -The router address is typically printed in the startup logs: +### Start Profiling + +```bash +cd /root/vime +python tools/profile_rollout.py \ + --router-url http://127.0.0.1:3521 \ + --action start ``` -Router launched at 127.0.0.1:3000 + +### Stop Profiling (optional) + +If `--vllm-profiler-config` sets `max_iterations`, the worker **auto-stops and flushes** after enough steps. In practice, traces often appear under `torch_profiler_dir` right after inference—you **do not** need to call `stop_profile` manually. Use this only to end collection early: + +```bash +python tools/profile_rollout.py \ + --router-url http://127.0.0.1:3521 \ + --action stop ``` -You can use `curl` to view the workers: +## 5. Send Inference Requests + +While `sleep_rollout` is waiting: + +1. `profile_rollout.py --action start` +2. Send a few completion requests to the router or **directly to a worker** (2–4 is enough; traces get large) +3. (Optional) `profile_rollout.py --action stop`; or wait for `max_iterations` to auto-flush +4. Inspect traces under `torch_profiler_dir` + +Example request (`model` is the HF checkpoint path): + ```bash -curl http://127.0.0.1:3000/workers +curl -X POST http://127.0.0.1:15000/v1/completions \ + -H "Content-Type: application/json" \ + -d '{"model":"/root/models/Qwen3-4B","prompt":"Hello","max_tokens":32}' ``` -## 4. Using Automated Profiling Tool -To simplify profiling across multiple engines simultaneously, we provide an automated script: `tools/profile_rollout.py`. +## 6. View Traces + +### Perfetto + +1. Open [https://ui.perfetto.dev/](https://ui.perfetto.dev/) +2. **Open trace file**, pick `*.trace.json.gz` +3. Inspect GPU kernels, CPU ops, and the timeline + +### Chrome Tracing -### Starting Profiling +Open `chrome://tracing` in the browser and **Load** a trace file. -By default, this tool starts profiling on all workers: +### Analysis Tool ```bash -python tools/profile_rollout.py --router-url http://127.0.0.1:3000 --action start +cd /root/vime +python tools/analyze_profile.py --profile-dir /root/logs/vllm_profile --all-ranks ``` -### Stopping Profiling Manually -If you set `max_iterations`, the worker auto-stops and flushes. To stop early: +## 7. Troubleshooting + +| Symptom | Fix | +|------|------| +| `POST /start_profile` 404 | Pass `--vllm-profiler-config` as JSON; restart the job | +| Start OK but empty output dir | Confirm curl hits a worker and returns 200; increase `max_iterations` or send more requests | +| Router 503 | Confirm the current job's router port; connect directly to a worker | +| Slow or timed-out stop | Increase `VLLM_RPC_TIMEOUT`; reduce request count | + +## 8. Full Runnable Example + +The script below assumes a **container** environment, vime at `/root/vime`, models and data under `/root/models` and `/root/data`. Two parts: + +1. **`launch_train_for_profiling`**: start train with the profiler (`sleep_rollout`, minimal single-GPU colocate example—adjust GPU layout for your machine). +2. **`run_profiling_session`**: run profiling from **another terminal** after train is ready. + +Save as `/root/vime/run_profiling_demo.sh` and run: ```bash -python tools/profile_rollout.py --router-url http://127.0.0.1:3000 --action stop +#!/usr/bin/env bash +# +# Full vime rollout profiling example +# Usage: +# bash /root/vime/run_profiling_demo.sh launch # terminal 1: start train +# bash /root/vime/run_profiling_demo.sh profile # terminal 2: capture traces after train is ready +# +set -euo pipefail + +VIME_ROOT="${VIME_ROOT:-/root/vime}" +HF_CKPT="${HF_CKPT:-/root/models/Qwen3-4B}" +REF_LOAD="${REF_LOAD:-/root/models/Qwen3-4B_torch_dist}" +PROMPT_DATA="${PROMPT_DATA:-/root/data/gsm8k/train.parquet}" +LOG_ROOT="${LOG_ROOT:-/root/logs/vime_profiling}" +PROFILE_DIR="${PROFILE_DIR:-/root/logs/vllm_profile}" +TRAIN_LOG="${LOG_ROOT}/train_profiling.log" +ROUTER_HOST="${ROUTER_HOST:-127.0.0.1}" + +mkdir -p "${LOG_ROOT}" "${PROFILE_DIR}" + +VLLM_PROFILER_CONFIG_JSON="$(printf \ + '{"profiler":"torch","torch_profiler_dir":"%s","max_iterations":3,"ignore_frontend":true}' \ + "${PROFILE_DIR}")" + +launch_train_for_profiling() { + cd "${VIME_ROOT}" + + # Clean up old Ray / vLLM processes (comment out if not needed) + ray stop --force || true + pkill -9 -f "vllm serve" || true + sleep 2 + + ray start --head --node-ip-address 127.0.0.1 --num-gpus 2 --disable-usage-stats + + source "${VIME_ROOT}/scripts/models/qwen3-4B.sh" + + export VLLM_RPC_TIMEOUT="${VLLM_RPC_TIMEOUT:-1800000}" + + RUNTIME_ENV_JSON="{ + \"env_vars\": { + \"PYTHONPATH\": \"${VIME_ROOT}:/root/Megatron-LM\", + \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", + \"VLLM_RPC_TIMEOUT\": \"${VLLM_RPC_TIMEOUT}\" + } + }" + + echo "=== Launching train; log: ${TRAIN_LOG} ===" + echo "=== After engines are up, run: bash $0 profile ===" + + ray job submit --address="http://127.0.0.1:8265" \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + -- python3 train.py \ + --train-backend megatron \ + --colocate \ + --actor-num-nodes 1 \ + --actor-num-gpus-per-node 1 \ + --rollout-num-gpus 1 \ + --rollout-num-gpus-per-engine 1 \ + --rollout-backend vllm \ + --rollout-function-path vime.rollout.sleep_rollout.sleep \ + --hf-checkpoint "${HF_CKPT}" \ + --ref-load "${REF_LOAD}" \ + --prompt-data "${PROMPT_DATA}" \ + --input-key question \ + --label-key label \ + --apply-chat-template \ + --rm-type deepscaler \ + --num-rollout 1 \ + --rollout-batch-size 4 \ + --n-samples-per-prompt 1 \ + --rollout-max-response-len 512 \ + --global-batch-size 4 \ + --vllm-gpu-memory-utilization 0.7 \ + --vllm-profiler-config "${VLLM_PROFILER_CONFIG_JSON}" \ + ${MODEL_ARGS[@]} \ + 2>&1 | tee "${TRAIN_LOG}" +} + +discover_router_url() { + local line port + line="$(grep -E 'Router launched at' "${TRAIN_LOG}" | tail -1 || true)" + if [[ -z "${line}" ]]; then + echo "ERROR: Router not found in ${TRAIN_LOG}. Is train still starting?" >&2 + exit 1 + fi + # Router launched at 127.0.0.1:3521, Prometheus port: ... + port="$(echo "${line}" | sed -n 's/.*Router launched at [^:]*:\([0-9]*\).*/\1/p')" + echo "http://${ROUTER_HOST}:${port}" +} + +discover_worker_url() { + local router_url="$1" + python3 - <<'PY' "${router_url}" +import json, sys, urllib.request +router = sys.argv[1] +with urllib.request.urlopen(f"{router}/workers", timeout=10) as r: + workers = json.load(r).get("workers", []) +if not workers: + raise SystemExit("No workers registered") +print(workers[0]["url"]) +PY +} + +run_profiling_session() { + cd "${VIME_ROOT}" + + local router_url worker_url model="${HF_CKPT}" + router_url="$(discover_router_url)" + worker_url="$(discover_worker_url "${router_url}")" + + echo "=== ROUTER=${router_url} WORKER=${worker_url} PROFILE_DIR=${PROFILE_DIR} ===" + + echo "=== 1/3 start_profile (all workers via router) ===" + python tools/profile_rollout.py --router-url "${router_url}" --action start + + echo "=== 2/3 send completions (direct to worker; 3 requests) ===" + for i in 1 2 3; do + curl -sS -X POST "${worker_url}/v1/completions" \ + -H "Content-Type: application/json" \ + -d "{\"model\":\"${model}\",\"prompt\":\"Hello ${i}\",\"max_tokens\":32}" \ + | head -c 400 + echo + done + + echo "=== 3/3 list trace files (max_iterations=3 auto-stop; add --action stop if needed) ===" + sleep 2 + find "${PROFILE_DIR}" -type f \( -name '*.json*' -o -name 'profiler_out_*' \) | sort + echo "Open *.trace.json.gz in https://ui.perfetto.dev/ or run:" + echo " python tools/analyze_profile.py --profile-dir ${PROFILE_DIR} --all-ranks" +} + +case "${1:-}" in + launch) launch_train_for_profiling ;; + profile) run_profiling_session ;; + *) + echo "Usage: $0 {launch|profile}" >&2 + exit 1 + ;; +esac ``` -## 5. Running Stress Tests +**Steps:** + +```bash +# Terminal 1: start train (wait until logs show Router launched at ...) +bash /root/vime/run_profiling_demo.sh launch + +# Terminal 2: capture traces +bash /root/vime/run_profiling_demo.sh profile +``` -While the Rollout process is in a waiting state via `sleep_rollout`, you can: -1. Start profiling using `tools/profile_rollout.py`. -2. Use stress testing tools to send requests to the router or directly to the engines. -3. Wait for profiling to complete (if `max_iterations` was set) or stop it manually. -4. Collect the `.json` trace files from the `torch_profiler_dir` and view them using `chrome://tracing` in Chrome or [Perfetto](https://ui.perfetto.dev/). +Adjust paths at the top of the script (`/root/models/...`, `/root/data/...`) and GPU layout (`actor-num-gpus-per-node`, `rollout-num-gpus`, etc.) as needed. diff --git a/docs/en/examples/qwen3-30B-A3B.md b/docs/en/examples/qwen3-30B-A3B.md index 9fd7ebd2b..bca56238d 100644 --- a/docs/en/examples/qwen3-30B-A3B.md +++ b/docs/en/examples/qwen3-30B-A3B.md @@ -74,15 +74,82 @@ Here, we will briefly introduce the MoE-related parts in the [run-qwen3-30B-A3B. ### Multi-Node Support -For a multi-node environment, the following modifications are necessary: +The following uses **two machines with 8 GPUs each (16 GPUs total)** as the starting example; scripts and parameters scale to **N nodes**. Key differences from single-node: - - Place the training model and data on a path accessible by all nodes. - - Set the `MASTER_ADDR` to an address that is accessible by all nodes. - - Remove configurations related to CPU Adam. This is because a distributed optimizer is used, which significantly reduces the optimizer's video memory (VRAM) usage in a multi-node setup. +- Place weights, checkpoints, and data on storage visible to every node (e.g. NFS). +- Set `MASTER_ADDR` to the head **LAN IP** (not `127.0.0.1`). +- Omit CPU Adam (multi-node uses a distributed optimizer; do not use `--optimizer-cpu-offload`). +- `global-batch-size` must equal `rollout-batch-size × n-samples-per-prompt`. -In addition, you can make the following changes: +#### Topology - - When the total number of GPUs is not a multiple or divisor of the total number of experts, you can enable vLLM's EPLB (Expert Parallelism Load Balancer) and configure redundant experts via `--vllm-eplb-config` to add redundant experts. For example, in a 24-GPU scenario, you can configure it as follows: +| Component | Dual-node defaults | +|-----------|-------------------| +| Cluster | `ACTOR_NUM_NODES=2`, `ACTOR_NUM_GPUS_PER_NODE=8` | +| Megatron training | TP=8, EP=8, CP=2 (experts sharded across nodes) | +| vLLM rollout | Cross-node TP=16 (`rollout-num-gpus-per-engine = nodes × GPUs per node`) | +| Scheduling | Ray cluster + `--colocate` mode | + +Convert checkpoints with Megatron parallelism matching training (dual-node: TP=8, EP=8). Checkpoint EP must match `--expert-model-parallel-size`, or `load_checkpoint` may hang or resharding may be extremely slow. + +#### Start the Ray Cluster + +Start Ray **outside** the training script on each node. Join all workers first; verify `ray status` reports the expected GPU count, then submit training from the head. Dual-node example: + +```bash +# === Head node === +export MASTER_ADDR= +ray start --head --node-ip-address="${MASTER_ADDR}" --num-gpus 8 --disable-usage-stats \ + --dashboard-host=0.0.0.0 --dashboard-port=8265 + +# === Each worker node === +export MASTER_ADDR= +ray start --address="${MASTER_ADDR}:6379" --node-ip-address= --num-gpus 8 +``` + +See [Quick Start — Multi-node training](../get_started/quick_start.md#multi-node-training-for-large-scale-moe-models) for more details. + +#### Run Training + +After the Ray cluster is ready, on the **head node** set multi-node env vars and run the **same script as single-node** (`ACTOR_NUM_NODES>1` skips Ray startup and applies multi-node defaults): + +```bash +export MASTER_ADDR= +export ACTOR_NUM_NODES=2 +export ACTOR_NUM_GPUS_PER_NODE=8 +cd /root/vime +bash scripts/run-qwen3-30B-A3B.sh +``` + +2-step smoke test: + +```bash +NUM_ROLLOUT=2 ENABLE_R3=0 bash scripts/run-qwen3-30B-A3B.sh +``` + +To scale to N nodes (e.g. 4×8), join all workers to Ray, set `ACTOR_NUM_NODES=4` on the head, and tune `MEGATRON_TP` / `MEGATRON_EP` / `MEGATRON_CP` / `ROLLOUT_NUM_GPUS_PER_ENGINE` for total GPU count. + +#### Key Multi-Node Parameters + +| Variable | Dual-node default | Description | +|----------|-------------------|-------------| +| `ACTOR_NUM_NODES` | 2 (default 1 for single-node) | Total nodes including head; script skips Ray startup when >1 | +| `ACTOR_NUM_GPUS_PER_NODE` | 8 | GPUs per node | +| `MEGATRON_TP` / `MEGATRON_EP` / `MEGATRON_CP` | 8 / 8 / 2 | Megatron parallelism | +| `ROLLOUT_NUM_GPUS_PER_ENGINE` | total GPUs | vLLM engine GPU count | +| `ENABLE_R3` | 1 | set to 0 to disable R3 | + +Default batch: `rollout-batch-size=4`, `n-samples-per-prompt=2`, `global-batch-size=8`; vLLM uses `--vllm-moe-backend triton`. + +#### Multi-Node Troubleshooting + +- **Worker cannot join Ray / NCCL failures**: check `MASTER_ADDR`, container `/etc/hosts` (hostname must not map to `127.0.0.1`), `NCCL_SOCKET_IFNAME` / `GLOO_SOCKET_IFNAME`. +- **`Not enough samples X for global_batch_size Y`**: keep `global-batch-size` equal to `rollout-batch-size × n-samples-per-prompt`. +- **GPU memory full but no processes**: restart the container or run `ray stop --force` to clear stale vLLM contexts. + +#### EPLB + +When the total number of GPUs is not a multiple or divisor of the total number of experts, enable vLLM's EPLB (Expert Parallelism Load Balancer) and configure redundant experts via `--vllm-eplb-config`. For example, in a 24-GPU scenario: ```bash VLLM_ARGS=( diff --git a/docs/en/examples/qwen3-4B.md b/docs/en/examples/qwen3-4B.md index 3068ea221..ddb6151b2 100644 --- a/docs/en/examples/qwen3-4B.md +++ b/docs/en/examples/qwen3-4B.md @@ -193,7 +193,7 @@ OPTIMIZER_ARGS=( #### VLLM\_ARGS -These are the parameters required by vLLM. Here, `--rollout-num-gpus-per-engine` basically corresponds to vLLM's `tensor_parallel_size`. Other vLLM parameters are passed to vime by adding the `--vllm-` prefix. +Parameters for vLLM inference. vime uses vLLM as the rollout backend by default (`rollout.py` launches `VLLMEngine`; the default rollout function is `vime.rollout.vllm_rollout.generate_rollout`), so no extra backend flag is needed. `--rollout-num-gpus-per-engine` corresponds to each vLLM engine's `tensor_parallel_size`. Other vLLM parameters are passed to vime with a `--vllm-` prefix (for example, `--vllm-max-model-len`). ```bash VLLM_ARGS=( @@ -202,7 +202,9 @@ VLLM_ARGS=( ) ``` -⚠️ vime uses the vLLM router to schedule multiple vLLM servers. `dp_size` is not supported when DP attention is disabled. +When rollout concurrency is high, tune the vLLM scheduler via the `--vllm-` prefix—for example, `--vllm-max-num-seqs` and `--vllm-max-num-batched-tokens`. Add `--vllm-enforce-eager` for debugging or to work around CUDA graph limits. + +⚠️ vime uses the vLLM router to schedule multiple vLLM servers. With co-located training and inference (`--colocate`), weights are synchronized via CUDA IPC; with decoupled training and inference, the trainer synchronizes weights with vLLM engines over NCCL. ### Dynamic Sampling @@ -276,21 +278,22 @@ ray job submit ... \ ... ``` -In this case, 2 GPUs will be allocated for training, and 6 GPUs will be allocated for inference. - -⚠️ If the concurrency on each vLLM server is too high, it may exceed vLLM's default CUDA graph concurrency limit, which will affect inference speed. You can adjust this in the following two ways: +In this case, 2 GPUs will be allocated for training, and 6 GPUs will be allocated for inference. Like `--actor-num-gpus-per-node`, `--rollout-num-gpus` is a **Ray resource argument** passed to `train.py`: the framework uses it to build the placement group and assign the first bundles to training actors and the remaining bundles to rollout engines (see `vime/ray/placement_group.py`). **Under co-located mode (`--colocate`), this argument is ignored** and is set automatically to `actor_num_gpus_per_node * actor_num_nodes`. Do not put `--rollout-num-gpus` in `VLLM_ARGS`. -1. Use `--vllm-server-concurrency` to limit the maximum number of concurrent requests sent to a single vLLM server. For example: +For decoupled training and inference, `VLLM_ARGS` only needs inference-backend settings, for example: - ```bash - --vllm-server-concurrency 160 - ``` +```bash +VLLM_ARGS=( + --rollout-num-gpus-per-engine 2 + --vllm-gpu-memory-utilization 0.9 + --vllm-max-num-seqs 256 + --vllm-max-num-batched-tokens 8192 +) +``` -2. Use `--vllm-cudagraph-capture-sizes` to increase the number of CUDA graphs initialized by vLLM. For example: +Add `--vllm-enforce-eager` when debugging or to work around CUDA graph limits. - ```bash - --vllm-cudagraph-capture-sizes 1 2 4 8 $(seq 16 8 256) - ``` +⚠️ When using co-located training and inference, Megatron will always occupy some GPU memory. Reduce vLLM's memory footprint with `--vllm-gpu-memory-utilization`, and reserve headroom for training with `--train-memory-margin-bytes`. ### Asynchronous Training diff --git a/docs/zh/developer_guide/debug.md b/docs/zh/developer_guide/debug.md index 21148b038..c41a3a95c 100644 --- a/docs/zh/developer_guide/debug.md +++ b/docs/zh/developer_guide/debug.md @@ -68,4 +68,45 @@ vime 支持将训练部分和推理部分分开进行调试，从而实现： ## 使用 Ray Distributed Debugger 单步调试 -请参考 verl 的 [Ray Debugging Tutorial](https://verl.readthedocs.io/en/latest/start/ray_debug_tutorial.html)。 +Ray 提供了基于 debugpy 的[分布式调试器](https://docs.ray.io/en/latest/ray-observability/ray-distributed-debugger.html)，可以在 driver 进程中设置断点并单步执行代码。 + +1. 安装 debugpy： + + ```bash + pip install debugpy==1.8.0 + ``` + +2. 在启动脚本中启用 `RAY_DEBUG_POSTMORTEM`： + + ```bash + export RAY_DEBUG_POSTMORTEM=1 + + RUNTIME_ENV_JSON="{ + \"env_vars\": { + ... + \"RAY_DEBUG_POSTMORTEM\": \"${RAY_DEBUG_POSTMORTEM:-0}\" + } + }" + + ray job submit --address="http://127.0.0.1:8265" \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + -- python3 train.py [args...] + ``` + +3. 在 `train.py` 中的 `breakpoint()` 前添加 `ray.init()`： + + ```python + if __name__ == "__main__": + ray.init() + breakpoint() + args = parse_args() + train(args) + ``` + + 必须先调用 `ray.init()`，因为分布式调试器依赖 `core_worker`，只有 Ray 初始化后才可用。否则 `breakpoint()` 会报错 `AttributeError: 'Worker' object has no attribute 'core_worker'`。 + +4. 通过 VS Code 连接调试器： + + 在 VS Code 中安装 [Ray Distributed Debugger](https://marketplace.visualstudio.com/items?itemName=ray-project.ray-distributed-debugger) 扩展。运行启动脚本提交 job 后，当 job 执行到 `breakpoint()` 暂停后，在 VS Code 的 Ray Dashboard 面板中点击活跃的断点即可 attach 调试器，之后可以直接在编辑器中单步执行、查看变量、设置新断点。 + +> **注意**：调试完成后务必移除 `ray.init()` 和 `breakpoint()`。不带参数的 `ray.init()` 在多节点训练中可能导致问题，因为 `ray job submit` 会注入特定的 namespace 和 runtime environment 配置。 diff --git a/docs/zh/developer_guide/profiling.md b/docs/zh/developer_guide/profiling.md index c74fbff27..2717c2218 100644 --- a/docs/zh/developer_guide/profiling.md +++ b/docs/zh/developer_guide/profiling.md @@ -1,12 +1,24 @@ -# 性能分析 (Profiling) +# 性能分析（Profiling） -在 vime 中，我们可以通过 vLLM 提供的 profiling 接口对 rollout 过程进行详细的性能分析。 +在vime中，我们可以通过vLLM提供的profiling接口对**rollout（vLLM推理）**过程做详细的性能分析。Profiling针对vLLM engine侧，不是Megatron训练侧。 -## 1. 使 Rollout 进程进入等待状态 (Sleep Rollout) +典型流程： -为了更自由地进行压力测试和性能分析，我们通常需要让 vime 的 rollout 进程在初始化完成后进入等待状态，而不是立即开始生成。 +- 启动train（sleep_rollout + vllm-profiler-config） +- 等待vLLM engine与router就绪 +- 从日志确认router/worker地址 +- start_profile +- 发送少量推理请求 +-（可选）stop_profile；或达到max_iterations后自动落盘 +- 在torch_profiler_dir查看trace文件 -你可以通过在启动参数中替换 `rollout_function_path` 来实现，而无需修改代码： + + +## 1. 使Rollout进入等待状态（sleep_rollout） + +为了更灵活地压测和profiling，通常让rollout在初始化完成后进入等待，而不是立即开始生成。 + +在 `train.py` 启动参数中替换 `rollout_function_path` 即可，无需改代码： ```bash python train.py \ @@ -14,60 +26,305 @@ python train.py \ ... (其他参数) ``` -该函数会让 rollout 进程进入无限循环等待状态，方便你手动发送请求或运行压测工具。 +该函数会让rollout进程进入无限循环等待，便于手动发HTTP请求或运行压测工具。 + +## 2. 启用vLLM Profiler（启动train时配置） -## 2. 启用 vLLM Profiler +vLLM只有在启动时配置了`--profiler-config`，才会注册`/start_profile`与`/stop_profile`路由。在vime中通过**`--vllm-profiler-config`**转发给`vllm serve`子进程。 -vLLM 只有在启动时配置了 profiler config 才会注册 `/start_profile` 与 `/stop_profile` 接口。在 vime 中通过 `--vllm-profiler-config` 转发给 `vllm serve` 子进程： +### 2.1 使用JSON整包传参 ```bash --vllm-profiler-config '{"profiler":"torch","torch_profiler_dir":"/root/logs/vllm_profile","max_iterations":3,"ignore_frontend":true}' ``` -**常用字段说明：** -* `profiler`: `"torch"` 或 `"cuda"`。 -* `torch_profiler_dir`: trace 输出目录（绝对路径）。 -* `max_iterations`: worker 记录该步数后自动 stop 并落盘。 -* `ignore_frontend`: 建议 `true`，仅 profile worker。 +常用JSON字段： + +| 字段 | 说明 | +|------|------| +| `profiler` | `"torch"` 或 `"cuda"` | +| `torch_profiler_dir` | trace输出目录（绝对路径） | +| `max_iterations` | worker记录超过N步后自动stop并落盘（条件为`> N`） | +| `ignore_frontend` | 建议`true`，仅profile worker，降低前端开销 | + +**防止`stop_profile`时RPC超时：** vLLM APIServer与EngineCore/worker之间通过内部RPC通信。手动调用`stop_profile`触发trace落盘可能耗时数分钟，而默认`VLLM_RPC_TIMEOUT`仅**10秒**（10000 ms），容易导致flush中断或trace不完整。Profiling时建议设为**30分钟**（1800000 ms）。 + +该变量须在**启动train、拉起vLLM之前**传入Ray worker环境（仅在本机shell `export`不一定会进入Ray job）。在`ray job submit`的`runtime-env-json`中写入，例如： + +```bash +export VLLM_RPC_TIMEOUT="${VLLM_RPC_TIMEOUT:-1800000}" + +RUNTIME_ENV_JSON="{ + \"env_vars\": { + \"PYTHONPATH\": \"/root/vime:/root/Megatron-LM\", + \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", + \"VLLM_RPC_TIMEOUT\": \"${VLLM_RPC_TIMEOUT}\" + } +}" + +ray job submit --address=\"http://127.0.0.1:8265\" \ + --runtime-env-json=\"${RUNTIME_ENV_JSON}\" \ + -- python3 train.py \ + ... \ + --vllm-profiler-config '{\"profiler\":\"torch\",\"torch_profiler_dir\":\"/root/logs/vllm_profile\",...}' +``` + + +### 2.2 验证是否生效 + +启动train后，在日志中确认以下三点（缺任一项说明profiler未正确启用）： + +1. **参数已解析**：出现`vllm_profiler_config ... profiler='torch'`（及`torch_profiler_dir`路径）。 +2. **已转发给vLLM子进程**：出现`Launching vLLM server: ... --profiler-config {"profiler":"torch",...}`。 +3. **HTTP路由已注册**：vLLM启动时的路由列表中包含`/start_profile`与`/stop_profile`（否则`POST /start_profile`会返回404）。 + +## 3. 获取Router与Worker地址 + +vLLM engine（workers）注册在vllm-router上。启动日志示例： + +```text +Router launched at 127.0.0.1:3521, Prometheus port: 4153 +Ports for engine 0: {'host': '127.0.0.1', 'port': 15000, ...} +Starting vLLM server on http://127.0.0.1:15000 +``` + +**注意：router端口每次job可能变化**（默认在3000–4000随机），不要沿用上次端口。可用curl验证： + +```bash +curl http://127.0.0.1:3521/workers +``` + +返回每个worker的`url`与`is_healthy`。 -## 3. 获取 vLLM 引擎列表 +## 4. 使用`tools/profile_rollout.py` -vLLM 引擎（workers）注册在 router 上。你可以通过访问 router 的 `/workers` 接口来获取所有活跃引擎的列表。 +脚本通过router的`/workers`列表，对所有worker调用`/start_profile`或`/stop_profile`。 -通常 router 地址会在启动日志中打印： +### 启动Profiling + +```bash +cd /root/vime +python tools/profile_rollout.py \ + --router-url http://127.0.0.1:3521 \ + --action start ``` -Router launched at 127.0.0.1:3000 + +### 停止Profiling（可选） + +若在`--vllm-profiler-config`中设置了`max_iterations`，worker在记录足够步数后会**自动stop并落盘**，实践中发完推理后常可直接在`torch_profiler_dir`看到trace，**不必**再手动`stop_profile`。需要提前结束采集时再执行： + +```bash +python tools/profile_rollout.py \ + --router-url http://127.0.0.1:3521 \ + --action stop ``` -你可以使用 `curl` 查看 workers： +## 5. 发送推理请求 + +在sleep_rollout等待期间，执行步骤如下： + +1. `profile_rollout.py --action start` +2. 向router或**直连worker**发送少量completion请求（2～4条即可，trace会很大） +3. （可选）`profile_rollout.py --action stop`；或等待`max_iterations`触发自动落盘 +4. 在`torch_profiler_dir`查看trace + +请求示例（`model`使用HF checkpoint路径）： + ```bash -curl http://127.0.0.1:3000/workers +curl -X POST http://127.0.0.1:15000/v1/completions \ + -H "Content-Type: application/json" \ + -d '{"model":"/root/models/Qwen3-4B","prompt":"Hello","max_tokens":32}' ``` -## 4. 使用自动化 Profiling 工具 -为了简化对多个引擎同时进行 profiling 的操作，我们提供了一个自动化脚本 `tools/profile_rollout.py`。 +## 6. 查看Trace + +### Perfetto + +1. 打开 [https://ui.perfetto.dev/](https://ui.perfetto.dev/) +2. **Open trace file**，选择`*.trace.json.gz` +3. 查看GPU kernel、CPU算子与时间线 + +### Chrome Tracing -### 启动 Profiling +浏览器访问`chrome://tracing`，Load加载trace文件。 -默认情况下，该工具会对所有 worker 启动 profiling： +### 分析工具 ```bash -python tools/profile_rollout.py --router-url http://127.0.0.1:3000 --action start +cd /root/vime +python tools/analyze_profile.py --profile-dir /root/logs/vllm_profile --all-ranks ``` -### 手动停止 Profiling -如果你设置了 `max_iterations`，worker 会自动 stop 并落盘。想要提前停止： +## 7. 常见问题 + +| 现象 | 处理 | +|------|------| +| `POST /start_profile` 404 | 用JSON传`--vllm-profiler-config`；重启job | +| start成功但目录为空 | 确认curl打到worker且返回200；适当增大`max_iterations`或补发推理 | +| router 503 | 确认当前job的router端口；改直连worker | +| stop很慢或超时 | 增大`VLLM_RPC_TIMEOUT`；减少请求条数 | + +## 8. 完整可运行示例 + +以下脚本假设在**容器内**、vime仓库位于`/root/vime`，模型与数据在`/root/models`、`/root/data`。分两段： + +1. **`launch_train_for_profiling`**：启动带profiler的train（sleep_rollout，单卡colocate最小示例，可按机器改GPU数）。 +2. **`run_profiling_session`**：train就绪后，在**另一个终端**执行profiling。 + +将脚本保存为 `/root/vime/run_profiling_demo.sh` 后执行。 ```bash -python tools/profile_rollout.py --router-url http://127.0.0.1:3000 --action stop +#!/usr/bin/env bash +# +# vime rollout profiling 完整示例 +# 用法: +# bash /root/vime/run_profiling_demo.sh launch # 终端1：启动train +# bash /root/vime/run_profiling_demo.sh profile # 终端2：train就绪后抓trace +# +set -euo pipefail + +VIME_ROOT="${VIME_ROOT:-/root/vime}" +HF_CKPT="${HF_CKPT:-/root/models/Qwen3-4B}" +REF_LOAD="${REF_LOAD:-/root/models/Qwen3-4B_torch_dist}" +PROMPT_DATA="${PROMPT_DATA:-/root/data/gsm8k/train.parquet}" +LOG_ROOT="${LOG_ROOT:-/root/logs/vime_profiling}" +PROFILE_DIR="${PROFILE_DIR:-/root/logs/vllm_profile}" +TRAIN_LOG="${LOG_ROOT}/train_profiling.log" +ROUTER_HOST="${ROUTER_HOST:-127.0.0.1}" + +mkdir -p "${LOG_ROOT}" "${PROFILE_DIR}" + +VLLM_PROFILER_CONFIG_JSON="$(printf \ + '{"profiler":"torch","torch_profiler_dir":"%s","max_iterations":3,"ignore_frontend":true}' \ + "${PROFILE_DIR}")" + +launch_train_for_profiling() { + cd "${VIME_ROOT}" + + # 清理旧 Ray / vLLM 进程（按需注释） + ray stop --force || true + pkill -9 -f "vllm serve" || true + sleep 2 + + ray start --head --node-ip-address 127.0.0.1 --num-gpus 2 --disable-usage-stats + + source "${VIME_ROOT}/scripts/models/qwen3-4B.sh" + + export VLLM_RPC_TIMEOUT="${VLLM_RPC_TIMEOUT:-1800000}" + + RUNTIME_ENV_JSON="{ + \"env_vars\": { + \"PYTHONPATH\": \"${VIME_ROOT}:/root/Megatron-LM\", + \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", + \"VLLM_RPC_TIMEOUT\": \"${VLLM_RPC_TIMEOUT}\" + } + }" + + echo "=== Launching train; log: ${TRAIN_LOG} ===" + echo "=== After engines are up, run: bash $0 profile ===" + + ray job submit --address="http://127.0.0.1:8265" \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + -- python3 train.py \ + --train-backend megatron \ + --colocate \ + --actor-num-nodes 1 \ + --actor-num-gpus-per-node 1 \ + --rollout-num-gpus 1 \ + --rollout-num-gpus-per-engine 1 \ + --rollout-backend vllm \ + --rollout-function-path vime.rollout.sleep_rollout.sleep \ + --hf-checkpoint "${HF_CKPT}" \ + --ref-load "${REF_LOAD}" \ + --prompt-data "${PROMPT_DATA}" \ + --input-key question \ + --label-key label \ + --apply-chat-template \ + --rm-type deepscaler \ + --num-rollout 1 \ + --rollout-batch-size 4 \ + --n-samples-per-prompt 1 \ + --rollout-max-response-len 512 \ + --global-batch-size 4 \ + --vllm-gpu-memory-utilization 0.7 \ + --vllm-profiler-config "${VLLM_PROFILER_CONFIG_JSON}" \ + ${MODEL_ARGS[@]} \ + 2>&1 | tee "${TRAIN_LOG}" +} + +discover_router_url() { + local line port + line="$(grep -E 'Router launched at' "${TRAIN_LOG}" | tail -1 || true)" + if [[ -z "${line}" ]]; then + echo "ERROR: Router not found in ${TRAIN_LOG}. Is train still starting?" >&2 + exit 1 + fi + # Router launched at 127.0.0.1:3521, Prometheus port: ... + port="$(echo "${line}" | sed -n 's/.*Router launched at [^:]*:\([0-9]*\).*/\1/p')" + echo "http://${ROUTER_HOST}:${port}" +} + +discover_worker_url() { + local router_url="$1" + python3 - <<'PY' "${router_url}" +import json, sys, urllib.request +router = sys.argv[1] +with urllib.request.urlopen(f"{router}/workers", timeout=10) as r: + workers = json.load(r).get("workers", []) +if not workers: + raise SystemExit("No workers registered") +print(workers[0]["url"]) +PY +} + +run_profiling_session() { + cd "${VIME_ROOT}" + + local router_url worker_url model="${HF_CKPT}" + router_url="$(discover_router_url)" + worker_url="$(discover_worker_url "${router_url}")" + + echo "=== ROUTER=${router_url} WORKER=${worker_url} PROFILE_DIR=${PROFILE_DIR} ===" + + echo "=== 1/3 start_profile (all workers via router) ===" + python tools/profile_rollout.py --router-url "${router_url}" --action start + + echo "=== 2/3 send completions (direct to worker; 3 requests) ===" + for i in 1 2 3; do + curl -sS -X POST "${worker_url}/v1/completions" \ + -H "Content-Type: application/json" \ + -d "{\"model\":\"${model}\",\"prompt\":\"Hello ${i}\",\"max_tokens\":32}" \ + | head -c 400 + echo + done + + echo "=== 3/3 list trace files (max_iterations=3 auto-stop; add --action stop if needed) ===" + sleep 2 + find "${PROFILE_DIR}" -type f \( -name '*.json*' -o -name 'profiler_out_*' \) | sort + echo "Open *.trace.json.gz in https://ui.perfetto.dev/ or run:" + echo " python tools/analyze_profile.py --profile-dir ${PROFILE_DIR} --all-ranks" +} + +case "${1:-}" in + launch) launch_train_for_profiling ;; + profile) run_profiling_session ;; + *) + echo "Usage: $0 {launch|profile}" >&2 + exit 1 + ;; +esac ``` -## 5. 进行压力测试 +**操作步骤：** + +```bash +# 终端1：启动train（等待vLLM与router就绪，日志出现Router launched at ...） +bash /root/vime/run_profiling_demo.sh launch + +# 终端2：抓trace +bash /root/vime/run_profiling_demo.sh profile +``` -在 Rollout 进程通过 `sleep_rollout` 处于等待状态时，你可以： -1. 使用 `tools/profile_rollout.py` 启动 profiling。 -2. 使用压测工具向 router 或直接向引擎发送请求。 -3. 等待 profiling 完成（如果设置了 `max_iterations`）或手动停止。 -4. 在 `torch_profiler_dir` 中获取 `.json` trace 文件，并使用 `chrome://tracing` 或 [Perfetto](https://ui.perfetto.dev/) 查看。 +按需修改脚本顶部的`/root/models/...`、`/root/data/...`与GPU布局（`actor-num-gpus-per-node`、`rollout-num-gpus`等）。 diff --git a/docs/zh/examples/qwen3-30B-A3B.md b/docs/zh/examples/qwen3-30B-A3B.md index 525478417..039b3920b 100644 --- a/docs/zh/examples/qwen3-30B-A3B.md +++ b/docs/zh/examples/qwen3-30B-A3B.md @@ -73,14 +73,82 @@ bash scripts/run-qwen3-30B-A3B.sh ### 多机支持 -对于多机环境，需要进行如下的几点修改： -- 将训练模型，数据放在所有机器都可以访问到的路径上； -- 设置各台机器都可以访问到的 `MASTER_ADDR` 之外； -- 去掉 CPU adam 相关的配置，因为使用了 distributed optimizer，所以多机环境下 optimizer 的显存占比会明显下降。 +以下以 **2台机器、每台8卡（共16GPU）** 为入门示例；脚本与参数可扩展到 **N节点**。多机与单机的主要差异： -除此之外，还可以进行如下的修改： +- 训练模型、数据放在所有节点均可访问的路径（如 NFS）； +- `MASTER_ADDR` 设为 head 节点的 **局域网 IP**（非 `127.0.0.1`）； +- 去掉 CPU Adam（多机使用 distributed optimizer，无需 `--optimizer-cpu-offload`）； +- `global-batch-size` 必须等于 `rollout-batch-size × n-samples-per-prompt`。 -- 当总卡数并不能被 expert 总数乘除时，可以开启 vLLM 的 EPLB（Expert Parallelism Load Balancer），通过 `--vllm-eplb-config` 来增加冗余的 expert，例如对于 24 卡的场景，可以配置： +#### 拓扑概览 + +| 组件 | 双机默认配置 | +|------|------| +| 集群 | `ACTOR_NUM_NODES=2`，`ACTOR_NUM_GPUS_PER_NODE=8` | +| Megatron训练 | TP=8, EP=8, CP=2（expert 分片跨节点） | +| vLLM Rollout | 跨节点 TP=16（`rollout-num-gpus-per-engine = 节点数 × 每节点GPU`） | +| 调度 | Ray 集群 + `--colocate` 共卡模式 | + +转换 checkpoint 时建议使用与训练一致的 Megatron 并行度（双机示例 TP=8, EP=8）。checkpoint 的 EP 需与 `--expert-model-parallel-size` 一致，否则 `load_checkpoint` 可能极慢或卡住。 + +#### 启动 Ray 集群 + +Ray 集群需在各节点上 **单独启动**，不在训练脚本内。先在所有 worker 节点加入集群，确认 `ray status` 显示预期 GPU 总数后，再在 head 提交训练。示例（双机）： + +```bash +# === Head 节点 === +export MASTER_ADDR= +ray start --head --node-ip-address="${MASTER_ADDR}" --num-gpus 8 --disable-usage-stats \ + --dashboard-host=0.0.0.0 --dashboard-port=8265 + +# === 各 Worker 节点 === +export MASTER_ADDR= +ray start --address="${MASTER_ADDR}:6379" --node-ip-address=<本机_局域网_IP> --num-gpus 8 +``` + +更多说明见 [快速开始 — 多机训练](../get_started/quick_start.md#multi-node-training-for-large-scale-moe-models)。 + +#### 执行训练 + +Ray 集群就绪后，在 **head 节点** 设置多机环境变量并运行 **与单机相同的脚本**（`ACTOR_NUM_NODES>1` 时脚本不会启动 Ray，并使用多机默认参数）： + +```bash +export MASTER_ADDR= +export ACTOR_NUM_NODES=2 +export ACTOR_NUM_GPUS_PER_NODE=8 +cd /root/vime +bash scripts/run-qwen3-30B-A3B.sh +``` + +2 step 冒烟示例： + +```bash +NUM_ROLLOUT=2 ENABLE_R3=0 bash scripts/run-qwen3-30B-A3B.sh +``` + +扩展到 N 节点（例如 4×8）时，在各 worker 加入 Ray 后，于 head 设置 `ACTOR_NUM_NODES=4` 并按总卡数调整 `MEGATRON_TP` / `MEGATRON_EP` / `MEGATRON_CP` / `ROLLOUT_NUM_GPUS_PER_ENGINE`。 + +#### 多机关键参数 + +| 变量 | 双机默认 | 说明 | +|------|----------|------| +| `ACTOR_NUM_NODES` | 2（单机默认为 1） | 训练节点总数（含 head）；>1 时脚本不启动 Ray | +| `ACTOR_NUM_GPUS_PER_NODE` | 8 | 每节点 GPU 数 | +| `MEGATRON_TP` / `MEGATRON_EP` / `MEGATRON_CP` | 8 / 8 / 2 | Megatron 并行 | +| `ROLLOUT_NUM_GPUS_PER_ENGINE` | 总 GPU 数 | vLLM engine 占用卡数 | +| `ENABLE_R3` | 1 | 设为 0 可关闭 R3 路径 | + +脚本默认 batch：`rollout-batch-size=4`，`n-samples-per-prompt=2`，`global-batch-size=8`；vLLM 使用 `--vllm-moe-backend triton`。 + +#### 多机常见问题 + +- **Worker 无法加入 Ray / NCCL 失败**：检查 `MASTER_ADDR`、容器 `/etc/hosts`（hostname 勿指向 `127.0.0.1`）、`NCCL_SOCKET_IFNAME` / `GLOO_SOCKET_IFNAME`。 +- **`Not enough samples X for global_batch_size Y`**：同步调整 `global-batch-size` 与 `rollout-batch-size × n-samples-per-prompt`。 +- **GPU 显存占满但无进程**：重启容器或 `ray stop --force` 清理残留 vLLM 上下文。 + +#### EPLB + +当总卡数并不能被 expert 总数整除时，可以开启 vLLM 的 EPLB（Expert Parallelism Load Balancer），通过 `--vllm-eplb-config` 配置冗余 expert。例如对于 24 卡的场景： ```bash VLLM_ARGS=( diff --git a/docs/zh/examples/qwen3-4B.md b/docs/zh/examples/qwen3-4B.md index 85720a35a..e0fac9059 100644 --- a/docs/zh/examples/qwen3-4B.md +++ b/docs/zh/examples/qwen3-4B.md @@ -193,7 +193,7 @@ OPTIMIZER_ARGS=( #### VLLM_ARGS -vLLM 所需的参数，这里 `--rollout-num-gpus-per-engine` 基本对应 vLLM 的 `tensor_parallel_size`，除此之外的 vLLM 参数均通过添加 `--vllm-` 的前缀来传给 vime。 +vLLM 推理所需的参数。vime 默认使用 vLLM 作为 rollout 后端（`rollout.py` 启动 `VLLMEngine`，默认 rollout 函数为 `vime.rollout.vllm_rollout.generate_rollout`），无需额外指定 backend。`--rollout-num-gpus-per-engine` 对应每个 vLLM engine 的 `tensor_parallel_size`；除此之外的 vLLM 参数均通过添加 `--vllm-` 前缀传给 vime（例如 `--vllm-max-model-len`）。 ```bash VLLM_ARGS=( @@ -202,7 +202,9 @@ VLLM_ARGS=( ) ``` -⚠️ vime 会用 vLLM router 来调度多个 vLLM server，在不开启 dp attention 的情况下不支持 `dp_size`。 +rollout 并发较高时，还可以通过 `--vllm-` 前缀调节 vLLM scheduler，例如 `--vllm-max-num-seqs`、`--vllm-max-num-batched-tokens`；调试或规避 CUDA graph 相关限制时可加 `--vllm-enforce-eager`。 + +⚠️ vime 会用 vLLM router 来调度多个 vLLM server。训推一体（`--colocate`）时，训练与推理权重经 CUDA IPC 同步；训推分离时，训练侧经 NCCL 与 vLLM engine 同步权重。 ### dynamic sampling @@ -276,21 +278,22 @@ ray job submit ... \ ... ``` -此时，就会分配 2 张卡给训练，6 张卡给推理。 - -⚠️ 在进行训推分离的时候，每个 vLLM server 上的并发度太大，超过了 vLLM 默认的 cuda graph 的并发度，影响推理速度。可以用以下 2 种方式进行调整： +此时，就会分配 2 张卡给训练，6 张卡给推理。`--rollout-num-gpus` 与 `--actor-num-gpus-per-node` 一样，是传给 `train.py` 的 **Ray 资源参数**：框架据此创建 placement group，并把前若干 bundle 分给训练 actor、后续 bundle 分给 rollout engine（见 `vime/ray/placement_group.py`）。**共卡模式（`--colocate`）下该参数会被忽略**，并自动设为 `actor_num_gpus_per_node * actor_num_nodes`。请勿把 `--rollout-num-gpus` 写在 `VLLM_ARGS` 中。 -1. 通过 `--vllm-server-concurrency` 限制发给一个 vLLM server 的最大并发量，例如： +训推分离时，`VLLM_ARGS` 仅需配置推理后端相关参数，例如： - ```bash - --vllm-server-concurrency 160 - ``` +```bash +VLLM_ARGS=( + --rollout-num-gpus-per-engine 2 + --vllm-gpu-memory-utilization 0.9 + --vllm-max-num-seqs 256 + --vllm-max-num-batched-tokens 8192 +) +``` -2. 使用 `--vllm-cudagraph-capture-sizes` 增大 vLLM 初始化的 cuda graph 数量，例如： +如需调试或规避 CUDA graph 相关限制，可额外加上 `--vllm-enforce-eager`。 - ```bash - --vllm-cudagraph-capture-sizes 1 2 4 8 $(seq 16 8 256) - ``` +⚠️ 在训推一体的训练时，megatron 始终会占据一些显存，需要通过 `--vllm-gpu-memory-utilization` 来降低 vLLM 占据的显存比例，并配合 `--train-memory-margin-bytes` 为训练侧预留空间。 ### 异步训练 From e92c4f8729383c39c1d8e002ce18b13ddc3b35d8 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Tue, 9 Jun 2026 02:39:06 +0000 Subject: [PATCH 3/5] examples/scripts: mirror slime for repro script + fully_async README Both files had drifted far from their slime originals; restore them to slime's content with only the sanctioned translations. run-qwen2.5-0.5B-reproducibility.sh: #180 re-created this file with vime-only scaffolding (NUM_GPUS var, VIME_ROOT on PYTHONPATH, --dashboard-* flags, MASTER_ADDR, a RUNTIME_ENV_JSON heredoc, --train-backend megatron) and dropped the engine-kill. Revert to slime's script verbatim, applying only: pkill -9 sglang -> pkill -9 -f "vllm serve"; slime-dev -> vime-dev; SGLANG_ARGS -> VLLM_ARGS (+ --vllm-gpu-memory-utilization / --vllm-attention-backend / --vllm-enable-deterministic-inference). --train-backend megatron is vime's only choice + default (redundant); VIME_ROOT is unneeded (run-qwen3-30B-A3B.sh and slime both omit it). Kept one fix vs slime: source "${SCRIPT_DIR}/models/..." (slime's "${SCRIPT_DIR}/scripts/models/..." doesn't resolve in this layout). fully_async/README.md: was a full vime-invented rewrite (even its "Creating new global async worker..." log lines don't match the code). Restore slime's README structure/prose, translated, pointing at vime's actual run-qwen3-4b-fully_async.sh and examples/coding_agent_rl, with args.sglang_server_concurrency -> args.vllm_server_concurrency. The "You should see" lines now match what vime/rollout/fully_async_rollout.py actually logs. Co-Authored-By: Claude Opus 4.8 (1M context) --- examples/fully_async/README.md | 88 ++++++++++++++------- scripts/run-qwen2.5-0.5B-reproducibility.sh | 32 +++----- 2 files changed, 73 insertions(+), 47 deletions(-) diff --git a/examples/fully_async/README.md b/examples/fully_async/README.md index 7fb271f6f..ac05e1779 100644 --- a/examples/fully_async/README.md +++ b/examples/fully_async/README.md @@ -1,45 +1,77 @@ -# 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 vime'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 `vime.rollout.fully_async_rollout`; this +directory is just the launch script. ## Files -The worker itself lives in `vime.rollout.fully_async_rollout`; this directory is just the launch script: -* `run-qwen3-4b-fully_async.sh`: example launch script with Qwen3‑4B. -## Prerequisite -First set up model & environment following the Qwen3-4B example. +* `run-qwen3-4b-fully_async.sh` — fully-async demo with Qwen3-4B on + dapo-math-17k. + +## Prerequisites + +``` +/root/Qwen3-4B/ # HF checkpoint +/root/Qwen3-4B_torch_dist/ # tools/convert_hf_to_torch_dist.py +/path/to/dapo-math-17k.jsonl # set PROMPT_SET in the script +``` + +## Run -## Quick Start ```bash cd vime bash examples/fully_async/run-qwen3-4b-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 vime.rollout.fully_async_rollout.generate_rollout_fully_async - ``` + ``` + --rollout-function-path vime.rollout.fully_async_rollout.generate_rollout_fully_async + ``` -Why is it still "fully" async although `train_async.py` itself schedules rollouts step‑by‑step? +For custom per-sample logic, use vime'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] +``` + +See `examples/coding_agent_rl/` for a non-trivial example that plugs in a +multi-turn agent 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.vllm_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/scripts/run-qwen2.5-0.5B-reproducibility.sh b/scripts/run-qwen2.5-0.5B-reproducibility.sh index 72dad7608..e8c6dedb5 100644 --- a/scripts/run-qwen2.5-0.5B-reproducibility.sh +++ b/scripts/run-qwen2.5-0.5B-reproducibility.sh @@ -1,6 +1,8 @@ #!/bin/bash # for rerun the task +pkill -9 -f "vllm serve" +sleep 3 ray stop --force pkill -9 ray pkill -9 python @@ -13,10 +15,7 @@ set -ex # will prevent ray from buffering stdout/stderr export PYTHONUNBUFFERED=1 -NUM_GPUS=8 - SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" -VIME_ROOT="$(cd -- "${SCRIPT_DIR}/.." &>/dev/null && pwd)" source "${SCRIPT_DIR}/models/qwen2.5-0.5B.sh" CKPT_ARGS=( @@ -110,26 +109,21 @@ MISC_ARGS=( ) # 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 ${NUM_GPUS} --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 - -# Build the runtime environment JSON with proper variable substitution. -RUNTIME_ENV_JSON="{ - \"env_vars\": { - \"PYTHONPATH\": \"${VIME_ROOT}:/root/Megatron-LM/\", - \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", - \"NCCL_ALGO\": \"Ring\", - \"NVTE_ALLOW_NONDETERMINISTIC_ALGO\": \"0\", - \"CUBLAS_WORKSPACE_CONFIG\": \":4096:8\" - } -}" +ray start --head --node-ip-address 127.0.0.1 --num-gpus 8 --disable-usage-stats ray job submit --address="http://127.0.0.1:8265" \ - --runtime-env-json="${RUNTIME_ENV_JSON}" \ + --runtime-env-json='{ + "env_vars": { + "PYTHONPATH": "/root/Megatron-LM", + "CUDA_DEVICE_MAX_CONNECTIONS": "1", + "NCCL_ALGO": "Ring", + "NVTE_ALLOW_NONDETERMINISTIC_ALGO": "0", + "CUBLAS_WORKSPACE_CONFIG": ":4096:8" + } + }' \ -- python3 train.py \ - --train-backend megatron \ --actor-num-nodes 1 \ - --actor-num-gpus-per-node ${NUM_GPUS} \ + --actor-num-gpus-per-node 8 \ --colocate \ --calculate-per-token-loss \ ${MODEL_ARGS[@]} \ From bb2d21b87a448bc38bdf1fb73a127a19bf9ada42 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Tue, 9 Jun 2026 02:57:08 +0000 Subject: [PATCH 4/5] scripts(repro): use command_utils pkill pattern for renamed vLLM children pkill -9 -f 'vllm serve' only catches the parent; vLLM renames its subprocesses (VLLM::EngineCore / Worker_TP*). Match the renamed children too, mirroring vime/utils/external_utils/command_utils.py. Co-Authored-By: Claude Opus 4.8 (1M context) --- scripts/run-qwen2.5-0.5B-reproducibility.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/scripts/run-qwen2.5-0.5B-reproducibility.sh b/scripts/run-qwen2.5-0.5B-reproducibility.sh index e8c6dedb5..fb753a97d 100644 --- a/scripts/run-qwen2.5-0.5B-reproducibility.sh +++ b/scripts/run-qwen2.5-0.5B-reproducibility.sh @@ -1,7 +1,7 @@ #!/bin/bash # for rerun the task -pkill -9 -f "vllm serve" +pkill -9 -f '[v]llm serve|VLL[M]::' sleep 3 ray stop --force pkill -9 ray From f3f01dee34a67c7fd13e269219918fc077d703f4 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Tue, 9 Jun 2026 11:17:23 +0800 Subject: [PATCH 5/5] Update README.md Signed-off-by: aoshen02 --- examples/geo3k_vlm/README.md | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/examples/geo3k_vlm/README.md b/examples/geo3k_vlm/README.md index 0d82487de..df88e7a46 100644 --- a/examples/geo3k_vlm/README.md +++ b/examples/geo3k_vlm/README.md @@ -115,4 +115,10 @@ Our initial geo3k-specific verifier produced "format scores" (**0 and 0.9**) ins We fixed this by switching to the default math RM with clean **binary 0/1 rewards**. If you encounter similar precision issues with non-binary rewards, you can change the reward tensor dtype from `torch.float` to `torch.float16` in `vime/ray/rollout.py` (`_post_process_rewards` method) to truncate precision artifacts. ## B200 -On Blackwell (SM100), vLLM automatically selects FlashAttention 4 for the ViT encoder, so no manual override is needed. +On Blackwell (SM100), vllm automatically dispatches the ViT encoder to +FlashAttention 4 (or FA2 fallback) — no manual override is needed +([vllm/v1/attention/backends/fa_utils.py:81](https://github.com/vllm-project/vllm/blob/main/vllm/v1/attention/backends/fa_utils.py#L81)). +If you hit a kernel issue on a specific model, you can force SDPA with +`--vllm-mm-encoder-attn-backend TORCH_SDPA`. The HF-side +`--attn-implementation flash_attention_2` flag is still relevant when the +model is loaded via Hugging Face Transformers.