diff --git a/docs/training/rlhf.md b/docs/training/rlhf.md index 2aae267e1d..ccce0f37dd 100644 --- a/docs/training/rlhf.md +++ b/docs/training/rlhf.md @@ -18,6 +18,8 @@ The following open-source RL libraries use vLLM for fast rollouts (sorted alphab For weight synchronization between training and inference, see the [Weight Transfer](weight_transfer/README.md) documentation, which covers the pluggable backend system with [NCCL](weight_transfer/nccl.md) (multi-GPU) and [IPC](weight_transfer/ipc.md) (same-GPU) engines. To verify that an update landed on every rank, see the [Weight Checker](weight_checker.md). +For MoE policies, see [Routed Experts Return](routed_experts.md) (`--enable-return-routed-experts`) to get the per-token expert routing of each rollout for routing replay (R3) in the trainer. + For pipelining generation and training to improve GPU utilization and throughput, see the [Async Reinforcement Learning](async_rl.md) guide, which covers the pause/resume API for safely updating weights mid-flight. See the following notebooks showing how to use vLLM for GRPO: diff --git a/docs/training/routed_experts.md b/docs/training/routed_experts.md new file mode 100644 index 0000000000..f531a0ebbd --- /dev/null +++ b/docs/training/routed_experts.md @@ -0,0 +1,54 @@ +# Routed Experts Return (MoE Routing Replay / R3) + +`--enable-return-routed-experts` makes vLLM return the routed experts (the expert IDs each MoE layer's router selected) for every token of a request. RL trainers use it for MoE routing replay (R3): the training forward pass reuses the rollout's routing. + +```bash +vllm serve Qwen/Qwen3-30B-A3B --enable-return-routed-experts +# offline: LLM("Qwen/Qwen3-30B-A3B", enable_return_routed_experts=True) +``` + +The flag sets `AuxOutputConfig.enable_return_routed_experts` (full form: `--aux-output-config '{"enable_return_routed_experts": true, "max_bytes": N}'`). There is no per-request on/off switch. + +## Requirements + +Startup fails with a `ValueError` unless: + +- **Model Runner V2** is used (`VLLM_USE_V2_MODEL_RUNNER=1` if vLLM falls back to V1); +- the runner is `generate` and the model is **MoE**; +- **prefix caching is enabled** (the default). `--no-enable-prefix-caching` fails; the old v0.21 advice to disable it is wrong. + +Rejected at startup: adaptive speculative verification, PP > 1, DCP/PCP > 1, and the `NixlConnector`, `NixlPullConnector`, `NixlPushConnector`, `MoRIIOConnector`, `MooncakeConnector` KV connectors. Regular speculative decoding works. + +### How it works: aux output connector + +Routing is stored by the aux output connector (`vllm/distributed/aux_output_connector/`, configured via `--aux-output-config`), keyed by KV block hash and shared with prefix caching, which is why prefix caching is required. `max_bytes` is its host-memory LRU capacity (default: derived from KV cache capacity). If a prefix-cache hit needs evicted rows, the worker raises `BlockObjectStoreError`; increase `max_bytes`. + +## Output format + +| API | Field | +| --- | --- | +| Offline (`LLM`, `AsyncLLM`) | `CompletionOutput.routed_experts`: `np.ndarray` or `None` | +| `/v1/completions`, `/v1/chat/completions` | `choices[i].routed_experts`: base64 `.npy`, **non-streaming only** | +| `/inference/v1/generate` | `routed_experts` per choice (base64 `.npy`); with `stream: true`, on the final chunk | + +- **Shape:** `(num_tokens - 1 - routed_experts_prompt_start, num_moe_layers, num_experts_per_tok)`, `num_tokens` = prompt + output. The last sampled token has no row. +- **Rows:** row `i` is position `routed_experts_prompt_start + i` of `prompt_token_ids + output_token_ids`. +- **dtype:** `uint8` for at most 256 experts, else `uint16`; logical expert IDs. +- Filled only when the request finishes. `None` if the flag is off or the request aborted before any forward pass. + +```python +import base64, io +import numpy as np + +experts = np.load(io.BytesIO(base64.b64decode(choice["routed_experts"]))) +``` + +## `routed_experts_prompt_start` + +Per-request, default `0`: drops the first N prompt-token rows (e.g. routing already received in earlier turns). Must be `0 <= value <= prompt_len`, else HTTP 400. Pass it top-level on `/v1/*completions`, inside `sampling_params` on `/inference/v1/generate`, or as `SamplingParams(routed_experts_prompt_start=...)`. + +## Limitations + +- The Rust HTTP frontend never returns `routed_experts` and always sends `routed_experts_prompt_start = 0`. Use the Python API server. + +Example: [examples/rl/routed_experts_e2e.py](../../examples/rl/routed_experts_e2e.py). diff --git a/mkdocs.yaml b/mkdocs.yaml index bed8318998..d591df7b14 100644 --- a/mkdocs.yaml +++ b/mkdocs.yaml @@ -106,6 +106,7 @@ plugins: examples/features/pause_resume.md: examples/rl/pause_resume.md serving/online_serving/decisions.md: models/structured_decisions_models.md serving/online_serving/structured_decisions.md: models/structured_decisions_models.md + training/routed_experts_replay.md: training/routed_experts.md markdown_extensions: - attr_list