From 08ce80bbc1853823520b10ae482f7e7d5553e449 Mon Sep 17 00:00:00 2001 From: SamitHuang <285365963@qq.com> Date: Mon, 2 Mar 2026 16:32:32 +0800 Subject: [PATCH 01/39] temp save rfc Signed-off-by: SamitHuang <285365963@qq.com> --- docs/en/advanced/rfc-vllm-rollout-backend.md | 360 +++++++++++++++++++ 1 file changed, 360 insertions(+) create mode 100644 docs/en/advanced/rfc-vllm-rollout-backend.md diff --git a/docs/en/advanced/rfc-vllm-rollout-backend.md b/docs/en/advanced/rfc-vllm-rollout-backend.md new file mode 100644 index 0000000000..5c3cbd20f5 --- /dev/null +++ b/docs/en/advanced/rfc-vllm-rollout-backend.md @@ -0,0 +1,360 @@ +# RFC: Add vLLM as a Rollout Backend in Slime + +- **Author**: \ +- **Status**: Draft +- **Audience**: Slime rollout/runtime maintainers, RL training maintainers +- **Last Updated**: 2026-03-02 + +## 1. Summary + +This RFC proposes adding **vLLM** as a first-class rollout backend in Slime while preserving current SGLang behavior and avoiding regressions in GRPO workflows. + +The design is based on: + +1. A backend-agnostic rollout request/response contract. +2. Backend adapters (`SGLangClient`, `VLLMClient`) that isolate protocol differences. +3. Capability-aware behavior for non-parity features (abort, routed experts, prompt logprobs). +4. Managed vLLM mode from day one (Slime manages vLLM process lifecycle inside Ray, same as SGLang path). +5. Weight sync via NCCL broadcast (GPU direct transfer, no disk I/O). + +## 2. Why this is needed + +Slime currently assumes SGLang behavior in multiple places: + +- rollout generation response schema (`meta_info.finish_reason.type`, `output_token_logprobs`, optional `routed_experts`) +- router control-plane interactions (`/workers`, `/list_workers`, `/abort_request`) +- rollout server startup flow in Ray + +Supporting vLLM is therefore not a URL replacement. It requires a compatibility layer across both data plane and control plane semantics. + +## 3. Goals and Non-goals + +### Goals + +- Add `--rollout-backend {sglang,vllm}`. +- Keep SGLang path unchanged by default. +- Keep trainer/algorithm interfaces stable. +- Support GRPO rollout with explicit compatibility semantics and observability. + +### Non-goals (initial phase) + +- Full parity for all SGLang-specific features. +- Colocate mode (training and rollout sharing same GPUs). +- R3 routed expert replay on vLLM. +- Multi-instance vLLM with router load balancing. + +## 4. Architecture and Interface Changes (Explicit) + +This section is the key communication point for Slime maintainers. + +### 4.1 Architecture changes + +#### End-to-end architecture (string diagram) + +```text + +--------------------+ + | TrainerLoop | + +---------+----------+ + | + v + +--------------------+ + | RolloutFunction | + +---------+----------+ + | + v + +-------------------------------+ + | CanonicalRolloutRequest | + | (input_ids, sampling_params, | + | return_logprob, prompt_text) | + +---------------+---------------+ + | + v + +--------------------------+ + | RolloutBackendClient | + +------------+-------------+ + | + +-----------------+-----------------+ + | | + v v + +------------------------+ +------------------------+ + | SGLangClient | | VLLMClient | + +-----------+------------+ +------------+-----------+ + | | + v v + +-------------------------------+ +------------------------------+ + | SGLangRouter or SlimeRouter | | SlimeRouter(generic) or | + | (SGLang control-plane aware) | | direct vLLM endpoint | + +---------------+---------------+ +--------------+---------------+ + | | + v v + +----------------------+ +----------------------+ + | SGLang workers | | vLLM workers | + +----------+-----------+ +----------+-----------+ + \ / + \ / + v v + +-----------------------------------+ + | CanonicalRolloutResponse | + | (text, token_ids, token_logprobs, | + | finish_reason, backend_raw) | + +----------------+------------------+ + | + v + +-----------------------------------+ + | SampleUpdate + Training Pipeline | + | (backend-agnostic consumption) | + +-----------------------------------+ +``` + +#### Component responsibility map + +| Component | Responsibility before RFC | Responsibility after RFC | +|---|---|---| +| `RolloutFunction` | Contains generic rollout + SGLang protocol details | Contains generic rollout orchestration only | +| Backend protocol layer | Implicit in rollout logic | Explicit via `RolloutBackendClient` adapters | +| `SGLangClient` | N/A (scattered logic) | Owns SGLang request/response/control-plane specifics | +| `VLLMClient` | N/A | Owns vLLM request/response mapping and retries | +| Trainer/sample pipeline | Consumes SGLang-shaped fields indirectly | Consumes canonical fields, backend-agnostic | +| Router integration | Mixed generic + SGLang-specific assumptions | SGLang-specific control-plane isolated to SGLang adapter | + +#### Control-plane behavior split + +| Control-plane behavior | SGLang path | vLLM initial path | +|---|---|---| +| Worker registration/discovery APIs | Supported | Not required in generic path | +| Worker-level abort | Supported (`/abort_request`) | Fallback to timeout/cancel semantics | +| Routed experts replay metadata | Supported | Explicitly unsupported (capability-gated) | +| Health/load-balance routing | Existing SGLang/SlimeRouter behavior | SlimeRouter generic mode or direct endpoint | + +#### A) Rollout backend abstraction layer (new) + +- Introduce a backend client interface: + - `RolloutBackendClient` + - backend capability descriptor +- Add concrete adapters: + - `SGLangClient` (existing behavior extraction) + - `VLLMClient` (new) + +**Impact**: rollout logic calls a unified backend interface instead of directly embedding SGLang HTTP semantics. + +#### B) Rollout execution path refactor + +- `sglang_rollout.generate` no longer directly depends on SGLang HTTP payload/response shape. +- It builds a canonical request and consumes a canonical response. + +**Impact**: trainer-side logic remains stable while backend-specific details move into adapters. + +#### C) Startup path split in RolloutManager + +- Existing path (SGLang): keep current managed startup behavior. +- vLLM path: managed mode -- Slime creates `VLLMEngine` Ray actors that launch and manage local vLLM server processes, just like SGLang path uses `SGLangEngine`. + +**Impact**: vLLM gets the same lifecycle management as SGLang (process startup, health check, shutdown). + +#### D) Weight sync: VLLMEngine with same interface as SGLangEngine + +Training-side weight updater (`UpdateWeightFromDistributed`) calls engine methods via Ray remote. The core call chain with source locations: + +```text +Training side (Megatron actor, UNCHANGED) + UpdateWeightFromDistributed.update_weights() + -> engine.pause_generation.remote() # Ray remote call + -> VLLMEngine.pause_generation() # Ray actor method + -> requests.post("http://localhost:8000/sleep?level=2") + -> requests.post("http://localhost:8000/wake_up?tags=weights") + + -> engine.flush_cache.remote() + -> VLLMEngine.flush_cache() # no-op, sleep level 2 covers this + + -> engine.init_weights_update_group.remote() + -> VLLMEngine.init_weights_update_group() + -> requests.post("http://localhost:8000/collective_rpc", + json={"method": "init_weight_update_group", + "master_address": ..., "master_port": ..., + "rank_offset": ..., "world_size": ...}) + + -> dist.broadcast(param, src=0, group=nccl_group) # training process NCCL broadcast + -> engine.update_weights_from_distributed.remote() # tell vLLM "receive these params" + -> VLLMEngine.update_weights_from_distributed() + -> requests.post("http://localhost:8000/collective_rpc", + json={"method": "update_weight", + "name": ..., "dtype_name": ..., "shape": ...}) + -> vLLM worker: NCCL broadcast recv + model.load_weights() + + -> engine.continue_generation.remote() + -> VLLMEngine.continue_generation() + -> requests.post("http://localhost:8000/wake_up?tags=kv_cache") + + +Rollout generation side + sglang_rollout.generate() + -> VLLMClient.generate() + -> httpx.post("http://localhost:8000/v1/completions") # direct to local vLLM +``` + +`VLLMEngine` exposes **the same method signatures** as `SGLangEngine`. Endpoint mapping: + +| Method (same signature) | SGLangEngine internal | VLLMEngine internal | +|---|---|---| +| `pause_generation()` | `POST /pause_generation` | `POST /sleep?level=2` + `POST /wake_up?tags=weights` | +| `flush_cache()` | `GET /flush_cache` | no-op (sleep level 2 covers this) | +| `init_weights_update_group(...)` | `POST /init_weights_update_group` | `POST /collective_rpc {"method":"init_weight_update_group",...}` | +| `update_weights_from_distributed(...)` | `POST /update_weights_from_distributed` | `POST /collective_rpc {"method":"update_weight",...}` | +| `continue_generation()` | `POST /continue_generation` | `POST /wake_up?tags=kv_cache` | + +**Impact**: training-side code (`UpdateWeightFromDistributed`, `train.py`, actor code) requires **zero changes**. Weight sync uses NCCL broadcast (GPU direct transfer), same efficiency as SGLang path. + +Source links: +- [train.py#L88](../../../train.py#L88), [actor_group.py#L126](../../../slime/ray/actor_group.py#L126), [actor.py#L532](../../../slime/backends/megatron_utils/actor.py#L532) +- [update_weight_from_distributed.py#L82](../../../slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py#L82), [#L89](../../../slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py#L89), [#L90](../../../slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py#L90), [#L280](../../../slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py#L280), [#L321](../../../slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py#L321), [#L139](../../../slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py#L139) +- [sglang_engine.py#L415](../../../slime/backends/sglang_utils/sglang_engine.py#L415), [#L296](../../../slime/backends/sglang_utils/sglang_engine.py#L296), [#L373](../../../slime/backends/sglang_utils/sglang_engine.py#L373), [#L398](../../../slime/backends/sglang_utils/sglang_engine.py#L398), [#L420](../../../slime/backends/sglang_utils/sglang_engine.py#L420) +- [sglang_rollout.py#L47](../../../slime/rollout/sglang_rollout.py#L47), [#L72](../../../slime/rollout/sglang_rollout.py#L72), [#L165](../../../slime/rollout/sglang_rollout.py#L165), [#L311](../../../slime/rollout/sglang_rollout.py#L311) +- [rollout.py#L477](../../../slime/ray/rollout.py#L477), [#L1028](../../../slime/ray/rollout.py#L1028), [#L1041](../../../slime/ray/rollout.py#L1041) + +#### E) Router decision: no router for vLLM initial phase + +- SGLang Model Gateway: only supports SGLang workers, not applicable. +- SlimeRouter: only needed for R3 / radix-tree caching; Qwen2.5-0.5B is not MoE and uses token-in/token-out. +- Single vLLM instance, `VLLMClient` connects directly to local vLLM server port. + +### 4.2 Interface changes + +#### A) New CLI interfaces + +- `--rollout-backend {sglang,vllm}` +- `--vllm-base-url` +- `--vllm-api-mode` (e.g. OpenAI-compatible completion mode) +- `--vllm-model` +- `--vllm-max-retries` + +#### B) New internal canonical interfaces + +Add canonical rollout contract types: + +- `RolloutBackendRequest` +- `RolloutBackendResponse` + +These carry backend-neutral fields such as: + +- input token ids +- sampling params +- output token ids/logprobs +- canonical finish reason (`stop|length|abort`) +- backend raw response for debugging + +#### C) Capability-gated interface behavior + +Backends declare capabilities (abort support, routed experts support, prompt logprobs support). Unsupported features are explicitly gated, logged, and/or failed fast. + +## 5. File-level Change Map (for maintainers) + +### New files + +- `slime/rollout/backends/base_client.py` -- backend client interface + capability model +- `slime/rollout/backends/sglang_client.py` -- extracted SGLang rollout client +- `slime/rollout/backends/vllm_client.py` -- vLLM rollout client +- `slime/rollout/backends/__init__.py` -- adapter exports +- `slime/backends/vllm_utils/vllm_engine.py` -- `VLLMEngine` Ray actor (analogous to `SGLangEngine`) +- `slime/backends/vllm_utils/worker_extension.py` -- vLLM WorkerExtension for NCCL weight sync + +### Modified files + +- `slime/rollout/base_types.py` -- add canonical backend request/response types +- `slime/rollout/sglang_rollout.py` -- use backend adapters instead of hardcoded SGLang protocol +- `slime/utils/arguments.py` -- add `--rollout-backend`, vLLM args; skip SGLang parse when vLLM +- `slime/ray/rollout.py` -- startup flow split: create `VLLMEngine` when `rollout_backend=vllm` + +### Unchanged files (by design) + +- `train.py` -- training loop does not know about backend +- `slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py` -- engine method signatures are identical +- `slime/backends/megatron_utils/actor.py` -- calls engine methods polymorphically + +## 6. Compatibility Matrix (initial) + +| Capability | SGLang | vLLM (initial) | Handling | +|---|---|---|---| +| Token-level response logprobs | Yes | Partial/endpoint-dependent | Adapter normalization | +| Finish reason mapping (`stop|length|abort`) | Native | Different enums likely | Canonical mapping | +| Worker-level abort | Yes | Typically no direct equivalent | timeout/cancel fallback strategy | +| Prompt logprobs (OPD-related) | Yes | Partial/unknown by endpoint | capability-gated | +| Routed experts replay (R3) | Yes | No | explicit unsupported gate | +| SGLang worker API dependency | Yes | No | isolate in SGLang adapter | + +## 7. Phased Plan + +### Phase 1: Managed vLLM GRPO training (small-scale validation) + +**Goal**: Qwen2.5-0.5B GRPO 8-GPU sync training on GSM8K, loss/reward convergence comparable to SGLang. + +Key deliverables: +- `VLLMEngine` Ray actor (process lifecycle, sleep/wake_up, weight sync via NCCL) +- `VLLMClient` rollout adapter (request/response mapping, logprob/finish_reason normalization) +- `RolloutBackendRequest/Response` canonical contract +- `SGLangClient` extraction (isolate SGLang protocol details) +- `start_rollout_servers` branching for vLLM path +- Argument parsing split (`--rollout-backend vllm`) + +Technical decisions: +- No router (single vLLM instance, direct connection) +- Weight sync: NCCL broadcast via `VLLMEngine` with same method signatures as `SGLangEngine` +- Non-colocate mode (separate training and rollout GPUs) +- Rollout function: reuse `sglang_rollout.py` (backend-agnostic orchestration) +- Training-side code: zero changes + +Acceptance criteria: +- `num_rollout=3` smoke passes without errors +- Weight sync correctness: rollout output changes with each training update +- reward mean delta < 5% vs SGLang path (same seed/config, 20 steps) +- Existing SGLang tests remain green + +### Phase 2: Performance, scale, and advanced features + +- Colocate mode support (GPU IPC weight transfer) +- Multi-instance vLLM with load balancing +- Abort/cancel strategy refinement +- Prompt logprob support (OPD scenarios) +- Deterministic computation verification +- Larger model validation (e.g. Qwen3-4B) + +## 8. Validation strategy + +### Unit + +- finish reason normalization +- token/logprob alignment +- capability-gated behavior + +### Integration + +- SGLang vs vLLM response schema comparisons +- timeout/retry/non-crash behavior checks + +### End-to-end GRPO + +- smoke (short rollout) +- stability (medium horizon) +- stress (higher concurrency / latency pressure) + +## 9. Risks and mitigations + +- **Logprob semantic mismatch**: strict adapter checks + canonicalization + tests. +- **Abort mismatch**: capability model + timeout/cancel fallback + explicit logs. +- **Training drift**: mandatory A/B runs with fixed seeds and aligned configs. +- **Middleware assumptions**: enforce capability checks and default-disable incompatible middleware paths. + +## 10. Open questions for Slime maintainers + +1. Should vLLM initial path be direct endpoint only, or standardize via SlimeRouter generic mode immediately? +2. What minimum logprob fidelity is required to claim GRPO support? +3. Should OPD prompt-logprob paths be blocked for vLLM until full parity? +4. What backend quality gates are required before default recommendation? + +## 11. Decision requested + +Approve incremental implementation with: + +- contract-first adapter architecture, +- external vLLM initial support, +- explicit capability-gated behavior, +- strict SGLang non-regression requirement. From 3af806ee731e77d3b9030cb57b6fc758280c8aca Mon Sep 17 00:00:00 2001 From: SamitHuang <285365963@qq.com> Date: Tue, 3 Mar 2026 17:03:34 +0800 Subject: [PATCH 02/39] add plan Signed-off-by: SamitHuang <285365963@qq.com> --- goal_plan.md | 27 ++ rfc-vllm-rollout-backend-en.md | 401 ++++++++++++++++++++++++++++++ rfc-vllm-rollout-backend.md | 436 +++++++++++++++++++++++++++++++++ 3 files changed, 864 insertions(+) create mode 100644 goal_plan.md create mode 100644 rfc-vllm-rollout-backend-en.md create mode 100644 rfc-vllm-rollout-backend.md diff --git a/goal_plan.md b/goal_plan.md new file mode 100644 index 0000000000..adb39bf047 --- /dev/null +++ b/goal_plan.md @@ -0,0 +1,27 @@ +### 阶段一：打通Qwen2.5-0.5B GRPO 8卡同步/异步训练（train.py和train_async.py），GSM8K 数据集，loss/reward 收敛与 SGLang backend 基本一致，且满足确定性计算，多次重复运行Loss曲线完全一致。 + +#### 初步方案： +- 对标SGLang，Slime 在 Ray 内管理 vLLM 的完整生命周期，包括进程拉起、权重同步、推理暂停/恢复 +- 暂不使用Router，SGLang Model Gateway仅只支持SGLang Worker，SlimeRouter仅在 R3 / radix-tree caching 时需要，Qwen2.5-0.5B 非 MoE 且用 token-in/token-out +- 单vLLM实例，无router，通过vLLMClient 直连本地 vLLM 进程端口 +- 若训推不共卡，权重同步采用NCCL broadcast，对标SGLang update_weights_from_distributed (默认，优先用于支持异步训推） +- 若colocate，权重同步采用GPU IPC（vLLM update_weights_from_ipc, update_weights_from_tensor），对标SGLang update_weights_from_tensor + +#### 风险： +- slime, sglang版本依赖，和vllm 0.16的版本依赖冲突(numpy, torch, transformers, etc) +- 算力 + +First Design and RFC by 03/06 + + +### 阶段二：接入vllm-project/router，支持多实例vLLM + +- vllm router forked from SGLang Model Gateway + +### 阶段三：多节点大规模验证，MoE模型，optional：验证MTP Speculative Decoding，FP8 rollout 等高级特性 + +- Model: Qwen/Qwen3-30B-A3B or GLM4.7 +- Parallel: 16卡 or 128卡, Train mixed EP+FSDP, Rollout EP+DP +- Verify more features: + - Bf16 train, FP8 rollout + - MTP Speculative Decoding diff --git a/rfc-vllm-rollout-backend-en.md b/rfc-vllm-rollout-backend-en.md new file mode 100644 index 0000000000..dad61a0731 --- /dev/null +++ b/rfc-vllm-rollout-backend-en.md @@ -0,0 +1,401 @@ +# RFC: Supporting vLLM as a Rollout Backend in Slime + +- **Author**: \ +- **Status**: Phase 1 Done +- **Audience**: Slime rollout/runtime maintainers, RL training maintainers +- **Last Updated**: 2026-03-03 + +## 1. Summary + +This RFC proposes adding **vLLM** as a first-class rollout backend in Slime while keeping the existing SGLang behavior unchanged and the GRPO training pipeline unaffected. + +Core design principles: + +1. Define backend-agnostic rollout request/response contracts (`RolloutBackendRequest`/`RolloutBackendResponse`). +2. Isolate protocol differences through backend adapters (`SGLangClient`, `VLLMClient`). +3. Apply explicit capability gating for non-equivalent features (abort, routed experts, prompt logprobs). +4. Managed mode -- Slime manages the vLLM process lifecycle within Ray, on par with the SGLang path. +5. Weight synchronization leverages vLLM's native weight transfer API, automatically selecting the backend based on deployment mode: + - **Colocate mode**: CUDA IPC (`IPCWeightTransferEngine`) -- zero-copy via shared GPU memory. + - **Non-colocate mode**: NCCL broadcast (`NCCLWeightTransferEngine`) -- direct GPU transfer. + +**Current status**: Phase 1 is complete and verified. GRPO training runs successfully with Qwen2.5-0.5B + GSM8K on 4 GPUs in colocate mode. + +## 2. Motivation + +Slime currently assumes SGLang behavior in multiple places: + +- Rollout generation response format (`meta_info.finish_reason.type`, `output_token_logprobs`, optional `routed_experts`) +- Router control-plane interactions (`/workers`, `/list_workers`, `/abort_request`) +- Ray-based rollout server startup flow + +Supporting vLLM is therefore not just "swapping a URL" -- it requires compatibility layers on both the data plane and control plane. + +## 3. Goals and Non-Goals + +### Goals + +- Add `--rollout-backend {sglang,vllm}` +- Keep SGLang as the default, unchanged +- Keep the training-side interface stable +- Support GRPO rollout with explicit compatibility semantics and observability +- **Support colocate mode** (training and inference share GPUs, weight sync via CUDA IPC) + +### Non-Goals (Current Phase) + +- Full feature parity with all SGLang capabilities +- R3 routing replay on vLLM +- Multi-instance vLLM + router load balancing + +## 4. Architecture and Interface Changes + +### 4.1 Architecture Changes + +#### End-to-End Architecture Diagram + +```text + +--------------------+ + | TrainerLoop | + +---------+----------+ + | + +---------------------+---------------------+ + | (generate) (update_weights) + v v + +--------------------+ +------------------------------------+ + | RolloutFunction | | weight updater (auto-selected) | + | (sglang_rollout) | | colocate → UpdateWeightFromTensor | + +---------+----------+ | otherwise → ...FromDistributed | + | +---------------+--------------------+ + v | + +-------------------------------+ | + | RolloutBackendRequest | | + +---------------+---------------+ | + | v + v +----------------------------+ + +--------------------------+ | VLLMEngine (Ray actor) | + | RolloutBackendClient | | weight transfer backend: | + +------------+-------------+ | colocate → IPC | + | | otherwise → NCCL | + +-------------+-------------+ +-------------+--------------+ + | | | + v v v ++------------+ +-----------+ +---------------+ +| SGLang | | VLLM | | vLLM server | +| Client | | Client | | /update_weights| ++-----+------+ +-----+-----+ | /pause /resume | + | | | /sleep /wake_up| + v v +---------------+ + +-----------+ +-------------+ + | SGLang | | vLLM server | + | Router | | /v1/compl. | + +-----------+ +-------------+ + \ / + \ / + v v + +-----------------------------------+ + | RolloutBackendResponse | + | (text, token_ids, token_logprobs, | + | finish_reason, backend_raw) | + +----------------+------------------+ + | + v + +-----------------------------------+ + | SampleUpdate + Training Pipeline | + | (backend-agnostic consumption) | + +-----------------------------------+ +``` + +#### Component Responsibility Comparison + +| Component | Before RFC | After RFC | +|---|---|---| +| `RolloutFunction` | Generic rollout + SGLang protocol details | Generic rollout orchestration only | +| Backend protocol layer | Implicit in rollout logic | Explicit `RolloutBackendClient` adapter | +| `SGLangClient` | Did not exist (logic scattered) | Owns all SGLang request/response/control-plane details | +| `VLLMClient` | Did not exist | Owns vLLM `/v1/completions` request/response mapping | +| Training/sample pipeline | Indirectly consumed SGLang-format fields | Consumes unified contract fields, backend-agnostic | +| Weight sync | SGLang IPC or NCCL only | Auto-adapts: SGLang IPC / vLLM IPC / vLLM NCCL | + +#### Control-Plane Behavior Split + +| Control-plane behavior | SGLang path | vLLM path | +|---|---|---| +| Worker registration/discovery API | Supported | Not needed | +| Worker-level abort | Supported (`/abort_request`) | Degraded to timeout/cancel strategy | +| Routed experts replay | Supported | Explicitly unsupported (capability-gated) | +| Health check | Existing SGLang/SlimeRouter | `GET /health` direct connection | +| Memory management (sleep/wake) | `release/resume_memory_occupation` | `POST /sleep` / `POST /wake_up` | + +#### A) Rollout Backend Abstraction Layer + +- `RolloutBackendClient` (`slime/rollout/backends/base_client.py`): defines `generate()` + `capabilities` abstract interface +- `BackendCapabilities` dataclass: declares `supports_abort`, `supports_routed_experts`, `supports_prompt_logprobs` +- `SGLangClient` (extracted from existing code), `VLLMClient` (new) + +**Impact**: Rollout logic calls a unified interface and no longer directly embeds SGLang HTTP semantics. + +#### B) Rollout Execution Path + +- `sglang_rollout.generate` constructs `RolloutBackendRequest` and consumes `RolloutBackendResponse` +- Selects `SGLangClient` or `VLLMClient` based on `--rollout-backend` +- `VLLMClient` calls vLLM's OpenAI-compatible `/v1/completions` endpoint + +**Impact**: Training-side logic remains stable; backend details are encapsulated in adapters. + +#### C) RolloutManager Startup Path + +- Existing path (SGLang): unchanged +- vLLM path: `_start_vllm_rollout_servers()` creates a `VLLMEngine` Ray actor that starts and manages a local vLLM server process + +**Impact**: vLLM gets the same lifecycle management as SGLang. + +#### D) Weight Synchronization + +Weight synchronization supports two modes, automatically selected based on deployment: + +**Selection logic** (`actor.py`): +```python +update_weight_cls = UpdateWeightFromTensor if args.colocate else UpdateWeightFromDistributed +``` +Identical to the SGLang selection logic -- backend-agnostic. + +##### D.1) Colocate Mode: CUDA IPC + +The vLLM server starts with `--weight-transfer-config '{"backend": "ipc"}'`. + +Call chain: +```text +UpdateWeightFromTensor.update_weights() + → engine.pause_generation.remote() # VLLMEngine → POST /pause?mode=abort + → _send_to_colocated_vllm_engine() + → each training rank: + reduce_tensor(tensor) # create CUDA IPC handle + {gpu_uuid: ipc_handle} # keyed by physical GPU UUID + → dist.gather_object (Gloo) # collect handles from all TP ranks + → pickle + base64 encode + → engine.update_weights_from_tensor.remote() + → VLLMEngine → POST /update_weights + { "update_info": { + "names": [...], + "dtype_names": [...], + "shapes": [...], + "ipc_handles_pickled": "base64..." + }} + → vLLM IPCWeightTransferEngine.receive_weights() + → each TP worker looks up its IPC handle by GPU UUID + → func(*args) reconstructs tensor → load_weights() + → engine.continue_generation.remote() # VLLMEngine → POST /resume +``` + +Key points: +- The training process and vLLM workers share the same physical GPU; CUDA IPC enables zero-copy weight transfer +- For TP>1, IPC handles from all training ranks are merged via Gloo gather; each parameter contains a mapping of all GPU UUIDs +- The training side must keep tensor references alive until `ray.get()` returns (vLLM has finished reading) + +##### D.2) Non-Colocate Mode: NCCL Broadcast + +The vLLM server starts with `--weight-transfer-config '{"backend": "nccl"}'`. + +Call chain: +```text +UpdateWeightFromDistributed.update_weights() + → engine.init_weights_update_group.remote() # VLLMEngine → POST /init_weight_transfer_engine + → vLLM NCCLWeightTransferEngine initializes StatelessProcessGroup + PyNcclCommunicator + → training side: NCCLWeightTransferEngine.trainer_init() + → establishes a matching NCCL communicator with vLLM + → PyNcclCommunicator.broadcast(tensor, src=0) # direct GPU transfer to vLLM workers + → engine.update_weights_from_distributed.remote() + → VLLMEngine → POST /update_weights + { "update_info": { "names": [...], "dtype_names": [...], "shapes": [...] }} + → vLLM NCCLWeightTransferEngine.receive_weights() +``` + +Key points: +- The training side uses vLLM's `NCCLWeightTransferEngine.trainer_init()` to initialize NCCL, compatible with vLLM's internal `StatelessProcessGroup` + `PyNcclCommunicator` +- Cannot use `torch.distributed.init_process_group` (incompatible with vLLM) + +#### VLLMEngine Endpoint Mapping + +`VLLMEngine` exposes method signatures compatible with `SGLangEngine`. Internal HTTP endpoint mapping: + +| Method | VLLMEngine internal HTTP call | +|---|---| +| `pause_generation()` | `POST /pause?mode=abort` | +| `flush_cache()` | no-op | +| `continue_generation()` | `POST /resume` | +| `init_weights_update_group(...)` | `POST /init_weight_transfer_engine` | +| `update_weights_from_distributed(...)` | `POST /update_weights` (NCCL mode) | +| `update_weights_from_tensor(...)` | `POST /update_weights` (IPC mode) | +| `release_memory_occupation()` | `POST /sleep?level=1&mode=abort` | +| `resume_memory_occupation()` | `POST /wake_up` | +| `health_generate()` | `GET /health` | +| `shutdown()` | `process.terminate()` | + +**Impact**: Training-side code (`UpdateWeightFromTensor`, `UpdateWeightFromDistributed`, actor code) uses a unified interface, unaware of the specific backend. + +#### E) Router Decision + +vLLM does not use a router in the current phase: +- Single vLLM instance; `VLLMClient` connects directly to the local vLLM server port +- SGLang Model Gateway only supports SGLang workers, not applicable +- SlimeRouter is only needed for R3 / radix-tree caching scenarios + +### 4.2 Interface Changes + +#### A) New CLI Arguments + +- `--rollout-backend {sglang,vllm}` -- select rollout backend +- `--vllm-base-url` -- manually specify vLLM server address (not needed when auto-managed) +- `--vllm-model` -- model path for vLLM to load (defaults to `--hf-checkpoint`) +- `--vllm-max-retries` -- max retries for generation requests +- `--vllm-enforce-eager` -- disable CUDA graph (default True) + +#### B) New Internal Unified Interface + +Added to `slime/rollout/base_types.py`: + +- `RolloutBackendRequest`: input_ids, sampling_params, return_logprob, return_routed_experts, image_data, session_id +- `RolloutBackendResponse`: text, output_token_ids, output_token_logprobs, finish_reason (`stop|length|abort`), prompt_tokens, completion_tokens, backend_raw, routed_experts + +#### C) Capability Gating + +`BackendCapabilities` dataclass declares capabilities per backend: + +| Capability | SGLangClient | VLLMClient | +|---|---|---| +| `supports_abort` | True | False | +| `supports_routed_experts` | True | False | +| `supports_prompt_logprobs` | True | False | + +Unsupported features are explicitly gated, logged, or fail fast at call time. + +## 5. File-Level Change List + +### New Files + +| File | Description | +|---|---| +| `slime/rollout/backends/base_client.py` | `RolloutBackendClient` abstract interface + `BackendCapabilities` | +| `slime/rollout/backends/sglang_client.py` | SGLang rollout client (extracted from existing code) | +| `slime/rollout/backends/vllm_client.py` | vLLM rollout client (`/v1/completions` adapter) | +| `slime/rollout/backends/__init__.py` | Adapter exports | +| `slime/backends/vllm_utils/vllm_engine.py` | `VLLMEngine` Ray actor (process management + weight sync) | +| `run-qwen2.5-0.5B-vllm.sh` | vLLM validation script (Qwen2.5-0.5B + GSM8K + colocate) | + +### Modified Files + +| File | Description | +|---|---| +| `slime/rollout/base_types.py` | Added `RolloutBackendRequest` / `RolloutBackendResponse` | +| `slime/rollout/sglang_rollout.py` | Uses backend adapter instead of hard-coded SGLang protocol | +| `slime/utils/arguments.py` | Added `--rollout-backend`, vLLM arguments; sets sglang alias defaults for vLLM | +| `slime/ray/rollout.py` | `_start_vllm_rollout_servers()` to start VLLMEngine | +| `slime/backends/megatron_utils/actor.py` | Weight updater selection: `colocate → UpdateWeightFromTensor` (backend-agnostic) | +| `slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py` | Added `_send_to_colocated_vllm_engine()` (CUDA IPC path) | +| `slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py` | vLLM NCCL compatibility: `NCCLWeightTransferEngine.trainer_init()` + `PyNcclCommunicator.broadcast()` | + +### Unchanged Files + +| File | Description | +|---|---| +| `train.py` | Training loop is backend-agnostic | +| SGLang-related files | SGLang path is completely unaffected | + +### Not Used (Differences from Initial RFC Draft) + +- **No `worker_extension.py`**: vLLM's native weight transfer API fully meets requirements; no custom WorkerExtension needed +- **No `collective_rpc`**: Uses vLLM's `/init_weight_transfer_engine` and `/update_weights` endpoints + +## 6. Capability Compatibility Matrix + +| Capability | SGLang | vLLM | Handling | +|---|---|---|---| +| Token-level response logprobs | Supported | Supported (`choice.logprobs.token_logprobs`) | Adapter normalization | +| Output token IDs | `meta_info` | `choice.token_ids` | Adapter normalization | +| Finish reason | `stop\|length\|abort` | `stop\|length` + others | Canonical mapping | +| Worker-level abort | Supported | Not supported | Timeout/cancel degradation | +| Prompt logprobs (OPD) | Supported | Partial | Capability-gated | +| Routed experts (R3) | Supported | Not supported | Explicit gate | +| Colocate (GPU sharing) | Supported (FlattenedTensorBucket IPC) | Supported (CUDA IPC handles) | Native IPC per backend | +| Memory management | `release/resume_memory_occupation` | `POST /sleep` / `POST /wake_up` | Unified interface | +| `include_stop_str_in_output` | `no_stop_trim` | `include_stop_str_in_output` | Parameter mapping | + +## 7. Phased Plan + +### Phase 1: Managed vLLM GRPO Training ✅ Done + +**Goal**: Qwen2.5-0.5B GRPO training on 4 GPUs in colocate mode with GSM8K dataset, running successfully. + +Completed deliverables: +- `VLLMEngine` Ray actor: process lifecycle management, `/pause`/`/resume`/`/sleep`/`/wake_up` +- `VLLMClient` rollout adapter: `/v1/completions` request mapping, logprob/finish_reason/token_ids normalization +- `RolloutBackendRequest`/`Response` unified contract +- `SGLangClient` extraction +- `_start_vllm_rollout_servers` startup path +- Argument parsing branching (`--rollout-backend vllm`, sglang alias defaults) +- **Colocate mode**: CUDA IPC weight transfer (`_send_to_colocated_vllm_engine` → `IPCWeightTransferEngine`) +- **Non-colocate mode**: NCCL weight transfer (`NCCLWeightTransferEngine.trainer_init()` → `PyNcclCommunicator.broadcast()`) +- Validation script: `run-qwen2.5-0.5B-vllm.sh` + +Technical decisions: +- No router (single vLLM instance, direct connection) +- Colocate weight sync: CUDA IPC (training and vLLM share GPU memory, zero-copy) +- Non-colocate weight sync: vLLM native NCCL (`StatelessProcessGroup` + `PyNcclCommunicator`) +- Rollout function: reuses `sglang_rollout.py` (backend-agnostic orchestration) +- Training-side weight updater selection is identical to SGLang (`colocate → UpdateWeightFromTensor`) + +Acceptance results: +- ✅ Qwen2.5-0.5B + GSM8K + 4 GPU colocate training runs successfully +- ✅ Weight sync correctness: rollout outputs change with each training update +- ✅ Existing SGLang path unaffected + +### Phase 2: Multi-Instance vLLM + Async Training + +**Goal**: Support multi-instance vLLM rollout with [vllm-project/router](https://github.com/vllm-project/router) for load balancing, and verify async training (`train_async.py`) on larger models. + +Key deliverables: +- Integrate [vllm-project/router](https://github.com/vllm-project/router) as the vLLM-side load balancer +- `start_rollout_servers` launches N `VLLMEngine` actors + one vllm-router process +- `VLLMClient` generation requests routed through vllm-router instead of direct connection +- Verify async training (`train_async.py`) correctness with the vLLM backend + +Acceptance criteria: +- Multi-instance (e.g., 2-4 vLLM workers) rollout stable over 20+ rollout steps +- Async training (`train_async.py`) with no hangs or weight sync race conditions +- Throughput improvement compared to single-instance baseline +- Larger model verification (e.g., Qwen3-4B, TP=2) + +### Phase 3 (Future): Advanced Features + +- Abort/cancel strategy refinement +- Prompt logprob support (OPD scenarios) +- Deterministic computation verification +- Performance benchmarks (vLLM vs SGLang throughput comparison) + +## 8. Validation Strategy + +### Completed (Phase 1) + +- ✅ End-to-end GRPO training: `run-qwen2.5-0.5B-vllm.sh` (Qwen2.5-0.5B, GSM8K, 4 GPUs, colocate) +- ✅ Weight sync correctness verification +- ✅ SGLang path non-regression + +### Remaining + +- Unit tests: finish reason normalization, token/logprob alignment, capability-gated behavior +- Integration tests: SGLang vs vLLM response schema comparison +- Stress tests: higher concurrency / latency pressure + +## 9. Key Implementation Challenges and Solutions + +### 9.1 vLLM NCCL Incompatibility + +**Problem**: vLLM internally uses `StatelessProcessGroup` + `PyNcclCommunicator` for NCCL communication, which is incompatible with `torch.distributed.init_process_group`. + +**Solution**: The training side uses `NCCLWeightTransferEngine.trainer_init()` to initialize NCCL, ensuring a matching communicator is established with the vLLM side. + +### 9.2 vLLM Logprobs Format Differences + +**Problem**: vLLM's `/v1/completions` logprobs format (`choice.token_ids`, `choice.logprobs.token_logprobs`) differs from SGLang's. + +**Solution**: `VLLMClient` performs explicit mapping, converting vLLM's format into the unified `RolloutBackendResponse` format. diff --git a/rfc-vllm-rollout-backend.md b/rfc-vllm-rollout-backend.md new file mode 100644 index 0000000000..fc8e569793 --- /dev/null +++ b/rfc-vllm-rollout-backend.md @@ -0,0 +1,436 @@ +# RFC: 在 Slime 中支持 vLLM 作为 Rollout Backend + +## 1. 概要 + +本 RFC 提议在 Slime 中增加 **vLLM** 作为一等 rollout backend，同时保持现有 SGLang 行为不变，不影响 GRPO 训练流程。 + +核心设计思路： + +1. 定义 backend 无关的 rollout 请求/响应契约（`RolloutBackendRequest`/`RolloutBackendResponse`）。 +2. 通过 backend adapter（`SGLangClient`、`VLLMClient`）隔离协议差异。 +3. 对不等价能力（abort、routed experts、prompt logprobs）做显式 capability gating。 +4. Managed 模式 —— Slime 在 Ray 内管理 vLLM 进程生命周期，与 SGLang 路径同等级别。 +5. 权重同步利用 vLLM 原生 weight transfer API，根据部署模式自动选择后端： + - **Colocate 模式**：CUDA IPC（`IPCWeightTransferEngine`），GPU 共享内存零拷贝。 + - **Non-colocate 模式**：NCCL broadcast（`NCCLWeightTransferEngine`），GPU 直传。 + +**当前状态**：Phase 1 已完成并验证通过。Qwen2.5-0.5B + GSM8K + 4 GPU colocate 模式下 GRPO 训练正常运行。 + +## 2. 为什么需要这个 + +Slime 当前在多处假设 SGLang 行为： + +- rollout 生成响应格式（`meta_info.finish_reason.type`、`output_token_logprobs`、可选 `routed_experts`） +- router 控制面交互（`/workers`、`/list_workers`、`/abort_request`） +- Ray 内 rollout server 启动流程 + +因此支持 vLLM 不是"换个 URL"，而是需要在数据面和控制面做兼容层。 + +## 3. 目标与非目标 + +### 目标 + +- 新增 `--rollout-backend {sglang,vllm}` +- SGLang 路径默认不变 +- 训练侧接口保持稳定 +- 支持 GRPO rollout，具有显式兼容性语义和可观测性 +- **支持 colocate 模式**（训练与推理共享 GPU，通过 CUDA IPC 同步权重） + +### 非目标（当前阶段） + +- 所有 SGLang 特性的完全对等 +- vLLM 上的 R3 路由回放 +- 多 vLLM 实例 + router 负载均衡 + +## 4. 架构和接口改动 + +### 4.1 架构改动 + +#### 端到端架构图 + +```text + +--------------------+ + | TrainerLoop | + +---------+----------+ + | + +---------------------+---------------------+ + | (generate) (update_weights) + v v + +--------------------+ +------------------------------------+ + | RolloutFunction | | weight updater (自动选择) | + | (sglang_rollout) | | colocate → UpdateWeightFromTensor | + +---------+----------+ | otherwise → ...FromDistributed | + | +---------------+--------------------+ + v | + +-------------------------------+ | + | RolloutBackendRequest | | + +---------------+---------------+ | + | v + v +----------------------------+ + +--------------------------+ | VLLMEngine (Ray actor) | + | RolloutBackendClient | | weight transfer backend: | + +------------+-------------+ | colocate → IPC | + | | otherwise → NCCL | + +-------------+-------------+ +-------------+--------------+ + | | | + v v v ++------------+ +-----------+ +---------------+ +| SGLang | | VLLM | | vLLM server | +| Client | | Client | | /update_weights| ++-----+------+ +-----+-----+ | /pause /resume | + | | | /sleep /wake_up| + v v +---------------+ + +-----------+ +-------------+ + | SGLang | | vLLM server | + | Router | | /v1/compl. | + +-----------+ +-------------+ + \ / + \ / + v v + +-----------------------------------+ + | RolloutBackendResponse | + | (text, token_ids, token_logprobs, | + | finish_reason, backend_raw) | + +----------------+------------------+ + | + v + +-----------------------------------+ + | SampleUpdate + Training Pipeline | + | (backend 无关消费) | + +-----------------------------------+ +``` + +#### 组件职责对照 + +| 组件 | RFC 前 | RFC 后 | +|---|---|---| +| `RolloutFunction` | 包含通用 rollout + SGLang 协议细节 | 只包含通用 rollout 编排 | +| Backend 协议层 | 隐含在 rollout 逻辑里 | 显式的 `RolloutBackendClient` adapter | +| `SGLangClient` | 不存在（逻辑分散） | 拥有 SGLang 请求/响应/控制面的全部细节 | +| `VLLMClient` | 不存在 | 拥有 vLLM `/v1/completions` 请求/响应映射 | +| 训练/sample 管线 | 间接消费 SGLang 格式字段 | 消费统一契约字段，backend 无关 | +| 权重同步 | 仅 SGLang IPC 或 NCCL | 自动适配：SGLang IPC / vLLM IPC / vLLM NCCL | + +#### 控制面行为拆分 + +| 控制面行为 | SGLang 路径 | vLLM 路径 | +|---|---|---| +| Worker 注册/发现 API | 支持 | 不需要 | +| Worker 级 abort | 支持（`/abort_request`） | 降级为超时/取消策略 | +| Routed experts 回放 | 支持 | 显式不支持（capability-gated） | +| 健康检查 | 现有 SGLang/SlimeRouter | `GET /health` 直连 | +| 内存管理 (sleep/wake) | `release/resume_memory_occupation` | `POST /sleep` / `POST /wake_up` | + +#### A) Rollout backend 抽象层 + +- `RolloutBackendClient`（`slime/rollout/backends/base_client.py`）：定义 `generate()` + `capabilities` 抽象接口 +- `BackendCapabilities` dataclass：声明 `supports_abort`、`supports_routed_experts`、`supports_prompt_logprobs` +- `SGLangClient`（从现有代码提取）、`VLLMClient`（新建） + +**影响**：rollout 逻辑调用统一接口，不再直接嵌入 SGLang HTTP 语义。 + +#### B) Rollout 执行路径 + +- `sglang_rollout.generate` 构造 `RolloutBackendRequest`，消费 `RolloutBackendResponse` +- 根据 `--rollout-backend` 选择 `SGLangClient` 或 `VLLMClient` +- `VLLMClient` 调用 vLLM OpenAI-compatible `/v1/completions` 端点 + +**影响**：训练侧逻辑保持稳定，backend 细节移入 adapter。 + +#### C) RolloutManager 启动路径 + +- 现有路径（SGLang）：保持不变 +- vLLM 路径：`_start_vllm_rollout_servers()` 创建 `VLLMEngine` Ray actor，启动管理本地 vLLM server 进程 + +**影响**：vLLM 获得与 SGLang 相同的生命周期管理。 + +#### D) 权重同步 + +权重同步支持两种模式，根据部署方式自动选择： + +**模式选择逻辑**（`actor.py`）： +```python +update_weight_cls = UpdateWeightFromTensor if args.colocate else UpdateWeightFromDistributed +``` +与 SGLang 使用完全相同的选择逻辑，backend 无关。 + +##### D.1) Colocate 模式：CUDA IPC + +vLLM server 启动时配置 `--weight-transfer-config '{"backend": "ipc"}'`。 + +调用链： +```text +UpdateWeightFromTensor.update_weights() + → engine.pause_generation.remote() # VLLMEngine → POST /pause?mode=abort + → _send_to_colocated_vllm_engine() + → 每个训练 rank: + reduce_tensor(tensor) # 创建 CUDA IPC handle + {gpu_uuid: ipc_handle} # 以物理 GPU UUID 为 key + → dist.gather_object (Gloo) # 收集所有 TP rank 的 handles + → pickle + base64 编码 + → engine.update_weights_from_tensor.remote() + → VLLMEngine → POST /update_weights + { "update_info": { + "names": [...], + "dtype_names": [...], + "shapes": [...], + "ipc_handles_pickled": "base64..." + }} + → vLLM IPCWeightTransferEngine.receive_weights() + → 每个 TP worker 根据自己的 GPU UUID 查找 IPC handle + → func(*args) 重建 tensor → load_weights() + → engine.continue_generation.remote() # VLLMEngine → POST /resume +``` + +关键点： +- 训练进程和 vLLM worker 共享同一物理 GPU，CUDA IPC 实现零拷贝权重传输 +- TP>1 时，各 training rank 的 IPC handles 通过 Gloo gather 合并，每个参数包含所有 GPU UUID 的映射 +- 训练侧必须保持 tensor 引用直到 `ray.get()` 返回（vLLM 读取完毕） + +##### D.2) Non-colocate 模式：NCCL broadcast + +vLLM server 启动时配置 `--weight-transfer-config '{"backend": "nccl"}'`。 + +调用链： +```text +UpdateWeightFromDistributed.update_weights() + → engine.init_weights_update_group.remote() # VLLMEngine → POST /init_weight_transfer_engine + → vLLM NCCLWeightTransferEngine 初始化 StatelessProcessGroup + PyNcclCommunicator + → 训练侧: NCCLWeightTransferEngine.trainer_init() + → 与 vLLM 建立匹配的 NCCL 通信组 + → PyNcclCommunicator.broadcast(tensor, src=0) # GPU 直传到 vLLM worker + → engine.update_weights_from_distributed.remote() + → VLLMEngine → POST /update_weights + { "update_info": { "names": [...], "dtype_names": [...], "shapes": [...] }} + → vLLM NCCLWeightTransferEngine.receive_weights() +``` + +关键点： +- 训练侧使用 vLLM 的 `NCCLWeightTransferEngine.trainer_init()` 初始化 NCCL，与 vLLM 内部的 `StatelessProcessGroup` + `PyNcclCommunicator` 兼容 +- 不能使用 `torch.distributed.init_process_group`（vLLM 不兼容） + +#### VLLMEngine 端点映射 + +`VLLMEngine` 暴露与 `SGLangEngine` 兼容的方法签名。内部 HTTP 端点映射： + +| 方法 | VLLMEngine 内部 HTTP 调用 | +|---|---| +| `pause_generation()` | `POST /pause?mode=abort` | +| `flush_cache()` | no-op | +| `continue_generation()` | `POST /resume` | +| `init_weights_update_group(...)` | `POST /init_weight_transfer_engine` | +| `update_weights_from_distributed(...)` | `POST /update_weights` (NCCL 模式) | +| `update_weights_from_tensor(...)` | `POST /update_weights` (IPC 模式) | +| `release_memory_occupation()` | `POST /sleep?level=1&mode=abort` | +| `resume_memory_occupation()` | `POST /wake_up` | +| `health_generate()` | `GET /health` | +| `shutdown()` | `process.terminate()` | + +**影响**：训练侧代码（`UpdateWeightFromTensor`、`UpdateWeightFromDistributed`、actor 代码）使用统一接口，不感知具体 backend。 + +#### E) Router 决策 + +vLLM 当前阶段不使用 router： +- 单 vLLM 实例，`VLLMClient` 直连本地 vLLM server 端口 +- SGLang Model Gateway 只支持 SGLang worker，不适用 +- SlimeRouter 仅在 R3 / radix-tree caching 时需要 + +### 4.2 接口改动 + +#### A) 新增 CLI 接口 + +- `--rollout-backend {sglang,vllm}` —— 选择 rollout backend +- `--vllm-base-url` —— 手动指定 vLLM server 地址（自动管理时无需设置） +- `--vllm-model` —— vLLM 加载的模型路径（默认同 `--hf-checkpoint`） +- `--vllm-max-retries` —— 生成请求最大重试次数 +- `--vllm-enforce-eager` —— 是否禁用 CUDA graph（默认 True） + +#### B) 新增内部统一接口 + +`slime/rollout/base_types.py` 中新增： + +- `RolloutBackendRequest`：input_ids、sampling_params、return_logprob、return_routed_experts、image_data、session_id +- `RolloutBackendResponse`：text、output_token_ids、output_token_logprobs、finish_reason（`stop|length|abort`）、prompt_tokens、completion_tokens、backend_raw、routed_experts + +#### C) 能力门控 + +`BackendCapabilities` dataclass 声明各 backend 支持的能力： + +| 能力 | SGLangClient | VLLMClient | +|---|---|---| +| `supports_abort` | True | False | +| `supports_routed_experts` | True | False | +| `supports_prompt_logprobs` | True | False | + +不支持的特性在调用时显式 gate、记录日志、或 fail fast。 + +## 5. 文件级改动清单 + +### 新增文件 + +| 文件 | 说明 | +|---|---| +| `slime/rollout/backends/base_client.py` | `RolloutBackendClient` 抽象接口 + `BackendCapabilities` | +| `slime/rollout/backends/sglang_client.py` | SGLang rollout client（从现有代码提取） | +| `slime/rollout/backends/vllm_client.py` | vLLM rollout client（`/v1/completions` 适配） | +| `slime/rollout/backends/__init__.py` | adapter 导出 | +| `slime/backends/vllm_utils/vllm_engine.py` | `VLLMEngine` Ray actor（进程管理 + 权重同步） | +| `run-qwen2.5-0.5B-vllm.sh` | vLLM 验证脚本（Qwen2.5-0.5B + GSM8K + colocate） | + +### 修改文件 + +| 文件 | 说明 | +|---|---| +| `slime/rollout/base_types.py` | 新增 `RolloutBackendRequest` / `RolloutBackendResponse` | +| `slime/rollout/sglang_rollout.py` | 使用 backend adapter 代替硬编码 SGLang 协议 | +| `slime/utils/arguments.py` | 新增 `--rollout-backend`、vLLM 参数；vLLM 时设置 sglang 别名默认值 | +| `slime/ray/rollout.py` | `_start_vllm_rollout_servers()` 启动 VLLMEngine | +| `slime/backends/megatron_utils/actor.py` | 权重更新器选择：`colocate → UpdateWeightFromTensor`（backend 无关） | +| `slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py` | 新增 `_send_to_colocated_vllm_engine()`（CUDA IPC 路径） | +| `slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py` | vLLM NCCL 兼容：`NCCLWeightTransferEngine.trainer_init()` + `PyNcclCommunicator.broadcast()` | + +### 不变文件 + +| 文件 | 说明 | +|---|---| +| `train.py` | 训练循环不感知 backend | +| SGLang 相关文件 | SGLang 路径完全不受影响 | + +### 未使用（与 RFC 初稿的差异） + +- **无 `worker_extension.py`**：vLLM 原生 weight transfer API 完全满足需求，无需自定义 WorkerExtension +- **无 `collective_rpc`**：使用 vLLM 的 `/init_weight_transfer_engine` 和 `/update_weights` 端点 + +## 6. 能力兼容矩阵 + +| 能力 | SGLang | vLLM | 处理方式 | +|---|---|---|---| +| Token 级响应 logprobs | 支持 | 支持（`choice.logprobs.token_logprobs`） | adapter 归一化 | +| Output token IDs | `meta_info` | `choice.token_ids` | adapter 归一化 | +| Finish reason | `stop\|length\|abort` | `stop\|length` + 其他 | canonical 映射 | +| Worker 级 abort | 支持 | 不支持 | 超时/取消降级 | +| Prompt logprobs（OPD） | 支持 | 部分 | capability-gated | +| Routed experts（R3） | 支持 | 不支持 | 显式 gate | +| Colocate（GPU 共享） | 支持（FlattenedTensorBucket IPC） | 支持（CUDA IPC handles） | 各自原生 IPC | +| 内存管理 | `release/resume_memory_occupation` | `POST /sleep` / `POST /wake_up` | 接口统一 | +| `include_stop_str_in_output` | `no_stop_trim` | `include_stop_str_in_output` | 参数映射 | + +## 7. 分阶段计划 + +### 第一阶段：Managed vLLM GRPO 训练 ✅ Done + +**目标**：Qwen2.5-0.5B GRPO 4 卡 colocate 训练，GSM8K 数据集，训练正常运行。 + +已完成交付物： +- `VLLMEngine` Ray actor：进程生命周期管理、`/pause`/`/resume`/`/sleep`/`/wake_up` +- `VLLMClient` rollout adapter：`/v1/completions` 请求映射、logprob/finish_reason/token_ids 归一化 +- `RolloutBackendRequest`/`Response` 统一契约 +- `SGLangClient` 提取 +- `_start_vllm_rollout_servers` 启动路径 +- 参数解析分流（`--rollout-backend vllm`，sglang 别名默认值） +- **Colocate 模式**：CUDA IPC 权重传输（`_send_to_colocated_vllm_engine` → `IPCWeightTransferEngine`） +- **Non-colocate 模式**：NCCL 权重传输（`NCCLWeightTransferEngine.trainer_init()` → `PyNcclCommunicator.broadcast()`） +- 验证脚本：`run-qwen2.5-0.5B-vllm.sh` + +技术决策： +- 不使用 router（单 vLLM 实例，直连） +- Colocate 权重同步：CUDA IPC（训练和 vLLM 共享 GPU 内存，零拷贝） +- Non-colocate 权重同步：vLLM 原生 NCCL（`StatelessProcessGroup` + `PyNcclCommunicator`） +- Rollout 函数：复用 `sglang_rollout.py`（backend 无关编排） +- 训练侧权重更新器选择与 SGLang 一致（`colocate → UpdateWeightFromTensor`） + +验收结果： +- ✅ Qwen2.5-0.5B + GSM8K + 4 GPU colocate 训练运行通过 +- ✅ 权重同步正确性：每轮 rollout 输出随训练更新变化 +- ✅ 现有 SGLang 路径不受影响 + +### 第二阶段：多 vLLM 实例 + 异步训推 + +**目标**：支持多 vLLM 实例 rollout，使用 [vllm-project/router](https://github.com/vllm-project/router) 做负载均衡，以异步训推（`train_async.py`）在更大模型上验证。 + +关键交付物： +- 集成 [vllm-project/router](https://github.com/vllm-project/router) 作为 vLLM 侧负载均衡器 +- `start_rollout_servers` 拉起 N 个 `VLLMEngine` actor + 一个 vllm-router 进程 +- `VLLMClient` 生成请求从直连改为走 vllm-router +- 验证异步训推（`train_async.py`）在 vLLM backend 下的正确性 + +验收标准： +- 多实例（如 2-4 个 vLLM worker）rollout 在 20+ rollout step 下稳定 +- 异步训推（`train_async.py`）无 hang、无权重同步竞态 +- 相比单实例基线有吞吐提升 +- 更大模型验证（如 Qwen3-4B，TP=2） + +### 第三阶段（未来）：高级特性 + +- Abort/cancel 策略完善 +- Prompt logprob 支持（OPD 场景） +- 确定性计算验证 +- 性能基准测试（vLLM vs SGLang 吞吐量对比） + +## 8. 验证策略 + +### 已完成（Phase 1） + +- ✅ 端到端 GRPO 训练：`run-qwen2.5-0.5B-vllm.sh`（Qwen2.5-0.5B, GSM8K, 4 GPU, colocate） +- ✅ 权重同步正确性验证 +- ✅ SGLang 路径非回归 + +### 待完成 + +- 单元测试：finish reason 归一化、token/logprob 对齐、capability-gated 行为 +- 集成测试：SGLang vs vLLM 响应 schema 对比 +- 压测：更高并发 / 延迟压力 + +## 9. 实现中遇到的关键问题及解决 + +### 9.1 vLLM NCCL 不兼容 + +**问题**：vLLM 内部使用 `StatelessProcessGroup` + `PyNcclCommunicator` 管理 NCCL 通信，与 `torch.distributed.init_process_group` 不兼容。 + +**解决**：训练侧使用 `NCCLWeightTransferEngine.trainer_init()` 初始化 NCCL，确保与 vLLM 端建立匹配的通信组。 + +### 9.2 vLLM logprobs 格式差异 + +**问题**：vLLM `/v1/completions` 的 logprobs 格式（`choice.token_ids`、`choice.logprobs.token_logprobs`）与 SGLang 不同。 + +**解决**：`VLLMClient` 中做显式映射，将 vLLM 格式转换为 `RolloutBackendResponse` 统一格式。对于缺失 `token_ids` 的情况，回退到 tokenizer。 + +### 9.3 Colocate 模式下的 IPC 格式差异 + +**问题**：SGLang 使用 `FlattenedTensorBucket` + `MultiprocessingSerializer` 做 IPC，vLLM 使用独立的 CUDA IPC handles（`reduce_tensor`）。两者格式不兼容。 + +**解决**：新增 `_send_to_colocated_vllm_engine()` 函数，直接使用 `torch.multiprocessing.reductions.reduce_tensor` 创建 CUDA IPC handles，以 GPU UUID 为 key。TP>1 时通过 Gloo gather 合并各 rank 的 handles，然后 pickle + base64 编码后通过 `VLLMEngine` 转发到 vLLM 的 `/update_weights` 端点。 + +### 9.4 缺失的 sglang 参数别名 + +**问题**：vLLM backend 跳过 `sglang_validate_args()`，导致 `sglang_dp_size` 等别名未设置，后续代码 `AttributeError`。 + +**解决**：在 `arguments.py` 中为 vLLM backend 添加条件分支，设置 `sglang_dp_size`、`sglang_pp_size`、`sglang_ep_size`、`sglang_tp_size` 等默认值。 + +## 10. 风险与缓解 + +- **Logprob 语义不匹配**：严格 adapter 检查 + 归一化 + 测试 +- **Abort 不匹配**：能力模型 + 超时/取消降级 + 显式日志 +- **训练漂移**：强制 A/B 对照运行，固定 seed 和配置 +- **IPC handle 生命周期**：训练侧通过 `kept_alive` 列表 + `ray.get()` 阻塞确保 tensor 在 vLLM 读取完成前不被回收 + +## 11. 向 Slime 维护者提出的开放问题 + +1. 是否需要在 Phase 2 中将 vLLM 多实例 + router 作为默认部署模式？ +2. 声称支持 GRPO 所需的最低 logprob 保真度是什么？ +3. 是否应在完全对等前阻止 vLLM 的 OPD prompt-logprob 路径？ +4. 推荐为默认 backend 之前需要满足什么质量门槛？ +5. vLLM colocate IPC 路径是否需要支持混合模式（部分 colocate + 部分 distributed）？ + +## 12. 请求决策 + +Phase 1 已实现并验证： + +- ✅ Contract-first adapter 架构 +- ✅ Managed vLLM 模式（与 SGLang 同等生命周期管理） +- ✅ 双模式权重同步（colocate: CUDA IPC / non-colocate: NCCL） +- ✅ 显式 capability-gated 行为 +- ✅ 严格 SGLang 非回归 + +请批准进入 Phase 2：多 vLLM 实例 + router + 异步训推。 + From 48fbde3c0ba7d1f1d68a427a99ca4baf3d4843ff Mon Sep 17 00:00:00 2001 From: SamitHuang <285365963@qq.com> Date: Tue, 3 Mar 2026 17:06:01 +0800 Subject: [PATCH 03/39] update Signed-off-by: SamitHuang <285365963@qq.com> --- goal_plan.md | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/goal_plan.md b/goal_plan.md index adb39bf047..d77618195b 100644 --- a/goal_plan.md +++ b/goal_plan.md @@ -1,17 +1,23 @@ ### 阶段一：打通Qwen2.5-0.5B GRPO 8卡同步/异步训练（train.py和train_async.py），GSM8K 数据集，loss/reward 收敛与 SGLang backend 基本一致，且满足确定性计算，多次重复运行Loss曲线完全一致。 +First Design and RFC by 03/06 + #### 初步方案： - 对标SGLang，Slime 在 Ray 内管理 vLLM 的完整生命周期，包括进程拉起、权重同步、推理暂停/恢复 - 暂不使用Router，SGLang Model Gateway仅只支持SGLang Worker，SlimeRouter仅在 R3 / radix-tree caching 时需要，Qwen2.5-0.5B 非 MoE 且用 token-in/token-out - 单vLLM实例，无router，通过vLLMClient 直连本地 vLLM 进程端口 -- 若训推不共卡，权重同步采用NCCL broadcast，对标SGLang update_weights_from_distributed (默认，优先用于支持异步训推） -- 若colocate，权重同步采用GPU IPC（vLLM update_weights_from_ipc, update_weights_from_tensor），对标SGLang update_weights_from_tensor +- 先支持和验证colocate，权重同步采用GPU IPC（vLLM update_weights_from_ipc, update_weights_from_tensor），对标SGLang update_weights_from_tensor，以验证Reproductivity +- 再支持训推不共卡，权重同步采用NCCL broadcast，对标SGLang update_weights_from_distributed (默认） #### 风险： - slime, sglang版本依赖，和vllm 0.16的版本依赖冲突(numpy, torch, transformers, etc) +- slime代码较挫，可靠性差，强依赖preset docker - 算力 -First Design and RFC by 03/06 + +#### Reference + +https://thudm.github.io/slime/advanced/reproducibility.html ### 阶段二：接入vllm-project/router，支持多实例vLLM From f8ceed67937f14bb7a2b95fedba0cf744077c317 Mon Sep 17 00:00:00 2001 From: samithuang <285365963@qq.com> Date: Wed, 4 Mar 2026 06:18:35 +0000 Subject: [PATCH 04/39] qwen2.5 0.5b non-colocate (first attempt ok, but nccl error later) Signed-off-by: samithuang <285365963@qq.com> --- ...wen2.5-0.5B-reproducibility-noncolocate.sh | 137 +++++++++++ run-qwen2.5-0.5B-vllm.sh | 137 +++++++++++ slime/backends/megatron_utils/actor.py | 3 +- .../update_weight_from_distributed.py | 73 ++++-- slime/backends/vllm_utils/__init__.py | 3 + slime/backends/vllm_utils/vllm_engine.py | 214 ++++++++++++++++++ slime/ray/rollout.py | 47 ++++ slime/rollout/backends/__init__.py | 10 + slime/rollout/backends/base_client.py | 31 +++ slime/rollout/backends/sglang_client.py | 82 +++++++ slime/rollout/backends/vllm_client.py | 91 ++++++++ slime/rollout/base_types.py | 26 +++ slime/rollout/sglang_rollout.py | 152 ++++++------- slime/utils/arguments.py | 29 ++- 14 files changed, 939 insertions(+), 96 deletions(-) create mode 100644 run-qwen2.5-0.5B-reproducibility-noncolocate.sh create mode 100644 run-qwen2.5-0.5B-vllm.sh create mode 100644 slime/backends/vllm_utils/__init__.py create mode 100644 slime/backends/vllm_utils/vllm_engine.py create mode 100644 slime/rollout/backends/__init__.py create mode 100644 slime/rollout/backends/base_client.py create mode 100644 slime/rollout/backends/sglang_client.py create mode 100644 slime/rollout/backends/vllm_client.py diff --git a/run-qwen2.5-0.5B-reproducibility-noncolocate.sh b/run-qwen2.5-0.5B-reproducibility-noncolocate.sh new file mode 100644 index 0000000000..5499949340 --- /dev/null +++ b/run-qwen2.5-0.5B-reproducibility-noncolocate.sh @@ -0,0 +1,137 @@ +#!/bin/bash +# Non-colocate version of run-qwen2.5-0.5B-reproducibility.sh +# 2 GPUs: 1 for training, 1 for SGLang rollout + +# for rerun the task +pkill -9 sglang +sleep 3 +ray stop --force +pkill -9 ray +pkill -9 python +sleep 3 +pkill -9 ray +pkill -9 python + +set -ex + +export PYTHONBUFFERED=16 + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +source "${SCRIPT_DIR}/scripts/models/qwen2.5-0.5B.sh" + +CKPT_ARGS=( + --hf-checkpoint /root/Qwen2.5-0.5B-Instruct/ + --ref-load /root/Qwen2.5-0.5B-Instruct_torch_dist/ +) + +ROLLOUT_ARGS=( + --prompt-data /root/gsm8k/train.parquet + --input-key messages + --label-key label + --apply-chat-template + --rollout-shuffle + --rm-type math + --num-rollout 100 + --rollout-batch-size 32 + --n-samples-per-prompt 8 + --rollout-max-response-len 1024 + --rollout-temperature 1 + + --global-batch-size 256 +) + +EVAL_ARGS=( + --eval-interval 20 + --eval-prompt-data gsm8k /root/gsm8k/test.parquet + --n-samples-per-eval-prompt 1 + --eval-max-response-len 1024 + --eval-top-k 1 +) + +PERF_ARGS=( + --tensor-model-parallel-size 1 + --sequence-parallel + --pipeline-model-parallel-size 1 + --context-parallel-size 1 + --expert-model-parallel-size 1 + --expert-tensor-parallel-size 1 + + --use-dynamic-batch-size + --max-tokens-per-gpu 9216 +) + +GRPO_ARGS=( + --advantage-estimator grpo + --use-kl-loss + --kl-loss-coef 0.00 + --kl-loss-type low_var_kl + --kl-coef 0.00 + --entropy-coef 0.00 + --eps-clip 0.2 + --eps-clip-high 0.28 +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-6 + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 +) + +WANDB_ARGS=( + --use-wandb + --wandb-host https://wandb.ai/ + --wandb-entity samithuang + --wandb-project slime-rl + --wandb-group qwen2.5-0.5B-gsm8k-noncolocate +) + +SGLANG_ARGS=( + --rollout-num-gpus-per-engine 1 + --sglang-mem-fraction-static 0.7 + + --sglang-enable-deterministic-inference + --sglang-attention-backend flashinfer + + --deterministic-mode +) + +MISC_ARGS=( + --attention-dropout 0.0 + --hidden-dropout 0.0 + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + --attention-backend flash +) + +ray start --head --node-ip-address 127.0.0.1 --num-gpus 2 --disable-usage-stats + +ray job submit --address="http://127.0.0.1:8265" \ + --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 \ + --actor-num-nodes 1 \ + --actor-num-gpus-per-node 1 \ + --num-gpus-per-node 2 \ + --rollout-num-gpus 1 \ + --calculate-per-token-loss \ + --use-slime-router \ + ${MODEL_ARGS[@]} \ + ${CKPT_ARGS[@]} \ + ${ROLLOUT_ARGS[@]} \ + ${OPTIMIZER_ARGS[@]} \ + ${GRPO_ARGS[@]} \ + ${WANDB_ARGS[@]} \ + ${PERF_ARGS[@]} \ + ${EVAL_ARGS[@]} \ + ${SGLANG_ARGS[@]} \ + ${MISC_ARGS[@]} diff --git a/run-qwen2.5-0.5B-vllm.sh b/run-qwen2.5-0.5B-vllm.sh new file mode 100644 index 0000000000..bcc6807c69 --- /dev/null +++ b/run-qwen2.5-0.5B-vllm.sh @@ -0,0 +1,137 @@ +#!/bin/bash +# vLLM rollout backend validation script (Phase 1) +# Based on run-qwen2.5-0.5B-reproducibility.sh + +# for rerun the task +pkill -9 vllm +pkill -9 sglang +sleep 3 +ray stop --force +pkill -9 ray +pkill -9 python +sleep 3 +pkill -9 ray +pkill -9 python + + +set -ex + +export PYTHONBUFFERED=16 + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +source "${SCRIPT_DIR}/scripts/models/qwen2.5-0.5B.sh" + +CKPT_ARGS=( + --hf-checkpoint /root/Qwen2.5-0.5B-Instruct/ + --ref-load /root/Qwen2.5-0.5B-Instruct_torch_dist/ +) + +ROLLOUT_ARGS=( + --prompt-data /root/gsm8k/train.parquet + --input-key messages + --label-key label + --apply-chat-template + --rollout-shuffle + --rm-type math + --num-rollout 100 + --rollout-batch-size 32 + --n-samples-per-prompt 8 + --rollout-max-response-len 1024 + --rollout-temperature 1 + + --global-batch-size 256 +) + +EVAL_ARGS=( + --eval-interval 20 + --eval-prompt-data gsm8k /root/gsm8k/test.parquet + --n-samples-per-eval-prompt 1 + --eval-max-response-len 1024 + --eval-top-k 1 +) + +PERF_ARGS=( + --tensor-model-parallel-size 1 + --sequence-parallel + --pipeline-model-parallel-size 1 + --context-parallel-size 1 + --expert-model-parallel-size 1 + --expert-tensor-parallel-size 1 + + --use-dynamic-batch-size + --max-tokens-per-gpu 9216 +) + +GRPO_ARGS=( + --advantage-estimator grpo + --use-kl-loss + --kl-loss-coef 0.00 + --kl-loss-type low_var_kl + --kl-coef 0.00 + --entropy-coef 0.00 + --eps-clip 0.2 + --eps-clip-high 0.28 +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-6 + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 +) + +WANDB_ARGS=( + --use-wandb + --wandb-host https://wandb.ai/ + --wandb-entity samithuang + --wandb-project slime-rl + --wandb-group qwen2.5-0.5B-gsm8k-vllm +) + +VLLM_ARGS=( + --rollout-backend vllm + --rollout-num-gpus-per-engine 1 + --sglang-server-concurrency 512 +) + +MISC_ARGS=( + --attention-dropout 0.0 + --hidden-dropout 0.0 + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + --attention-backend flash + --deterministic-mode +) + +ray start --head --node-ip-address 127.0.0.1 --num-gpus 2 --disable-usage-stats + +ray job submit --address="http://127.0.0.1:8265" \ + --runtime-env-json='{ + "env_vars": { + "PYTHONPATH": "/root/Megatron-LM", + "CUDA_DEVICE_MAX_CONNECTIONS": "1", + "NCCL_ALGO": "Ring", + "NCCL_P2P_DISABLE": "1", + "NCCL_DEBUG": "INFO", + "NVTE_ALLOW_NONDETERMINISTIC_ALGO": "0", + "CUBLAS_WORKSPACE_CONFIG": ":4096:8" + } + }' \ + -- python3 train.py \ + --actor-num-nodes 1 \ + --actor-num-gpus-per-node 1 \ + --num-gpus-per-node 2 \ + --rollout-num-gpus 1 \ + --calculate-per-token-loss \ + ${MODEL_ARGS[@]} \ + ${CKPT_ARGS[@]} \ + ${ROLLOUT_ARGS[@]} \ + ${OPTIMIZER_ARGS[@]} \ + ${GRPO_ARGS[@]} \ + ${WANDB_ARGS[@]} \ + ${PERF_ARGS[@]} \ + ${EVAL_ARGS[@]} \ + ${VLLM_ARGS[@]} \ + ${MISC_ARGS[@]} diff --git a/slime/backends/megatron_utils/actor.py b/slime/backends/megatron_utils/actor.py index 9abb2f96ad..1f45e5fc94 100644 --- a/slime/backends/megatron_utils/actor.py +++ b/slime/backends/megatron_utils/actor.py @@ -128,7 +128,8 @@ def init( if self.args.vocab_size is None: self.args.vocab_size = self.tokenizer.vocab_size - update_weight_cls = UpdateWeightFromTensor if self.args.colocate else UpdateWeightFromDistributed + use_tensor_update = self.args.colocate and getattr(self.args, "rollout_backend", "sglang") != "vllm" + update_weight_cls = UpdateWeightFromTensor if use_tensor_update else UpdateWeightFromDistributed self.weight_updater = update_weight_cls( self.args, self.model, diff --git a/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py b/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py index 7b5b7817f1..44649f87f4 100644 --- a/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py +++ b/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py @@ -1,3 +1,4 @@ +import os import socket import time from argparse import Namespace @@ -241,6 +242,7 @@ def _update_bucket_weights_from_distributed( self.weight_version, self.rollout_engines, converted_named_tensors, + use_vllm=_is_vllm_backend(self.args), ) ray.get(refs) @@ -249,6 +251,10 @@ def _update_bucket_weights_from_distributed( pbar.update(1) +def _is_vllm_backend(args: Namespace) -> bool: + return getattr(args, "rollout_backend", "sglang") == "vllm" + + def connect_rollout_engines_from_distributed( args: Namespace, group_name: str, @@ -261,6 +267,10 @@ def connect_rollout_engines_from_distributed( ``engine_gpu_counts`` gives the number of GPUs per engine. When engines have heterogeneous TP sizes (e.g. prefill TP=2, decode TP=4), each engine occupies a different number of ranks in the NCCL group. + + For vLLM backend, uses vLLM's StatelessProcessGroup + PyNcclCommunicator + instead of torch.distributed, because vLLM's weight transfer engine uses + its own NCCL initialization protocol. """ if engine_gpu_counts is None: engine_gpu_counts = [args.rollout_num_gpus_per_engine] * len(rollout_engines) @@ -287,13 +297,39 @@ def connect_rollout_engines_from_distributed( ) for i, engine in enumerate(rollout_engines) ] - model_update_groups = init_process_group( - backend="nccl", - init_method=f"tcp://{master_address}:{master_port}", - world_size=world_size, - rank=0, - group_name=group_name, - ) + + if _is_vllm_backend(args): + # vLLM uses StatelessProcessGroup + PyNcclCommunicator for weight transfer. + # The training side must use the same mechanism for NCCL compatibility. + # + # Disable P2P transport: in colocate mode the trainer and vLLM server + # share the same physical GPU but have different CUDA_VISIBLE_DEVICES, + # which causes NCCL P2P (cudaIpc*) to fail with "invalid argument". + # SHM transport works correctly in this scenario. + from vllm.distributed.weight_transfer.nccl_engine import NCCLWeightTransferEngine + + old_p2p = os.environ.get("NCCL_P2P_DISABLE") + os.environ["NCCL_P2P_DISABLE"] = "1" + try: + model_update_groups = NCCLWeightTransferEngine.trainer_init({ + "master_address": master_address, + "master_port": master_port, + "world_size": world_size, + }) + finally: + if old_p2p is None: + os.environ.pop("NCCL_P2P_DISABLE", None) + else: + os.environ["NCCL_P2P_DISABLE"] = old_p2p + else: + model_update_groups = init_process_group( + backend="nccl", + init_method=f"tcp://{master_address}:{master_port}", + world_size=world_size, + rank=0, + group_name=group_name, + ) + ray.get(refs) return model_update_groups @@ -303,19 +339,24 @@ def disconnect_rollout_engines_from_distributed(args, group_name, model_update_g Destroy NCCL on training and engines. """ refs = [engine.destroy_weights_update_group.remote(group_name) for engine in rollout_engines] - dist.destroy_process_group(model_update_groups) + if _is_vllm_backend(args): + model_update_groups = None + else: + dist.destroy_process_group(model_update_groups) ray.get(refs) def update_weights_from_distributed( group_name: str, - group: dist.ProcessGroup, + group, weight_version: int, rollout_engines: Sequence[ActorHandle], converted_named_tensors: Sequence[tuple[str, torch.Tensor]], + use_vllm: bool = False, ) -> list[ObjectRef]: """ Send metadata (Ray), broadcast tensors (NCCL rank 0 → engines). + For vLLM, uses PyNcclCommunicator.broadcast instead of dist.broadcast. """ refs = [ engine.update_weights_from_distributed.remote( @@ -328,11 +369,15 @@ def update_weights_from_distributed( for engine in rollout_engines ] - handles = [] - for _, param in converted_named_tensors: - handles.append(dist.broadcast(param.data, 0, group=group, async_op=True)) - for handle in handles: - handle.wait() + if use_vllm: + for _, param in converted_named_tensors: + group.broadcast(param.data, src=0, stream=torch.cuda.current_stream()) + else: + handles = [] + for _, param in converted_named_tensors: + handles.append(dist.broadcast(param.data, 0, group=group, async_op=True)) + for handle in handles: + handle.wait() return refs diff --git a/slime/backends/vllm_utils/__init__.py b/slime/backends/vllm_utils/__init__.py new file mode 100644 index 0000000000..ea4e7311c3 --- /dev/null +++ b/slime/backends/vllm_utils/__init__.py @@ -0,0 +1,3 @@ +from slime.backends.vllm_utils.vllm_engine import VLLMEngine + +__all__ = ["VLLMEngine"] diff --git a/slime/backends/vllm_utils/vllm_engine.py b/slime/backends/vllm_utils/vllm_engine.py new file mode 100644 index 0000000000..12283de21e --- /dev/null +++ b/slime/backends/vllm_utils/vllm_engine.py @@ -0,0 +1,214 @@ +"""VLLMEngine: Ray actor that launches and manages a vLLM server.""" + +import logging +import os +import subprocess +import tempfile +import time + +import requests + +from slime.ray.ray_actor import RayActor +from slime.utils.http_utils import get_host_info +from slime.utils.misc import get_free_port + +logger = logging.getLogger(__name__) + + +class VLLMEngine(RayActor): + """Ray actor that runs vLLM server with same interface as SGLangEngine for weight sync.""" + + def __init__(self, args, rank: int, base_gpu_id: int | None = None, gpu_ids: list[int] | None = None, **kwargs): + self.args = args + self.rank = rank + self.base_gpu_id = base_gpu_id or 0 + self.gpu_ids = gpu_ids or [self.base_gpu_id] + self.server_host = None + self.server_port = None + self.process = None + self._log_file = None + + def init(self, port=None, host=None, **kwargs): + self.server_host = host or get_host_info()[1] + self.server_port = port or get_free_port(15000) + + model = getattr(self.args, "vllm_model", None) or self.args.hf_checkpoint + tp = self.args.rollout_num_gpus_per_engine + gpu_ids = self.gpu_ids[:tp] + cvd = os.environ.get("CUDA_VISIBLE_DEVICES", "") + if cvd: + visible = [int(x) for x in cvd.split(",") if x.strip()] + dev_str = ",".join(str(visible[gid]) if gid < len(visible) else str(gid) for gid in gpu_ids) + else: + dev_str = ",".join(str(g) for g in gpu_ids) + + cmd = [ + "vllm", "serve", model, + "--tensor-parallel-size", str(tp), + "--port", str(self.server_port), + "--host", "0.0.0.0", + "--weight-transfer-config", '{"backend": "nccl"}', + ] + if getattr(self.args, "offload_rollout", False): + cmd.append("--enable-sleep-mode") + if getattr(self.args, "vllm_enforce_eager", True): + cmd.append("--enforce-eager") + if getattr(self.args, "fp16", False): + cmd.extend(["--dtype", "float16"]) + + env = os.environ.copy() + env["VLLM_SERVER_DEV_MODE"] = "1" + env["CUDA_VISIBLE_DEVICES"] = dev_str + env.setdefault("NCCL_DEBUG", "INFO") + env.setdefault("NCCL_DEBUG_SUBSYS", "ALL") + env["NCCL_P2P_DISABLE"] = "1" + + self._log_file = tempfile.NamedTemporaryFile( + prefix="vllm_engine_", suffix=".log", delete=False, mode="w" + ) + logger.info("Launching vLLM: cmd=%s, CUDA_VISIBLE_DEVICES=%s, log=%s", + " ".join(cmd), dev_str, self._log_file.name) + self.process = subprocess.Popen( + cmd, + env=env, + stdout=self._log_file, + stderr=subprocess.STDOUT, + ) + self._wait_healthy() + + def _wait_healthy(self, timeout=300): + base = f"http://{self.server_host}:{self.server_port}" + start = time.time() + while time.time() - start < timeout: + try: + r = requests.get(f"{base}/health", timeout=5) + if r.status_code == 200: + logger.info("vLLM server healthy at %s:%s", self.server_host, self.server_port) + return + except Exception: + pass + if self.process and self.process.poll() is not None: + log_tail = self._read_log_tail() + raise RuntimeError(f"vLLM process exited with code {self.process.returncode}.\n{log_tail}") + time.sleep(2) + log_tail = self._read_log_tail() + raise TimeoutError(f"vLLM server failed to become healthy within {timeout}s.\n{log_tail}") + + def _read_log_tail(self, n=200): + if not self._log_file: + return "" + try: + self._log_file.flush() + with open(self._log_file.name) as f: + lines = f.readlines() + return "".join(lines[-n:]) + except Exception: + return "" + + def _post(self, path: str, json_data: dict | None = None, params: dict | None = None): + url = f"http://{self.server_host}:{self.server_port}{path}" + kwargs = {"timeout": 120} + if json_data is not None: + kwargs["json"] = json_data + if params is not None: + kwargs["params"] = params + r = requests.post(url, **kwargs) + if not r.ok: + body = r.text[:2000] if r.text else "(empty)" + log_tail = self._read_log_tail(50) + logger.error( + "vLLM %s returned %s: %s\n--- vLLM log tail ---\n%s", + path, r.status_code, body, log_tail, + ) + r.raise_for_status() + try: + return r.json() + except Exception: + return None + + def health_generate(self, timeout: float = 5.0) -> bool: + try: + r = requests.get( + f"http://{self.server_host}:{self.server_port}/health", + timeout=timeout, + ) + r.raise_for_status() + return True + except requests.RequestException: + return False + + def pause_generation(self): + self._post("/pause", params={"mode": "abort"}) + + def flush_cache(self): + pass + + def init_weights_update_group(self, master_address, master_port, rank_offset, world_size, group_name=None, backend=None): + logger.info( + "Initializing NCCL weight transfer: master=%s:%s, rank_offset=%d, world_size=%d", + master_address, master_port, rank_offset, world_size, + ) + self._post("/init_weight_transfer_engine", json_data={ + "init_info": { + "master_address": master_address, + "master_port": master_port, + "rank_offset": rank_offset, + "world_size": world_size, + } + }) + + def update_weights_from_distributed(self, names, dtypes, shapes, group_name=None, flush_cache=False, weight_version=None): + dtype_names = [str(d).replace("torch.", "") for d in dtypes] + shape_lists = [list(s) for s in shapes] + self._post("/update_weights", json_data={ + "update_info": { + "names": names, + "dtype_names": dtype_names, + "shapes": shape_lists, + "packed": False, + } + }) + + def continue_generation(self): + self._post("/resume") + + def destroy_weights_update_group(self, group_name): + pass + + def release_memory_occupation(self): + try: + self._post("/sleep", params={"level": "1", "mode": "abort"}) + except requests.RequestException as e: + logger.warning("vLLM sleep failed (need --enable-sleep-mode?): %s", e) + + def resume_memory_occupation(self, tags: list[str] | None = None): + try: + params = {} + if tags: + params["tags"] = tags + self._post("/wake_up", params=params) + except requests.RequestException as e: + logger.warning("vLLM wake_up failed: %s", e) + + def get_weight_version(self): + return None + + def check_weights(self, action: str): + pass + + def post_process_weights(self, **kwargs): + pass + + def shutdown(self): + if self.process: + self.process.terminate() + try: + self.process.wait(timeout=10) + except subprocess.TimeoutExpired: + self.process.kill() + self.process = None + if self._log_file: + try: + self._log_file.close() + except Exception: + pass diff --git a/slime/ray/rollout.py b/slime/ray/rollout.py index 8f919c1172..ba18fb9c47 100644 --- a/slime/ray/rollout.py +++ b/slime/ray/rollout.py @@ -16,6 +16,7 @@ from sglang.srt.constants import GPU_MEMORY_TYPE_CUDA_GRAPH, GPU_MEMORY_TYPE_KV_CACHE, GPU_MEMORY_TYPE_WEIGHTS from slime.backends.sglang_utils.sglang_engine import SGLangEngine +from slime.backends.vllm_utils.vllm_engine import VLLMEngine from slime.rollout.base_types import call_rollout_fn from slime.utils import logging_utils from slime.utils.health_monitor import RolloutHealthMonitor @@ -1020,6 +1021,49 @@ def _start_router(args, *, has_pd_disaggregation: bool = False, force_new: bool return router_ip, router_port +def _start_vllm_rollout_servers(args, pg) -> dict[str, RolloutServer]: + """Start vLLM rollout server (single instance, no router).""" + pg_obj, reordered_bundle_indices, reordered_gpu_ids = pg + tp = args.rollout_num_gpus_per_engine + gpu_ids = [int(reordered_gpu_ids[i]) for i in range(tp)] + scheduling_strategy = PlacementGroupSchedulingStrategy( + placement_group=pg_obj, + placement_group_capture_child_tasks=True, + placement_group_bundle_index=reordered_bundle_indices[0], + ) + + env_vars = {name: "1" for name in NOSET_VISIBLE_DEVICES_ENV_VARS_LIST} + + VLLMRayActor = ray.remote(VLLMEngine) + engine = VLLMRayActor.options( + num_cpus=0.2, + num_gpus=0.2, + scheduling_strategy=scheduling_strategy, + runtime_env={"env_vars": env_vars}, + ).remote(args, rank=0, base_gpu_id=gpu_ids[0], gpu_ids=gpu_ids) + + host, port = ray.get(engine._get_current_node_ip_and_free_port.remote(start_port=15000)) + ray.get(engine.init.remote(port=port, host=host)) + args.vllm_base_url = f"http://{host}:{port}" + args.sglang_router_ip = host + args.sglang_router_port = port + + group = EngineGroup( + args=args, + pg=pg, + all_engines=[engine], + num_gpus_per_engine=args.rollout_num_gpus_per_engine, + num_new_engines=1, + worker_type="regular", + rank_offset=0, + gpu_offset=0, + sglang_overrides={}, + router_ip=host, + router_port=port, + ) + return {"default": RolloutServer(engine_groups=[group], router_ip=host, router_port=port, model_name="default")} + + def start_rollout_servers(args, pg) -> dict[str, RolloutServer]: """Start rollout servers: one per model, each with its own router. @@ -1033,6 +1077,9 @@ def start_rollout_servers(args, pg) -> dict[str, RolloutServer]: Note: ``init_http_client`` should be called separately before this, as the HTTP client is shared across all servers. """ + if getattr(args, "rollout_backend", "sglang") == "vllm": + return _start_vllm_rollout_servers(args, pg) + config = _resolve_sglang_config(args) servers: dict[str, RolloutServer] = {} diff --git a/slime/rollout/backends/__init__.py b/slime/rollout/backends/__init__.py new file mode 100644 index 0000000000..49c4a86eef --- /dev/null +++ b/slime/rollout/backends/__init__.py @@ -0,0 +1,10 @@ +from slime.rollout.backends.base_client import BackendCapabilities, RolloutBackendClient +from slime.rollout.backends.sglang_client import SGLangClient +from slime.rollout.backends.vllm_client import VLLMClient + +__all__ = [ + "BackendCapabilities", + "RolloutBackendClient", + "SGLangClient", + "VLLMClient", +] diff --git a/slime/rollout/backends/base_client.py b/slime/rollout/backends/base_client.py new file mode 100644 index 0000000000..1730d6cbac --- /dev/null +++ b/slime/rollout/backends/base_client.py @@ -0,0 +1,31 @@ +from abc import ABC, abstractmethod +from dataclasses import dataclass + +from slime.rollout.base_types import RolloutBackendRequest, RolloutBackendResponse + + +@dataclass +class BackendCapabilities: + supports_abort: bool + supports_routed_experts: bool + supports_prompt_logprobs: bool + + +class RolloutBackendClient(ABC): + @property + @abstractmethod + def capabilities(self) -> BackendCapabilities: + ... + + @abstractmethod + async def generate( + self, + request: RolloutBackendRequest, + base_url: str, + headers: dict | None = None, + ) -> RolloutBackendResponse: + ... + + async def abort(self) -> list[str]: + """Return worker URLs for abort. Empty for backends without worker-level abort.""" + return [] diff --git a/slime/rollout/backends/sglang_client.py b/slime/rollout/backends/sglang_client.py new file mode 100644 index 0000000000..505070c37d --- /dev/null +++ b/slime/rollout/backends/sglang_client.py @@ -0,0 +1,82 @@ +import logging + +import numpy as np +import pybase64 +import sglang_router +from packaging.version import parse + +from slime.rollout.backends.base_client import BackendCapabilities, RolloutBackendClient +from slime.rollout.base_types import RolloutBackendRequest, RolloutBackendResponse +from slime.utils.http_utils import get, post + +logger = logging.getLogger(__name__) + + +class SGLangClient(RolloutBackendClient): + def __init__(self, args): + self.args = args + + @property + def capabilities(self) -> BackendCapabilities: + return BackendCapabilities( + supports_abort=True, + supports_routed_experts=bool(getattr(self.args, "use_rollout_routing_replay", False)), + supports_prompt_logprobs=True, + ) + + async def generate( + self, + request: RolloutBackendRequest, + base_url: str, + headers: dict | None = None, + ) -> RolloutBackendResponse: + payload = { + "input_ids": request.input_ids, + "sampling_params": request.sampling_params, + "return_logprob": request.return_logprob, + "return_routed_experts": request.return_routed_experts, + } + if request.image_data: + payload["image_data"] = request.image_data + + url = f"{base_url.rstrip('/')}/generate" + output = await post(url, payload, headers=headers) + + meta = output.get("meta_info", {}) + logprobs = meta.get("output_token_logprobs", []) + output_token_ids = [item[1] for item in logprobs] + output_token_logprobs = [item[0] for item in logprobs] + + finish_reason = meta.get("finish_reason", {}).get("type", "stop") + routed_experts = None + if "routed_experts" in meta and self.capabilities.supports_routed_experts: + num_layers = getattr(self.args, "num_layers", 0) + moe_topk = getattr(self.args, "moe_router_topk", 1) + if num_layers and moe_topk: + routed_experts = np.frombuffer( + pybase64.b64decode(meta["routed_experts"].encode("ascii")), + dtype=np.int32, + ).reshape( + len(request.input_ids) + len(output_token_ids) - 1, + num_layers, + moe_topk, + ) + + return RolloutBackendResponse( + text=output.get("text", ""), + output_token_ids=output_token_ids, + output_token_logprobs=output_token_logprobs, + finish_reason=finish_reason, + prompt_tokens=meta.get("prompt_tokens", len(request.input_ids)), + completion_tokens=meta.get("completion_tokens", len(output_token_ids)), + backend_raw=output, + routed_experts=routed_experts, + ) + + async def abort(self) -> list[str]: + base = f"http://{self.args.sglang_router_ip}:{self.args.sglang_router_port}" + if parse(sglang_router.__version__) <= parse("0.2.1") or getattr(self.args, "use_slime_router", False): + r = await get(f"{base}/list_workers") + return r.get("urls", []) + r = await get(f"{base}/workers") + return [w["url"] for w in r.get("workers", [])] diff --git a/slime/rollout/backends/vllm_client.py b/slime/rollout/backends/vllm_client.py new file mode 100644 index 0000000000..537569499d --- /dev/null +++ b/slime/rollout/backends/vllm_client.py @@ -0,0 +1,91 @@ +import logging + +from slime.rollout.backends.base_client import BackendCapabilities, RolloutBackendClient +from slime.rollout.base_types import RolloutBackendRequest, RolloutBackendResponse +from slime.utils.http_utils import post + +logger = logging.getLogger(__name__) + +_FINISH_REASON_MAP = { + "stop": "stop", + "length": "length", + "abort": "abort", + "end_turn": "stop", + "max_tokens": "length", +} + + +class VLLMClient(RolloutBackendClient): + def __init__(self, args): + self.args = args + self._max_retries = getattr(args, "vllm_max_retries", 3) + + @property + def capabilities(self) -> BackendCapabilities: + return BackendCapabilities( + supports_abort=False, + supports_routed_experts=False, + supports_prompt_logprobs=False, + ) + + async def generate( + self, + request: RolloutBackendRequest, + base_url: str, + headers: dict | None = None, + ) -> RolloutBackendResponse: + sp = request.sampling_params + payload = { + "prompt": request.input_ids, + "max_tokens": sp.get("max_new_tokens", 1024), + "temperature": sp.get("temperature", 1.0), + "top_p": sp.get("top_p", 1.0), + "top_k": sp.get("top_k", -1), + "stop": sp.get("stop"), + "stop_token_ids": sp.get("stop_token_ids"), + "skip_special_tokens": sp.get("skip_special_tokens", True), + "logprobs": 1 if request.return_logprob else None, + "include_stop_str_in_output": sp.get("no_stop_trim", False), + "return_token_ids": True, + } + payload = {k: v for k, v in payload.items() if v is not None} + + url = f"{base_url.rstrip('/')}/v1/completions" + output = await post(url, payload, headers=headers) + + choice = output.get("choices", [{}])[0] + text = choice.get("text", "") + finish = choice.get("finish_reason", "stop") + finish_reason = _FINISH_REASON_MAP.get(finish, "stop") + + # vLLM /v1/completions response format: + # choice.token_ids: list[int] (output token IDs) + # choice.logprobs.token_logprobs: list[float|None] + # choice.logprobs.tokens: list[str] (token text, not IDs) + output_token_ids = choice.get("token_ids") or [] + logprobs_obj = choice.get("logprobs") or {} + raw_logprobs = logprobs_obj.get("token_logprobs") or [] + output_token_logprobs = [float(lp) if lp is not None else 0.0 for lp in raw_logprobs] + + # If token_ids not in response, fall back to tokenizer + if not output_token_ids and text: + logger.warning("vLLM response missing token_ids, falling back to tokenizer") + from slime.utils.processing_utils import load_tokenizer + tokenizer = load_tokenizer(self.args.hf_checkpoint, trust_remote_code=True) + output_token_ids = tokenizer.encode(text, add_special_tokens=False) + + # Ensure logprobs list matches token count + if len(output_token_logprobs) < len(output_token_ids): + output_token_logprobs.extend([0.0] * (len(output_token_ids) - len(output_token_logprobs))) + + usage = output.get("usage", {}) + return RolloutBackendResponse( + text=text, + output_token_ids=output_token_ids, + output_token_logprobs=output_token_logprobs, + finish_reason=finish_reason, + prompt_tokens=usage.get("prompt_tokens", len(request.input_ids)), + completion_tokens=usage.get("completion_tokens", len(output_token_ids)), + backend_raw=output, + routed_experts=None, + ) diff --git a/slime/rollout/base_types.py b/slime/rollout/base_types.py index f0eb0d96ce..75c140778a 100644 --- a/slime/rollout/base_types.py +++ b/slime/rollout/base_types.py @@ -4,6 +4,32 @@ from slime.utils.types import Sample +@dataclass +class RolloutBackendRequest: + """Backend-agnostic rollout request.""" + + input_ids: list[int] + sampling_params: dict[str, Any] + return_logprob: bool = True + return_routed_experts: bool = False + image_data: list[str] | None = None + session_id: str | None = None + + +@dataclass +class RolloutBackendResponse: + """Backend-agnostic rollout response.""" + + text: str + output_token_ids: list[int] + output_token_logprobs: list[float] + finish_reason: str # "stop" | "length" | "abort" + prompt_tokens: int + completion_tokens: int + backend_raw: dict + routed_experts: Any = None + + @dataclass class RolloutFnTrainOutput: samples: list[list[Sample]] diff --git a/slime/rollout/sglang_rollout.py b/slime/rollout/sglang_rollout.py index 042fdab315..0b59be4477 100644 --- a/slime/rollout/sglang_rollout.py +++ b/slime/rollout/sglang_rollout.py @@ -9,17 +9,14 @@ from typing import Any import numpy as np -import pybase64 -import sglang_router -from packaging.version import parse from tqdm import tqdm -from slime.rollout.base_types import RolloutFnEvalOutput, RolloutFnTrainOutput +from slime.rollout.base_types import RolloutBackendRequest, RolloutBackendResponse, RolloutFnEvalOutput, RolloutFnTrainOutput from slime.rollout.filter_hub.base_types import MetricGatherer, call_dynamic_filter from slime.utils.async_utils import run from slime.utils.data import Dataset from slime.utils.eval_config import EvalDatasetConfig -from slime.utils.http_utils import get, post +from slime.utils.http_utils import post from slime.utils.misc import SingletonMeta, load_function from slime.utils.processing_utils import ( build_processor_kwargs, @@ -33,6 +30,37 @@ __all__ = ["generate_rollout"] + +def _get_backend_client(args): + backend = getattr(args, "rollout_backend", "sglang") + if backend == "vllm": + from .backends.vllm_client import VLLMClient + return VLLMClient(args) + from .backends.sglang_client import SGLangClient + return SGLangClient(args) + + +def _apply_backend_response(sample, resp: RolloutBackendResponse, args): + """Apply RolloutBackendResponse to sample.""" + sample.tokens = sample.tokens + resp.output_token_ids + sample.response_length += len(resp.output_token_ids) + sample.response += resp.text + if sample.loss_mask is not None: + assert args.partial_rollout and args.mask_offpolicy_in_partial_rollout + sample.loss_mask += [1] * len(resp.output_token_ids) + if sample.rollout_log_probs is None: + sample.rollout_log_probs = [] + sample.rollout_log_probs += resp.output_token_logprobs + if resp.routed_experts is not None: + sample.rollout_routed_experts = resp.routed_experts + meta = { + "finish_reason": {"type": resp.finish_reason}, + "prompt_tokens": resp.prompt_tokens, + "completion_tokens": resp.completion_tokens, + **{k: v for k, v in resp.backend_raw.get("meta_info", {}).items() if k not in ("finish_reason",)}, + } + sample.update_from_meta_info(args, meta) + logger = logging.getLogger(__name__) @@ -106,13 +134,11 @@ def submit_generate_tasks(self, samples: list[list[Sample]]) -> None: async def generate(args: Namespace, sample: Sample, sampling_params: dict[str, Any]) -> Sample: - """Generate using traditional SGLang router with token-based workflow""" + """Generate using backend client (SGLang or vLLM).""" if args.ci_test: assert isinstance(sample.prompt, str) state = GenerateState(args) - url = f"http://{args.sglang_router_ip}:{args.sglang_router_port}/generate" - assert ( sample.status == Sample.Status.PENDING or sample.status == Sample.Status.ABORTED ), f"Sample status is {sample.status}" @@ -137,70 +163,41 @@ async def generate(args: Namespace, sample: Sample, sampling_params: dict[str, A sample.status = Sample.Status.TRUNCATED return sample - # Prepare payload for sglang server - payload = { - "sampling_params": sampling_params, - "return_logprob": True, - } + input_ids = sample.tokens if len(sample.response) > 0 else prompt_ids + if not sample.tokens: + sample.tokens = list(input_ids) - if args.use_rollout_routing_replay: - payload["return_routed_experts"] = True - - if sample.multimodal_inputs and sample.multimodal_inputs["images"]: - image_data = sample.multimodal_inputs["images"] - payload["image_data"] = [encode_image_for_rollout_engine(image) for image in image_data] - - # Use existing tokens for multi-turn or tokenize the new prompt - if len(sample.response) > 0: - payload["input_ids"] = sample.tokens - else: - payload["input_ids"] = prompt_ids - if not sample.tokens: # Initialize sample.tokens for the first turn - sample.tokens = prompt_ids - - # Use session_id for consistent hashing routing if router uses consistent_hashing policy - headers = None - if args.sglang_router_policy == "consistent_hashing" and sample.session_id: - headers = {"X-SMG-Routing-Key": sample.session_id} - - output = await post(url, payload, headers=headers) - - if args.use_slime_router and "RadixTreeMiddleware" in args.slime_router_middleware_paths: + use_radix = args.use_slime_router and "RadixTreeMiddleware" in getattr(args, "slime_router_middleware_paths", []) + if use_radix: + # RadixTree path: SGLang-specific, keep direct post from slime.router.middleware_hub.radix_tree_middleware import postprocess_sample_with_radix_tree + url = f"http://{args.sglang_router_ip}:{args.sglang_router_port}/generate" + payload = { + "input_ids": input_ids, + "sampling_params": sampling_params, + "return_logprob": True, + "return_routed_experts": args.use_rollout_routing_replay, + } + if sample.multimodal_inputs and sample.multimodal_inputs.get("images"): + payload["image_data"] = [encode_image_for_rollout_engine(img) for img in sample.multimodal_inputs["images"]] + headers = {"X-SMG-Routing-Key": sample.session_id} if getattr(args, "sglang_router_policy") == "consistent_hashing" and sample.session_id else None + output = await post(url, payload, headers=headers) sample = await postprocess_sample_with_radix_tree(args, sample, output) else: - if "output_token_logprobs" in output["meta_info"]: - new_response_tokens = [item[1] for item in output["meta_info"]["output_token_logprobs"]] - new_response_log_probs = [item[0] for item in output["meta_info"]["output_token_logprobs"]] - else: - new_response_tokens, new_response_log_probs = [], [] - - # Update sample with tokens directly - avoiding re-tokenization - sample.tokens = sample.tokens + new_response_tokens - sample.response_length += len(new_response_tokens) - sample.response += output["text"] - - # When partial rollout and masking off policy is enabled, update the loss mask - if sample.loss_mask is not None: - assert args.partial_rollout and args.mask_offpolicy_in_partial_rollout - sample.loss_mask += [1] * len(new_response_tokens) - - if sample.rollout_log_probs is None: - sample.rollout_log_probs = [] - sample.rollout_log_probs += new_response_log_probs - - if "routed_experts" in output["meta_info"]: - sample.rollout_routed_experts = np.frombuffer( - pybase64.b64decode(output["meta_info"]["routed_experts"].encode("ascii")), - dtype=np.int32, - ).reshape( - len(sample.tokens) - 1, - args.num_layers, - args.moe_router_topk, + backend = _get_backend_client(args) + base_url = getattr(args, "vllm_base_url", None) or f"http://{args.sglang_router_ip}:{args.sglang_router_port}" + headers = {"X-SMG-Routing-Key": sample.session_id} if getattr(args, "sglang_router_policy", None) == "consistent_hashing" and sample.session_id else None + req = RolloutBackendRequest( + input_ids=input_ids, + sampling_params=sampling_params, + return_logprob=True, + return_routed_experts=args.use_rollout_routing_replay, + image_data=[encode_image_for_rollout_engine(img) for img in sample.multimodal_inputs["images"]] if (sample.multimodal_inputs and sample.multimodal_inputs.get("images")) else None, + session_id=sample.session_id, ) - - sample.update_from_meta_info(args, output["meta_info"]) + resp = await backend.generate(req, base_url, headers=headers) + _apply_backend_response(sample, resp, args) return sample @@ -306,24 +303,19 @@ async def generate_and_rm_group( async def abort(args: Namespace, rollout_id: int) -> list[list[Sample]]: aborted_samples = [] - state = GenerateState(args) assert not state.aborted state.aborted = True - if parse(sglang_router.__version__) <= parse("0.2.1") or args.use_slime_router: - response = await get(f"http://{args.sglang_router_ip}:{args.sglang_router_port}/list_workers") - urls = response["urls"] - else: - response = await get(f"http://{args.sglang_router_ip}:{args.sglang_router_port}/workers") - urls = [worker["url"] for worker in response["workers"]] - - logger.info(f"Abort request for {urls}") - abort_tasks = [post(f"{url}/abort_request", {"abort_all": True}) for url in urls] - abort_results = await asyncio.gather(*abort_tasks, return_exceptions=True) - for url, result in zip(urls, abort_results, strict=False): - if isinstance(result, Exception): - logger.warning(f"Failed to abort worker at {url}: {result}") + backend = _get_backend_client(args) + urls = await backend.abort() + if urls: + logger.info(f"Abort request for {urls}") + abort_tasks = [post(f"{url}/abort_request", {"abort_all": True}) for url in urls] + abort_results = await asyncio.gather(*abort_tasks, return_exceptions=True) + for url, result in zip(urls, abort_results, strict=False): + if isinstance(result, Exception): + logger.warning(f"Failed to abort worker at {url}: {result}") # make sure all the pending tasks are finished count = 0 diff --git a/slime/utils/arguments.py b/slime/utils/arguments.py index fb2d305a63..76e681a87f 100644 --- a/slime/utils/arguments.py +++ b/slime/utils/arguments.py @@ -238,6 +238,25 @@ def add_rollout_arguments(parser): "Also, sometimes this will help alleviate the bug that transformers cannot find certain model." ), ) + parser.add_argument( + "--rollout-backend", + type=str, + choices=["sglang", "vllm"], + default="sglang", + help="Rollout inference backend: sglang (default) or vllm.", + ) + parser.add_argument( + "--vllm-base-url", + type=str, + default=None, + help="vLLM server URL (set automatically when using managed vLLM).", + ) + parser.add_argument( + "--vllm-max-retries", + type=int, + default=3, + help="Max HTTP retries for vLLM requests.", + ) parser.add_argument( "--rollout-function-path", type=str, @@ -1484,8 +1503,16 @@ def parse_args(add_custom_arguments=None): if pre.train_backend == "megatron" and not args.debug_rollout_only: megatron_validate_args(args) - if not args.debug_train_only: + if not args.debug_train_only and getattr(args, "rollout_backend", "sglang") == "sglang": sglang_validate_args(args) + elif getattr(args, "rollout_backend", "sglang") == "vllm": + # Set sglang aliases that the rest of the codebase expects + args.sglang_dp_size = getattr(args, "sglang_data_parallel_size", 1) or 1 + args.sglang_pp_size = getattr(args, "sglang_pipeline_parallel_size", 1) or 1 + args.sglang_ep_size = getattr(args, "sglang_expert_parallel_size", 1) or 1 + args.sglang_tp_size = args.rollout_num_gpus_per_engine + if not hasattr(args, "sglang_speculative_algorithm"): + args.sglang_speculative_algorithm = None return args From 2caa4a067220d22ca24b782d642d63128c50ed60 Mon Sep 17 00:00:00 2001 From: samithuang <285365963@qq.com> Date: Wed, 4 Mar 2026 06:27:17 +0000 Subject: [PATCH 05/39] add convert script --- convert_qwen2.5_ckpt.sh | 5 +++++ 1 file changed, 5 insertions(+) create mode 100644 convert_qwen2.5_ckpt.sh diff --git a/convert_qwen2.5_ckpt.sh b/convert_qwen2.5_ckpt.sh new file mode 100644 index 0000000000..fae03587dc --- /dev/null +++ b/convert_qwen2.5_ckpt.sh @@ -0,0 +1,5 @@ +source scripts/models/qwen2.5-0.5B.sh +PYTHONPATH=/root/Megatron-LM python tools/convert_hf_to_torch_dist.py \ + ${MODEL_ARGS[@]} \ + --hf-checkpoint /root/Qwen2.5-0.5B-Instruct \ + --save /root/Qwen2.5-0.5B-Instruct_torch_dist/ From 8caa8bae5f36615a238a6f3f1189544f691cd16b Mon Sep 17 00:00:00 2001 From: samithuang <285365963@qq.com> Date: Wed, 4 Mar 2026 06:39:11 +0000 Subject: [PATCH 06/39] add setup doc --- setup_for_vllm.md | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) create mode 100644 setup_for_vllm.md diff --git a/setup_for_vllm.md b/setup_for_vllm.md new file mode 100644 index 0000000000..ccdc8fd593 --- /dev/null +++ b/setup_for_vllm.md @@ -0,0 +1,21 @@ +``` +docker pull slimerl/slime:latest +``` + +``` +docker run -itd --gpus all --ipc=host --shm-size=128g --net=host --privileged=true --restart=always \ +--ulimit memlock=-1 --ulimit stack=67108864 \ +--ulimit nofile=65536:65536 \ +--name DNAME \ +-it slimerl/slime:latest /bin/bash \ + +``` +docker exec -it --user root DNAME bash +``` + +``` +pip install vllm=0.16 + +# for compatibility +pip install numpy==1.26.4 +``` From 25ee00543d07cacad5fb8e57ecaee28315e9d6c7 Mon Sep 17 00:00:00 2001 From: samithuang <285365963@qq.com> Date: Thu, 5 Mar 2026 03:45:03 +0000 Subject: [PATCH 07/39] fix nccl error by NcclBridge subprocess --- run-qwen2.5-0.5B-vllm.sh | 1 + .../update_weight_from_distributed.py | 335 +++++++++++++++--- slime/backends/vllm_utils/vllm_engine.py | 23 +- slime/ray/actor_group.py | 6 + slime/utils/arguments.py | 6 + 5 files changed, 308 insertions(+), 63 deletions(-) diff --git a/run-qwen2.5-0.5B-vllm.sh b/run-qwen2.5-0.5B-vllm.sh index bcc6807c69..2746098478 100644 --- a/run-qwen2.5-0.5B-vllm.sh +++ b/run-qwen2.5-0.5B-vllm.sh @@ -113,6 +113,7 @@ ray job submit --address="http://127.0.0.1:8265" \ "PYTHONPATH": "/root/Megatron-LM", "CUDA_DEVICE_MAX_CONNECTIONS": "1", "NCCL_ALGO": "Ring", + "NCCL_IB_DISABLE": "1", "NCCL_P2P_DISABLE": "1", "NCCL_DEBUG": "INFO", "NVTE_ALLOW_NONDETERMINISTIC_ALGO": "0", diff --git a/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py b/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py index 44649f87f4..94fcecc203 100644 --- a/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py +++ b/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py @@ -1,6 +1,9 @@ +import logging +import multiprocessing as mp import os import socket import time +import traceback from argparse import Namespace from collections.abc import Callable, Mapping, Sequence @@ -17,6 +20,155 @@ from ..megatron_to_hf import convert_to_hf from .common import all_gather_param, named_params_and_buffers +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# NcclBridge: isolate vLLM's PyNcclCommunicator in a subprocess so that it +# never coexists with torch.distributed NCCL groups in the Megatron trainer. +# +# vLLM's weight transfer uses raw NCCL (PyNcclCommunicator) which conflicts +# with torch.distributed's NCCL backend when both exist in the same process +# (see https://github.com/vllm-project/vllm/issues/5477). sglang avoids +# this because it uses torch.distributed process groups for weight sync. +# --------------------------------------------------------------------------- + + +def _nccl_bridge_worker(conn, master_address, master_port, world_size, device, cvd, env_snapshot): + """Subprocess entry-point: creates PyNcclCommunicator and serves requests. + + Protocol over *conn* (multiprocessing.Connection): + parent → child: + {"op": "broadcast", "tensors": [cpu_tensor, ...]} + {"op": "send_packed", "named_tensors": [(name, cpu_tensor), ...]} + None → shutdown + child → parent: + "ready" (after init) + "ok" (after each op) + "error: " + """ + try: + os.environ.update(env_snapshot) + if cvd: + os.environ["CUDA_VISIBLE_DEVICES"] = cvd + + import torch + torch.cuda.set_device(device) + + from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator + from vllm.distributed.utils import StatelessProcessGroup + + pg = StatelessProcessGroup.create( + host=master_address, port=master_port, rank=0, world_size=world_size, + ) + comm = PyNcclCommunicator(pg, device=device) + + conn.send("ready") + + while True: + cmd = conn.recv() + if cmd is None: + break + + op = cmd["op"] + if op == "broadcast": + for cpu_t in cmd["tensors"]: + gpu_t = cpu_t.cuda(device) + comm.broadcast(gpu_t, src=0, stream=torch.cuda.current_stream()) + del gpu_t + torch.cuda.synchronize() + conn.send("ok") + + elif op == "send_packed": + from vllm.distributed.weight_transfer.nccl_engine import ( + NCCLWeightTransferEngine, + ) + + named_tensors = cmd["named_tensors"] + + def _gpu_iter(): + for name, cpu_t in named_tensors: + yield (name, cpu_t.cuda(device).contiguous()) + + NCCLWeightTransferEngine.trainer_send_weights( + iterator=_gpu_iter(), + group=comm, + packed=True, + ) + torch.cuda.synchronize() + conn.send("ok") + + except Exception as e: + try: + conn.send(f"error: {e}") + except Exception: + pass + traceback.print_exc() + + +class _NcclBridge: + """Runs vLLM's PyNcclCommunicator in a separate subprocess. + + This prevents NCCL communicator conflicts with torch.distributed groups + that already exist in the Megatron trainer process. Communication between + the trainer and the bridge uses a multiprocessing Pipe; GPU tensors are + staged through CPU (pinned memory when possible) for the transfer. + """ + + def __init__(self, master_address: str, master_port: int, world_size: int, device: int): + ctx = mp.get_context("spawn") + self._parent_conn, child_conn = ctx.Pipe() + + # Pass the full environment so the subprocess inherits all NCCL, CUDA, + # and networking settings from the trainer (set by Ray runtime_env). + env_snapshot = dict(os.environ) + cvd = os.environ.get("CUDA_VISIBLE_DEVICES", "") + + self._process = ctx.Process( + target=_nccl_bridge_worker, + args=(child_conn, master_address, master_port, world_size, device, cvd, env_snapshot), + daemon=True, + ) + self._process.start() + + msg = self._parent_conn.recv() + if isinstance(msg, str) and msg.startswith("error:"): + raise RuntimeError(f"NcclBridge init failed: {msg}") + if msg != "ready": + raise RuntimeError(f"NcclBridge init unexpected response: {msg}") + logger.info("NcclBridge ready (pid=%d, device=%d)", self._process.pid, device) + + def broadcast_tensors(self, tensors: list[torch.Tensor]) -> None: + """Broadcast a list of tensors (one-by-one) via the bridge subprocess.""" + cpu_tensors = [t.cpu().contiguous() for t in tensors] + self._parent_conn.send({"op": "broadcast", "tensors": cpu_tensors}) + self._wait_ok("broadcast_tensors") + + def send_weights_packed(self, named_tensors: list[tuple[str, torch.Tensor]]) -> None: + """Send weights using vLLM's packed broadcast protocol.""" + cpu_pairs = [] + for name, t in named_tensors: + data = t.data if hasattr(t, "data") else t + cpu_pairs.append((name, data.cpu().contiguous())) + self._parent_conn.send({"op": "send_packed", "named_tensors": cpu_pairs}) + self._wait_ok("send_weights_packed") + + def _wait_ok(self, label: str, timeout: float = 600.0) -> None: + if not self._parent_conn.poll(timeout): + raise TimeoutError(f"NcclBridge {label} timed out after {timeout}s") + msg = self._parent_conn.recv() + if msg != "ok": + raise RuntimeError(f"NcclBridge {label} failed: {msg}") + + def shutdown(self) -> None: + try: + self._parent_conn.send(None) + self._process.join(timeout=30) + except Exception: + pass + if self._process.is_alive(): + self._process.terminate() + class UpdateWeightFromDistributed: """ @@ -99,34 +251,52 @@ def update_weights(self) -> None: ) dist.barrier(group=get_gloo_group()) - buffer_size = 0 - converted_named_tensors = [] - # non expert params - pbar = tqdm(desc=f"[{self._group_name}] Update weights", total=0) if self._is_pp_src_rank else None - - for name, param in named_params_and_buffers(self.args, self.model): - if ".experts." in name: - continue - buffer_size = self._update_weight_from_distributed( - name, param, converted_named_tensors, buffer_size, pbar=pbar - ) + use_vllm_packed = self._use_vllm_packed() + if use_vllm_packed and self._is_pp_src_rank: + logger.info("Using vLLM packed weight sync (one-shot metadata + trainer_send_weights)") + if use_vllm_packed: + # vLLM packed path: gather all non-expert params, one-shot update (aligned with vLLM New Weight Syncing) + converted_named_tensors = [] + for name, param in named_params_and_buffers(self.args, self.model): + if ".experts." in name: + continue + param = all_gather_param(name, param) + if self._is_pp_src_rank: + converted_named_tensors += convert_to_hf( + self.args, self.model_name, name, param, self.quantization_config + ) + if converted_named_tensors and self._is_pp_src_rank: + self._update_weights_vllm_packed(converted_named_tensors) + else: + buffer_size = 0 + converted_named_tensors = [] + pbar = tqdm(desc=f"[{self._group_name}] Update weights", total=0) if self._is_pp_src_rank else None + + for name, param in named_params_and_buffers(self.args, self.model): + if ".experts." in name: + continue + buffer_size = self._update_weight_from_distributed( + name, param, converted_named_tensors, buffer_size, pbar=pbar + ) - if converted_named_tensors: - self._update_bucket_weights_from_distributed(converted_named_tensors, pbar=pbar) + if converted_named_tensors: + self._update_bucket_weights_from_distributed(converted_named_tensors, pbar=pbar) dist.barrier(group=get_gloo_group()) - buffer_size = 0 - named_tensors = [] - for name, param in named_params_and_buffers(self.args, self.model): - if ".experts." not in name: - continue - buffer_size = self._update_expert_weight_from_distributed( - name, param, named_tensors, buffer_size, pbar=pbar - ) + if not use_vllm_packed: + buffer_size = 0 + named_tensors = [] + pbar = tqdm(desc=f"[{self._group_name}] Update weights (experts)", total=0) if self._is_pp_src_rank else None + for name, param in named_params_and_buffers(self.args, self.model): + if ".experts." not in name: + continue + buffer_size = self._update_expert_weight_from_distributed( + name, param, named_tensors, buffer_size, pbar=pbar + ) - if named_tensors: - self._update_expert_bucket_weights_from_distributed(named_tensors, pbar=pbar) + if named_tensors: + self._update_expert_bucket_weights_from_distributed(named_tensors, pbar=pbar) dist.barrier(group=get_gloo_group()) if dist.get_rank() == 0: @@ -164,6 +334,41 @@ def _update_weight_from_distributed( buffer_size += param_size return buffer_size + def _use_vllm_packed(self) -> bool: + """Use vLLM packed weight transfer (one-shot metadata + trainer_send_weights).""" + if not _is_vllm_backend(self.args): + return False + if not getattr(self.args, "vllm_weight_sync_packed", True): + return False + # MoE models need expert path, skip packed + if any(".experts." in name for name, _ in named_params_and_buffers(self.args, self.model)): + return False + # compressed-tensors needs pre/post process + if self.quantization_config and self.quantization_config.get("quant_method") == "compressed-tensors": + return False + return True + + def _update_weights_vllm_packed( + self, converted_named_tensors: list[tuple[str, torch.Tensor]] + ) -> None: + """Single-shot vLLM weight update using packed broadcast (aligned with vLLM New Weight Syncing).""" + while not ray.get(self.rollout_engine_lock.acquire.remote()): + time.sleep(0.1) + + try: + refs = update_weights_from_distributed( + self._group_name, + self._model_update_groups, + self.weight_version, + self.rollout_engines, + converted_named_tensors, + use_vllm=True, + packed=True, + ) + ray.get(refs) + finally: + ray.get(self.rollout_engine_lock.release.remote()) + def _update_expert_weight_from_distributed( self, name: str, @@ -268,9 +473,10 @@ def connect_rollout_engines_from_distributed( have heterogeneous TP sizes (e.g. prefill TP=2, decode TP=4), each engine occupies a different number of ranks in the NCCL group. - For vLLM backend, uses vLLM's StatelessProcessGroup + PyNcclCommunicator - instead of torch.distributed, because vLLM's weight transfer engine uses - its own NCCL initialization protocol. + For vLLM backend, the trainer-side NCCL communicator is created inside a + separate subprocess (_NcclBridge) to avoid conflicts between vLLM's raw + NCCL (PyNcclCommunicator) and the torch.distributed NCCL groups that + Megatron already holds in this process. """ if engine_gpu_counts is None: engine_gpu_counts = [args.rollout_num_gpus_per_engine] * len(rollout_engines) @@ -286,6 +492,7 @@ def connect_rollout_engines_from_distributed( for c in engine_gpu_counts: cumulative.append(cumulative[-1] + c) + # Fire engine init remotes first (non-blocking Ray calls). refs = [ engine.init_weights_update_group.remote( master_address, @@ -298,29 +505,28 @@ def connect_rollout_engines_from_distributed( for i, engine in enumerate(rollout_engines) ] + torch.cuda.synchronize() + torch.cuda.empty_cache() + if _is_vllm_backend(args): - # vLLM uses StatelessProcessGroup + PyNcclCommunicator for weight transfer. - # The training side must use the same mechanism for NCCL compatibility. - # - # Disable P2P transport: in colocate mode the trainer and vLLM server - # share the same physical GPU but have different CUDA_VISIBLE_DEVICES, - # which causes NCCL P2P (cudaIpc*) to fail with "invalid argument". - # SHM transport works correctly in this scenario. - from vllm.distributed.weight_transfer.nccl_engine import NCCLWeightTransferEngine - - old_p2p = os.environ.get("NCCL_P2P_DISABLE") - os.environ["NCCL_P2P_DISABLE"] = "1" - try: - model_update_groups = NCCLWeightTransferEngine.trainer_init({ - "master_address": master_address, - "master_port": master_port, - "world_size": world_size, - }) - finally: - if old_p2p is None: - os.environ.pop("NCCL_P2P_DISABLE", None) - else: - os.environ["NCCL_P2P_DISABLE"] = old_p2p + # vLLM uses StatelessProcessGroup + PyNcclCommunicator (raw NCCL). + # Creating PyNcclCommunicator in the Megatron trainer process would + # conflict with torch.distributed's NCCL groups (issue #5477). + # Instead, we spawn a bridge subprocess that owns the PyNcclCommunicator, + # keeping the trainer process free of raw NCCL communicators. + device = torch.cuda.current_device() + logger.info( + "vLLM weight transfer via NcclBridge: addr=%s port=%d " + "world_size=%d device=%d CVD=%s", + master_address, master_port, world_size, device, + os.environ.get("CUDA_VISIBLE_DEVICES", ""), + ) + model_update_groups = _NcclBridge( + master_address=master_address, + master_port=master_port, + world_size=world_size, + device=device, + ) else: model_update_groups = init_process_group( backend="nccl", @@ -340,6 +546,8 @@ def disconnect_rollout_engines_from_distributed(args, group_name, model_update_g """ refs = [engine.destroy_weights_update_group.remote(group_name) for engine in rollout_engines] if _is_vllm_backend(args): + if isinstance(model_update_groups, _NcclBridge): + model_update_groups.shutdown() model_update_groups = None else: dist.destroy_process_group(model_update_groups) @@ -353,25 +561,34 @@ def update_weights_from_distributed( rollout_engines: Sequence[ActorHandle], converted_named_tensors: Sequence[tuple[str, torch.Tensor]], use_vllm: bool = False, + packed: bool = False, ) -> list[ObjectRef]: """ Send metadata (Ray), broadcast tensors (NCCL rank 0 → engines). - For vLLM, uses PyNcclCommunicator.broadcast instead of dist.broadcast. + + For vLLM the *group* is an ``_NcclBridge`` instance (subprocess) so that + raw NCCL never runs inside the Megatron trainer process. + For sglang the *group* is a ``torch.distributed.ProcessGroup``. """ + kwargs = { + "names": [name for name, _ in converted_named_tensors], + "dtypes": [param.dtype for _, param in converted_named_tensors], + "shapes": [param.shape for _, param in converted_named_tensors], + "group_name": group_name, + "weight_version": str(weight_version), + } + if use_vllm: + kwargs["packed"] = packed + refs = [ - engine.update_weights_from_distributed.remote( - names=[name for name, _ in converted_named_tensors], - dtypes=[param.dtype for _, param in converted_named_tensors], - shapes=[param.shape for _, param in converted_named_tensors], - group_name=group_name, - weight_version=str(weight_version), - ) + engine.update_weights_from_distributed.remote(**kwargs) for engine in rollout_engines ] - if use_vllm: - for _, param in converted_named_tensors: - group.broadcast(param.data, src=0, stream=torch.cuda.current_stream()) + if use_vllm and packed: + group.send_weights_packed(converted_named_tensors) + elif use_vllm: + group.broadcast_tensors([param.data for _, param in converted_named_tensors]) else: handles = [] for _, param in converted_named_tensors: diff --git a/slime/backends/vllm_utils/vllm_engine.py b/slime/backends/vllm_utils/vllm_engine.py index 12283de21e..22ff897ee5 100644 --- a/slime/backends/vllm_utils/vllm_engine.py +++ b/slime/backends/vllm_utils/vllm_engine.py @@ -62,6 +62,7 @@ def init(self, port=None, host=None, **kwargs): env.setdefault("NCCL_DEBUG", "INFO") env.setdefault("NCCL_DEBUG_SUBSYS", "ALL") env["NCCL_P2P_DISABLE"] = "1" + env.setdefault("NCCL_IB_DISABLE", "1") self._log_file = tempfile.NamedTemporaryFile( prefix="vllm_engine_", suffix=".log", delete=False, mode="w" @@ -145,8 +146,11 @@ def flush_cache(self): def init_weights_update_group(self, master_address, master_port, rank_offset, world_size, group_name=None, backend=None): logger.info( - "Initializing NCCL weight transfer: master=%s:%s, rank_offset=%d, world_size=%d", + "Initializing NCCL weight transfer: master=%s:%s, rank_offset=%d, " + "world_size=%d, vllm_url=http://%s:%s, vllm_log=%s", master_address, master_port, rank_offset, world_size, + self.server_host, self.server_port, + self._log_file.name if self._log_file else "", ) self._post("/init_weight_transfer_engine", json_data={ "init_info": { @@ -156,8 +160,19 @@ def init_weights_update_group(self, master_address, master_port, rank_offset, wo "world_size": world_size, } }) - - def update_weights_from_distributed(self, names, dtypes, shapes, group_name=None, flush_cache=False, weight_version=None): + log_tail = self._read_log_tail(30) + logger.info("vLLM log after init_weight_transfer_engine:\n%s", log_tail) + + def update_weights_from_distributed( + self, + names, + dtypes, + shapes, + group_name=None, + flush_cache=False, + weight_version=None, + packed: bool = True, + ): dtype_names = [str(d).replace("torch.", "") for d in dtypes] shape_lists = [list(s) for s in shapes] self._post("/update_weights", json_data={ @@ -165,7 +180,7 @@ def update_weights_from_distributed(self, names, dtypes, shapes, group_name=None "names": names, "dtype_names": dtype_names, "shapes": shape_lists, - "packed": False, + "packed": packed, } }) diff --git a/slime/ray/actor_group.py b/slime/ray/actor_group.py index 5542835313..1150e9ed35 100644 --- a/slime/ray/actor_group.py +++ b/slime/ray/actor_group.py @@ -50,12 +50,18 @@ def _allocate_gpus_for_actor(self, pg, num_gpus_per_actor): assert pg is not None pg, reordered_bundle_indices, _reordered_gpu_ids = pg + # Restrict CUDA_VISIBLE_DEVICES to only this group's GPUs so that + # NCCL / PyTorch do not allocate memory on rollout GPUs. + trainer_gpu_ids = [_reordered_gpu_ids[rank] for rank in range(world_size)] + trainer_cvd = ",".join(str(g) for g in trainer_gpu_ids) + env_vars = { # because sglang will always set NCCL_CUMEM_ENABLE to 0 # we need also set it to 0 to prevent nccl error. "NCCL_CUMEM_ENABLE": os.environ.get("NCCL_CUMEM_ENABLE", "0"), "NVTE_FP8_BLOCK_SCALING_FP32_SCALES": os.environ.get("NVTE_FP8_BLOCK_SCALING_FP32_SCALES", "1"), **{name: "1" for name in NOSET_VISIBLE_DEVICES_ENV_VARS_LIST}, + "CUDA_VISIBLE_DEVICES": trainer_cvd, **self.args.train_env_vars, } diff --git a/slime/utils/arguments.py b/slime/utils/arguments.py index 76e681a87f..7ef8e70b71 100644 --- a/slime/utils/arguments.py +++ b/slime/utils/arguments.py @@ -257,6 +257,12 @@ def add_rollout_arguments(parser): default=3, help="Max HTTP retries for vLLM requests.", ) + parser.add_argument( + "--vllm-weight-sync-packed", + action=argparse.BooleanOptionalAction, + default=True, + help="Use vLLM packed weight transfer for non-colocate (default: True). Disable for per-bucket mode.", + ) parser.add_argument( "--rollout-function-path", type=str, From ab7eb0bc11412fb1ff232c1f9a4560f8f1e86df0 Mon Sep 17 00:00:00 2001 From: samithuang <285365963@qq.com> Date: Thu, 5 Mar 2026 08:31:05 +0000 Subject: [PATCH 08/39] eliminate gpu to cpu weight transfer Signed-off-by: samithuang <285365963@qq.com> --- run-qwen2.5-0.5B-vllm.sh | 3 +- .../update_weight_from_distributed.py | 40 +++++++++---------- slime/backends/vllm_utils/vllm_engine.py | 3 ++ slime/rollout/backends/vllm_client.py | 2 + 4 files changed, 25 insertions(+), 23 deletions(-) diff --git a/run-qwen2.5-0.5B-vllm.sh b/run-qwen2.5-0.5B-vllm.sh index 2746098478..208ec60c91 100644 --- a/run-qwen2.5-0.5B-vllm.sh +++ b/run-qwen2.5-0.5B-vllm.sh @@ -26,6 +26,7 @@ CKPT_ARGS=( --ref-load /root/Qwen2.5-0.5B-Instruct_torch_dist/ ) +# num-rollout:100 ROLLOUT_ARGS=( --prompt-data /root/gsm8k/train.parquet --input-key messages @@ -33,7 +34,7 @@ ROLLOUT_ARGS=( --apply-chat-template --rollout-shuffle --rm-type math - --num-rollout 100 + --num-rollout 500 --rollout-batch-size 32 --n-samples-per-prompt 8 --rollout-max-response-len 1024 diff --git a/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py b/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py index 94fcecc203..03d480801f 100644 --- a/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py +++ b/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py @@ -1,5 +1,5 @@ import logging -import multiprocessing as mp +import torch.multiprocessing as mp import os import socket import time @@ -37,10 +37,13 @@ def _nccl_bridge_worker(conn, master_address, master_port, world_size, device, cvd, env_snapshot): """Subprocess entry-point: creates PyNcclCommunicator and serves requests. + GPU tensors are shared from the parent via CUDA IPC (torch.multiprocessing + handles this transparently). No GPU→CPU→GPU copies are needed. + Protocol over *conn* (multiprocessing.Connection): parent → child: - {"op": "broadcast", "tensors": [cpu_tensor, ...]} - {"op": "send_packed", "named_tensors": [(name, cpu_tensor), ...]} + {"op": "broadcast", "tensors": [gpu_tensor, ...]} + {"op": "send_packed", "named_tensors": [(name, gpu_tensor), ...]} None → shutdown child → parent: "ready" (after init) @@ -53,6 +56,7 @@ def _nccl_bridge_worker(conn, master_address, master_port, world_size, device, c os.environ["CUDA_VISIBLE_DEVICES"] = cvd import torch + import torch.multiprocessing # ensure CUDA IPC reducers are registered torch.cuda.set_device(device) from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator @@ -72,10 +76,8 @@ def _nccl_bridge_worker(conn, master_address, master_port, world_size, device, c op = cmd["op"] if op == "broadcast": - for cpu_t in cmd["tensors"]: - gpu_t = cpu_t.cuda(device) - comm.broadcast(gpu_t, src=0, stream=torch.cuda.current_stream()) - del gpu_t + for t in cmd["tensors"]: + comm.broadcast(t, src=0, stream=torch.cuda.current_stream()) torch.cuda.synchronize() conn.send("ok") @@ -84,14 +86,8 @@ def _nccl_bridge_worker(conn, master_address, master_port, world_size, device, c NCCLWeightTransferEngine, ) - named_tensors = cmd["named_tensors"] - - def _gpu_iter(): - for name, cpu_t in named_tensors: - yield (name, cpu_t.cuda(device).contiguous()) - NCCLWeightTransferEngine.trainer_send_weights( - iterator=_gpu_iter(), + iterator=iter(cmd["named_tensors"]), group=comm, packed=True, ) @@ -110,9 +106,9 @@ class _NcclBridge: """Runs vLLM's PyNcclCommunicator in a separate subprocess. This prevents NCCL communicator conflicts with torch.distributed groups - that already exist in the Megatron trainer process. Communication between - the trainer and the bridge uses a multiprocessing Pipe; GPU tensors are - staged through CPU (pinned memory when possible) for the transfer. + that already exist in the Megatron trainer process. GPU tensors are shared + with the subprocess via CUDA IPC (handled transparently by + torch.multiprocessing), avoiding any GPU→CPU→GPU copies. """ def __init__(self, master_address: str, master_port: int, world_size: int, device: int): @@ -140,17 +136,17 @@ def __init__(self, master_address: str, master_port: int, world_size: int, devic def broadcast_tensors(self, tensors: list[torch.Tensor]) -> None: """Broadcast a list of tensors (one-by-one) via the bridge subprocess.""" - cpu_tensors = [t.cpu().contiguous() for t in tensors] - self._parent_conn.send({"op": "broadcast", "tensors": cpu_tensors}) + gpu_tensors = [t.contiguous() for t in tensors] + self._parent_conn.send({"op": "broadcast", "tensors": gpu_tensors}) self._wait_ok("broadcast_tensors") def send_weights_packed(self, named_tensors: list[tuple[str, torch.Tensor]]) -> None: """Send weights using vLLM's packed broadcast protocol.""" - cpu_pairs = [] + gpu_pairs = [] for name, t in named_tensors: data = t.data if hasattr(t, "data") else t - cpu_pairs.append((name, data.cpu().contiguous())) - self._parent_conn.send({"op": "send_packed", "named_tensors": cpu_pairs}) + gpu_pairs.append((name, data.contiguous())) + self._parent_conn.send({"op": "send_packed", "named_tensors": gpu_pairs}) self._wait_ok("send_weights_packed") def _wait_ok(self, label: str, timeout: float = 600.0) -> None: diff --git a/slime/backends/vllm_utils/vllm_engine.py b/slime/backends/vllm_utils/vllm_engine.py index 22ff897ee5..bf645cd0bd 100644 --- a/slime/backends/vllm_utils/vllm_engine.py +++ b/slime/backends/vllm_utils/vllm_engine.py @@ -42,12 +42,15 @@ def init(self, port=None, host=None, **kwargs): else: dev_str = ",".join(str(g) for g in gpu_ids) + seed = getattr(self.args, "seed", 1234) + self.rank cmd = [ "vllm", "serve", model, "--tensor-parallel-size", str(tp), "--port", str(self.server_port), "--host", "0.0.0.0", "--weight-transfer-config", '{"backend": "nccl"}', + "--seed", str(seed), + "--trust-remote-code", ] if getattr(self.args, "offload_rollout", False): cmd.append("--enable-sleep-mode") diff --git a/slime/rollout/backends/vllm_client.py b/slime/rollout/backends/vllm_client.py index 537569499d..32be01d66c 100644 --- a/slime/rollout/backends/vllm_client.py +++ b/slime/rollout/backends/vllm_client.py @@ -44,9 +44,11 @@ async def generate( "stop": sp.get("stop"), "stop_token_ids": sp.get("stop_token_ids"), "skip_special_tokens": sp.get("skip_special_tokens", True), + "spaces_between_special_tokens": sp.get("spaces_between_special_tokens", False), "logprobs": 1 if request.return_logprob else None, "include_stop_str_in_output": sp.get("no_stop_trim", False), "return_token_ids": True, + "seed": sp.get("sampling_seed"), } payload = {k: v for k, v in payload.items() if v is not None} From 546d2ad734a28e887457ce9e124398df150ded86 Mon Sep 17 00:00:00 2001 From: Samit <285365963@qq.com> Date: Fri, 6 Mar 2026 00:19:19 +0800 Subject: [PATCH 09/39] Revise weight synchronization strategy in goal plan Reorder weight synchronization support for colocate and non-colocate scenarios in the goal plan. --- goal_plan.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/goal_plan.md b/goal_plan.md index d77618195b..f7b8827c8b 100644 --- a/goal_plan.md +++ b/goal_plan.md @@ -6,8 +6,8 @@ First Design and RFC by 03/06 - 对标SGLang，Slime 在 Ray 内管理 vLLM 的完整生命周期，包括进程拉起、权重同步、推理暂停/恢复 - 暂不使用Router，SGLang Model Gateway仅只支持SGLang Worker，SlimeRouter仅在 R3 / radix-tree caching 时需要，Qwen2.5-0.5B 非 MoE 且用 token-in/token-out - 单vLLM实例，无router，通过vLLMClient 直连本地 vLLM 进程端口 -- 先支持和验证colocate，权重同步采用GPU IPC（vLLM update_weights_from_ipc, update_weights_from_tensor），对标SGLang update_weights_from_tensor，以验证Reproductivity -- 再支持训推不共卡，权重同步采用NCCL broadcast，对标SGLang update_weights_from_distributed (默认） +- 先支持训推不共卡(non-colocate)，权重同步采用NCCL broadcast，对标SGLang update_weights_from_distributed (默认） +- 再支持和验证colocate，权重同步采用GPU IPC（vLLM update_weights_from_ipc, update_weights_from_tensor），对标SGLang update_weights_from_tensor，以验证Reproductivity。**IPC 依赖vllm 0.17** #### 风险： - slime, sglang版本依赖，和vllm 0.16的版本依赖冲突(numpy, torch, transformers, etc) From d480da0da4ef1228b4f47b3bd5de88dafff6cffc Mon Sep 17 00:00:00 2001 From: knlnguyen1802 Date: Mon, 23 Mar 2026 15:55:22 +0700 Subject: [PATCH 10/39] Router for vllm (#5) * Draft router design Signed-off-by: knlnguyen1802 * Add vllm router Signed-off-by: knlnguyen1802 * Add router to script Signed-off-by: knlnguyen1802 * Fix gpu memory utilization Signed-off-by: knlnguyen1802 * Fix output token ids Signed-off-by: knlnguyen1802 * Add more nccl flag Signed-off-by: knlnguyen1802 * Fix bug Signed-off-by: knlnguyen1802 --------- Signed-off-by: knlnguyen1802 --- docs/en/vllm/ROUTER_DESIGN.md | 462 ++++++++++++++++++ run-qwen2.5-0.5B-vllm.sh | 4 + slime/backends/vllm_utils/__init__.py | 3 +- slime/backends/vllm_utils/vllm_engine.py | 130 ++++- .../vllm_utils/vllm_translation_sidecar.py | 454 +++++++++++++++++ slime/ray/rollout.py | 91 +++- slime/rollout/backends/vllm_client.py | 92 +++- slime/utils/arguments.py | 7 + 8 files changed, 1212 insertions(+), 31 deletions(-) create mode 100644 docs/en/vllm/ROUTER_DESIGN.md create mode 100644 slime/backends/vllm_utils/vllm_translation_sidecar.py diff --git a/docs/en/vllm/ROUTER_DESIGN.md b/docs/en/vllm/ROUTER_DESIGN.md new file mode 100644 index 0000000000..c18a29c997 --- /dev/null +++ b/docs/en/vllm/ROUTER_DESIGN.md @@ -0,0 +1,462 @@ +# RFC: Replace SGLang Backend with vLLM — Router Integration + +--- + +## Summary + +Replace the SGLang inference backend behind **SlimeRouter** with **vLLM** while keeping the existing router and middleware stack completely unchanged. +This RFC covers **only the router layer** — what APIs the vLLM backend must expose, how the existing SlimeRouter is reused, and what translation is needed between the two formats. + +**Key design decision:** Reuse vLLM's built-in [OpenAI-compatible API server](https://docs.vllm.ai/en/stable/serving/openai_compatible_server/) (`vllm serve`) + + +--- + +## 1. Target Architecture + +``` + Rollout Workers SlimeRouter (NO CHANGE) vLLM Engines (NEW) + ────────────── ──────────────────────── ────────────────── + ┌──────────────────────┐ + POST /generate ──────────────────▶│ RadixTreeMiddleware │ + │ • prefix cache │ + │ • retry on abort │ + │ • token/logprob cache│ + └──────────┬───────────┘ + │ + ┌──────────▼───────────┐ + │ SlimeRouter.proxy() │ ┌─────────────────────┐ + │ • least-connections │────────▶│ vLLM Translation │ + │ load balancer │ │ Sidecar (per engine) │ + │ • health check loop │ │ │ + └──────────────────────┘ │ POST /generate │ + │ ↓ translate │ + │ POST /v1/completions │ + │ ↓ translate back │ + │ → SGLang-format JSON │ + └─────────┬───────────┘ + │ + ┌─────────▼───────────┐ + │ vLLM Server │ + │ (vllm serve) │ + │ • /v1/completions │ + │ • /health │ + │ • /sleep, /wake_up │ + │ • /pause, /resume │ + │ • /update_weights │ + └─────────────────────┘ +``` + +### What stays the same + +| Component | Change | Reason | +|---|---|---| +| `SlimeRouter` ([router.py](slime/router/router.py)) | **None** | Engine-agnostic HTTP proxy; only reads JSON responses | +| `RadixTreeMiddleware` ([radix_tree_middleware.py](slime/router/middleware_hub/radix_tree_middleware.py)) | **None** | Operates on request/response JSON; has no engine-specific code | +| `StringRadixTrie` ([radix_tree.py](slime/router/middleware_hub/radix_tree.py)) | **None** | Pure data structure, no engine coupling | +| Middleware loading (`--slime-router-middleware-paths`) | **None** | Dynamic import via `load_function()` | + +### What is new + +| Component | Description | +|---|---| +| `vllm_translation_sidecar.py` | Lightweight FastAPI process co-located with each vLLM engine. Receives SGLang-format `/generate` requests, translates to vLLM's `/v1/completions`, translates responses back. Also proxies lifecycle endpoints (`/abort_request`, `/health_generate`, etc.). | +| `vllm_engine.py` | Ray actor that manages the vLLM server process lifecycle (via `vllm serve`), the translation sidecar, weight updates, and registration with the router. | + +--- + +## 2. Reusing SlimeRouter — Zero Modification + +The SlimeRouter communicates with backends through **five interaction points**. All are already engine-agnostic: + +### 2.1 Worker Registration + +**Flow:** Engine starts → engine calls `POST /add_worker?url=http://{host}:{port}` → router adds to pool. + +``` +Router state after registration: + worker_request_counts["http://10.0.0.1:10090"] = 0 + worker_failure_counts["http://10.0.0.1:10090"] = 0 +``` + +**vLLM action:** The `VLLMEngine` Ray actor calls this endpoint after verifying the vLLM server + translation sidecar are healthy. The registered URL points to the **sidecar**, not the raw vLLM server. No router change needed. + +### 2.2 Request Proxying + +**Flow:** `POST /generate` → middleware pipeline → `SlimeRouter.proxy()` → `httpx` forwards to backend (sidecar). + +The router selects a backend via **least-connections** (`_use_url()`), forwards the raw request body as-is, and returns the response as-is. It never inspects or transforms the request/response payload. + +**vLLM action:** The sidecar receives the forwarded request, translates it to `/v1/completions`, calls the co-located vLLM server, translates the response back to SGLang format, and returns it. + +### 2.3 Health Check + +**Flow:** Background loop calls `GET {worker_url}/health` every N seconds. + +- 200 → healthy, reset failure count +- Non-200 or timeout → increment failure count +- Failures ≥ threshold (default 3) → quarantine worker permanently + +**vLLM action:** The sidecar's `/health` proxies to vLLM's built-in `/health` endpoint (returns 200 when ready). Compatible out of the box. + +### 2.4 Worker Listing + +**Flow:** `GET /list_workers` → returns `{"urls": [...]}` + +Used by the rollout to discover engines for direct abort calls. No engine involvement. + +### 2.5 Retrieve from Text (Radix Tree) + +**Flow:** `POST /retrieve_from_text` → router looks up the radix tree cache → returns tokens/logprobs. + +Fully router-internal. Never reaches the engine. + +--- + +## 3. API Contract — What the Translation Sidecar Must Expose + +The translation sidecar sits between SlimeRouter and the vLLM server. It receives SGLang-format requests and returns SGLang-format responses. + +### 3.1 `POST /generate` — Generation + +This is the primary endpoint. The sidecar translates between Slime's format and vLLM's `/v1/completions`. + +#### Incoming Request (from router) + +```json +{ + "input_ids": [128000, 2610, 553, 264, 11190, 18328, 13], + "input_tokens": [128000, 2610, 553, 264, 11190, 18328, 13], + "sampling_params": { + "temperature": 0.7, + "top_p": 0.9, + "top_k": -1, + "max_new_tokens": 1024, + "stop": ["<|endoftext|>"], + "stop_token_ids": [128001], + "skip_special_tokens": false, + "no_stop_trim": true, + "spaces_between_special_tokens": false + }, + "return_logprob": true, + "stream": false +} +``` + +#### Translated Request (to vLLM `/v1/completions`) + +```json +{ + "model": "", + "prompt": [128000, 2610, 553, 264, 11190, 18328, 13], + "max_tokens": 1024, + "temperature": 0.7, + "top_p": 0.9, + "top_k": -1, + "stop": ["<|endoftext|>"], + "stop_token_ids": [128001], + "skip_special_tokens": false, + "include_stop_str_in_output": true, + "spaces_between_special_tokens": false, + "logprobs": 1, + "stream": false, + "extra_body": { + "return_token_ids": true + } +} +``` + +**Key translations:** +- `input_ids` → `prompt` (vLLM accepts `list[int]` as pre-tokenized prompt) +- `max_new_tokens` → `max_tokens` +- `no_stop_trim: true` → `include_stop_str_in_output: true` +- `return_logprob: true` → `logprobs: 1` + `extra_body.return_token_ids: true` + +#### vLLM Response (from `/v1/completions`) + +```json +{ + "id": "cmpl-abc123", + "choices": [{ + "text": "I'll help you with that. The answer is 42.", + "logprobs": { + "token_logprobs": [-0.152, -0.089, -0.203], + "tokens": ["I", "'ll", " help"] + }, + "token_ids": [40, 3358, 1520], + "finish_reason": "stop" + }], + "usage": { + "prompt_tokens": 7, + "completion_tokens": 3, + "total_tokens": 10 + } +} +``` + +#### Translated Response (returned to router) + +```json +{ + "text": "I'll help you with that. The answer is 42.", + "output_ids": [40, 3358, 1520], + "meta_info": { + "output_token_logprobs": [ + [-0.152, 40], + [-0.089, 3358], + [-0.203, 1520] + ], + "finish_reason": { + "type": "stop" + }, + "weight_version": 3, + "prompt_tokens": 7, + "cached_tokens": 0 + } +} +``` + +##### Field-by-field contract + +| Field | Type | Required | Consumer | Description | +|---|---|---|---|---| +| `text` | `str` | **Yes** | Rollout, Middleware | Generated text (output only, not including prompt) | +| `output_ids` | `list[int]` | **Yes** | Middleware | Generated token IDs. Middleware checks existence as a gate for caching. | +| `meta_info.output_token_logprobs` | `list[[float, int]]` | **Yes** (if `return_logprob`) | Rollout, Middleware | Each element is `[logprob, token_id]`. Used for RL policy ratio calculation. | +| `meta_info.finish_reason` | `{"type": str}` | **Yes** | Rollout, Middleware | Must be `{"type": "stop"}`, `{"type": "length"}`, or `{"type": "abort"}`. **Not** a plain string. | +| `meta_info.weight_version` | `int` | **Yes** | Middleware, Rollout | Current model weight version. Tracked by the sidecar (incremented on each weight update). | +| `meta_info.prompt_tokens` | `int` | Nice-to-have | Rollout (stats) | From `usage.prompt_tokens`. | +| `meta_info.cached_tokens` | `int` | Nice-to-have | Rollout (stats) | vLLM doesn't expose this directly; default to `0`. | + +### 3.2 `GET /health` — Health Check + +``` +GET /health +→ Sidecar proxies to vLLM's GET /health +→ 200 OK (engine ready) +→ 503 or timeout (engine not ready / overloaded) +``` + +vLLM already provides this endpoint. **Passthrough — no translation needed.** + +### 3.3 `POST /abort_request` — Cancel Generation + +``` +POST /abort_request +Body: {"abort_all": true} +→ 200 OK +``` + +Called **directly** by the rollout to each engine (bypasses the router). The rollout discovers engine URLs via `GET /list_workers`, then sends abort to each. + +**vLLM approach:** vLLM uses **HTTP connection close** for abort (via its `@with_cancellation` decorator). When a client disconnects, the in-flight request is automatically cancelled. + +**Implementation options:** +1. **Track active connections.** The sidecar maintains a set of active `httpx` connections to the vLLM server. On `POST /abort_request`, close all of them — triggering vLLM's cancellation. +2. **Use vLLM's `/pause` endpoint.** Call `POST /pause` to block new requests, then `POST /resume` after the RL training step completes. This is semantically closer to how Slime uses abort (clearing the decks between training generations). + +> **Note:** vLLM has `POST /abort_requests` only in disaggregated mode. For standard mode, HTTP disconnect is the canonical abort mechanism. + +### 3.4 `GET /health_generate` — Startup Readiness Probe + +``` +GET /health_generate +→ 200 OK (model loaded, engine ready for generation) +``` + +Called by `VLLMEngine.init()` during startup to block until the engine is fully ready. The sidecar implements this by calling vLLM's `GET /health` and optionally performing a dummy `/v1/completions` call with `max_tokens=1` to verify end-to-end readiness. + +### 3.5 Sampling Params Translation + +The request uses SGLang-format parameter names. The sidecar translates to vLLM's `/v1/completions` format: + +| SGLang field (in request) | vLLM `/v1/completions` field | Notes | +|---|---|---| +| `input_ids` | `prompt` | Direct — vLLM accepts `list[int]` as pre-tokenized prompt | +| `temperature` | `temperature` | Direct | +| `top_p` | `top_p` | Direct | +| `top_k` | `top_k` | Both use `-1` for disabled | +| `max_new_tokens` | `max_tokens` | **Name change** | +| `stop` | `stop` | Direct (list of strings) | +| `stop_token_ids` | `stop_token_ids` | Direct | +| `skip_special_tokens` | `skip_special_tokens` | Direct | +| `no_stop_trim` | `include_stop_str_in_output` | **Same semantics, different name** | +| `spaces_between_special_tokens` | `spaces_between_special_tokens` | Direct | +| `return_logprob` | `logprobs` (set to `1`) | Also add `extra_body.return_token_ids = true` | +| `sampling_seed` | `seed` | Optional | +| — | `model` | Must be set to the model name served by vLLM | + +### 3.6 Response Translation Pseudocode + +```python +def translate_vllm_response(vllm_resp: dict, weight_version: int) -> dict: + """Translate vLLM /v1/completions response to SGLang format.""" + choice = vllm_resp["choices"][0] + usage = vllm_resp.get("usage", {}) + + # Build output_token_logprobs: zip logprobs with token IDs + output_token_logprobs = None + if choice.get("logprobs") and choice.get("token_ids"): + output_token_logprobs = [ + [logprob, token_id] + for logprob, token_id in zip( + choice["logprobs"]["token_logprobs"], + choice["token_ids"] + ) + ] + + # Translate finish_reason: plain string → {"type": str} + raw_reason = choice.get("finish_reason") + finish_reason = {"type": raw_reason if raw_reason else "abort"} + + return { + "text": choice["text"], + "output_ids": choice.get("token_ids", []), + "meta_info": { + "output_token_logprobs": output_token_logprobs, + "finish_reason": finish_reason, + "weight_version": weight_version, + "prompt_tokens": usage.get("prompt_tokens", 0), + "cached_tokens": 0, + } + } +``` + +### 3.7 `finish_reason` Translation Table + +| vLLM returns | Translate to | Notes | +|---|---|---| +| `"stop"` | `{"type": "stop"}` | Normal completion | +| `"length"` | `{"type": "length"}` | Hit `max_tokens` | +| `None` (aborted/incomplete) | `{"type": "abort"}` | Triggers middleware retry logic (sleep 30s, up to 5 retries) | + +--- + +## 4. Server Launch Configuration + +The `VLLMEngine` Ray actor should launch vLLM as follows: + +```bash +# Environment +export VLLM_SERVER_DEV_MODE=1 + +# Launch vLLM server +vllm serve \ + --host 0.0.0.0 \ + --port \ + --tensor-parallel-size \ + --enable-sleep-mode \ + --enforce-eager \ + --gpu-memory-utilization 0.9 \ + --disable-log-requests +``` + +The translation sidecar runs on a separate port (``) and is the URL registered with the router via `POST /add_worker?url=http://{host}:{sidecar_port}`. + +``` + Router + │ + ▼ + ┌─────────────────────────┐ + │ Translation Sidecar │ ◄── registered with router + │ port: sidecar_port │ + │ │ + │ /generate ──translate──▶│──┐ + │ /health ──passthrough──▶│ │ + │ /abort_request │ │ + │ /health_generate │ │ + └─────────────────────────┘ │ + │ + ┌─────────────────────────┐ │ + │ vLLM Server │◄─┘ + │ port: engine_port │ + │ │ + │ /v1/completions │ + │ /health │ + │ /sleep, /wake_up │ + │ /pause, /resume │ + │ /update_weights │ + │ /init_weight_transfer │ + └─────────────────────────┘ +``` + +--- + +## 5. Abort Strategy — Detailed Design + +vLLM's abort mechanism differs fundamentally from SGLang's: + +| Aspect | SGLang | vLLM | +|---|---|---| +| Abort granularity | Per-request via `POST /abort_request` with `rid` | Per-connection via HTTP disconnect | +| Bulk abort | `{"abort_all": true}` | No built-in equivalent | +| Mechanism | Engine tracks `request_id`, explicit `abort()` | `@with_cancellation` decorator; request cancelled when client disconnects | +| Between-generation abort | Abort + restart | `POST /pause` → training → `POST /resume` | + +### Recommended implementation + +For the Slime RL use case, the rollout calls `abort_all` between generation rounds (to clear the engine before the next batch). The best vLLM equivalent is: + +```python +# In the translation sidecar +@app.post("/abort_request") +async def abort_request(request: Request): + body = await request.json() + if body.get("abort_all"): + # Option 1: Close all tracked httpx connections → triggers vLLM cancellation + for conn in active_connections: + await conn.aclose() + active_connections.clear() + + # Option 2: Use pause/resume (cleaner) + await httpx.post(f"{vllm_url}/pause") + await httpx.post(f"{vllm_url}/resume") + + return {"status": "ok"} +``` + +--- + +## 6. Endpoints Summary — Gap Analysis + +### Engine-side endpoints (vLLM built-in vs. needs implementation) + +| Endpoint | SGLang | vLLM Built-in | Action | +|---|---|---|---| +| `POST /v1/completions` | — | ✅ | **Reuse** — target for translation | +| `GET /health` | ✅ | ✅ | **Reuse** as-is (passthrough) | +| `POST /pause` | — | ✅ (dev mode) | **Reuse** for abort/weight-update | +| `POST /resume` | — | ✅ (dev mode) | **Reuse** for abort/weight-update | +| `POST /sleep` | — | ✅ (dev mode) | **Reuse** for weight updates | +| `POST /wake_up` | — | ✅ (dev mode) | **Reuse** for weight updates | +| `POST /collective_rpc` | — | ✅ (dev mode) | **Reuse** for weight reload | +| `GET /is_sleeping` | — | ✅ (dev mode) | **Reuse** for state checks | +| `POST /init_weight_transfer_engine` | — | ✅ (dev mode) | **Reuse** for NCCL setup | +| `POST /update_weights` | — | ✅ (dev mode) | **Reuse** for NCCL weight apply | +| `GET /get_world_size` | — | ✅ (dev mode) | **Reuse** for TP world size | + +### Translation sidecar endpoints (to implement) + +| Endpoint | Description | Complexity | +|---|---|---| +| `POST /generate` | Translate SGLang → `/v1/completions` → SGLang | **Medium** — main logic | +| `GET /health` | Proxy to vLLM `/health` | **Trivial** | +| `GET /health_generate` | Health + optional dummy completion | **Low** | +| `POST /abort_request` | Close connections or pause/resume | **Low** | +| `GET /flush_cache` | `POST /sleep?level=1` + `POST /wake_up?tags=kv_cache` | **Low** | +| `GET /get_weight_version` | Return sidecar-tracked version counter | **Trivial** | + +### Router endpoints (no change needed) + +| Endpoint | Action | +|---|---| +| `POST /add_worker` | No change | +| `GET /list_workers` | No change | +| `POST /retrieve_from_text` | No change | +| Catch-all proxy | No change | + +--- + + + + diff --git a/run-qwen2.5-0.5B-vllm.sh b/run-qwen2.5-0.5B-vllm.sh index 208ec60c91..a315590808 100644 --- a/run-qwen2.5-0.5B-vllm.sh +++ b/run-qwen2.5-0.5B-vllm.sh @@ -95,6 +95,8 @@ VLLM_ARGS=( --rollout-backend vllm --rollout-num-gpus-per-engine 1 --sglang-server-concurrency 512 + --use-slime-router + --slime-router-middleware-paths slime.router.middleware_hub.radix_tree_middleware.RadixTreeMiddleware ) MISC_ARGS=( @@ -116,6 +118,8 @@ ray job submit --address="http://127.0.0.1:8265" \ "NCCL_ALGO": "Ring", "NCCL_IB_DISABLE": "1", "NCCL_P2P_DISABLE": "1", + "NCCL_SHM_DISABLE": "1", + "NCCL_NET_GDR_LEVEL": "0", "NCCL_DEBUG": "INFO", "NVTE_ALLOW_NONDETERMINISTIC_ALGO": "0", "CUBLAS_WORKSPACE_CONFIG": ":4096:8" diff --git a/slime/backends/vllm_utils/__init__.py b/slime/backends/vllm_utils/__init__.py index ea4e7311c3..41b5f9909c 100644 --- a/slime/backends/vllm_utils/__init__.py +++ b/slime/backends/vllm_utils/__init__.py @@ -1,3 +1,4 @@ from slime.backends.vllm_utils.vllm_engine import VLLMEngine +from slime.backends.vllm_utils.vllm_translation_sidecar import TranslationSidecar, run_sidecar -__all__ = ["VLLMEngine"] +__all__ = ["VLLMEngine", "TranslationSidecar", "run_sidecar"] diff --git a/slime/backends/vllm_utils/vllm_engine.py b/slime/backends/vllm_utils/vllm_engine.py index bf645cd0bd..713793b07c 100644 --- a/slime/backends/vllm_utils/vllm_engine.py +++ b/slime/backends/vllm_utils/vllm_engine.py @@ -1,6 +1,7 @@ -"""VLLMEngine: Ray actor that launches and manages a vLLM server.""" +"""VLLMEngine: Ray actor that launches and manages a vLLM server + translation sidecar.""" import logging +import multiprocessing import os import subprocess import tempfile @@ -25,14 +26,32 @@ def __init__(self, args, rank: int, base_gpu_id: int | None = None, gpu_ids: lis self.gpu_ids = gpu_ids or [self.base_gpu_id] self.server_host = None self.server_port = None + self.sidecar_port = None self.process = None + self.sidecar_process = None self._log_file = None + self._sidecar_log_file = None + self._weight_version: int = 0 - def init(self, port=None, host=None, **kwargs): + @property + def sidecar_url(self) -> str: + """URL of the translation sidecar (registered with the router).""" + return f"http://{self.server_host}:{self.sidecar_port}" + + @property + def vllm_url(self) -> str: + """URL of the raw vLLM server.""" + return f"http://{self.server_host}:{self.server_port}" + + def init(self, port=None, host=None, router_ip=None, router_port=None, **kwargs): self.server_host = host or get_host_info()[1] self.server_port = port or get_free_port(15000) + self.sidecar_port = get_free_port(self.server_port + 100) + self.router_ip = router_ip or getattr(self.args, "sglang_router_ip", None) + self.router_port = router_port or getattr(self.args, "sglang_router_port", None) model = getattr(self.args, "vllm_model", None) or self.args.hf_checkpoint + self._model_name = model tp = self.args.rollout_num_gpus_per_engine gpu_ids = self.gpu_ids[:tp] cvd = os.environ.get("CUDA_VISIBLE_DEVICES", "") @@ -52,6 +71,8 @@ def init(self, port=None, host=None, **kwargs): "--seed", str(seed), "--trust-remote-code", ] + gpu_mem_util = getattr(self.args, "vllm_gpu_memory_utilization", 0.4) + cmd.extend(["--gpu-memory-utilization", str(gpu_mem_util)]) if getattr(self.args, "offload_rollout", False): cmd.append("--enable-sleep-mode") if getattr(self.args, "vllm_enforce_eager", True): @@ -80,6 +101,14 @@ def init(self, port=None, host=None, **kwargs): ) self._wait_healthy() + # Launch the translation sidecar + self._launch_sidecar() + self._wait_sidecar_healthy() + + # Register the sidecar URL with the router + if self.router_ip and self.router_port: + self._register_with_router() + def _wait_healthy(self, timeout=300): base = f"http://{self.server_host}:{self.server_port}" start = time.time() @@ -98,6 +127,78 @@ def _wait_healthy(self, timeout=300): log_tail = self._read_log_tail() raise TimeoutError(f"vLLM server failed to become healthy within {timeout}s.\n{log_tail}") + def _launch_sidecar(self): + """Launch the translation sidecar as a subprocess.""" + from slime.backends.vllm_utils.vllm_translation_sidecar import run_sidecar + + self._sidecar_log_file = tempfile.NamedTemporaryFile( + prefix="vllm_sidecar_", suffix=".log", delete=False, mode="w" + ) + + def _target(): + run_sidecar( + vllm_host="127.0.0.1", + vllm_port=self.server_port, + sidecar_host="0.0.0.0", + sidecar_port=self.sidecar_port, + model_name=self._model_name, + log_level="info", + ) + + self.sidecar_process = multiprocessing.Process(target=_target, daemon=True) + self.sidecar_process.start() + logger.info( + "Launched translation sidecar on port %s (vLLM → %s:%s), log=%s", + self.sidecar_port, + self.server_host, + self.server_port, + self._sidecar_log_file.name, + ) + + def _wait_sidecar_healthy(self, timeout: float = 60.0): + """Block until the sidecar /health endpoint responds 200.""" + url = f"{self.sidecar_url}/health" + start = time.time() + while time.time() - start < timeout: + try: + r = requests.get(url, timeout=5) + if r.status_code == 200: + logger.info("Translation sidecar healthy at %s", self.sidecar_url) + return + except Exception: + pass + if self.sidecar_process and not self.sidecar_process.is_alive(): + raise RuntimeError( + f"Sidecar process exited with code {self.sidecar_process.exitcode}" + ) + time.sleep(1) + raise TimeoutError(f"Sidecar failed to become healthy within {timeout}s") + + def _register_with_router(self): + """Register the sidecar URL with the SlimeRouter.""" + router_url = f"http://{self.router_ip}:{self.router_port}" + response = requests.post( + f"{router_url}/add_worker", + params={"url": self.sidecar_url}, + ) + response.raise_for_status() + logger.info( + "Registered sidecar %s with router at %s", + self.sidecar_url, + router_url, + ) + + def _bump_weight_version(self, version: int | None = None): + """Notify the sidecar to increment (or set) its weight version counter.""" + url = f"{self.sidecar_url}/set_weight_version" + payload = {"weight_version": version} if version is not None else {} + try: + r = requests.post(url, json=payload, timeout=10) + r.raise_for_status() + self._weight_version = r.json().get("weight_version", self._weight_version) + except Exception as exc: + logger.warning("Failed to bump sidecar weight version: %s", exc) + def _read_log_tail(self, n=200): if not self._log_file: return "" @@ -186,6 +287,8 @@ def update_weights_from_distributed( "packed": packed, } }) + # Notify the sidecar about the new weight version + self._bump_weight_version(weight_version) def continue_generation(self): self._post("/resume") @@ -209,7 +312,14 @@ def resume_memory_occupation(self, tags: list[str] | None = None): logger.warning("vLLM wake_up failed: %s", e) def get_weight_version(self): - return None + if self.sidecar_port: + try: + r = requests.get(f"{self.sidecar_url}/get_weight_version", timeout=5) + r.raise_for_status() + return r.json().get("weight_version", self._weight_version) + except Exception: + pass + return self._weight_version def check_weights(self, action: str): pass @@ -218,6 +328,20 @@ def post_process_weights(self, **kwargs): pass def shutdown(self): + # Shutdown translation sidecar first + if self.sidecar_process and self.sidecar_process.is_alive(): + self.sidecar_process.terminate() + self.sidecar_process.join(timeout=10) + if self.sidecar_process.is_alive(): + self.sidecar_process.kill() + self.sidecar_process = None + if self._sidecar_log_file: + try: + self._sidecar_log_file.close() + except Exception: + pass + + # Shutdown vLLM server if self.process: self.process.terminate() try: diff --git a/slime/backends/vllm_utils/vllm_translation_sidecar.py b/slime/backends/vllm_utils/vllm_translation_sidecar.py new file mode 100644 index 0000000000..73998d8bf4 --- /dev/null +++ b/slime/backends/vllm_utils/vllm_translation_sidecar.py @@ -0,0 +1,454 @@ +""" +vLLM Translation Sidecar +======================== + +Lightweight FastAPI process co-located with each vLLM engine. +Receives SGLang-format ``/generate`` requests from the SlimeRouter, +translates them to vLLM ``/v1/completions``, and translates responses back. + +Also proxies lifecycle endpoints: + /health, /health_generate, /abort_request, /flush_cache, /get_weight_version + +See docs/en/vllm/ROUTER_DESIGN.md for the full specification. +""" + +from __future__ import annotations + +import argparse +import asyncio +import logging +import signal +import sys +from contextlib import asynccontextmanager +from typing import Any + +import httpx +import uvicorn +from fastapi import FastAPI, Request +from fastapi.responses import JSONResponse + +logger = logging.getLogger(__name__) + +# --------------------------------------------------------------------------- +# Sampling-param translation tables +# --------------------------------------------------------------------------- + +# SGLang name → vLLM /v1/completions name (only entries that differ) +_PARAM_RENAME = { + "max_new_tokens": "max_tokens", + "no_stop_trim": "include_stop_str_in_output", + "sampling_seed": "seed", +} + +# Parameters passed through unchanged +_PARAM_DIRECT = frozenset( + { + "temperature", + "top_p", + "top_k", + "stop", + "stop_token_ids", + "skip_special_tokens", + "spaces_between_special_tokens", + } +) + +# finish_reason vLLM → SGLang-style {"type": ...} +_FINISH_REASON_MAP = { + "stop": "stop", + "length": "length", + None: "abort", +} + + +# --------------------------------------------------------------------------- +# Request / response translation helpers +# --------------------------------------------------------------------------- + + +def translate_generate_request( + body: dict[str, Any], + model_name: str, +) -> dict[str, Any]: + """Translate an SGLang-format /generate request → vLLM /v1/completions payload.""" + + sp: dict = body.get("sampling_params", {}) + + vllm_payload: dict[str, Any] = { + "model": model_name, + # vLLM accepts list[int] as a pre-tokenized prompt + "prompt": body.get("input_ids") or body.get("input_tokens", []), + "stream": False, + } + + # --- direct-copy params --- + for key in _PARAM_DIRECT: + if key in sp: + vllm_payload[key] = sp[key] + + # --- renamed params --- + for src, dst in _PARAM_RENAME.items(): + if src in sp: + vllm_payload[dst] = sp[src] + + # --- logprob handling --- + if body.get("return_logprob", False): + vllm_payload["logprobs"] = 1 + # request token IDs alongside logprobs + # NOTE: must be a top-level param; "extra_body" is an OpenAI SDK + # client concept and is ignored by vLLM's raw HTTP API. + vllm_payload["return_token_ids"] = True + + return vllm_payload + + +def translate_vllm_response( + vllm_resp: dict[str, Any], + weight_version: int, +) -> dict[str, Any]: + """Translate a vLLM /v1/completions response → SGLang-format JSON.""" + + choice: dict = vllm_resp.get("choices", [{}])[0] + usage: dict = vllm_resp.get("usage", {}) + + # --- output token IDs --- + output_ids: list[int] = choice.get("token_ids", []) + + # --- logprobs: zip(logprob, token_id) --- + output_token_logprobs: list[list[float | int]] = [] + logprobs_obj = choice.get("logprobs") + if logprobs_obj and output_ids: + raw_lp: list[float | None] = logprobs_obj.get("token_logprobs", []) + output_token_logprobs = [ + [float(lp) if lp is not None else 0.0, tid] + for lp, tid in zip(raw_lp, output_ids) + ] + + # --- finish reason --- + raw_reason = choice.get("finish_reason") + mapped = _FINISH_REASON_MAP.get(raw_reason, raw_reason or "abort") + finish_reason = {"type": mapped} + + meta_info: dict[str, Any] = { + "finish_reason": finish_reason, + "weight_version": weight_version, + "prompt_tokens": usage.get("prompt_tokens", 0), + "completion_tokens": usage.get("completion_tokens", len(output_ids)), + "cached_tokens": 0, + } + # Only include output_token_logprobs when we have valid paired data; + # a None value causes RadixTreeMiddleware to silently fail when iterating. + if output_token_logprobs: + meta_info["output_token_logprobs"] = output_token_logprobs + + return { + "text": choice.get("text", ""), + "output_ids": output_ids, + "meta_info": meta_info, + } + + +# --------------------------------------------------------------------------- +# Sidecar application +# --------------------------------------------------------------------------- + + +class TranslationSidecar: + """Manages state and provides the FastAPI app for the translation sidecar.""" + + def __init__( + self, + vllm_base_url: str, + model_name: str, + *, + timeout: float = 600.0, + max_connections: int = 256, + ): + self.vllm_base_url = vllm_base_url.rstrip("/") + self.model_name = model_name + self._weight_version: int = 0 + self._active_connections: set[httpx.Response] = set() + self._lock = asyncio.Lock() + + self._client: httpx.AsyncClient | None = None + self._timeout = timeout + self._max_connections = max_connections + + self.app = self._build_app() + + # ---- lifecycle ----------------------------------------------------------- + + async def startup(self): + self._client = httpx.AsyncClient( + limits=httpx.Limits( + max_connections=self._max_connections, + max_keepalive_connections=self._max_connections, + ), + timeout=httpx.Timeout(self._timeout), + ) + + async def shutdown(self): + if self._client: + await self._client.aclose() + self._client = None + + # ---- app factory --------------------------------------------------------- + + def _build_app(self) -> FastAPI: + + @asynccontextmanager + async def lifespan(app: FastAPI): + await self.startup() + yield + await self.shutdown() + + app = FastAPI(title="vLLM Translation Sidecar", lifespan=lifespan) + + app.post("/generate")(self.generate) + app.get("/health")(self.health) + app.get("/health_generate")(self.health_generate) + app.post("/abort_request")(self.abort_request) + app.get("/flush_cache")(self.flush_cache) + app.get("/get_weight_version")(self.get_weight_version) + app.post("/set_weight_version")(self.set_weight_version) + + return app + + # ---- endpoints ----------------------------------------------------------- + + async def generate(self, request: Request): + """Translate SGLang /generate → vLLM /v1/completions → SGLang response.""" + + body = await request.json() + vllm_payload = translate_generate_request(body, self.model_name) + + url = f"{self.vllm_base_url}/v1/completions" + + resp: httpx.Response | None = None + try: + async with self._lock: + # We don't actually hold the lock during the request, + # just use it to safely add to the tracking set. + pass + + resp = await self._client.post(url, json=vllm_payload) + self._active_connections.add(resp) + + resp.raise_for_status() + vllm_data = resp.json() + except httpx.HTTPStatusError as exc: + logger.error( + "vLLM /v1/completions returned %s: %s", + exc.response.status_code, + exc.response.text[:2000], + ) + # Return an abort-style response so the middleware retries + return JSONResponse( + content={ + "text": "", + "output_ids": [], + "meta_info": { + "finish_reason": {"type": "abort"}, + "weight_version": self._weight_version, + "prompt_tokens": 0, + "completion_tokens": 0, + "cached_tokens": 0, + }, + }, + status_code=200, + ) + except Exception as exc: + logger.error("Error calling vLLM: %s", exc, exc_info=True) + return JSONResponse( + content={ + "text": "", + "output_ids": [], + "meta_info": { + "finish_reason": {"type": "abort"}, + "weight_version": self._weight_version, + "prompt_tokens": 0, + "completion_tokens": 0, + "cached_tokens": 0, + }, + }, + status_code=200, + ) + finally: + if resp is not None: + self._active_connections.discard(resp) + await resp.aclose() + + translated = translate_vllm_response(vllm_data, self._weight_version) + return JSONResponse(content=translated) + + async def health(self): + """Proxy to vLLM's built-in /health endpoint.""" + try: + resp = await self._client.get(f"{self.vllm_base_url}/health", timeout=5.0) + return JSONResponse(content={"status": "ok"}, status_code=resp.status_code) + except Exception: + return JSONResponse(content={"status": "unhealthy"}, status_code=503) + + async def health_generate(self): + """ + Startup readiness probe. + + Checks vLLM /health and optionally fires a dummy /v1/completions + with max_tokens=1 to verify end-to-end readiness. + """ + try: + resp = await self._client.get(f"{self.vllm_base_url}/health", timeout=10.0) + if resp.status_code != 200: + return JSONResponse(content={"status": "not_ready"}, status_code=503) + + # Lightweight smoke test: single-token completion + dummy_payload = { + "model": self.model_name, + "prompt": "hi", + "max_tokens": 1, + "stream": False, + } + resp2 = await self._client.post( + f"{self.vllm_base_url}/v1/completions", + json=dummy_payload, + timeout=30.0, + ) + if resp2.status_code == 200: + return JSONResponse(content={"status": "ready"}, status_code=200) + else: + return JSONResponse(content={"status": "not_ready"}, status_code=503) + except Exception as exc: + logger.debug("health_generate check failed: %s", exc) + return JSONResponse(content={"status": "not_ready"}, status_code=503) + + async def abort_request(self, request: Request): + """ + Handle abort requests. + + vLLM uses HTTP disconnect for cancellation. We close all active + connections to the vLLM backend, which triggers vLLM's + ``@with_cancellation`` decorator to abort in-flight requests. + + Alternatively, use pause/resume for a cleaner between-generation abort. + """ + body = await request.json() + + if body.get("abort_all", False): + # Strategy: pause + resume clears the pipeline + try: + await self._client.post( + f"{self.vllm_base_url}/pause", + params={"mode": "abort"}, + timeout=30.0, + ) + await self._client.post( + f"{self.vllm_base_url}/resume", + timeout=30.0, + ) + except Exception as exc: + logger.warning("pause/resume abort failed, falling back to connection close: %s", exc) + # Fallback: close all tracked connections + conns = list(self._active_connections) + self._active_connections.clear() + for conn in conns: + try: + await conn.aclose() + except Exception: + pass + + return JSONResponse(content={"status": "ok"}) + + async def flush_cache(self): + """ + Flush the KV cache. + + vLLM equivalent: sleep(level=1) + wake_up(tags=kv_cache). + """ + try: + await self._client.post( + f"{self.vllm_base_url}/sleep", + params={"level": "1", "mode": "abort"}, + timeout=30.0, + ) + await self._client.post( + f"{self.vllm_base_url}/wake_up", + params={"tags": "kv_cache"}, + timeout=30.0, + ) + return JSONResponse(content={"status": "ok"}) + except Exception as exc: + logger.warning("flush_cache failed: %s", exc) + return JSONResponse(content={"status": "ok"}) + + async def get_weight_version(self): + """Return the sidecar-tracked weight version counter.""" + return JSONResponse(content={"weight_version": self._weight_version}) + + async def set_weight_version(self, request: Request): + """Increment or set the weight version (called by VLLMEngine after weight update).""" + body = await request.json() + if "weight_version" in body: + self._weight_version = int(body["weight_version"]) + else: + self._weight_version += 1 + return JSONResponse(content={"weight_version": self._weight_version}) + + +# --------------------------------------------------------------------------- +# Standalone entry point +# --------------------------------------------------------------------------- + + +def run_sidecar( + vllm_host: str = "127.0.0.1", + vllm_port: int = 8000, + sidecar_host: str = "0.0.0.0", + sidecar_port: int = 8100, + model_name: str = "default", + timeout: float = 600.0, + max_connections: int = 256, + log_level: str = "info", +): + """Launch the translation sidecar as a standalone uvicorn process.""" + + vllm_base_url = f"http://{vllm_host}:{vllm_port}" + sidecar = TranslationSidecar( + vllm_base_url=vllm_base_url, + model_name=model_name, + timeout=timeout, + max_connections=max_connections, + ) + uvicorn.run( + sidecar.app, + host=sidecar_host, + port=sidecar_port, + log_level=log_level, + ) + + +def main(): + parser = argparse.ArgumentParser(description="vLLM Translation Sidecar") + parser.add_argument("--vllm-host", type=str, default="127.0.0.1") + parser.add_argument("--vllm-port", type=int, default=8000) + parser.add_argument("--sidecar-host", type=str, default="0.0.0.0") + parser.add_argument("--sidecar-port", type=int, default=8100) + parser.add_argument("--model-name", type=str, default="default") + parser.add_argument("--timeout", type=float, default=600.0) + parser.add_argument("--max-connections", type=int, default=256) + parser.add_argument("--log-level", type=str, default="info") + args = parser.parse_args() + + run_sidecar( + vllm_host=args.vllm_host, + vllm_port=args.vllm_port, + sidecar_host=args.sidecar_host, + sidecar_port=args.sidecar_port, + model_name=args.model_name, + timeout=args.timeout, + max_connections=args.max_connections, + log_level=args.log_level, + ) + + +if __name__ == "__main__": + main() diff --git a/slime/ray/rollout.py b/slime/ray/rollout.py index ba18fb9c47..9d1a786ae0 100644 --- a/slime/ray/rollout.py +++ b/slime/ray/rollout.py @@ -1022,46 +1022,87 @@ def _start_router(args, *, has_pd_disaggregation: bool = False, force_new: bool def _start_vllm_rollout_servers(args, pg) -> dict[str, RolloutServer]: - """Start vLLM rollout server (single instance, no router).""" + """Start vLLM rollout server(s) with optional SlimeRouter. + + When ``args.use_slime_router`` is True, a SlimeRouter is launched first + and each vLLM engine registers its translation sidecar with the router. + Otherwise, a single engine is started and used directly (no router). + """ pg_obj, reordered_bundle_indices, reordered_gpu_ids = pg tp = args.rollout_num_gpus_per_engine - gpu_ids = [int(reordered_gpu_ids[i]) for i in range(tp)] - scheduling_strategy = PlacementGroupSchedulingStrategy( - placement_group=pg_obj, - placement_group_capture_child_tasks=True, - placement_group_bundle_index=reordered_bundle_indices[0], - ) + num_gpu_per_engine_local = min(tp, args.num_gpus_per_node) + num_engines = max(1, args.rollout_num_gpus // num_gpu_per_engine_local) - env_vars = {name: "1" for name in NOSET_VISIBLE_DEVICES_ENV_VARS_LIST} + use_router = getattr(args, "use_slime_router", False) + + # --- Launch the SlimeRouter if requested --- + if use_router: + router_ip, router_port = _start_router(args, has_pd_disaggregation=False) + args.sglang_router_ip = router_ip + args.sglang_router_port = router_port + else: + router_ip = None + router_port = None + env_vars = {name: "1" for name in NOSET_VISIBLE_DEVICES_ENV_VARS_LIST} VLLMRayActor = ray.remote(VLLMEngine) - engine = VLLMRayActor.options( - num_cpus=0.2, - num_gpus=0.2, - scheduling_strategy=scheduling_strategy, - runtime_env={"env_vars": env_vars}, - ).remote(args, rank=0, base_gpu_id=gpu_ids[0], gpu_ids=gpu_ids) - - host, port = ray.get(engine._get_current_node_ip_and_free_port.remote(start_port=15000)) - ray.get(engine.init.remote(port=port, host=host)) - args.vllm_base_url = f"http://{host}:{port}" - args.sglang_router_ip = host - args.sglang_router_port = port + + engines = [] + init_handles = [] + for i in range(num_engines): + gpu_index = i * num_gpu_per_engine_local + gpu_ids = [int(reordered_gpu_ids[gpu_index + j]) for j in range(num_gpu_per_engine_local)] + scheduling_strategy = PlacementGroupSchedulingStrategy( + placement_group=pg_obj, + placement_group_capture_child_tasks=True, + placement_group_bundle_index=reordered_bundle_indices[gpu_index], + ) + + engine = VLLMRayActor.options( + num_cpus=0.2, + num_gpus=0.2, + scheduling_strategy=scheduling_strategy, + runtime_env={"env_vars": env_vars}, + ).remote(args, rank=i, base_gpu_id=gpu_ids[0], gpu_ids=gpu_ids) + + host, port = ray.get(engine._get_current_node_ip_and_free_port.remote(start_port=15000 + i * 10)) + + init_handles.append( + engine.init.remote( + port=port, + host=host, + router_ip=router_ip, + router_port=router_port, + ) + ) + engines.append(engine) + + # Use the first engine's host for args if no router + if i == 0 and not use_router: + args.vllm_base_url = f"http://{host}:{port}" + args.sglang_router_ip = host + args.sglang_router_port = port + + # Wait for all engines to be healthy + registered with router + ray.get(init_handles) + + first_host = args.sglang_router_ip + first_port = args.sglang_router_port group = EngineGroup( args=args, pg=pg, - all_engines=[engine], + all_engines=engines, num_gpus_per_engine=args.rollout_num_gpus_per_engine, - num_new_engines=1, + num_new_engines=num_engines, worker_type="regular", rank_offset=0, gpu_offset=0, sglang_overrides={}, - router_ip=host, - router_port=port, + router_ip=first_host, + router_port=first_port, ) - return {"default": RolloutServer(engine_groups=[group], router_ip=host, router_port=port, model_name="default")} + return {"default": RolloutServer(engine_groups=[group], router_ip=first_host, router_port=first_port, model_name="default")} def start_rollout_servers(args, pg) -> dict[str, RolloutServer]: diff --git a/slime/rollout/backends/vllm_client.py b/slime/rollout/backends/vllm_client.py index 32be01d66c..d5e2ed6bdf 100644 --- a/slime/rollout/backends/vllm_client.py +++ b/slime/rollout/backends/vllm_client.py @@ -2,7 +2,7 @@ from slime.rollout.backends.base_client import BackendCapabilities, RolloutBackendClient from slime.rollout.base_types import RolloutBackendRequest, RolloutBackendResponse -from slime.utils.http_utils import post +from slime.utils.http_utils import get, post logger = logging.getLogger(__name__) @@ -16,14 +16,27 @@ class VLLMClient(RolloutBackendClient): + """Rollout backend client for vLLM. + + Supports two modes: + + 1. **Direct mode** (default): sends requests directly to vLLM's + ``/v1/completions`` endpoint and parses the native vLLM response. + 2. **Router mode** (``use_slime_router=True``): sends SGLang-format + requests to ``/generate`` through the SlimeRouter, which forwards + to the translation sidecar. The sidecar handles the vLLM translation + and returns SGLang-format responses. + """ + def __init__(self, args): self.args = args self._max_retries = getattr(args, "vllm_max_retries", 3) + self._use_router = getattr(args, "use_slime_router", False) @property def capabilities(self) -> BackendCapabilities: return BackendCapabilities( - supports_abort=False, + supports_abort=self._use_router, # abort is supported through the sidecar supports_routed_experts=False, supports_prompt_logprobs=False, ) @@ -33,6 +46,69 @@ async def generate( request: RolloutBackendRequest, base_url: str, headers: dict | None = None, + ) -> RolloutBackendResponse: + if self._use_router: + return await self._generate_via_router(request, base_url, headers) + return await self._generate_direct(request, base_url, headers) + + # ------------------------------------------------------------------ + # Router mode: SGLang-format /generate → sidecar → vLLM + # ------------------------------------------------------------------ + + async def _generate_via_router( + self, + request: RolloutBackendRequest, + base_url: str, + headers: dict | None = None, + ) -> RolloutBackendResponse: + """Send SGLang-format request through the SlimeRouter → sidecar pipeline.""" + + payload = { + "input_ids": request.input_ids, + "sampling_params": request.sampling_params, + "return_logprob": request.return_logprob, + "stream": False, + } + if request.image_data: + payload["image_data"] = request.image_data + + url = f"{base_url.rstrip('/')}/generate" + output = await post(url, payload, headers=headers) + + # Parse SGLang-format response (produced by translation sidecar) + meta = output.get("meta_info", {}) + logprobs_data = meta.get("output_token_logprobs", []) + + # output_token_logprobs is list of [logprob, token_id] + output_token_ids = [item[1] for item in logprobs_data] if logprobs_data else [] + output_token_logprobs = [item[0] for item in logprobs_data] if logprobs_data else [] + + # Fall back to output_ids if logprobs not available + if not output_token_ids: + output_token_ids = output.get("output_ids") or [] + + finish_reason = meta.get("finish_reason", {}).get("type", "stop") + + return RolloutBackendResponse( + text=output.get("text", ""), + output_token_ids=output_token_ids, + output_token_logprobs=output_token_logprobs, + finish_reason=finish_reason, + prompt_tokens=meta.get("prompt_tokens", len(request.input_ids)), + completion_tokens=len(output_token_ids), + backend_raw=output, + routed_experts=None, + ) + + # ------------------------------------------------------------------ + # Direct mode: vLLM /v1/completions (no router) + # ------------------------------------------------------------------ + + async def _generate_direct( + self, + request: RolloutBackendRequest, + base_url: str, + headers: dict | None = None, ) -> RolloutBackendResponse: sp = request.sampling_params payload = { @@ -91,3 +167,15 @@ async def generate( backend_raw=output, routed_experts=None, ) + + # ------------------------------------------------------------------ + # Abort (router mode only) + # ------------------------------------------------------------------ + + async def abort(self) -> list[str]: + """Return worker URLs for abort. Only works in router mode.""" + if not self._use_router: + return [] + base = f"http://{self.args.sglang_router_ip}:{self.args.sglang_router_port}" + r = await get(f"{base}/list_workers") + return r.get("urls", []) diff --git a/slime/utils/arguments.py b/slime/utils/arguments.py index 7ef8e70b71..4bb5d3506d 100644 --- a/slime/utils/arguments.py +++ b/slime/utils/arguments.py @@ -263,6 +263,13 @@ def add_rollout_arguments(parser): default=True, help="Use vLLM packed weight transfer for non-colocate (default: True). Disable for per-bucket mode.", ) + parser.add_argument( + "--vllm-gpu-memory-utilization", + type=float, + default=0.4, + help="Fraction of GPU memory for vLLM KV cache (default: 0.85). " + "Lower this if the training model leaves insufficient free memory.", + ) parser.add_argument( "--rollout-function-path", type=str, From c15407873cf697e81d51fe31108635322d6da576 Mon Sep 17 00:00:00 2001 From: knlnguyen1802 Date: Mon, 30 Mar 2026 16:52:45 +0800 Subject: [PATCH 11/39] Plan refactor vllm/sglang Signed-off-by: knlnguyen1802 --- rfc-rollout-backend-separation-plan.md | 271 +++++++++++++++++++++++++ 1 file changed, 271 insertions(+) create mode 100644 rfc-rollout-backend-separation-plan.md diff --git a/rfc-rollout-backend-separation-plan.md b/rfc-rollout-backend-separation-plan.md new file mode 100644 index 0000000000..ef715c2a6d --- /dev/null +++ b/rfc-rollout-backend-separation-plan.md @@ -0,0 +1,271 @@ +# RFC: Rollout Separation Plan (EngineGroup Generalization + Executor Cleanup) + +## 1. Summary + +This RFC proposes backend separation with minimal churn in runtime orchestration. + +Scope is four items: + +1. Refactor [slime/ray/rollout.py](slime/ray/rollout.py) by generalizing `EngineGroup.start_engines()` and abstracting engine/server creation (no new runtime manager class hierarchy). +2. Refactor [slime/rollout/sglang_rollout.py](slime/rollout/sglang_rollout.py) in-place: extract the one SGLang-specific code path (RadixTree) into a strategy hook, rename SGLang-prefixed args to generic names. No class hierarchy, no new files. +3. Refactor [slime/utils/arguments.py](slime/utils/arguments.py) into shared args + backend arg groups/finalizers. +4. Decouple [slime/backends/fsdp_utils/update_weight_utils.py](slime/backends/fsdp_utils/update_weight_utils.py) from SGLang internals so FSDP weight sync works with both SGLang and vLLM engines. + +## 2. Already Done (Reuse, Do Not Rewrite) + +- Unified rollout contracts in [slime/rollout/base_types.py](slime/rollout/base_types.py). +- Backend client abstraction in [slime/rollout/backends/base_client.py](slime/rollout/backends/base_client.py). +- Backend adapters in [slime/rollout/backends/sglang_client.py](slime/rollout/backends/sglang_client.py) and [slime/rollout/backends/vllm_client.py](slime/rollout/backends/vllm_client.py). +- Managed vLLM engine actor in [slime/backends/vllm_utils/vllm_engine.py](slime/backends/vllm_utils/vllm_engine.py). +- vLLM translation sidecar in [slime/backends/vllm_utils/vllm_translation_sidecar.py](slime/backends/vllm_utils/vllm_translation_sidecar.py). + +## 3. Problem + +- [slime/ray/rollout.py](slime/ray/rollout.py) mixes shared and backend-specific engine creation paths. +- [slime/rollout/sglang_rollout.py](slime/rollout/sglang_rollout.py) is ~95% backend-agnostic but has one inlined SGLang-specific code path (RadixTree, ~14 lines) and uses SGLang-prefixed arg names for generic rollout concepts. +- [slime/utils/arguments.py](slime/utils/arguments.py) still has SGLang alias behavior in vLLM path. +- [slime/backends/fsdp_utils/update_weight_utils.py](slime/backends/fsdp_utils/update_weight_utils.py) hard-imports SGLang internals (`FlattenedTensorBucket`, `MultiprocessingSerializer`, `monkey_patch_torch_reductions`) and calls SGLang-specific engine RPC names (`update_weights_from_tensor`, `update_weights_from_distributed`), making FSDP weight sync unusable with vLLM engines. + +## 4. Goals and Non-Goals + +### Goals + +- Keep runtime refactor minimal and localized to `EngineGroup` + creation abstraction. +- Isolate the one SGLang-specific executor code path behind a strategy hook; keep functions as functions. +- Rename SGLang-prefixed arg names to generic rollout names to eliminate naming coupling. +- Reduce backend leakage in argument finalization. +- Preserve current external behavior, call sites, and import paths. + +### Non-Goals + +- No rewrite of `SGLangEngine` or `VLLMEngine` internals. +- No algorithmic changes to GRPO/PPO. +- No mandatory feature parity for unsupported backend capabilities. + +## 5. Design + +### 5.1 Runtime: Generalize `EngineGroup` (No New Runtime Manager Classes) + +Keep [slime/ray/rollout.py](slime/ray/rollout.py) as the orchestration entry file. + +Refactor focus: + +1. Generalize `EngineGroup.start_engines()` to call backend-aware creation hooks. +2. Abstract engine creation and rollout-server assembly helpers. +3. Keep existing startup function API (`start_rollout_servers`) and return shape. + +Proposed helper abstraction points: + +- `create_engine_actor_cls(args, worker_type)` + - returns `ray.remote(SGLangEngine)` or `ray.remote(VLLMEngine)`. +- `create_engine_remote(args, actor_cls, scheduling_strategy, ...)` + - encapsulates `.options(...).remote(...)` with backend-specific init kwargs. +- `build_rollout_server(...)` + - standardizes `RolloutServer` construction from engine groups. + +`EngineGroup` and `RolloutServer` remain the main shared dataclasses. + +### 5.2 Executor: Isolate Backend Logic In-Place (No Class Hierarchy) + +#### Current state analysis + +[slime/rollout/sglang_rollout.py](slime/rollout/sglang_rollout.py) (577 lines) is **already ~95% backend-agnostic**: + +| Function / Class | Lines | Backend-specific? | Notes | +|---|---|---|---| +| `_get_backend_client()` | 6 | Factory only | Delegates to existing `RolloutBackendClient` subclasses | +| `_apply_backend_response()` | 20 | No | Uses `RolloutBackendResponse` contract | +| `GenerateState` | 37 | **Naming only** | References `sglang_server_concurrency`, `sglang_dp_size`, `sglang_enable_deterministic_inference` — all are generic rollout concepts with SGLang-prefixed names | +| `generate()` | 58 | **14 lines** | RadixTree middleware path (L170-183) is 100% SGLang-specific; the else branch (L185-195) already uses `RolloutBackendClient` | +| `generate_and_rm()` | 60 | No | Shared orchestration (semaphore, custom func, reward) | +| `generate_and_rm_group()` | 37 | No | Group parallelism + deterministic seeds | +| `abort()` | 38 | No | Already uses `backend.abort()` | +| `generate_rollout_async()` | 71 | No | Main loop, filtering, metrics | +| `eval_rollout()` / `eval_rollout_single_dataset()` | 118 | **Naming only** | `sglang_enable_deterministic_inference` reference | +| `generate_rollout()` | 41 | No | Sync entry point | + +**Conclusion**: the actual backend logic that needs isolation is **one code path** (~14 lines) inside `generate()`. Everything else is either already abstracted through `RolloutBackendClient` or is a naming-only coupling (SGLang-prefixed arg names for generic concepts). + +#### Approach + +1. **Extract the RadixTree path into a strategy hook** that `generate()` calls conditionally. +2. **Rename SGLang-prefixed args** to generic names (coordinated with Phase 3 args refactor). +3. **Keep functions as functions** — they compose well and callers (`train.py`, OPD, multi-agent) import them directly. + +#### Concrete changes + +**Step 1 — Extract RadixTree strategy from `generate()`** + +Current `generate()` has an `if use_radix: ... else: backend.generate(...)` branch. +Refactor into: + +```python +# slime/rollout/sglang_rollout.py — generate() simplified + +async def generate(args, sample, sampling_params): + ... + input_ids = ... # shared prompt encoding (unchanged) + + strategy = _get_generate_strategy(args) + resp = await strategy(args, sample, input_ids, sampling_params) + _apply_backend_response(sample, resp, args) + return sample +``` + +Two strategies: + +```python +# Still in sglang_rollout.py (no new file needed) + +def _get_generate_strategy(args): + """Return the generate coroutine to use.""" + if _is_radix_tree_enabled(args): + return _generate_radix_tree # SGLang-only path + return _generate_via_backend_client # Generic path (SGLang or vLLM) + +def _is_radix_tree_enabled(args) -> bool: + return ( + args.use_slime_router + and "RadixTreeMiddleware" in getattr(args, "slime_router_middleware_paths", []) + ) + +async def _generate_radix_tree(args, sample, input_ids, sampling_params) -> RolloutBackendResponse: + """SGLang RadixTree middleware path — returns normalized response.""" + from slime.router.middleware_hub.radix_tree_middleware import postprocess_sample_with_radix_tree + url = f"http://{args.rollout_router_ip}:{args.rollout_router_port}/generate" + payload = { ... } # existing payload construction + output = await post(url, payload, headers=headers) + sample = await postprocess_sample_with_radix_tree(args, sample, output) + return _extract_response_from_sample(sample) # normalize to RolloutBackendResponse + +async def _generate_via_backend_client(args, sample, input_ids, sampling_params) -> RolloutBackendResponse: + """Generic backend client path — works for SGLang and vLLM.""" + backend = _get_backend_client(args) + base_url = f"http://{args.rollout_router_ip}:{args.rollout_router_port}" + req = RolloutBackendRequest(...) + return await backend.generate(req, base_url, headers=headers) +``` + +**Step 2 — Rename SGLang-prefixed args to generic names** + +| Current name | New name | Reason | +|---|---|---| +| `sglang_server_concurrency` | `rollout_concurrency` | Controls request parallelism for any backend | +| `sglang_dp_size` | `rollout_dp_size` | Data-parallel sharding, not SGLang-specific | +| `sglang_router_ip` / `sglang_router_port` | `rollout_router_ip` / `rollout_router_port` | Router endpoint, backend-agnostic | +| `sglang_router_policy` | `rollout_router_policy` | Routing strategy | +| `sglang_enable_deterministic_inference` | `rollout_deterministic_inference` | Seed-based determinism | +| `vllm_base_url` | (remove) | Folded into `rollout_router_ip:port`, no special case | + +Legacy aliases kept in [slime/utils/arguments.py](slime/utils/arguments.py) for one release cycle (coordinated with Phase 3). + +**Step 3 — No new files** + +The file stays as [slime/rollout/sglang_rollout.py](slime/rollout/sglang_rollout.py) during this phase. +Optionally rename to `slime/rollout/rollout.py` in Phase 4 cleanup, since the file is backend-agnostic after the refactor. + +### 5.3 Arguments/Config Refactor + +In [slime/utils/arguments.py](slime/utils/arguments.py), split into: + +1. Shared rollout args. +2. SGLang backend args/validation. +3. vLLM backend args/validation. + +Add backend finalizers: + +- `finalize_sglang_args(args)` +- `finalize_vllm_args(args)` + +Move SGLang alias fallback out of shared finalize flow. + +### 5.4 Weight Sync: Decouple FSDP `update_weight_utils.py` from SGLang Internals + +#### Current state analysis + +[slime/backends/fsdp_utils/update_weight_utils.py](slime/backends/fsdp_utils/update_weight_utils.py) (287 lines) has two concrete classes: + +| Class | Weight-push method | SGLang coupling | +|---|---|---| +| `UpdateWeightFromTensor` | IPC via Gloo gather → `engine.update_weights_from_tensor.remote()` | Imports `FlattenedTensorBucket`, `MultiprocessingSerializer`, `monkey_patch_torch_reductions` directly from `sglang.srt.*` | +| `UpdateWeightFromDistributed` | NCCL broadcast → `engine.update_weights_from_distributed.remote()` | Calls `engine.init_weights_update_group.remote()` — SGLang engine API | + +The abstract base `UpdateWeight` itself is clean (only PyTorch + Ray). + +The Megatron side already solved this: [slime/backends/megatron_utils/sglang.py](slime/backends/megatron_utils/sglang.py) centralizes all SGLang imports into one shim. The FSDP side duplicates these imports inline. + +#### Coupling points + +1. **SGLang utility imports (L13-26)** — `monkey_patch_torch_reductions`, `MultiprocessingSerializer`, `FlattenedTensorBucket` are imported directly from `sglang.srt.*` with try/except version fallbacks. +2. **`UpdateWeightFromTensor.update_bucket_weights()`** — uses `FlattenedTensorBucket` to flatten tensors, `MultiprocessingSerializer` to serialize, then calls `engine.update_weights_from_tensor.remote()`. +3. **`UpdateWeightFromDistributed.connect_rollout_engines()`** — calls `engine.init_weights_update_group.remote()` which is an SGLang engine method. +4. **`UpdateWeightFromDistributed.update_bucket_weights()`** — calls `engine.update_weights_from_distributed.remote()` which is an SGLang engine method. + +All four points assume the engine actor exposes SGLang's RPC interface. vLLM engines expose different method names. + +#### Approach + +**Step 1 — Centralize SGLang imports via a shim (same pattern as Megatron)** + +Reuse or mirror the existing [slime/backends/megatron_utils/sglang.py](slime/backends/megatron_utils/sglang.py) pattern: + +```python +# slime/backends/fsdp_utils/sglang_compat.py (new, ~15 lines) +try: + from sglang.srt.utils.patch_torch import monkey_patch_torch_reductions +except ImportError: + from sglang.srt.patch_torch import monkey_patch_torch_reductions + +from sglang.srt.utils import MultiprocessingSerializer + +try: + from sglang.srt.weight_sync.tensor_bucket import FlattenedTensorBucket +except ImportError: + from sglang.srt.model_executor.model_runner import FlattenedTensorBucket +``` + +Then `update_weight_utils.py` imports from `sglang_compat` — one import line instead of five, and only loaded when SGLang is the active backend. + +**Step 2 — Abstract engine RPC calls behind a protocol** + +The engine actors (`SGLangEngine`, `VLLMEngine`) already expose weight-sync methods but with different names/signatures. Introduce a lightweight protocol or adapter: + +```python +# In update_weight_utils.py or a small helper + +def _call_engine_update_tensor(engine, backend: str, **kwargs): + """Dispatch IPC weight update to the correct engine method.""" + if backend == "vllm": + return engine.update_weights_from_tensor.remote(**kwargs) # vLLM uses same name via NcclBridge + return engine.update_weights_from_tensor.remote(**kwargs) # SGLang native + +def _call_engine_update_distributed(engine, backend: str, **kwargs): + if backend == "vllm": + return engine.update_weights_from_distributed.remote(**kwargs) + return engine.update_weights_from_distributed.remote(**kwargs) + +def _call_engine_init_weight_group(engine, backend: str, **kwargs): + if backend == "vllm": + return engine.init_weights_update_group.remote(**kwargs) + return engine.init_weights_update_group.remote(**kwargs) +``` + +> Note: Currently `VLLMEngine` already mirrors these method names (it wraps them via `NcclBridge`), so the dispatch functions may initially be identical. The value is making the indirection explicit so future method-name divergence is handled in one place. + +**Step 3 — Lazy-import SGLang utilities only when backend is SGLang** + +Move the `FlattenedTensorBucket` / `MultiprocessingSerializer` imports inside `UpdateWeightFromTensor.update_bucket_weights()` behind a lazy import, so the module can be loaded in a vLLM-only environment without SGLang installed. + +#### What changes + +| File | Change | +|---|---| +| `slime/backends/fsdp_utils/sglang_compat.py` | New shim file (~15 lines) centralizing SGLang imports | +| `slime/backends/fsdp_utils/update_weight_utils.py` | Replace 5 inline SGLang imports with one `from .sglang_compat import ...`; add backend-aware engine dispatch helpers | + +#### What does NOT change + +- `UpdateWeight` abstract base class — already clean. +- `UpdateWeightFromDistributed` NCCL logic — the broadcast itself is pure PyTorch; only the engine RPC dispatch gets a thin wrapper. +- Megatron-side weight sync — already has its own shim, not touched. + From 1a2dcf55acd2e4d3e408e5ce534020d3a25ed4bd Mon Sep 17 00:00:00 2001 From: knlnguyen1802 Date: Tue, 31 Mar 2026 09:43:21 +0800 Subject: [PATCH 12/39] Code implemented Signed-off-by: knlnguyen1802 --- .../fsdp_utils/update_weight_utils.py | 34 +++++++++++++------ slime/ray/rollout.py | 12 +++++-- slime/rollout/backends/__init__.py | 10 +++++- slime/utils/arguments.py | 23 +++++++++---- 4 files changed, 58 insertions(+), 21 deletions(-) diff --git a/slime/backends/fsdp_utils/update_weight_utils.py b/slime/backends/fsdp_utils/update_weight_utils.py index 75e217beb9..64a6770c06 100644 --- a/slime/backends/fsdp_utils/update_weight_utils.py +++ b/slime/backends/fsdp_utils/update_weight_utils.py @@ -10,23 +10,31 @@ from ray.actor import ActorHandle from torch.distributed.tensor import DTensor, Replicate -try: - from sglang.srt.utils.patch_torch import monkey_patch_torch_reductions # type: ignore[import] -except ImportError: - from sglang.srt.patch_torch import monkey_patch_torch_reductions # type: ignore[import] +from slime.utils.distributed_utils import init_process_group -from sglang.srt.utils import MultiprocessingSerializer -from slime.utils.distributed_utils import init_process_group +logger = logging.getLogger(__name__) -try: - from sglang.srt.weight_sync.tensor_bucket import FlattenedTensorBucket # type: ignore[import] -except ImportError: - from sglang.srt.model_executor.model_runner import FlattenedTensorBucket # type: ignore[import] +def _import_sglang_weight_sync_utils(): + """Lazy-import SGLang serialization utilities. + Centralizes the try/except version fallbacks so callers get a clean tuple. + Raises ImportError if sglang is not installed at all. + """ + try: + from sglang.srt.utils.patch_torch import monkey_patch_torch_reductions # type: ignore[import] + except ImportError: + from sglang.srt.patch_torch import monkey_patch_torch_reductions # type: ignore[import] -logger = logging.getLogger(__name__) + from sglang.srt.utils import MultiprocessingSerializer + + try: + from sglang.srt.weight_sync.tensor_bucket import FlattenedTensorBucket # type: ignore[import] + except ImportError: + from sglang.srt.model_executor.model_runner import FlattenedTensorBucket # type: ignore[import] + + return monkey_patch_torch_reductions, MultiprocessingSerializer, FlattenedTensorBucket class UpdateWeight(abc.ABC): @@ -135,6 +143,10 @@ def update_bucket_weights(self, named_tensors, weight_version=None) -> None: if self._ipc_gather_group is None: return + monkey_patch_torch_reductions, MultiprocessingSerializer, FlattenedTensorBucket = ( + _import_sglang_weight_sync_utils() + ) + monkey_patch_torch_reductions() # Use flattened bucket approach similar to Megatron logger.info("Using flattened tensor bucket") diff --git a/slime/ray/rollout.py b/slime/ray/rollout.py index 9d1a786ae0..ee6d2da50f 100644 --- a/slime/ray/rollout.py +++ b/slime/ray/rollout.py @@ -13,9 +13,13 @@ import torch import yaml from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy -from sglang.srt.constants import GPU_MEMORY_TYPE_CUDA_GRAPH, GPU_MEMORY_TYPE_KV_CACHE, GPU_MEMORY_TYPE_WEIGHTS -from slime.backends.sglang_utils.sglang_engine import SGLangEngine +# GPU memory-tag constants (originally from sglang.srt.constants). +# Duplicated here to avoid a hard sglang dependency when running the vLLM backend. +GPU_MEMORY_TYPE_KV_CACHE = "kv_cache" +GPU_MEMORY_TYPE_WEIGHTS = "weights" +GPU_MEMORY_TYPE_CUDA_GRAPH = "cuda_graph" + from slime.backends.vllm_utils.vllm_engine import VLLMEngine from slime.rollout.base_types import call_rollout_fn from slime.utils import logging_utils @@ -234,6 +238,8 @@ def start_engines(self, port_cursors: dict[int, int] | None = None) -> tuple[lis pg, reordered_bundle_indices, reordered_gpu_ids = self.pg + from slime.backends.sglang_utils.sglang_engine import SGLangEngine + RolloutRayActor = ray.remote(SGLangEngine) rollout_engines = [] @@ -989,7 +995,7 @@ def _start_router(args, *, has_pd_disaggregation: bool = False, force_new: bool router_args.sglang_router_port = router_port else: - from sglang_router.launch_router import RouterArgs + from sglang_router.launch_router import RouterArgs # noqa: delayed import — only needed for sglang router from slime.utils.http_utils import run_router diff --git a/slime/rollout/backends/__init__.py b/slime/rollout/backends/__init__.py index 49c4a86eef..ea9ccd8a9c 100644 --- a/slime/rollout/backends/__init__.py +++ b/slime/rollout/backends/__init__.py @@ -1,7 +1,15 @@ from slime.rollout.backends.base_client import BackendCapabilities, RolloutBackendClient -from slime.rollout.backends.sglang_client import SGLangClient from slime.rollout.backends.vllm_client import VLLMClient + +def __getattr__(name): + if name == "SGLangClient": + from slime.rollout.backends.sglang_client import SGLangClient + + return SGLangClient + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + __all__ = [ "BackendCapabilities", "RolloutBackendClient", diff --git a/slime/utils/arguments.py b/slime/utils/arguments.py index 4bb5d3506d..7874237abc 100644 --- a/slime/utils/arguments.py +++ b/slime/utils/arguments.py @@ -5,10 +5,7 @@ from typing import Any import yaml -from sglang_router.launch_router import RouterArgs -from slime.backends.sglang_utils.arguments import sglang_parse_args -from slime.backends.sglang_utils.arguments import validate_args as sglang_validate_args from slime.utils.eval_config import EvalDatasetConfig, build_eval_dataset_configs, ensure_dataset_list from slime.utils.logging_utils import configure_logger @@ -1069,7 +1066,12 @@ def add_router_arguments(parser): default=3, help="Number of consecutive failures before marking a worker as unhealthy.", ) - RouterArgs.add_cli_args(parser, use_router_prefix=True, exclude_host_port=True) + try: + from sglang_router.launch_router import RouterArgs + + RouterArgs.add_cli_args(parser, use_router_prefix=True, exclude_host_port=True) + except ImportError: + pass # sglang not installed — router args unavailable (vllm-only mode) return parser # wandb @@ -1463,6 +1465,7 @@ def _pre_parse_mode(): temp_parser.add_argument("--debug-rollout-only", action="store_true", default=False) temp_parser.add_argument("--debug-train-only", action="store_true", default=False) temp_parser.add_argument("--load-debug-rollout-data", type=str, default=None) + temp_parser.add_argument("--rollout-backend", type=str, choices=["sglang", "vllm"], default="sglang") temp_args, _ = temp_parser.parse_known_args() return temp_args @@ -1474,12 +1477,18 @@ def parse_args(add_custom_arguments=None): add_slime_arguments = get_slime_extra_args_provider(add_custom_arguments) pre = _pre_parse_mode() - skip_sglang = pre.debug_train_only or pre.load_debug_rollout_data is not None + skip_sglang = ( + pre.debug_train_only + or pre.load_debug_rollout_data is not None + or pre.rollout_backend == "vllm" + ) # Phase 1: Parse sglang args independently (separate parser, parse_known_args). - # Skipped when sglang servers are not needed. + # Skipped when sglang servers are not needed or when using vLLM backend. sglang_ns = None if not skip_sglang: + from slime.backends.sglang_utils.arguments import sglang_parse_args + sglang_ns = sglang_parse_args() # Phase 2: Parse megatron/fsdp + slime args. @@ -1517,6 +1526,8 @@ def parse_args(add_custom_arguments=None): megatron_validate_args(args) if not args.debug_train_only and getattr(args, "rollout_backend", "sglang") == "sglang": + from slime.backends.sglang_utils.arguments import validate_args as sglang_validate_args + sglang_validate_args(args) elif getattr(args, "rollout_backend", "sglang") == "vllm": # Set sglang aliases that the rest of the codebase expects From 8addb37546e33e55f22ee977e6d2c958bf393646 Mon Sep 17 00:00:00 2001 From: knlnguyen1802 Date: Tue, 31 Mar 2026 14:42:28 +0800 Subject: [PATCH 13/39] Fix bug Signed-off-by: knlnguyen1802 --- slime/backends/megatron_utils/sglang.py | 25 +- .../megatron_utils/weight_sync_utils.py | 268 ++++++++++++++++++ 2 files changed, 288 insertions(+), 5 deletions(-) create mode 100644 slime/backends/megatron_utils/weight_sync_utils.py diff --git a/slime/backends/megatron_utils/sglang.py b/slime/backends/megatron_utils/sglang.py index 97c82a31cd..4c045c4817 100644 --- a/slime/backends/megatron_utils/sglang.py +++ b/slime/backends/megatron_utils/sglang.py @@ -1,4 +1,9 @@ # the file to manage all sglang deps in the megatron actor +# When sglang is installed we prefer its implementations; otherwise we +# fall back to API-compatible reimplementations in +# slime.backends.megatron_utils.weight_sync_utils. + +# ── FP8 quantisation helpers (sglang-only, no local fallback) ─────── try: from sglang.srt.layers.quantization.fp8_utils import quant_weight_ue8m0, transform_scale_ue8m0 from sglang.srt.model_loader.utils import should_deepgemm_weight_requant_ue8m0 @@ -7,19 +12,29 @@ transform_scale_ue8m0 = None should_deepgemm_weight_requant_ue8m0 = None +# ── monkey_patch_torch_reductions ─────────────────────────────────── try: from sglang.srt.utils.patch_torch import monkey_patch_torch_reductions except ImportError: - from sglang.srt.patch_torch import monkey_patch_torch_reductions - - -from sglang.srt.utils import MultiprocessingSerializer + try: + from sglang.srt.patch_torch import monkey_patch_torch_reductions + except ImportError: + from .weight_sync_utils import monkey_patch_torch_reductions +# ── MultiprocessingSerializer ─────────────────────────────────────── +try: + from sglang.srt.utils import MultiprocessingSerializer +except ImportError: + from .weight_sync_utils import MultiprocessingSerializer +# ── FlattenedTensorBucket ─────────────────────────────────────────── try: from sglang.srt.weight_sync.tensor_bucket import FlattenedTensorBucket # type: ignore[import] except ImportError: - from sglang.srt.model_executor.model_runner import FlattenedTensorBucket # type: ignore[import] + try: + from sglang.srt.model_executor.model_runner import FlattenedTensorBucket # type: ignore[import] + except ImportError: + from .weight_sync_utils import FlattenedTensorBucket __all__ = [ "quant_weight_ue8m0", diff --git a/slime/backends/megatron_utils/weight_sync_utils.py b/slime/backends/megatron_utils/weight_sync_utils.py new file mode 100644 index 0000000000..5fef12e58f --- /dev/null +++ b/slime/backends/megatron_utils/weight_sync_utils.py @@ -0,0 +1,268 @@ +""" +Local reimplementations of sglang weight-sync utilities. + +Used as fallback when sglang is not installed (e.g. vLLM-only mode). +The three classes/functions here are API-compatible with their sglang +counterparts so that the rest of the megatron weight-update code can +work unchanged. + +Origin (sglang): + - FlattenedTensorBucket → sglang.srt.weight_sync.tensor_bucket + - MultiprocessingSerializer / SafeUnpickler → sglang.srt.utils.common + - monkey_patch_torch_reductions → sglang.srt.utils.patch_torch +""" + +from __future__ import annotations + +import base64 +import io +import pickle +from dataclasses import dataclass +from multiprocessing.reduction import ForkingPickler +from typing import Callable, Union + +import torch +from torch.multiprocessing import reductions + +# ── FlattenedTensorBucket ─────────────────────────────────────────── + + +@dataclass +class FlattenedTensorMetadata: + """Metadata for a tensor in a flattened bucket.""" + + name: str + shape: torch.Size + dtype: torch.dtype + start_idx: int + end_idx: int + numel: int + + +class FlattenedTensorBucket: + """ + A bucket that flattens multiple tensors into a single uint8 tensor + for efficient serialisation, while preserving all metadata needed + for reconstruction. + + API-compatible with ``sglang.srt.weight_sync.tensor_bucket.FlattenedTensorBucket``. + """ + + # Checked by callers to decide whether to group tensors by dtype. + supports_multi_dtypes = True + + def __init__( + self, + named_tensors: list[tuple[str, torch.Tensor]] | None = None, + flattened_tensor: torch.Tensor | None = None, + metadata: list[FlattenedTensorMetadata] | None = None, + ): + if named_tensors is not None: + if not named_tensors: + raise ValueError("Cannot create empty tensor bucket") + + self.metadata: list[FlattenedTensorMetadata] = [None] * len(named_tensors) + current_idx = 0 + flat_parts: list[torch.Tensor] = [None] * len(named_tensors) + + for i, (name, tensor) in enumerate(named_tensors): + flat = tensor.flatten().view(torch.uint8) + numel = flat.numel() + flat_parts[i] = flat + self.metadata[i] = FlattenedTensorMetadata( + name=name, + shape=tensor.shape, + dtype=tensor.dtype, + start_idx=current_idx, + end_idx=current_idx + numel, + numel=numel, + ) + current_idx += numel + + self.flattened_tensor: torch.Tensor = torch.cat(flat_parts, dim=0) + else: + if flattened_tensor is None or metadata is None: + raise ValueError( + "Must provide either named_tensors or both flattened_tensor and metadata" + ) + self.flattened_tensor = flattened_tensor + self.metadata = metadata + + def get_flattened_tensor(self) -> torch.Tensor: + """Return the single flat uint8 tensor.""" + return self.flattened_tensor + + def get_metadata(self) -> list[FlattenedTensorMetadata]: + """Return per-tensor metadata list.""" + return self.metadata + + def reconstruct_tensors(self) -> list[tuple[str, torch.Tensor]]: + """Reconstruct the original named tensors from the flat representation.""" + reconstructed = [None] * len(self.metadata) + for i, meta in enumerate(self.metadata): + tensor = ( + self.flattened_tensor[meta.start_idx : meta.end_idx] + .view(meta.dtype) + .reshape(meta.shape) + ) + reconstructed[i] = (meta.name, tensor) + return reconstructed + + +# ── SafeUnpickler / MultiprocessingSerializer ─────────────────────── + + +class SafeUnpickler(pickle.Unpickler): + """ + Unpickler with an allow-list to prevent arbitrary code execution. + + API-compatible with the ``SafeUnpickler`` in ``sglang.srt.utils.common``. + """ + + ALLOWED_MODULE_PREFIXES = { + # Python builtins + "builtins.", + "collections.", + "copyreg.", + "functools.", + "itertools.", + "operator.", + "types.", + "weakref.", + # PyTorch + "torch.", + "torch._tensor.", + "torch.storage.", + "torch.nn.parameter.", + "torch.autograd.function.", + # torch.distributed + "torch.distributed.", + "torch.distributed._shard.", + "torch.distributed._composable.", + "torch._C._distributed_c10d.", + "torch._C._distributed_fsdp.", + "torch.distributed.optim.", + # multiprocessing + "multiprocessing.resource_sharer.", + "multiprocessing.reduction.", + "pickletools.", + # HuggingFace / PEFT + "peft.", + "transformers.", + "huggingface_hub.", + # slime local reimplementation + "slime.backends.megatron_utils.weight_sync_utils.", + # sglang (if installed alongside) + "sglang.srt.weight_sync.tensor_bucket.", + "sglang.srt.model_executor.model_runner.", + "sglang.srt.layers.", + "sglang.srt.utils.", + # NPU + "torch_npu.", + } + + DENY_CLASSES = { + ("builtins", "eval"), + ("builtins", "exec"), + ("builtins", "compile"), + ("os", "system"), + ("subprocess", "Popen"), + ("subprocess", "run"), + ("codecs", "decode"), + ("types", "CodeType"), + ("types", "FunctionType"), + } + + def find_class(self, module: str, name: str): + if (module, name) in self.DENY_CLASSES: + raise RuntimeError( + f"Blocked unsafe class loading ({module}.{name}), " + f"to prevent exploitation of CVE-2025-10164" + ) + if any((module + ".").startswith(prefix) for prefix in self.ALLOWED_MODULE_PREFIXES): + return super().find_class(module, name) + raise RuntimeError( + f"Blocked unsafe class loading ({module}.{name}), " + f"to prevent exploitation of CVE-2025-10164" + ) + + +class MultiprocessingSerializer: + """ + Serialize / deserialize Python objects via ``ForkingPickler`` so that + CUDA tensors are transferred through shared memory (IPC handles). + + API-compatible with ``sglang.srt.utils.common.MultiprocessingSerializer``. + + Uses stdlib ``base64`` instead of ``pybase64`` to avoid adding a dependency. + """ + + @staticmethod + def serialize(obj, output_str: bool = False): + buf = io.BytesIO() + ForkingPickler(buf).dump(obj) + buf.seek(0) + output = buf.read() + if output_str: + output = base64.b64encode(output).decode("utf-8") + return output + + @staticmethod + def deserialize(data): + if isinstance(data, str): + data = base64.b64decode(data, validate=True) + return SafeUnpickler(io.BytesIO(data)).load() + + +# ── monkey_patch_torch_reductions ─────────────────────────────────── + +_REDUCE_TENSOR_ARG_DEVICE_INDEX = 6 + + +def _device_to_uuid(device: int) -> str: + return str(torch.cuda.get_device_properties(device).uuid) + + +def _device_from_maybe_uuid(device_maybe_uuid: Union[int, str]) -> int: + if isinstance(device_maybe_uuid, int): + return device_maybe_uuid + if isinstance(device_maybe_uuid, str): + for device in range(torch.cuda.device_count()): + if str(torch.cuda.get_device_properties(device).uuid) == device_maybe_uuid: + return device + raise RuntimeError("Invalid device_uuid=" + device_maybe_uuid) + raise RuntimeError(f"Unknown type: {device_maybe_uuid=}") + + +def _modify_tuple(t, index: int, modifier: Callable): + return (*t[:index], modifier(t[index]), *t[index + 1 :]) + + +def _reduce_tensor_modified(*args, **kwargs): + output_fn, output_args = reductions._reduce_tensor_original(*args, **kwargs) + output_args = _modify_tuple(output_args, _REDUCE_TENSOR_ARG_DEVICE_INDEX, _device_to_uuid) + return output_fn, output_args + + +def _rebuild_cuda_tensor_modified(*args): + args = _modify_tuple(args, _REDUCE_TENSOR_ARG_DEVICE_INDEX, _device_from_maybe_uuid) + return reductions._rebuild_cuda_tensor_original(*args) + + +def monkey_patch_torch_reductions(): + """ + Monkey-patch ``torch.multiprocessing.reductions`` so that CUDA tensors + are identified by device UUID rather than ordinal index. + + This works around https://github.com/pytorch/pytorch/pull/149248. + + API-compatible with ``sglang.srt.utils.patch_torch.monkey_patch_torch_reductions``. + """ + if hasattr(reductions, "_reduce_tensor_original"): + return # already patched + reductions._reduce_tensor_original = reductions.reduce_tensor + reductions._rebuild_cuda_tensor_original = reductions.rebuild_cuda_tensor + + reductions.reduce_tensor = _reduce_tensor_modified + reductions.rebuild_cuda_tensor = _rebuild_cuda_tensor_modified + reductions.init_reductions() From 9689a970eb3174f26203454722cd8321103b9490 Mon Sep 17 00:00:00 2001 From: knlnguyen1802 Date: Tue, 31 Mar 2026 14:43:08 +0800 Subject: [PATCH 14/39] Fix bug Signed-off-by: knlnguyen1802 --- slime/utils/arguments.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/slime/utils/arguments.py b/slime/utils/arguments.py index 7874237abc..5ef8424067 100644 --- a/slime/utils/arguments.py +++ b/slime/utils/arguments.py @@ -1500,7 +1500,7 @@ def parse_args(add_custom_arguments=None): args = megatron_parse_args( extra_args_provider=add_slime_arguments, - skip_hf_validate=pre.debug_rollout_only, + skip_hf_validate=True, #pre.debug_rollout_only, ) else: logger.warning( From 3753432aaea1d70402e05f5343188a1c6972a8e3 Mon Sep 17 00:00:00 2001 From: knlnguyen1802 Date: Tue, 31 Mar 2026 14:59:43 +0800 Subject: [PATCH 15/39] Fix bug Signed-off-by: knlnguyen1802 --- slime/backends/megatron_utils/model_provider.py | 12 ++++++++---- slime/utils/http_utils.py | 3 ++- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/slime/backends/megatron_utils/model_provider.py b/slime/backends/megatron_utils/model_provider.py index 31db8b0da8..cc7da48133 100644 --- a/slime/backends/megatron_utils/model_provider.py +++ b/slime/backends/megatron_utils/model_provider.py @@ -209,12 +209,16 @@ def model_provider(pre_process: bool = True, post_process: bool = True, vp_stage def wrap_model_provider_with_freeze(original_provider, args): - def wrapped_provider(pre_process=True, post_process=True, vp_stage=None): + def wrapped_provider(pre_process=True, post_process=True, vp_stage=None, **kwargs): sig = inspect.signature(original_provider) + call_kwargs = {"pre_process": pre_process, "post_process": post_process} if "vp_stage" in sig.parameters: - model = original_provider(pre_process=pre_process, post_process=post_process, vp_stage=vp_stage) - else: - model = original_provider(pre_process=pre_process, post_process=post_process) + call_kwargs["vp_stage"] = vp_stage + # Forward any extra kwargs (e.g. config) accepted by the provider + for k, v in kwargs.items(): + if k in sig.parameters: + call_kwargs[k] = v + model = original_provider(**call_kwargs) freeze_model_params(model, args) diff --git a/slime/utils/http_utils.py b/slime/utils/http_utils.py index d7807f3b7a..f393cbb50b 100644 --- a/slime/utils/http_utils.py +++ b/slime/utils/http_utils.py @@ -204,7 +204,8 @@ def init_http_client(args): if not args.rollout_num_gpus: return - _client_concurrency = args.sglang_server_concurrency * args.rollout_num_gpus // args.rollout_num_gpus_per_engine + server_concurrency = getattr(args, "sglang_server_concurrency", 256) + _client_concurrency = server_concurrency * args.rollout_num_gpus // args.rollout_num_gpus_per_engine if _http_client is None: _http_client = httpx.AsyncClient( limits=httpx.Limits(max_connections=_client_concurrency), From f1e75542f292bcfdcddb977b5688c45fe9a5c2cf Mon Sep 17 00:00:00 2001 From: knlnguyen1802 Date: Tue, 31 Mar 2026 15:10:46 +0800 Subject: [PATCH 16/39] Fix port Signed-off-by: knlnguyen1802 --- slime/ray/rollout.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/slime/ray/rollout.py b/slime/ray/rollout.py index ee6d2da50f..149780e188 100644 --- a/slime/ray/rollout.py +++ b/slime/ray/rollout.py @@ -973,14 +973,14 @@ def _start_router(args, *, has_pd_disaggregation: bool = False, force_new: bool ``force_new`` is False, skip launching and return the existing values. When ``force_new`` is True (multi-model), always allocate a fresh port. """ - if not force_new and args.sglang_router_ip is not None: + if not force_new and getattr(args, "sglang_router_ip", None) is not None: return args.sglang_router_ip, args.sglang_router_port router_ip = _wrap_ipv6(get_host_info()[1]) if force_new: router_port = find_available_port(random.randint(3000, 4000)) else: - router_port = args.sglang_router_port + router_port = getattr(args, "sglang_router_port", None) if router_port is None: router_port = find_available_port(random.randint(3000, 4000)) From 91cc780f764a775fc5ded42e6695cccf166b9827 Mon Sep 17 00:00:00 2001 From: knlnguyen1802 Date: Tue, 31 Mar 2026 17:50:07 +0800 Subject: [PATCH 17/39] Fix config Signed-off-by: knlnguyen1802 --- docs/en/examples/deepseek-r1.md | 4 ++-- docs/en/examples/qwen3-4B.md | 4 ++-- docs/zh/examples/deepseek-r1.md | 4 ++-- docs/zh/examples/qwen3-4B.md | 4 ++-- examples/fully_async/fully_async_rollout.py | 2 +- examples/tau-bench/run_qwen3_4B.sh | 2 +- run-qwen2.5-0.5B-vllm.sh | 2 +- scripts/low_precision/run-kimi-k2-Thinking-int4.sh | 2 +- scripts/run-deepseek-r1.sh | 2 +- scripts/run-kimi-k2-Instruct.sh | 2 +- scripts/run-kimi-k2-Thinking.sh | 2 +- slime/backends/sglang_utils/arguments.py | 1 - slime/rollout/sglang_rollout.py | 2 +- slime/router/router.py | 2 +- slime/utils/arguments.py | 7 +++++++ slime/utils/http_utils.py | 3 +-- 16 files changed, 25 insertions(+), 20 deletions(-) diff --git a/docs/en/examples/deepseek-r1.md b/docs/en/examples/deepseek-r1.md index ab10502eec..48cca18d31 100644 --- a/docs/en/examples/deepseek-r1.md +++ b/docs/en/examples/deepseek-r1.md @@ -171,7 +171,7 @@ OPTIMIZER_ARGS=( These are the parameters required by sglang. Here, `--rollout-num-gpus-per-engine` basically corresponds to sglang's `tp_size`. Other sglang parameters are passed to slime by adding a `--sglang-` prefix. To fully leverage sglang's large EP inference capabilities, we have added configurations like ep64, dp\_attention dp8, and deepep mode auto. -The final `--sglang-server-concurrency` is a parameter specific to slime. It is used to prevent the sglang server's concurrent requests from becoming too large and crashing the HTTP server. The default is 512. However, since we now have one server for 8 nodes, we have adjusted it to 1024 to ensure that each dp rank can have a concurrency of 128. +The final `--server-concurrency` is a parameter specific to slime. It is used to prevent the sglang server's concurrent requests from becoming too large and crashing the HTTP server. The default is 512. However, since we now have one server for 8 nodes, we have adjusted it to 1024 to ensure that each dp rank can have a concurrency of 128. ```bash SGLANG_ARGS=( @@ -190,7 +190,7 @@ SGLANG_ARGS=( --sglang-deepep-mode auto # make every dp rank have 128 concurrency - --sglang-server-concurrency 1024 + --server-concurrency 1024 ) ``` diff --git a/docs/en/examples/qwen3-4B.md b/docs/en/examples/qwen3-4B.md index 56711878f6..8a4022ab92 100644 --- a/docs/en/examples/qwen3-4B.md +++ b/docs/en/examples/qwen3-4B.md @@ -280,10 +280,10 @@ In this case, 2 GPUs will be allocated for training, and 6 GPUs will be allocate ⚠️ If the concurrency on each sglang server is too high, it may exceed sglang's default CUDA graph concurrency limit (the default maximum is 160), which will affect inference speed. You can adjust this in the following two ways: -1. Use `--sglang-server-concurrency` to limit the maximum number of concurrent requests sent to a single sglang server. For example: +1. Use `--server-concurrency` to limit the maximum number of concurrent requests sent to a single sglang server. For example: ```bash - --sglang-server-concurrency 160 + --server-concurrency 160 ``` 2. Use `--sglang-cuda-graph-bs` (which corresponds to sglang's native `--cuda-graph-bs` argument) to increase the number of CUDA graphs initialized by sglang. For example: diff --git a/docs/zh/examples/deepseek-r1.md b/docs/zh/examples/deepseek-r1.md index 368653844d..fba2e736a3 100644 --- a/docs/zh/examples/deepseek-r1.md +++ b/docs/zh/examples/deepseek-r1.md @@ -171,7 +171,7 @@ OPTIMIZER_ARGS=( sglang 所需的参数，这里 `--rollout-num-gpus-per-engine` 基本对应 sglang 的 `tp_size`，除此之外的 sglang 参数均通过添加 `--sglang-` 的前缀来传给 slime。为了充分利用 sglang 的大 EP 推理能力，我们加上了 ep64、dp_attention dp8、deepep mode auto 等配置。 -最后的 `--sglang-server-concurrency` 是 slime 的特有参数，是为了方式同时发给 sglang server 的并发太大打爆 http server，默认为 512。但是我们现在是 8 机一个 server，为了保证每个 dp rank 能有 128 的并发，我们调整为 1024。 +最后的 `--server-concurrency` 是 slime 的特有参数，是为了方式同时发给 sglang server 的并发太大打爆 http server，默认为 512。但是我们现在是 8 机一个 server，为了保证每个 dp rank 能有 128 的并发，我们调整为 1024。 ```bash SGLANG_ARGS=( @@ -190,7 +190,7 @@ SGLANG_ARGS=( --sglang-deepep-mode auto # make every dp rank has 128 concurrency - --sglang-server-concurrency 1024 + --server-concurrency 1024 ) ``` diff --git a/docs/zh/examples/qwen3-4B.md b/docs/zh/examples/qwen3-4B.md index 2afc70bfbe..ac0fad5b60 100644 --- a/docs/zh/examples/qwen3-4B.md +++ b/docs/zh/examples/qwen3-4B.md @@ -280,10 +280,10 @@ ray job submit ... \ ⚠️ 在进行训推分离的时候，每个 sglang server 上的并发度太大，超过了 sglang 默认的 cuda graph 的并发度（默认最大 160），影响推理速度。可以用以下 2 种方式进行调整： -1. 通过 `--sglang-server-concurrency` 限制发给一个 sglang server 的最大并发量，例如： +1. 通过 `--server-concurrency` 限制发给一个 sglang server 的最大并发量，例如： ```bash - --sglang-server-concurrency 160 + --server-concurrency 160 ``` 2. 使用 `--sglang-cuda-graph-bs`，即 sglang 原生的 `--cuda-graph-bs`, 增大 sglang 初始化的 cuda graph 数量，例如： diff --git a/examples/fully_async/fully_async_rollout.py b/examples/fully_async/fully_async_rollout.py index 7208365c18..25b832709d 100644 --- a/examples/fully_async/fully_async_rollout.py +++ b/examples/fully_async/fully_async_rollout.py @@ -20,7 +20,7 @@ def get_global_worker(args, data_buffer): with _worker_lock: if _global_worker is None or not _global_worker.worker_thread.is_alive(): print("Creating new global async worker...") - _global_worker = AsyncRolloutWorker(args, data_buffer, concurrency=args.sglang_server_concurrency) + _global_worker = AsyncRolloutWorker(args, data_buffer, concurrency=args.server_concurrency) _global_worker.start() return _global_worker diff --git a/examples/tau-bench/run_qwen3_4B.sh b/examples/tau-bench/run_qwen3_4B.sh index a821734012..81ad6f63fc 100644 --- a/examples/tau-bench/run_qwen3_4B.sh +++ b/examples/tau-bench/run_qwen3_4B.sh @@ -100,7 +100,7 @@ SGLANG_ARGS=( --rollout-num-gpus-per-engine 1 --sglang-mem-fraction-static 0.7 # If gemini API reports concurrency limit error, set this parameter to reduce the concurrency - # --sglang-server-concurrency 32 + # --server-concurrency 32 ) MISC_ARGS=( diff --git a/run-qwen2.5-0.5B-vllm.sh b/run-qwen2.5-0.5B-vllm.sh index a315590808..3fb6c896a9 100644 --- a/run-qwen2.5-0.5B-vllm.sh +++ b/run-qwen2.5-0.5B-vllm.sh @@ -94,7 +94,7 @@ WANDB_ARGS=( VLLM_ARGS=( --rollout-backend vllm --rollout-num-gpus-per-engine 1 - --sglang-server-concurrency 512 + --server-concurrency 512 --use-slime-router --slime-router-middleware-paths slime.router.middleware_hub.radix_tree_middleware.RadixTreeMiddleware ) diff --git a/scripts/low_precision/run-kimi-k2-Thinking-int4.sh b/scripts/low_precision/run-kimi-k2-Thinking-int4.sh index c41ea3df82..d09a8f3bfa 100644 --- a/scripts/low_precision/run-kimi-k2-Thinking-int4.sh +++ b/scripts/low_precision/run-kimi-k2-Thinking-int4.sh @@ -134,7 +134,7 @@ SGLANG_ARGS=( #--sglang-deepep-mode auto # make every dp rank has 128 concurrency - --sglang-server-concurrency 1024 + --server-concurrency 1024 --use-slime-router ) diff --git a/scripts/run-deepseek-r1.sh b/scripts/run-deepseek-r1.sh index c307d110ec..9b8bc64b08 100644 --- a/scripts/run-deepseek-r1.sh +++ b/scripts/run-deepseek-r1.sh @@ -126,7 +126,7 @@ SGLANG_ARGS=( --sglang-deepep-mode auto # make every dp rank has 128 concurrency - --sglang-server-concurrency 1024 + --server-concurrency 1024 ) MISC_ARGS=( diff --git a/scripts/run-kimi-k2-Instruct.sh b/scripts/run-kimi-k2-Instruct.sh index 3a591b923a..506b576d2d 100644 --- a/scripts/run-kimi-k2-Instruct.sh +++ b/scripts/run-kimi-k2-Instruct.sh @@ -132,7 +132,7 @@ SGLANG_ARGS=( # --sglang-deepep-mode auto # make every dp rank has 128 concurrency - --sglang-server-concurrency 1024 + --server-concurrency 1024 ) diff --git a/scripts/run-kimi-k2-Thinking.sh b/scripts/run-kimi-k2-Thinking.sh index 25cc3c475f..0883177616 100644 --- a/scripts/run-kimi-k2-Thinking.sh +++ b/scripts/run-kimi-k2-Thinking.sh @@ -134,7 +134,7 @@ SGLANG_ARGS=( # --sglang-deepep-mode auto # make every dp rank has 128 concurrency - --sglang-server-concurrency 1024 + --server-concurrency 1024 ) diff --git a/slime/backends/sglang_utils/arguments.py b/slime/backends/sglang_utils/arguments.py index 71543e02d3..9185f53365 100644 --- a/slime/backends/sglang_utils/arguments.py +++ b/slime/backends/sglang_utils/arguments.py @@ -42,7 +42,6 @@ def add_sglang_arguments(parser): """ parser = add_sglang_router_arguments(parser) parser.set_defaults(router_balance_abs_threshold=10, router_balance_rel_threshold=1.2) - parser.add_argument("--sglang-server-concurrency", type=int, default=512) old_add_argument = parser.add_argument diff --git a/slime/rollout/sglang_rollout.py b/slime/rollout/sglang_rollout.py index 0b59be4477..416d384072 100644 --- a/slime/rollout/sglang_rollout.py +++ b/slime/rollout/sglang_rollout.py @@ -76,7 +76,7 @@ def __init__(self, args: Namespace) -> None: self.processor = load_processor(args.hf_checkpoint, trust_remote_code=True) self.semaphore = asyncio.Semaphore( - args.sglang_server_concurrency * args.rollout_num_gpus // args.rollout_num_gpus_per_engine + args.server_concurrency * args.rollout_num_gpus // args.rollout_num_gpus_per_engine ) self.sampling_params: dict[str, Any] = dict( temperature=args.rollout_temperature, diff --git a/slime/router/router.py b/slime/router/router.py index 669094e442..9f660e24b5 100644 --- a/slime/router/router.py +++ b/slime/router/router.py @@ -45,7 +45,7 @@ def __init__(self, args, verbose=False): max_connections = getattr(args, "slime_router_max_connections", None) if max_connections is None: max_connections = ( - args.sglang_server_concurrency * args.rollout_num_gpus // args.rollout_num_gpus_per_engine + args.server_concurrency * args.rollout_num_gpus // args.rollout_num_gpus_per_engine ) timeout = getattr(args, "slime_router_timeout", None) diff --git a/slime/utils/arguments.py b/slime/utils/arguments.py index 5ef8424067..2699072f6a 100644 --- a/slime/utils/arguments.py +++ b/slime/utils/arguments.py @@ -267,6 +267,13 @@ def add_rollout_arguments(parser): help="Fraction of GPU memory for vLLM KV cache (default: 0.85). " "Lower this if the training model leaves insufficient free memory.", ) + parser.add_argument( + "--server-concurrency", + type=int, + default=512, + help="Maximum number of concurrent requests sent to the rollout server. " + "Controls request parallelism for any backend (sglang or vllm).", + ) parser.add_argument( "--rollout-function-path", type=str, diff --git a/slime/utils/http_utils.py b/slime/utils/http_utils.py index f393cbb50b..d8b92dc696 100644 --- a/slime/utils/http_utils.py +++ b/slime/utils/http_utils.py @@ -204,8 +204,7 @@ def init_http_client(args): if not args.rollout_num_gpus: return - server_concurrency = getattr(args, "sglang_server_concurrency", 256) - _client_concurrency = server_concurrency * args.rollout_num_gpus // args.rollout_num_gpus_per_engine + _client_concurrency = args.server_concurrency * args.rollout_num_gpus // args.rollout_num_gpus_per_engine if _http_client is None: _http_client = httpx.AsyncClient( limits=httpx.Limits(max_connections=_client_concurrency), From be1ecd490a77bf59b897b9929c18da72f4f6a2de Mon Sep 17 00:00:00 2001 From: knlnguyen1802 Date: Wed, 1 Apr 2026 14:14:15 +0800 Subject: [PATCH 18/39] Fix bug MOE weight sync --- .../megatron_utils/update_weight/common.py | 46 +++++++++++++++---- 1 file changed, 38 insertions(+), 8 deletions(-) diff --git a/slime/backends/megatron_utils/update_weight/common.py b/slime/backends/megatron_utils/update_weight/common.py index a2e4e129bc..e0a24cf1c9 100644 --- a/slime/backends/megatron_utils/update_weight/common.py +++ b/slime/backends/megatron_utils/update_weight/common.py @@ -12,6 +12,36 @@ from slime.utils.types import ParamInfo +def _cat_partitions( + partitions: list[torch.Tensor], partition_dim: int, stride: int +) -> torch.Tensor: + """Concatenate TP-gathered partitions, handling interleaved (stride > 1) layouts. + + When ``stride == 1`` each rank holds a single contiguous shard and a plain + ``torch.cat`` along *partition_dim* is sufficient. + + When ``stride > 1`` (e.g. Megatron-Core GroupedMLP expert weights) each + rank's shard contains multiple interleaved groups of ``stride`` consecutive + elements. We split each shard into those groups and interleave across + ranks to reconstruct the original full tensor. + """ + if stride == 1: + return torch.cat(partitions, dim=partition_dim) + + # Number of contiguous groups stored on each rank + num_groups = partitions[0].shape[partition_dim] // stride + if num_groups <= 0: + # Fallback: stride >= shard size — simple concat should still be correct + return torch.cat(partitions, dim=partition_dim) + + rank_groups = [p.chunk(num_groups, dim=partition_dim) for p in partitions] + interleaved: list[torch.Tensor] = [] + for g in range(num_groups): + for r in range(len(partitions)): + interleaved.append(rank_groups[r][g]) + return torch.cat(interleaved, dim=partition_dim) + + def all_gather_param(name: str, param: torch.nn.Parameter) -> torch.Tensor: """ All-gather TP-sharded param to full tensor. expert_bias→param, non-TP/duplicated→param.data. @@ -34,7 +64,7 @@ def all_gather_param(name: str, param: torch.nn.Parameter) -> torch.Tensor: param_partitions = [torch.empty_like(param.data) for _ in range(tp_size)] dist.all_gather(param_partitions, param.data, group=tp_group) partition_dim = param.partition_dim - assert param.partition_stride == 1, "partition_stride != 1 is not supported" + stride = getattr(param, "partition_stride", 1) # TODO: here we did an extra copy during concat, maybe merge this with convert_to_hf is better? # TODO: check only GLU is used. if "linear_fc1.weight" in name: @@ -44,7 +74,7 @@ def all_gather_param(name: str, param: torch.nn.Parameter) -> torch.Tensor: if "linear_fc2.weight" in name: if partition_dim == 0: partition_dim = 1 - param = torch.cat(param_partitions, dim=partition_dim) + param = _cat_partitions(param_partitions, partition_dim, stride) return param @@ -63,10 +93,10 @@ def all_gather_params_async( for info, param in param_infos_and_params: # Prepare async all_gather if "expert_bias" in info.name: - gather_tasks.append((info, param, None, None, None)) + gather_tasks.append((info, param, None, None, None, 1)) handles.append(None) elif not param.tensor_model_parallel or getattr(param, "parallel_mode", None) == "duplicated": - gather_tasks.append((info, param.data, None, None, None)) + gather_tasks.append((info, param.data, None, None, None, 1)) handles.append(None) else: # Start async all_gather @@ -79,7 +109,8 @@ def all_gather_params_async( param_partitions = [torch.empty_like(param.data) for _ in range(tp_size)] handle = dist.all_gather(param_partitions, param.data, group=tp_group, async_op=True) - gather_tasks.append((info, None, handle, param_partitions, param.partition_dim)) + partition_stride = getattr(param, "partition_stride", 1) + gather_tasks.append((info, None, handle, param_partitions, param.partition_dim, partition_stride)) handles.append(handle) # Phase 2: Wait for ALL async operations to complete at once @@ -90,13 +121,12 @@ def all_gather_params_async( # Phase 3: Process all results after all communications are done gathered_params = [] - for info, direct_param, handle, param_partitions, partition_dim in gather_tasks: + for info, direct_param, handle, param_partitions, partition_dim, partition_stride in gather_tasks: if handle is None: # No all_gather needed param = direct_param else: # Process the gathered partitions (same logic as original all_gather_param) - assert partition_dim is not None, "partition_stride != 1 is not supported" # TODO: here we did an extra copy during concat, maybe merge this with convert_to_hf is better? # TODO: check only GLU is used. if "linear_fc1.weight" in info.name: @@ -106,7 +136,7 @@ def all_gather_params_async( if "linear_fc2.weight" in info.name: if partition_dim == 0: partition_dim = 1 - param = torch.cat(param_partitions, dim=partition_dim) + param = _cat_partitions(param_partitions, partition_dim, partition_stride) gathered_params.append(param) From e7216d8356614479db66693818263312f6823195 Mon Sep 17 00:00:00 2001 From: knlnguyen1802 Date: Wed, 1 Apr 2026 14:29:17 +0800 Subject: [PATCH 19/39] Fix bug vllm transfer weight Signed-off-by: knlnguyen1802 --- .../update_weight/update_weight_from_distributed.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py b/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py index 03d480801f..e827fac6ba 100644 --- a/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py +++ b/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py @@ -83,14 +83,18 @@ def _nccl_bridge_worker(conn, master_address, master_port, world_size, device, c elif op == "send_packed": from vllm.distributed.weight_transfer.nccl_engine import ( + NCCLTrainerSendWeightsArgs, NCCLWeightTransferEngine, ) - NCCLWeightTransferEngine.trainer_send_weights( - iterator=iter(cmd["named_tensors"]), + trainer_args = NCCLTrainerSendWeightsArgs( group=comm, packed=True, ) + NCCLWeightTransferEngine.trainer_send_weights( + iterator=iter(cmd["named_tensors"]), + trainer_args=trainer_args, + ) torch.cuda.synchronize() conn.send("ok") From 2343647735876472553fdfee9954f8bc0d96dbed Mon Sep 17 00:00:00 2001 From: knlnguyen1802 Date: Wed, 1 Apr 2026 15:07:07 +0800 Subject: [PATCH 20/39] Fix weight sync Signed-off-by: knlnguyen1802 --- .../megatron_utils/update_weight/common.py | 36 ++++++++++--------- 1 file changed, 19 insertions(+), 17 deletions(-) diff --git a/slime/backends/megatron_utils/update_weight/common.py b/slime/backends/megatron_utils/update_weight/common.py index e0a24cf1c9..35faf97c13 100644 --- a/slime/backends/megatron_utils/update_weight/common.py +++ b/slime/backends/megatron_utils/update_weight/common.py @@ -15,31 +15,33 @@ def _cat_partitions( partitions: list[torch.Tensor], partition_dim: int, stride: int ) -> torch.Tensor: - """Concatenate TP-gathered partitions, handling interleaved (stride > 1) layouts. + """Concatenate TP-gathered partitions, reversing Megatron-Core's strided layout. When ``stride == 1`` each rank holds a single contiguous shard and a plain ``torch.cat`` along *partition_dim* is sufficient. - When ``stride > 1`` (e.g. Megatron-Core GroupedMLP expert weights) each - rank's shard contains multiple interleaved groups of ``stride`` consecutive - elements. We split each shard into those groups and interleave across - ranks to reconstruct the original full tensor. + When ``stride > 1``, Megatron-Core's ``_initialize_affine_weight_cpu`` + splits the master weight into ``world_size * stride`` equal sub-chunks and + assigns rank *r* the sub-chunks at indices ``r, r + world_size, + r + 2 * world_size, …`` (i.e. ``weight_list[rank::world_size]``). Each + rank therefore holds ``stride`` sub-chunks concatenated together. + + To reconstruct the original full tensor we split each rank's shard back + into ``stride`` sub-chunks and interleave them: + ``for i in range(stride): for r in range(world_size): append chunk[r][i]`` """ if stride == 1: return torch.cat(partitions, dim=partition_dim) - # Number of contiguous groups stored on each rank - num_groups = partitions[0].shape[partition_dim] // stride - if num_groups <= 0: - # Fallback: stride >= shard size — simple concat should still be correct - return torch.cat(partitions, dim=partition_dim) - - rank_groups = [p.chunk(num_groups, dim=partition_dim) for p in partitions] - interleaved: list[torch.Tensor] = [] - for g in range(num_groups): - for r in range(len(partitions)): - interleaved.append(rank_groups[r][g]) - return torch.cat(interleaved, dim=partition_dim) + world_size = len(partitions) + # Each rank holds exactly `stride` sub-chunks concatenated along partition_dim. + rank_chunks = [p.chunk(stride, dim=partition_dim) for p in partitions] + # Reconstruct original ordering: chunk index (i * world_size + r) + ordered: list[torch.Tensor] = [] + for i in range(stride): + for r in range(world_size): + ordered.append(rank_chunks[r][i]) + return torch.cat(ordered, dim=partition_dim) def all_gather_param(name: str, param: torch.nn.Parameter) -> torch.Tensor: From 8498b7b497026bb78426a8daa951f882105e600a Mon Sep 17 00:00:00 2001 From: knlnguyen1802 Date: Wed, 1 Apr 2026 15:24:12 +0800 Subject: [PATCH 21/39] Fix Signed-off-by: knlnguyen1802 --- slime/backends/megatron_utils/update_weight/common.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/slime/backends/megatron_utils/update_weight/common.py b/slime/backends/megatron_utils/update_weight/common.py index 35faf97c13..529bd793d6 100644 --- a/slime/backends/megatron_utils/update_weight/common.py +++ b/slime/backends/megatron_utils/update_weight/common.py @@ -72,6 +72,9 @@ def all_gather_param(name: str, param: torch.nn.Parameter) -> torch.Tensor: if "linear_fc1.weight" in name: param_partitions = [p.chunk(2, dim=0) for p in param_partitions] param_partitions = [p[0] for p in param_partitions] + [p[1] for p in param_partitions] + # GLU rechunking already reverses the stride=2 interleaving, so use + # plain concatenation from here on. + stride = 1 # this is bug in megatron's grouped moe. if "linear_fc2.weight" in name: if partition_dim == 0: @@ -134,6 +137,8 @@ def all_gather_params_async( if "linear_fc1.weight" in info.name: param_partitions = [p.chunk(2, dim=0) for p in param_partitions] param_partitions = [p[0] for p in param_partitions] + [p[1] for p in param_partitions] + # GLU rechunking already reverses the stride=2 interleaving. + partition_stride = 1 # this is bug in megatron's grouped moe. if "linear_fc2.weight" in info.name: if partition_dim == 0: From 8a41184c35c03c51720b8b67cd2dbbbfdc7e3e9b Mon Sep 17 00:00:00 2001 From: knlnguyen1802 Date: Thu, 2 Apr 2026 13:26:38 +0000 Subject: [PATCH 22/39] Fix config Signed-off-by: knlnguyen1802 --- slime/backends/megatron_utils/arguments.py | 6 ++++++ slime/utils/arguments.py | 2 +- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/slime/backends/megatron_utils/arguments.py b/slime/backends/megatron_utils/arguments.py index 3e6e7a0d19..eb4fad6e60 100644 --- a/slime/backends/megatron_utils/arguments.py +++ b/slime/backends/megatron_utils/arguments.py @@ -50,6 +50,12 @@ def equal(x, y): ("rope_theta", "rotary_base", equal), ]: if hasattr(hf_config, hf_config_name): + if not hasattr(args, megatron_config_name): + logger.warning( + f"Megatron args missing '{megatron_config_name}' (mapped from HF '{hf_config_name}') , " + f"Skip validate" + ) + continue if not compare_fn(getattr(hf_config, hf_config_name), getattr(args, megatron_config_name)): errors.append( f"{hf_config_name} in hf config {getattr(hf_config, hf_config_name)} is not equal to " diff --git a/slime/utils/arguments.py b/slime/utils/arguments.py index 2699072f6a..7d61297d87 100644 --- a/slime/utils/arguments.py +++ b/slime/utils/arguments.py @@ -1507,7 +1507,7 @@ def parse_args(add_custom_arguments=None): args = megatron_parse_args( extra_args_provider=add_slime_arguments, - skip_hf_validate=True, #pre.debug_rollout_only, + skip_hf_validate=pre.debug_rollout_only, ) else: logger.warning( From acc969034bd5fe2ebcf3213b89a63ae33a2295c2 Mon Sep 17 00:00:00 2001 From: knlnguyen1802 Date: Thu, 2 Apr 2026 13:45:58 +0800 Subject: [PATCH 23/39] Change name config Signed-off-by: knlnguyen1802 --- docs/en/examples/deepseek-r1.md | 4 ++-- docs/en/examples/qwen3-4B.md | 4 ++-- docs/zh/examples/deepseek-r1.md | 4 ++-- docs/zh/examples/qwen3-4B.md | 4 ++-- examples/fully_async/fully_async_rollout.py | 2 +- examples/tau-bench/run_qwen3_4B.sh | 2 +- run-qwen2.5-0.5B-vllm.sh | 2 +- scripts/low_precision/run-kimi-k2-Thinking-int4.sh | 2 +- scripts/run-deepseek-r1.sh | 2 +- scripts/run-kimi-k2-Instruct.sh | 2 +- scripts/run-kimi-k2-Thinking.sh | 2 +- slime/rollout/sglang_rollout.py | 2 +- slime/router/router.py | 2 +- slime/utils/arguments.py | 2 +- slime/utils/http_utils.py | 2 +- 15 files changed, 19 insertions(+), 19 deletions(-) diff --git a/docs/en/examples/deepseek-r1.md b/docs/en/examples/deepseek-r1.md index 48cca18d31..22e3fd814d 100644 --- a/docs/en/examples/deepseek-r1.md +++ b/docs/en/examples/deepseek-r1.md @@ -171,7 +171,7 @@ OPTIMIZER_ARGS=( These are the parameters required by sglang. Here, `--rollout-num-gpus-per-engine` basically corresponds to sglang's `tp_size`. Other sglang parameters are passed to slime by adding a `--sglang-` prefix. To fully leverage sglang's large EP inference capabilities, we have added configurations like ep64, dp\_attention dp8, and deepep mode auto. -The final `--server-concurrency` is a parameter specific to slime. It is used to prevent the sglang server's concurrent requests from becoming too large and crashing the HTTP server. The default is 512. However, since we now have one server for 8 nodes, we have adjusted it to 1024 to ensure that each dp rank can have a concurrency of 128. +The final `--rollout-server-concurrency` is a parameter specific to slime. It is used to prevent the sglang server's concurrent requests from becoming too large and crashing the HTTP server. The default is 512. However, since we now have one server for 8 nodes, we have adjusted it to 1024 to ensure that each dp rank can have a concurrency of 128. ```bash SGLANG_ARGS=( @@ -190,7 +190,7 @@ SGLANG_ARGS=( --sglang-deepep-mode auto # make every dp rank have 128 concurrency - --server-concurrency 1024 + --rollout-server-concurrency 1024 ) ``` diff --git a/docs/en/examples/qwen3-4B.md b/docs/en/examples/qwen3-4B.md index 8a4022ab92..688b30b03c 100644 --- a/docs/en/examples/qwen3-4B.md +++ b/docs/en/examples/qwen3-4B.md @@ -280,10 +280,10 @@ In this case, 2 GPUs will be allocated for training, and 6 GPUs will be allocate ⚠️ If the concurrency on each sglang server is too high, it may exceed sglang's default CUDA graph concurrency limit (the default maximum is 160), which will affect inference speed. You can adjust this in the following two ways: -1. Use `--server-concurrency` to limit the maximum number of concurrent requests sent to a single sglang server. For example: +1. Use `--rollout-server-concurrency` to limit the maximum number of concurrent requests sent to a single sglang server. For example: ```bash - --server-concurrency 160 + --rollout-server-concurrency 160 ``` 2. Use `--sglang-cuda-graph-bs` (which corresponds to sglang's native `--cuda-graph-bs` argument) to increase the number of CUDA graphs initialized by sglang. For example: diff --git a/docs/zh/examples/deepseek-r1.md b/docs/zh/examples/deepseek-r1.md index fba2e736a3..703a7730e4 100644 --- a/docs/zh/examples/deepseek-r1.md +++ b/docs/zh/examples/deepseek-r1.md @@ -171,7 +171,7 @@ OPTIMIZER_ARGS=( sglang 所需的参数，这里 `--rollout-num-gpus-per-engine` 基本对应 sglang 的 `tp_size`，除此之外的 sglang 参数均通过添加 `--sglang-` 的前缀来传给 slime。为了充分利用 sglang 的大 EP 推理能力，我们加上了 ep64、dp_attention dp8、deepep mode auto 等配置。 -最后的 `--server-concurrency` 是 slime 的特有参数，是为了方式同时发给 sglang server 的并发太大打爆 http server，默认为 512。但是我们现在是 8 机一个 server，为了保证每个 dp rank 能有 128 的并发，我们调整为 1024。 +最后的 `--rollout-server-concurrency` 是 slime 的特有参数，是为了方式同时发给 sglang server 的并发太大打爆 http server，默认为 512。但是我们现在是 8 机一个 server，为了保证每个 dp rank 能有 128 的并发，我们调整为 1024。 ```bash SGLANG_ARGS=( @@ -190,7 +190,7 @@ SGLANG_ARGS=( --sglang-deepep-mode auto # make every dp rank has 128 concurrency - --server-concurrency 1024 + --rollout-server-concurrency 1024 ) ``` diff --git a/docs/zh/examples/qwen3-4B.md b/docs/zh/examples/qwen3-4B.md index ac0fad5b60..1f3adc9c51 100644 --- a/docs/zh/examples/qwen3-4B.md +++ b/docs/zh/examples/qwen3-4B.md @@ -280,10 +280,10 @@ ray job submit ... \ ⚠️ 在进行训推分离的时候，每个 sglang server 上的并发度太大，超过了 sglang 默认的 cuda graph 的并发度（默认最大 160），影响推理速度。可以用以下 2 种方式进行调整： -1. 通过 `--server-concurrency` 限制发给一个 sglang server 的最大并发量，例如： +1. 通过 `--rollout-server-concurrency` 限制发给一个 sglang server 的最大并发量，例如： ```bash - --server-concurrency 160 + --rollout-server-concurrency 160 ``` 2. 使用 `--sglang-cuda-graph-bs`，即 sglang 原生的 `--cuda-graph-bs`, 增大 sglang 初始化的 cuda graph 数量，例如： diff --git a/examples/fully_async/fully_async_rollout.py b/examples/fully_async/fully_async_rollout.py index 25b832709d..1df005ed02 100644 --- a/examples/fully_async/fully_async_rollout.py +++ b/examples/fully_async/fully_async_rollout.py @@ -20,7 +20,7 @@ def get_global_worker(args, data_buffer): with _worker_lock: if _global_worker is None or not _global_worker.worker_thread.is_alive(): print("Creating new global async worker...") - _global_worker = AsyncRolloutWorker(args, data_buffer, concurrency=args.server_concurrency) + _global_worker = AsyncRolloutWorker(args, data_buffer, concurrency=args.rollout_server_concurrency) _global_worker.start() return _global_worker diff --git a/examples/tau-bench/run_qwen3_4B.sh b/examples/tau-bench/run_qwen3_4B.sh index 81ad6f63fc..fa36af5526 100644 --- a/examples/tau-bench/run_qwen3_4B.sh +++ b/examples/tau-bench/run_qwen3_4B.sh @@ -100,7 +100,7 @@ SGLANG_ARGS=( --rollout-num-gpus-per-engine 1 --sglang-mem-fraction-static 0.7 # If gemini API reports concurrency limit error, set this parameter to reduce the concurrency - # --server-concurrency 32 + # --rollout-server-concurrency 32 ) MISC_ARGS=( diff --git a/run-qwen2.5-0.5B-vllm.sh b/run-qwen2.5-0.5B-vllm.sh index 3fb6c896a9..b352fa37d0 100644 --- a/run-qwen2.5-0.5B-vllm.sh +++ b/run-qwen2.5-0.5B-vllm.sh @@ -94,7 +94,7 @@ WANDB_ARGS=( VLLM_ARGS=( --rollout-backend vllm --rollout-num-gpus-per-engine 1 - --server-concurrency 512 + --rollout-server-concurrency 512 --use-slime-router --slime-router-middleware-paths slime.router.middleware_hub.radix_tree_middleware.RadixTreeMiddleware ) diff --git a/scripts/low_precision/run-kimi-k2-Thinking-int4.sh b/scripts/low_precision/run-kimi-k2-Thinking-int4.sh index d09a8f3bfa..16bc65eaee 100644 --- a/scripts/low_precision/run-kimi-k2-Thinking-int4.sh +++ b/scripts/low_precision/run-kimi-k2-Thinking-int4.sh @@ -134,7 +134,7 @@ SGLANG_ARGS=( #--sglang-deepep-mode auto # make every dp rank has 128 concurrency - --server-concurrency 1024 + --rollout-server-concurrency 1024 --use-slime-router ) diff --git a/scripts/run-deepseek-r1.sh b/scripts/run-deepseek-r1.sh index 9b8bc64b08..dc6d6cb744 100644 --- a/scripts/run-deepseek-r1.sh +++ b/scripts/run-deepseek-r1.sh @@ -126,7 +126,7 @@ SGLANG_ARGS=( --sglang-deepep-mode auto # make every dp rank has 128 concurrency - --server-concurrency 1024 + --rollout-server-concurrency 1024 ) MISC_ARGS=( diff --git a/scripts/run-kimi-k2-Instruct.sh b/scripts/run-kimi-k2-Instruct.sh index 506b576d2d..0c9fc2c3d4 100644 --- a/scripts/run-kimi-k2-Instruct.sh +++ b/scripts/run-kimi-k2-Instruct.sh @@ -132,7 +132,7 @@ SGLANG_ARGS=( # --sglang-deepep-mode auto # make every dp rank has 128 concurrency - --server-concurrency 1024 + --rollout-server-concurrency 1024 ) diff --git a/scripts/run-kimi-k2-Thinking.sh b/scripts/run-kimi-k2-Thinking.sh index 0883177616..533019d697 100644 --- a/scripts/run-kimi-k2-Thinking.sh +++ b/scripts/run-kimi-k2-Thinking.sh @@ -134,7 +134,7 @@ SGLANG_ARGS=( # --sglang-deepep-mode auto # make every dp rank has 128 concurrency - --server-concurrency 1024 + --rollout-server-concurrency 1024 ) diff --git a/slime/rollout/sglang_rollout.py b/slime/rollout/sglang_rollout.py index 416d384072..f3aec2d55b 100644 --- a/slime/rollout/sglang_rollout.py +++ b/slime/rollout/sglang_rollout.py @@ -76,7 +76,7 @@ def __init__(self, args: Namespace) -> None: self.processor = load_processor(args.hf_checkpoint, trust_remote_code=True) self.semaphore = asyncio.Semaphore( - args.server_concurrency * args.rollout_num_gpus // args.rollout_num_gpus_per_engine + args.rollout_server_concurrency * args.rollout_num_gpus // args.rollout_num_gpus_per_engine ) self.sampling_params: dict[str, Any] = dict( temperature=args.rollout_temperature, diff --git a/slime/router/router.py b/slime/router/router.py index 9f660e24b5..b77dd975f4 100644 --- a/slime/router/router.py +++ b/slime/router/router.py @@ -45,7 +45,7 @@ def __init__(self, args, verbose=False): max_connections = getattr(args, "slime_router_max_connections", None) if max_connections is None: max_connections = ( - args.server_concurrency * args.rollout_num_gpus // args.rollout_num_gpus_per_engine + args.rollout_server_concurrency * args.rollout_num_gpus // args.rollout_num_gpus_per_engine ) timeout = getattr(args, "slime_router_timeout", None) diff --git a/slime/utils/arguments.py b/slime/utils/arguments.py index 7d61297d87..b7abc68b98 100644 --- a/slime/utils/arguments.py +++ b/slime/utils/arguments.py @@ -268,7 +268,7 @@ def add_rollout_arguments(parser): "Lower this if the training model leaves insufficient free memory.", ) parser.add_argument( - "--server-concurrency", + "--rollout-server-concurrency", type=int, default=512, help="Maximum number of concurrent requests sent to the rollout server. " diff --git a/slime/utils/http_utils.py b/slime/utils/http_utils.py index d8b92dc696..e6ebe6bf83 100644 --- a/slime/utils/http_utils.py +++ b/slime/utils/http_utils.py @@ -204,7 +204,7 @@ def init_http_client(args): if not args.rollout_num_gpus: return - _client_concurrency = args.server_concurrency * args.rollout_num_gpus // args.rollout_num_gpus_per_engine + _client_concurrency = args.rollout_server_concurrency * args.rollout_num_gpus // args.rollout_num_gpus_per_engine if _http_client is None: _http_client = httpx.AsyncClient( limits=httpx.Limits(max_connections=_client_concurrency), From 90934e624557d79f0599b9056d8eac1e113f5564 Mon Sep 17 00:00:00 2001 From: knlnguyen1802 Date: Mon, 13 Apr 2026 14:45:12 +0800 Subject: [PATCH 24/39] Resolve review Signed-off-by: knlnguyen1802 --- .../fsdp_utils/update_weight_utils.py | 22 +++++++-- .../megatron_utils/weight_sync_utils.py | 48 ++++++++++++++++++- slime/ray/rollout.py | 13 +++-- 3 files changed, 74 insertions(+), 9 deletions(-) diff --git a/slime/backends/fsdp_utils/update_weight_utils.py b/slime/backends/fsdp_utils/update_weight_utils.py index 64a6770c06..f56f6ad847 100644 --- a/slime/backends/fsdp_utils/update_weight_utils.py +++ b/slime/backends/fsdp_utils/update_weight_utils.py @@ -20,19 +20,33 @@ def _import_sglang_weight_sync_utils(): """Lazy-import SGLang serialization utilities. Centralizes the try/except version fallbacks so callers get a clean tuple. - Raises ImportError if sglang is not installed at all. + Falls back to local reimplementations in + ``slime.backends.megatron_utils.weight_sync_utils`` when sglang is not + installed (e.g. vLLM-only environments). """ + # ── monkey_patch_torch_reductions ── try: from sglang.srt.utils.patch_torch import monkey_patch_torch_reductions # type: ignore[import] except ImportError: - from sglang.srt.patch_torch import monkey_patch_torch_reductions # type: ignore[import] + try: + from sglang.srt.patch_torch import monkey_patch_torch_reductions # type: ignore[import] + except ImportError: + from slime.backends.megatron_utils.weight_sync_utils import monkey_patch_torch_reductions - from sglang.srt.utils import MultiprocessingSerializer + # ── MultiprocessingSerializer ── + try: + from sglang.srt.utils import MultiprocessingSerializer + except ImportError: + from slime.backends.megatron_utils.weight_sync_utils import MultiprocessingSerializer + # ── FlattenedTensorBucket ── try: from sglang.srt.weight_sync.tensor_bucket import FlattenedTensorBucket # type: ignore[import] except ImportError: - from sglang.srt.model_executor.model_runner import FlattenedTensorBucket # type: ignore[import] + try: + from sglang.srt.model_executor.model_runner import FlattenedTensorBucket # type: ignore[import] + except ImportError: + from slime.backends.megatron_utils.weight_sync_utils import FlattenedTensorBucket return monkey_patch_torch_reductions, MultiprocessingSerializer, FlattenedTensorBucket diff --git a/slime/backends/megatron_utils/weight_sync_utils.py b/slime/backends/megatron_utils/weight_sync_utils.py index 5fef12e58f..1626444ad8 100644 --- a/slime/backends/megatron_utils/weight_sync_utils.py +++ b/slime/backends/megatron_utils/weight_sync_utils.py @@ -120,8 +120,7 @@ class SafeUnpickler(pickle.Unpickler): """ ALLOWED_MODULE_PREFIXES = { - # Python builtins - "builtins.", + # Python builtins (specific safe classes only – see ALLOW_CLASSES) "collections.", "copyreg.", "functools.", @@ -161,10 +160,47 @@ class SafeUnpickler(pickle.Unpickler): "torch_npu.", } + # Specific builtins classes that are safe to unpickle. + ALLOW_CLASSES = { + ("builtins", "True"), + ("builtins", "False"), + ("builtins", "None"), + ("builtins", "dict"), + ("builtins", "list"), + ("builtins", "tuple"), + ("builtins", "set"), + ("builtins", "frozenset"), + ("builtins", "int"), + ("builtins", "float"), + ("builtins", "complex"), + ("builtins", "str"), + ("builtins", "bytes"), + ("builtins", "bytearray"), + ("builtins", "bool"), + ("builtins", "slice"), + ("builtins", "range"), + ("builtins", "enumerate"), + ("builtins", "map"), + ("builtins", "zip"), + ("builtins", "filter"), + ("builtins", "reversed"), + ("builtins", "sorted"), + } + DENY_CLASSES = { ("builtins", "eval"), ("builtins", "exec"), ("builtins", "compile"), + ("builtins", "getattr"), + ("builtins", "setattr"), + ("builtins", "delattr"), + ("builtins", "__import__"), + ("builtins", "globals"), + ("builtins", "locals"), + ("builtins", "open"), + ("builtins", "breakpoint"), + ("builtins", "input"), + ("builtins", "memoryview"), ("os", "system"), ("subprocess", "Popen"), ("subprocess", "run"), @@ -179,6 +215,14 @@ def find_class(self, module: str, name: str): f"Blocked unsafe class loading ({module}.{name}), " f"to prevent exploitation of CVE-2025-10164" ) + # Check explicit allow-list for builtins (strict whitelist) + if module == "builtins": + if (module, name) in self.ALLOW_CLASSES: + return super().find_class(module, name) + raise RuntimeError( + f"Blocked unsafe class loading ({module}.{name}), " + f"to prevent exploitation of CVE-2025-10164" + ) if any((module + ".").startswith(prefix) for prefix in self.ALLOWED_MODULE_PREFIXES): return super().find_class(module, name) raise RuntimeError( diff --git a/slime/ray/rollout.py b/slime/ray/rollout.py index 149780e188..ebdc97d59a 100644 --- a/slime/ray/rollout.py +++ b/slime/ray/rollout.py @@ -20,7 +20,6 @@ GPU_MEMORY_TYPE_WEIGHTS = "weights" GPU_MEMORY_TYPE_CUDA_GRAPH = "cuda_graph" -from slime.backends.vllm_utils.vllm_engine import VLLMEngine from slime.rollout.base_types import call_rollout_fn from slime.utils import logging_utils from slime.utils.health_monitor import RolloutHealthMonitor @@ -238,9 +237,14 @@ def start_engines(self, port_cursors: dict[int, int] | None = None) -> tuple[lis pg, reordered_bundle_indices, reordered_gpu_ids = self.pg - from slime.backends.sglang_utils.sglang_engine import SGLangEngine + if getattr(self.args, "rollout_backend", "sglang") == "vllm": + from slime.backends.vllm_utils.vllm_engine import VLLMEngine - RolloutRayActor = ray.remote(SGLangEngine) + RolloutRayActor = ray.remote(VLLMEngine) + else: + from slime.backends.sglang_utils.sglang_engine import SGLangEngine + + RolloutRayActor = ray.remote(SGLangEngine) rollout_engines = [] for i in range(len(self.all_engines)): @@ -1051,6 +1055,9 @@ def _start_vllm_rollout_servers(args, pg) -> dict[str, RolloutServer]: router_port = None env_vars = {name: "1" for name in NOSET_VISIBLE_DEVICES_ENV_VARS_LIST} + + from slime.backends.vllm_utils.vllm_engine import VLLMEngine + VLLMRayActor = ray.remote(VLLMEngine) engines = [] From 6b4c37388cf715e85e84026f2f3b3275751aa7bb Mon Sep 17 00:00:00 2001 From: knlnguyen1802 Date: Wed, 15 Apr 2026 11:28:17 +0800 Subject: [PATCH 25/39] Try colocated vllm weight Signed-off-by: knlnguyen1802 --- run-qwen2.5-0.5B-vllm.sh | 2 + slime/backends/megatron_utils/actor.py | 9 +- .../update_weight_from_tensor_vllm.py | 281 ++++++++++++++++++ slime/backends/vllm_utils/vllm_engine.py | 20 +- 4 files changed, 310 insertions(+), 2 deletions(-) create mode 100644 slime/backends/megatron_utils/update_weight/update_weight_from_tensor_vllm.py diff --git a/run-qwen2.5-0.5B-vllm.sh b/run-qwen2.5-0.5B-vllm.sh index b352fa37d0..e56de9860d 100644 --- a/run-qwen2.5-0.5B-vllm.sh +++ b/run-qwen2.5-0.5B-vllm.sh @@ -97,6 +97,8 @@ VLLM_ARGS=( --rollout-server-concurrency 512 --use-slime-router --slime-router-middleware-paths slime.router.middleware_hub.radix_tree_middleware.RadixTreeMiddleware + # Uncomment below to enable colocate mode (IPC weight transfer, same GPU) + --colocate ) MISC_ARGS=( diff --git a/slime/backends/megatron_utils/actor.py b/slime/backends/megatron_utils/actor.py index 1f45e5fc94..ad66204414 100644 --- a/slime/backends/megatron_utils/actor.py +++ b/slime/backends/megatron_utils/actor.py @@ -36,6 +36,7 @@ from .update_weight.common import named_params_and_buffers from .update_weight.update_weight_from_distributed import UpdateWeightFromDistributed from .update_weight.update_weight_from_tensor import UpdateWeightFromTensor +from .update_weight.update_weight_from_tensor_vllm import UpdateVLLMWeightFromTensor logging.getLogger("megatron").setLevel(logging.WARNING) @@ -128,8 +129,14 @@ def init( if self.args.vocab_size is None: self.args.vocab_size = self.tokenizer.vocab_size + use_vllm_colocate = self.args.colocate and getattr(self.args, "rollout_backend", "sglang") == "vllm" use_tensor_update = self.args.colocate and getattr(self.args, "rollout_backend", "sglang") != "vllm" - update_weight_cls = UpdateWeightFromTensor if use_tensor_update else UpdateWeightFromDistributed + if use_vllm_colocate: + update_weight_cls = UpdateVLLMWeightFromTensor + elif use_tensor_update: + update_weight_cls = UpdateWeightFromTensor + else: + update_weight_cls = UpdateWeightFromDistributed self.weight_updater = update_weight_cls( self.args, self.model, diff --git a/slime/backends/megatron_utils/update_weight/update_weight_from_tensor_vllm.py b/slime/backends/megatron_utils/update_weight/update_weight_from_tensor_vllm.py new file mode 100644 index 0000000000..5a6a206ae6 --- /dev/null +++ b/slime/backends/megatron_utils/update_weight/update_weight_from_tensor_vllm.py @@ -0,0 +1,281 @@ +""" +UpdateVLLMWeightFromTensor +========================== + +Update vLLM rollout engines using CUDA IPC (HTTP mode) when colocated on the +same GPU(s) as the trainer. This follows the vLLM RLHF HTTP IPC example: +https://docs.vllm.ai/en/stable/examples/rl/rlhf_http_ipc/ + +The flow: +1. Megatron params → TP all-gather → HF conversion (via HfWeightIteratorBase) +2. Send HF-named tensors to the vLLM server using + ``IPCWeightTransferEngine.trainer_send_weights()`` with + ``IPCTrainerSendWeightsArgs(mode="http", url=)``. +3. For any overflow (non-colocated) engines, fall back to NCCL distributed + broadcast identical to ``UpdateWeightFromDistributed``. + +The API is compatible with ``MegatronTrainRayActor`` — same ``__init__``, +``connect_rollout_engines``, and ``update_weights`` signatures as +``UpdateWeightFromTensor`` / ``UpdateWeightFromDistributed``. +""" + +from __future__ import annotations + +import logging +import os +import time +from argparse import Namespace +from collections.abc import Callable, Mapping, Sequence + +import ray +import torch +import torch.distributed as dist +from megatron.core import mpu +from ray.actor import ActorHandle + +from slime.utils.distributed_utils import get_gloo_group + +from .common import all_gather_param, named_params_and_buffers +from .hf_weight_iterator_base import HfWeightIteratorBase +from .update_weight_from_distributed import ( + connect_rollout_engines_from_distributed, + disconnect_rollout_engines_from_distributed, + post_process_weights, + update_weights_from_distributed, +) + +logger = logging.getLogger(__name__) + + +class UpdateVLLMWeightFromTensor: + """ + Update colocated vLLM engines from tensor via CUDA IPC (HTTP mode). + + Colocated path: + Megatron weights → TP all-gather → HF conversion → CUDA IPC to vLLM + server via ``IPCWeightTransferEngine.trainer_send_weights()``. + + Non-colocated overflow path (optional): + Falls back to NCCL distributed broadcast via ``_NcclBridge``. + """ + + def __init__( + self, + args: Namespace, + model: Sequence[torch.nn.Module], + weights_getter: Callable[[], Mapping[str, torch.Tensor]], + *, + model_name: str, + quantization_config: dict[str, int | str | list[str]] | None, + ) -> None: + self.args = args + self.model = model + self.weights_getter = weights_getter + self.model_name = model_name + self.quantization_config = quantization_config + self.weight_version = 0 + + self._hf_weight_iterator = HfWeightIteratorBase.create( + args=args, model=model, model_name=model_name, quantization_config=quantization_config, + ) + + # Populated by connect_rollout_engines + self.rollout_engines: list[ActorHandle] = [] + self._colocated_engines: list[ActorHandle] = [] + self._colocated_vllm_urls: list[str] = [] + self._distributed_engines: list[ActorHandle] = [] + self._model_update_groups = None + self._is_distributed_src_rank = False + self._group_name = "slime" + self._ipc_initialized = False + + # ------------------------------------------------------------------ + # connect / disconnect + # ------------------------------------------------------------------ + + def connect_rollout_engines( + self, + rollout_engines: Sequence[ActorHandle], + rollout_engine_lock: ActorHandle, + engine_gpu_counts: Sequence[int] | None = None, + engine_gpu_offsets: Sequence[int] | None = None, + ) -> None: + """ + Split colocated / distributed engines. + + For colocated engines we resolve their vLLM base URLs (needed by the + IPC weight transfer HTTP mode). For overflow distributed engines we + create the NCCL bridge just like ``UpdateWeightFromDistributed``. + """ + self.rollout_engines = list(rollout_engines) + self.rollout_engine_lock = rollout_engine_lock + + if engine_gpu_counts is None: + engine_gpu_counts = [self.args.rollout_num_gpus_per_engine] * len(rollout_engines) + if engine_gpu_offsets is None: + engine_gpu_offsets = [] + offset = 0 + for c in engine_gpu_counts: + engine_gpu_offsets.append(offset) + offset += c + + # Determine colocated vs distributed engines + total_actor_gpus = self.args.actor_num_nodes * self.args.actor_num_gpus_per_node + colocate_engine_nums = 0 + for gpu_offset, gpu_count in zip(engine_gpu_offsets, engine_gpu_counts, strict=True): + if gpu_offset + gpu_count > total_actor_gpus: + break + colocate_engine_nums += 1 + + self._colocated_engines = list(rollout_engines[:colocate_engine_nums]) + self._distributed_engines = list(rollout_engines[colocate_engine_nums:]) + + # Resolve vLLM base URLs for colocated engines (blocking Ray call) + if dist.get_rank() == 0 and self._colocated_engines: + url_refs = [engine.get_vllm_url.remote() for engine in self._colocated_engines] + self._colocated_vllm_urls = ray.get(url_refs) + logger.info("Colocated vLLM URLs for IPC weight transfer: %s", self._colocated_vllm_urls) + + # Set up NCCL bridge for distributed overflow engines + if self._distributed_engines: + distributed_gpu_counts = engine_gpu_counts[colocate_engine_nums:] + self._is_distributed_src_rank = ( + mpu.get_data_parallel_rank(with_context_parallel=True) == 0 + and mpu.get_tensor_model_parallel_rank() == 0 + and mpu.get_pipeline_model_parallel_rank() == 0 + ) + if self._is_distributed_src_rank: + if self._model_update_groups is not None: + disconnect_rollout_engines_from_distributed( + self.args, self._group_name, self._model_update_groups, self._distributed_engines, + ) + self._model_update_groups = connect_rollout_engines_from_distributed( + self.args, + self._group_name, + self._distributed_engines, + engine_gpu_counts=distributed_gpu_counts, + ) + + # ------------------------------------------------------------------ + # weight update + # ------------------------------------------------------------------ + + @torch.no_grad() + def update_weights(self) -> None: + """ + Main entry-point called by ``MegatronTrainRayActor.update_weights()``. + + Pause → flush → init IPC (once) → send weights via IPC → resume. + """ + self.weight_version += 1 + rank = dist.get_rank() + + # Pause generation and flush KV cache on all engines + if rank == 0: + all_engines = self._colocated_engines + self._distributed_engines + ray.get([engine.pause_generation.remote() for engine in all_engines]) + ray.get([engine.flush_cache.remote() for engine in all_engines]) + if self.quantization_config and self.quantization_config.get("quant_method") in ["compressed-tensors"]: + post_process_weights( + restore_weights_before_load=True, + post_process_quantization=False, + rollout_engines=all_engines, + ) + dist.barrier(group=get_gloo_group()) + + # Initialize IPC weight transfer engines (first time only) + if rank == 0 and self._colocated_engines and not self._ipc_initialized: + self._init_ipc_weight_transfer() + self._ipc_initialized = True + + # Get megatron weights and iterate HF chunks + megatron_local_weights = self.weights_getter() + + for hf_named_tensors in self._hf_weight_iterator.get_hf_weight_chunks(megatron_local_weights): + # Send to colocated engines via IPC + if rank == 0 and self._colocated_engines: + self._send_via_ipc(hf_named_tensors) + + # Send to distributed engines via NCCL + if self._distributed_engines and self._is_distributed_src_rank: + refs = update_weights_from_distributed( + self._group_name, + self._model_update_groups, + self.weight_version, + self._distributed_engines, + hf_named_tensors, + use_vllm=True, + packed=True, + ) + if refs: + ray.get(refs) + + dist.barrier(group=get_gloo_group()) + + # Post-process and resume + if rank == 0: + all_engines = self._colocated_engines + self._distributed_engines + if self.quantization_config and self.quantization_config.get("quant_method") in ["compressed-tensors"]: + post_process_weights( + restore_weights_before_load=False, + post_process_quantization=True, + rollout_engines=all_engines, + ) + # Bump weight version on sidecar for colocated engines + for engine in self._colocated_engines: + try: + ray.get(engine.set_weight_version.remote(self.weight_version)) + except Exception as exc: + logger.warning("Failed to set weight version on engine: %s", exc) + + ray.get([engine.continue_generation.remote() for engine in all_engines]) + dist.barrier(group=get_gloo_group()) + + # ------------------------------------------------------------------ + # IPC helpers + # ------------------------------------------------------------------ + + def _init_ipc_weight_transfer(self) -> None: + """ + Call ``/init_weight_transfer_engine`` on each colocated vLLM server. + For IPC backend this is a no-op on the server side but is still required + by vLLM's weight transfer protocol. + """ + import requests + + for url in self._colocated_vllm_urls: + try: + resp = requests.post( + f"{url}/init_weight_transfer_engine", + json={"init_info": {}}, + timeout=60, + ) + resp.raise_for_status() + logger.info("Initialized IPC weight transfer on %s", url) + except Exception as exc: + logger.error("Failed to init IPC weight transfer on %s: %s", url, exc) + raise + + def _send_via_ipc(self, hf_named_tensors: list[tuple[str, torch.Tensor]]) -> None: + """ + Send HF-named tensors to all colocated vLLM engines via CUDA IPC. + + Uses ``vllm.distributed.weight_transfer.ipc_engine.IPCWeightTransferEngine`` + with ``mode="http"`` so the IPC handles are sent via HTTP to the vLLM + server's ``/update_weights`` endpoint. + """ + # Allow insecure serialization for IPC handle serialization + os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1" + + from vllm.distributed.weight_transfer.ipc_engine import ( + IPCTrainerSendWeightsArgs, + IPCWeightTransferEngine, + ) + + for url in self._colocated_vllm_urls: + trainer_args = IPCTrainerSendWeightsArgs(mode="http", url=url) + IPCWeightTransferEngine.trainer_send_weights( + iterator=iter(hf_named_tensors), + trainer_args=trainer_args, + ) + logger.debug("IPC weight transfer completed to %s", url) diff --git a/slime/backends/vllm_utils/vllm_engine.py b/slime/backends/vllm_utils/vllm_engine.py index 713793b07c..372425b9f8 100644 --- a/slime/backends/vllm_utils/vllm_engine.py +++ b/slime/backends/vllm_utils/vllm_engine.py @@ -32,6 +32,7 @@ def __init__(self, args, rank: int, base_gpu_id: int | None = None, gpu_ids: lis self._log_file = None self._sidecar_log_file = None self._weight_version: int = 0 + self._colocate: bool = getattr(args, "colocate", False) @property def sidecar_url(self) -> str: @@ -62,12 +63,18 @@ def init(self, port=None, host=None, router_ip=None, router_port=None, **kwargs) dev_str = ",".join(str(g) for g in gpu_ids) seed = getattr(self.args, "seed", 1234) + self.rank + # Use IPC backend when colocated (same GPU), NCCL otherwise + if self._colocate: + weight_transfer_backend = '{"backend": "ipc"}' + else: + weight_transfer_backend = '{"backend": "nccl"}' + cmd = [ "vllm", "serve", model, "--tensor-parallel-size", str(tp), "--port", str(self.server_port), "--host", "0.0.0.0", - "--weight-transfer-config", '{"backend": "nccl"}', + "--weight-transfer-config", weight_transfer_backend, "--seed", str(seed), "--trust-remote-code", ] @@ -82,6 +89,8 @@ def init(self, port=None, host=None, router_ip=None, router_port=None, **kwargs) env = os.environ.copy() env["VLLM_SERVER_DEV_MODE"] = "1" + if self._colocate: + env["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1" env["CUDA_VISIBLE_DEVICES"] = dev_str env.setdefault("NCCL_DEBUG", "INFO") env.setdefault("NCCL_DEBUG_SUBSYS", "ALL") @@ -321,6 +330,15 @@ def get_weight_version(self): pass return self._weight_version + def get_vllm_url(self) -> str: + """Return the raw vLLM server URL (used by IPC weight transfer).""" + return self.vllm_url + + def set_weight_version(self, version: int) -> None: + """Set weight version on both the engine and the sidecar.""" + self._weight_version = version + self._bump_weight_version(version) + def check_weights(self, action: str): pass From 49d760f2cc7c1c41f49c5b7bfb7cc20cbf4821d0 Mon Sep 17 00:00:00 2001 From: knlnguyen1802 Date: Fri, 17 Apr 2026 10:50:48 +0000 Subject: [PATCH 26/39] [Draft] Local runable dev Signed-off-by: knlnguyen1802 --- run-qwen2.5-0.5B-vllm.sh | 12 +++++----- slime/backends/megatron_utils/actor.py | 2 +- .../update_weight_from_tensor_vllm.py | 1 + slime/backends/vllm_utils/vllm_engine.py | 10 ++++---- slime/ray/actor_group.py | 24 +++++++++---------- slime/utils/arguments.py | 2 +- 6 files changed, 27 insertions(+), 24 deletions(-) diff --git a/run-qwen2.5-0.5B-vllm.sh b/run-qwen2.5-0.5B-vllm.sh index e56de9860d..ea010b7e47 100644 --- a/run-qwen2.5-0.5B-vllm.sh +++ b/run-qwen2.5-0.5B-vllm.sh @@ -84,11 +84,11 @@ OPTIMIZER_ARGS=( ) WANDB_ARGS=( - --use-wandb - --wandb-host https://wandb.ai/ - --wandb-entity samithuang - --wandb-project slime-rl - --wandb-group qwen2.5-0.5B-gsm8k-vllm + # --use-wandb + # --wandb-host https://wandb.ai/ + # --wandb-entity samithuang + # --wandb-project slime-rl + # --wandb-group qwen2.5-0.5B-gsm8k-vllm ) VLLM_ARGS=( @@ -115,7 +115,7 @@ ray start --head --node-ip-address 127.0.0.1 --num-gpus 2 --disable-usage-stats ray job submit --address="http://127.0.0.1:8265" \ --runtime-env-json='{ "env_vars": { - "PYTHONPATH": "/root/Megatron-LM", + "PYTHONPATH": "/root/Megatron-LM:/data/n0090/SLIME_PJ/new_version/local_package", "CUDA_DEVICE_MAX_CONNECTIONS": "1", "NCCL_ALGO": "Ring", "NCCL_IB_DISABLE": "1", diff --git a/slime/backends/megatron_utils/actor.py b/slime/backends/megatron_utils/actor.py index ad66204414..f7563116fe 100644 --- a/slime/backends/megatron_utils/actor.py +++ b/slime/backends/megatron_utils/actor.py @@ -564,7 +564,7 @@ def update_weights(self) -> None: if dist.get_rank() == 0: ray.get(self.rollout_manager.clear_num_new_engines.remote()) - with torch_memory_saver.disable() if self.args.offload_train else nullcontext(): + with nullcontext(): #torch_memory_saver.disable() if self.args.offload_train else nullcontext(): print_memory("before update_weights") self.weight_updater.update_weights() print_memory("after update_weights") diff --git a/slime/backends/megatron_utils/update_weight/update_weight_from_tensor_vllm.py b/slime/backends/megatron_utils/update_weight/update_weight_from_tensor_vllm.py index 5a6a206ae6..d58d0377e7 100644 --- a/slime/backends/megatron_utils/update_weight/update_weight_from_tensor_vllm.py +++ b/slime/backends/megatron_utils/update_weight/update_weight_from_tensor_vllm.py @@ -167,6 +167,7 @@ def update_weights(self) -> None: Pause → flush → init IPC (once) → send weights via IPC → resume. """ + self.weight_version += 1 rank = dist.get_rank() diff --git a/slime/backends/vllm_utils/vllm_engine.py b/slime/backends/vllm_utils/vllm_engine.py index 372425b9f8..f431922545 100644 --- a/slime/backends/vllm_utils/vllm_engine.py +++ b/slime/backends/vllm_utils/vllm_engine.py @@ -78,7 +78,7 @@ def init(self, port=None, host=None, router_ip=None, router_port=None, **kwargs) "--seed", str(seed), "--trust-remote-code", ] - gpu_mem_util = getattr(self.args, "vllm_gpu_memory_utilization", 0.4) + gpu_mem_util = getattr(self.args, "vllm_gpu_memory_utilization", 0.25) cmd.extend(["--gpu-memory-utilization", str(gpu_mem_util)]) if getattr(self.args, "offload_rollout", False): cmd.append("--enable-sleep-mode") @@ -97,9 +97,11 @@ def init(self, port=None, host=None, router_ip=None, router_port=None, **kwargs) env["NCCL_P2P_DISABLE"] = "1" env.setdefault("NCCL_IB_DISABLE", "1") - self._log_file = tempfile.NamedTemporaryFile( - prefix="vllm_engine_", suffix=".log", delete=False, mode="w" - ) + # self._log_file = tempfile.NamedTemporaryFile( + # prefix="vllm_engine_", suffix=".log", delete=False, mode="w" + # ) + log_path = "/data/n0090/SLIME_PJ/new_version/slime/vllm_server_log.log" + self._log_file = open(log_path,"w") logger.info("Launching vLLM: cmd=%s, CUDA_VISIBLE_DEVICES=%s, log=%s", " ".join(cmd), dev_str, self._log_file.name) self.process = subprocess.Popen( diff --git a/slime/ray/actor_group.py b/slime/ray/actor_group.py index 1150e9ed35..7cdf639278 100644 --- a/slime/ray/actor_group.py +++ b/slime/ray/actor_group.py @@ -65,18 +65,18 @@ def _allocate_gpus_for_actor(self, pg, num_gpus_per_actor): **self.args.train_env_vars, } - if self.args.offload_train and self.args.train_backend == "megatron": - import torch_memory_saver - - dynlib_path = os.path.join( - os.path.dirname(os.path.dirname(torch_memory_saver.__file__)), - "torch_memory_saver_hook_mode_preload.abi3.so", - ) - assert os.path.exists(dynlib_path), f"LD_PRELOAD so file {dynlib_path} does not exist." - - env_vars["LD_PRELOAD"] = dynlib_path - env_vars["TMS_INIT_ENABLE"] = "1" - env_vars["TMS_INIT_ENABLE_CPU_BACKUP"] = "1" + # if self.args.offload_train and self.args.train_backend == "megatron": + # import torch_memory_saver + + # dynlib_path = os.path.join( + # os.path.dirname(os.path.dirname(torch_memory_saver.__file__)), + # "torch_memory_saver_hook_mode_preload.abi3.so", + # ) + # assert os.path.exists(dynlib_path), f"LD_PRELOAD so file {dynlib_path} does not exist." + + # env_vars["LD_PRELOAD"] = dynlib_path + # env_vars["TMS_INIT_ENABLE"] = "1" + # env_vars["TMS_INIT_ENABLE_CPU_BACKUP"] = "1" # We cannot do routing replay for critic. if self.args.use_routing_replay and self.role == "actor": diff --git a/slime/utils/arguments.py b/slime/utils/arguments.py index b7abc68b98..ad5f61d0cf 100644 --- a/slime/utils/arguments.py +++ b/slime/utils/arguments.py @@ -1762,7 +1762,7 @@ def slime_validate_args(args): args.offload_train = False if args.offload_rollout is None: args.offload_rollout = False - + args.offload_rollout = False if args.eval_function_path is None: args.eval_function_path = args.rollout_function_path From ee2702b568b21c7797ef4785881e979f14605196 Mon Sep 17 00:00:00 2001 From: knlnguyen1802 Date: Tue, 21 Apr 2026 14:44:16 +0800 Subject: [PATCH 27/39] Fix offload train Signed-off-by: knlnguyen1802 --- slime/backends/megatron_utils/actor.py | 86 +++++++++++++++++-- .../megatron_utils/update_weight/common.py | 5 +- slime/ray/actor_group.py | 23 +++++ 3 files changed, 107 insertions(+), 7 deletions(-) diff --git a/slime/backends/megatron_utils/actor.py b/slime/backends/megatron_utils/actor.py index f7563116fe..efc0909087 100644 --- a/slime/backends/megatron_utils/actor.py +++ b/slime/backends/megatron_utils/actor.py @@ -10,9 +10,17 @@ import torch.distributed as dist from megatron.core import mpu from ray.actor import ActorHandle -from torch_memory_saver import torch_memory_saver from transformers import AutoConfig, AutoTokenizer +# torch_memory_saver is only used when the rollout backend is sglang (the +# default). For the vLLM rollout backend we intentionally avoid importing it +# (the preload .so is not loaded), so we resolve it lazily and tolerate the +# import being unavailable. +try: # pragma: no cover - import guard + from torch_memory_saver import torch_memory_saver # type: ignore +except Exception: # noqa: BLE001 - missing native lib should not crash import + torch_memory_saver = None # type: ignore[assignment] + from slime.ray.train_actor import TrainRayActor from slime.utils import train_dump_utils from slime.utils.data import process_rollout_data @@ -79,9 +87,15 @@ def init( dist.barrier(group=get_gloo_group()) if args.offload_train: - if (x := args.train_memory_margin_bytes) > 0: - logger.info(f"Set torch_memory_saver.memory_margin_bytes to {x}") - torch_memory_saver.memory_margin_bytes = x + if self._uses_torch_memory_saver(): + if (x := args.train_memory_margin_bytes) > 0: + logger.info(f"Set torch_memory_saver.memory_margin_bytes to {x}") + torch_memory_saver.memory_margin_bytes = x + else: + logger.info( + "offload_train enabled with rollout_backend=vllm: using CPU offload " + "instead of torch_memory_saver." + ) if role == "critic": self.args.load = self.args.critic_load @@ -165,6 +179,59 @@ def init( return start_rollout_id + def _uses_torch_memory_saver(self) -> bool: + """Whether GPU memory should be released via torch_memory_saver. + + torch_memory_saver relies on an LD_PRELOAD-injected allocator hook that + is only set up for the sglang rollout backend (see ``RayTrainGroup``). + For the vLLM rollout backend we fall back to an explicit CPU-offload + path that moves model parameters and optimizer state to CPU. + """ + if not self.args.offload_train: + return False + if torch_memory_saver is None: + return False + return getattr(self.args, "rollout_backend", "sglang") != "vllm" + + @torch.no_grad() + def _move_optimizer_state_to_device(self, device: str | torch.device) -> None: + """Move all tensor entries in ``self.optimizer.state`` to ``device``. + + Mirrors the helper used by the FSDP backend; works for any optimizer + that exposes the standard ``param_groups`` / ``state`` interface, + including Megatron's DistributedOptimizer. + """ + optimizer = self.optimizer + if optimizer is None or not getattr(optimizer, "state", None): + return + for param_group in optimizer.param_groups: + for param in param_group["params"]: + state = optimizer.state.get(param, None) + if not state: + continue + for key, value in state.items(): + if isinstance(value, torch.Tensor): + state[key] = value.to(device, non_blocking=True) + torch.cuda.synchronize() + + def _cpu_offload_sleep(self) -> None: + """CPU-offload alternative to ``torch_memory_saver.pause()`` for vLLM.""" + if self.model is not None: + for module in self.model: + module.to("cpu", non_blocking=True) + if getattr(self, "optimizer", None) is not None: + self._move_optimizer_state_to_device("cpu") + + def _cpu_offload_wake_up(self) -> None: + """CPU-offload counterpart to ``_cpu_offload_sleep``.""" + device = torch.cuda.current_device() + if self.model is not None: + for module in self.model: + module.to(device, non_blocking=True) + if getattr(self, "optimizer", None) is not None: + self._move_optimizer_state_to_device(device) + torch.cuda.synchronize() + @timer def sleep(self) -> None: assert self.args.offload_train @@ -173,7 +240,11 @@ def sleep(self) -> None: print_memory("before offload model") destroy_process_groups() - torch_memory_saver.pause() + if self._uses_torch_memory_saver(): + torch_memory_saver.pause() + else: + self._cpu_offload_sleep() + clear_memory() print_memory("after offload model") @@ -182,7 +253,10 @@ def wake_up(self) -> None: assert self.args.offload_train print_memory("before wake_up model") - torch_memory_saver.resume() + if self._uses_torch_memory_saver(): + torch_memory_saver.resume() + else: + self._cpu_offload_wake_up() clear_memory() reload_process_groups() diff --git a/slime/backends/megatron_utils/update_weight/common.py b/slime/backends/megatron_utils/update_weight/common.py index 529bd793d6..b37b32a96e 100644 --- a/slime/backends/megatron_utils/update_weight/common.py +++ b/slime/backends/megatron_utils/update_weight/common.py @@ -168,7 +168,10 @@ def named_params_and_buffers( def _maybe_get_cpu_backup(x: torch.Tensor): - from torch_memory_saver import torch_memory_saver + try: + from torch_memory_saver import torch_memory_saver + except Exception: # noqa: BLE001 - native lib may be unavailable (e.g. vLLM backend) + return x if (cpu_tensor := torch_memory_saver.get_cpu_backup(x)) is not None: return cpu_tensor diff --git a/slime/ray/actor_group.py b/slime/ray/actor_group.py index 7cdf639278..ff18f95cdf 100644 --- a/slime/ray/actor_group.py +++ b/slime/ray/actor_group.py @@ -78,6 +78,29 @@ def _allocate_gpus_for_actor(self, pg, num_gpus_per_actor): # env_vars["TMS_INIT_ENABLE"] = "1" # env_vars["TMS_INIT_ENABLE_CPU_BACKUP"] = "1" + # torch_memory_saver requires an LD_PRELOAD'd allocator hook. It is only + # used when offloading the trainer for the sglang rollout backend. With + # the vLLM rollout backend we instead fall back to an explicit CPU + # offload path inside the Megatron actor (see ``MegatronTrainRayActor`` + # ``sleep`` / ``wake_up``), so we must NOT preload the .so here. + rollout_backend = getattr(self.args, "rollout_backend", "sglang") + if ( + self.args.offload_train + and self.args.train_backend == "megatron" + and rollout_backend != "vllm" + ): + import torch_memory_saver + + dynlib_path = os.path.join( + os.path.dirname(os.path.dirname(torch_memory_saver.__file__)), + "torch_memory_saver_hook_mode_preload.abi3.so", + ) + assert os.path.exists(dynlib_path), f"LD_PRELOAD so file {dynlib_path} does not exist." + + env_vars["LD_PRELOAD"] = dynlib_path + env_vars["TMS_INIT_ENABLE"] = "1" + env_vars["TMS_INIT_ENABLE_CPU_BACKUP"] = "1" + # We cannot do routing replay for critic. if self.args.use_routing_replay and self.role == "actor": env_vars["ENABLE_ROUTING_REPLAY"] = "1" From cb77fc5f3f36a9670db83ea29a308131cf73f26b Mon Sep 17 00:00:00 2001 From: knlnguyen1802 Date: Wed, 22 Apr 2026 10:39:44 +0800 Subject: [PATCH 28/39] Fix offload train Signed-off-by: knlnguyen1802 --- slime/backends/megatron_utils/actor.py | 32 ++++++++++++++++---------- 1 file changed, 20 insertions(+), 12 deletions(-) diff --git a/slime/backends/megatron_utils/actor.py b/slime/backends/megatron_utils/actor.py index efc0909087..6bb3335a84 100644 --- a/slime/backends/megatron_utils/actor.py +++ b/slime/backends/megatron_utils/actor.py @@ -197,21 +197,29 @@ def _uses_torch_memory_saver(self) -> bool: def _move_optimizer_state_to_device(self, device: str | torch.device) -> None: """Move all tensor entries in ``self.optimizer.state`` to ``device``. - Mirrors the helper used by the FSDP backend; works for any optimizer - that exposes the standard ``param_groups`` / ``state`` interface, - including Megatron's DistributedOptimizer. + Handles both the standard ``dict[param, state_dict]`` layout used by + stock torch optimizers and the ``ProxyDict`` layout used by Megatron's + ``DistributedOptimizer``, whose ``items()`` yields + ``((idx, inner_key), tensor)`` pairs directly (and which has no + ``.get()`` method). """ optimizer = self.optimizer - if optimizer is None or not getattr(optimizer, "state", None): + if optimizer is None: return - for param_group in optimizer.param_groups: - for param in param_group["params"]: - state = optimizer.state.get(param, None) - if not state: - continue - for key, value in state.items(): - if isinstance(value, torch.Tensor): - state[key] = value.to(device, non_blocking=True) + state = getattr(optimizer, "state", None) + if not state: + return + + # Snapshot items first so we can safely reassign via __setitem__. + for key, value in list(state.items()): + if isinstance(value, torch.Tensor): + # Megatron ProxyDict: value is the tensor itself. + state[key] = value.to(device, non_blocking=True) + elif isinstance(value, dict): + # Standard torch optimizer: value is a per-param state dict. + for sub_key, sub_value in value.items(): + if isinstance(sub_value, torch.Tensor): + value[sub_key] = sub_value.to(device, non_blocking=True) torch.cuda.synchronize() def _cpu_offload_sleep(self) -> None: From ba06919c3336b997ee36afedec907dc97bd1dcf3 Mon Sep 17 00:00:00 2001 From: knlnguyen1802 Date: Wed, 22 Apr 2026 14:46:52 +0800 Subject: [PATCH 29/39] Fix offload_rollout Signed-off-by: knlnguyen1802 --- slime/ray/rollout.py | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/slime/ray/rollout.py b/slime/ray/rollout.py index ebdc97d59a..25876be55a 100644 --- a/slime/ray/rollout.py +++ b/slime/ray/rollout.py @@ -457,6 +457,7 @@ def __init__(self, args, pg): self.pg = pg self.args = args + self.rollout_backend = getattr(args, "rollout_backend", "sglang") init_tracking(args, primary=False) @@ -589,9 +590,21 @@ def onload(self, tags: list[str] | None = None): srv.onload(tags) def onload_weights(self): - self.onload(tags=[GPU_MEMORY_TYPE_WEIGHTS]) + # vLLM does not expose a separate weights-only resume: a single + # resume_memory_occupation reloads both weights and KV cache. To keep + # the public ``onload_weights`` / ``onload_kv`` contract on the train + # side unchanged, we load everything here and make ``onload_kv`` a + # no-op for the vLLM backend. + if self.rollout_backend == "vllm": + self.onload(tags=[GPU_MEMORY_TYPE_WEIGHTS, GPU_MEMORY_TYPE_KV_CACHE, GPU_MEMORY_TYPE_CUDA_GRAPH]) + else: + self.onload(tags=[GPU_MEMORY_TYPE_WEIGHTS]) def onload_kv(self): + # See ``onload_weights``: for vLLM, weights+KV are already resumed + # together, so this is a no-op to avoid a double-resume. + if self.rollout_backend == "vllm": + return self.onload(tags=[GPU_MEMORY_TYPE_KV_CACHE, GPU_MEMORY_TYPE_CUDA_GRAPH]) def recover_rollout_engines(self, model_name: str | None = None): From 3491aae266a73791e1253d864206242ab5b3706a Mon Sep 17 00:00:00 2001 From: knlnguyen1802 Date: Fri, 24 Apr 2026 14:42:29 +0800 Subject: [PATCH 30/39] Fix vllm offload Signed-off-by: knlnguyen1802 --- slime/ray/rollout.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/slime/ray/rollout.py b/slime/ray/rollout.py index 25876be55a..08d2406e86 100644 --- a/slime/ray/rollout.py +++ b/slime/ray/rollout.py @@ -596,7 +596,7 @@ def onload_weights(self): # side unchanged, we load everything here and make ``onload_kv`` a # no-op for the vLLM backend. if self.rollout_backend == "vllm": - self.onload(tags=[GPU_MEMORY_TYPE_WEIGHTS, GPU_MEMORY_TYPE_KV_CACHE, GPU_MEMORY_TYPE_CUDA_GRAPH]) + self.onload(tags=[GPU_MEMORY_TYPE_WEIGHTS, GPU_MEMORY_TYPE_KV_CACHE]) else: self.onload(tags=[GPU_MEMORY_TYPE_WEIGHTS]) From af2bfb41b92c333460666e54d88623f55a5aeea6 Mon Sep 17 00:00:00 2001 From: knlnguyen1802 Date: Fri, 24 Apr 2026 14:50:10 +0800 Subject: [PATCH 31/39] Fix offload traing Signed-off-by: knlnguyen1802 --- slime/backends/megatron_utils/actor.py | 36 +++++++++++++++++++++----- 1 file changed, 29 insertions(+), 7 deletions(-) diff --git a/slime/backends/megatron_utils/actor.py b/slime/backends/megatron_utils/actor.py index 6bb3335a84..733415f4ac 100644 --- a/slime/backends/megatron_utils/actor.py +++ b/slime/backends/megatron_utils/actor.py @@ -222,20 +222,40 @@ def _move_optimizer_state_to_device(self, device: str | torch.device) -> None: value[sub_key] = sub_value.to(device, non_blocking=True) torch.cuda.synchronize() + @torch.no_grad() + def _move_module_tensors_to_device(self, device: str | torch.device) -> None: + """Move params/grads/buffers to ``device`` in-place. + + ``nn.Module.to()`` replaces ``nn.Parameter`` objects whenever the + device actually changes, which silently breaks any external references + the optimizer (or our ``weights_backuper``) still holds to the old + tensors -- gradients then accumulate on GPU tensors that the optimizer + no longer knows about, so ``optimizer.step()`` becomes a no-op on the + live parameters. Mutating ``p.data`` (and ``p.grad.data``) in place + keeps the ``Parameter`` identity intact, so every consumer continues + to see the moved storage. + """ + if self.model is None: + return + for module in self.model: + for p in module.parameters(recurse=True): + p.data = p.data.to(device, non_blocking=True) + if p.grad is not None: + p.grad.data = p.grad.data.to(device, non_blocking=True) + for b in module.buffers(recurse=True): + b.data = b.data.to(device, non_blocking=True) + def _cpu_offload_sleep(self) -> None: """CPU-offload alternative to ``torch_memory_saver.pause()`` for vLLM.""" - if self.model is not None: - for module in self.model: - module.to("cpu", non_blocking=True) + self._move_module_tensors_to_device("cpu") if getattr(self, "optimizer", None) is not None: self._move_optimizer_state_to_device("cpu") + torch.cuda.synchronize() def _cpu_offload_wake_up(self) -> None: """CPU-offload counterpart to ``_cpu_offload_sleep``.""" device = torch.cuda.current_device() - if self.model is not None: - for module in self.model: - module.to(device, non_blocking=True) + self._move_module_tensors_to_device(device) if getattr(self, "optimizer", None) is not None: self._move_optimizer_state_to_device(device) torch.cuda.synchronize() @@ -646,7 +666,9 @@ def update_weights(self) -> None: if dist.get_rank() == 0: ray.get(self.rollout_manager.clear_num_new_engines.remote()) - with nullcontext(): #torch_memory_saver.disable() if self.args.offload_train else nullcontext(): + weight_update_ctx = torch_memory_saver.disable() if self._uses_torch_memory_saver() else nullcontext() + + with weight_update_ctx: print_memory("before update_weights") self.weight_updater.update_weights() print_memory("after update_weights") From 14aa09ad7ca076ec0813621c4e856e07af4948fb Mon Sep 17 00:00:00 2001 From: knlnguyen1802 Date: Fri, 24 Apr 2026 17:27:38 +0800 Subject: [PATCH 32/39] Fix offload weight Signed-off-by: knlnguyen1802 --- slime/backends/megatron_utils/actor.py | 85 +++++++++++++++++++------- 1 file changed, 63 insertions(+), 22 deletions(-) diff --git a/slime/backends/megatron_utils/actor.py b/slime/backends/megatron_utils/actor.py index 733415f4ac..ca3164b212 100644 --- a/slime/backends/megatron_utils/actor.py +++ b/slime/backends/megatron_utils/actor.py @@ -195,31 +195,39 @@ def _uses_torch_memory_saver(self) -> bool: @torch.no_grad() def _move_optimizer_state_to_device(self, device: str | torch.device) -> None: - """Move all tensor entries in ``self.optimizer.state`` to ``device``. - - Handles both the standard ``dict[param, state_dict]`` layout used by - stock torch optimizers and the ``ProxyDict`` layout used by Megatron's - ``DistributedOptimizer``, whose ``items()`` yields - ``((idx, inner_key), tensor)`` pairs directly (and which has no - ``.get()`` method). + """Move all tensor entries in optimizer state to ``device``. + + Handles three layouts we see in practice: + * ``ChainedOptimizer`` used when there are multiple model chunks + (virtual pipeline) -- we must recurse into ``chained_optimizers`` + because the outer object's ``.state`` is empty. + * Megatron ``DistributedOptimizer`` whose ``state`` is a ``ProxyDict`` + that yields ``(key, tensor)`` directly (no ``.get``). + * Stock torch optimizer with ``dict[param, state_dict]`` layout. """ - optimizer = self.optimizer - if optimizer is None: - return - state = getattr(optimizer, "state", None) - if not state: + if self.optimizer is None: return - # Snapshot items first so we can safely reassign via __setitem__. - for key, value in list(state.items()): - if isinstance(value, torch.Tensor): - # Megatron ProxyDict: value is the tensor itself. - state[key] = value.to(device, non_blocking=True) - elif isinstance(value, dict): - # Standard torch optimizer: value is a per-param state dict. - for sub_key, sub_value in value.items(): - if isinstance(sub_value, torch.Tensor): - value[sub_key] = sub_value.to(device, non_blocking=True) + def _walk(opt) -> None: + # ChainedOptimizer (Megatron VP): dispatch to inner optimizers. + inner = getattr(opt, "chained_optimizers", None) + if inner: + for sub in inner: + _walk(sub) + return + + state = getattr(opt, "state", None) + if not state: + return + for key, value in list(state.items()): + if isinstance(value, torch.Tensor): + state[key] = value.to(device, non_blocking=True) + elif isinstance(value, dict): + for sub_key, sub_value in value.items(): + if isinstance(sub_value, torch.Tensor): + value[sub_key] = sub_value.to(device, non_blocking=True) + + _walk(self.optimizer) torch.cuda.synchronize() @torch.no_grad() @@ -234,14 +242,47 @@ def _move_module_tensors_to_device(self, device: str | torch.device) -> None: live parameters. Mutating ``p.data`` (and ``p.grad.data``) in place keeps the ``Parameter`` identity intact, so every consumer continues to see the moved storage. + + Additional Megatron specifics handled here: + * ``p.main_grad`` (used with gradient-accumulation fusion / DDP) is + moved alongside ``p.grad``; the optimizer reads from + ``main_grad`` when present. + * Megatron ``DistributedDataParallel`` stores parameters/grads in a + few big contiguous ``ParamAndGradBuffer`` tensors; individual + ``p.data`` / ``p.main_grad`` are views into them. Moving only the + views leaves the big buckets allocated on GPU and leaves the + views pointing at stale storage, so we also migrate each + buffer's underlying ``param_data`` / ``grad_data`` in place. """ if self.model is None: return + + def _move_tensor_attr(owner, attr: str) -> None: + t = getattr(owner, attr, None) + if isinstance(t, torch.Tensor): + t.data = t.data.to(device, non_blocking=True) + for module in self.model: + # 1) Move DDP bucket buffers first so subsequent param views land + # on the new device's storage without orphaning the buckets. + for bucket_attr in ("buffers", "expert_parallel_buffers"): + bucket_list = getattr(module, bucket_attr, None) + # ``buffers`` on nn.Module is a *method*; only treat it as our + # target when it is the Megatron list of ParamAndGradBuffer. + if isinstance(bucket_list, (list, tuple)): + for buf in bucket_list: + for t_attr in ("param_data", "grad_data"): + _move_tensor_attr(buf, t_attr) + + # 2) Move parameter / grad views in place (Parameter identity + # preserved so optimizer + weights_backuper stay consistent). for p in module.parameters(recurse=True): p.data = p.data.to(device, non_blocking=True) if p.grad is not None: p.grad.data = p.grad.data.to(device, non_blocking=True) + _move_tensor_attr(p, "main_grad") + + # 3) Registered buffers (e.g. RoPE caches, norm running stats). for b in module.buffers(recurse=True): b.data = b.data.to(device, non_blocking=True) From d82d1d0d67a0b66ce41590e9f9eb362cfe3a8d37 Mon Sep 17 00:00:00 2001 From: knlnguyen1802 Date: Fri, 24 Apr 2026 17:38:01 +0800 Subject: [PATCH 33/39] Fix offload weight Signed-off-by: knlnguyen1802 --- slime/backends/megatron_utils/actor.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/slime/backends/megatron_utils/actor.py b/slime/backends/megatron_utils/actor.py index ca3164b212..5b6d2384f6 100644 --- a/slime/backends/megatron_utils/actor.py +++ b/slime/backends/megatron_utils/actor.py @@ -283,7 +283,9 @@ def _move_tensor_attr(owner, attr: str) -> None: _move_tensor_attr(p, "main_grad") # 3) Registered buffers (e.g. RoPE caches, norm running stats). - for b in module.buffers(recurse=True): + # Use the unbound ``nn.Module.buffers`` because Megatron's DDP + # shadows the method with a ``list`` of ParamAndGradBuffer. + for b in torch.nn.Module.buffers(module, recurse=True): b.data = b.data.to(device, non_blocking=True) def _cpu_offload_sleep(self) -> None: From c8aaf018629686c4e3752f0777d828d8645d29c6 Mon Sep 17 00:00:00 2001 From: jingshenghang <48083555+jingshenghang@users.noreply.github.com> Date: Mon, 11 May 2026 16:05:32 +0800 Subject: [PATCH 34/39] Fix CI: update rollout_data_postprocess plugin contract for new call site (#1902) Co-authored-by: jingshenghang --- .../test_plugin_runtime_hook_contracts.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/plugin_contracts/test_plugin_runtime_hook_contracts.py b/tests/plugin_contracts/test_plugin_runtime_hook_contracts.py index 9c74876d96..a8380feecc 100644 --- a/tests/plugin_contracts/test_plugin_runtime_hook_contracts.py +++ b/tests/plugin_contracts/test_plugin_runtime_hook_contracts.py @@ -69,7 +69,7 @@ def reference_convert_samples_to_train_data(args, samples): } -def reference_rollout_data_postprocess(args) -> None: +def reference_rollout_data_postprocess(args, rollout_id, rollout_data) -> None: args.rollout_data_postprocess_called = True @@ -124,7 +124,7 @@ def invoke_convert_samples_to_train_data(fn): def invoke_rollout_data_postprocess(fn): args = type("Args", (), {})() - assert fn(args) is None + assert fn(args, 0, {}) is None assert args.rollout_data_postprocess_called is True @@ -170,8 +170,8 @@ def invoke_rollout_data_postprocess(fn): "ROLLOUT_DATA_POSTPROCESS_PATH", "plugin_contracts.test_plugin_runtime_hook_contracts.reference_rollout_data_postprocess", "slime/backends/megatron_utils/actor.py", - "self.rollout_data_postprocess(self.args)", - ("args",), + "self.rollout_data_postprocess(self.args, rollout_id, rollout_data)", + ("args", "rollout_id", "rollout_data"), invoke_rollout_data_postprocess, ), ] From 5b326e6587044c886b592183b4173903e47542af Mon Sep 17 00:00:00 2001 From: jingshenghang <48083555+jingshenghang@users.noreply.github.com> Date: Mon, 11 May 2026 16:26:49 +0800 Subject: [PATCH 35/39] Patch Megatron TP grad coalesce to chunked all-reduce (#1899) --- slime/backends/megatron_utils/__init__.py | 2 + .../megatron_utils/megatron_patch/__init__.py | 1 + .../megatron_chunked_grad_coalesce_patch.py | 146 ++++++++++++++++++ 3 files changed, 149 insertions(+) create mode 100644 slime/backends/megatron_utils/megatron_patch/__init__.py create mode 100644 slime/backends/megatron_utils/megatron_patch/megatron_chunked_grad_coalesce_patch.py diff --git a/slime/backends/megatron_utils/__init__.py b/slime/backends/megatron_utils/__init__.py index a4666fbeb9..b1936ae692 100644 --- a/slime/backends/megatron_utils/__init__.py +++ b/slime/backends/megatron_utils/__init__.py @@ -40,3 +40,5 @@ def _patched_forward(self, *args, packed_seq_params=None, **kwargs): pass logging.getLogger("megatron").setLevel(logging.WARNING) + +from . import megatron_patch # noqa: F401, E402 diff --git a/slime/backends/megatron_utils/megatron_patch/__init__.py b/slime/backends/megatron_utils/megatron_patch/__init__.py new file mode 100644 index 0000000000..8693ff85a0 --- /dev/null +++ b/slime/backends/megatron_utils/megatron_patch/__init__.py @@ -0,0 +1 @@ +from . import megatron_chunked_grad_coalesce_patch # noqa: F401 diff --git a/slime/backends/megatron_utils/megatron_patch/megatron_chunked_grad_coalesce_patch.py b/slime/backends/megatron_utils/megatron_patch/megatron_chunked_grad_coalesce_patch.py new file mode 100644 index 0000000000..f7ecceb2dd --- /dev/null +++ b/slime/backends/megatron_utils/megatron_patch/megatron_chunked_grad_coalesce_patch.py @@ -0,0 +1,146 @@ +# Patch _allreduce_non_tensor_model_parallel_grads (and its legacy alias +# _allreduce_layernorm_grads) in megatron.core.distributed.finalize_model_grads +# to coalesce/all_reduce TP-side grads in size-bounded chunks instead of one +# large _flatten_dense_tensors(grads). Lowers the peak contiguous-memory +# allocation during TP-side grad sync, avoiding OOM under allocator +# fragmentation when the combined grad buffer would otherwise be very large. +# SUM/AVG are element-wise, so chunking is mathematically equivalent. +# Chunk size: SLIME_GRAD_COALESCE_CHUNK_BYTES, default 1 GiB. +# +# Cross-compatible across the Megatron-LM versions slime is run against: +# the core_v0.13.0 line (DDP config exposes `use_custom_fsdp`, `_get_main_grad_attr` +# takes `(param, use_custom_fsdp)`, target function takes `(model, config)`) and +# the post-core_v0.15.0rc7 dev line (`use_megatron_fsdp`, single-arg +# `_get_main_grad_attr`, `(model, config, tp_group)`). API differences are +# resolved at runtime — no version-conditional imports. + +import inspect +import logging +import os +import sys +import warnings + +logger = logging.getLogger(__name__) + +try: + import torch + from megatron.core import parallel_state + from megatron.core.distributed.finalize_model_grads import ( + _flatten_dense_tensors, + _get_main_grad_attr, + _reshard_if_dtensor, + _unflatten_dense_tensors, + _unshard_if_dtensor, + get_attr_wrapped_model, + ) + + # post-core_v0.15.0rc7 dev takes (param); core_v0.13.0 line takes + # (param, use_custom_fsdp=False). + _gma_takes_fsdp_arg = len(inspect.signature(_get_main_grad_attr).parameters) >= 2 + + def _grad_attr(param, fsdp_on): + if _gma_takes_fsdp_arg: + return _get_main_grad_attr(param, fsdp_on) + return _get_main_grad_attr(param) + + def _fsdp_flag(ddp_config): + return bool(getattr(ddp_config, "use_megatron_fsdp", False) or getattr(ddp_config, "use_custom_fsdp", False)) + + _chunk_bytes = int(os.environ.get("SLIME_GRAD_COALESCE_CHUNK_BYTES") or (1 << 30)) + + def _split_into_chunks(params, grads, target_bytes): + """Greedy split keeping params/grads aligned. A single grad larger + than target_bytes is placed alone in its own chunk.""" + chunks, cur_p, cur_g, cur_b = [], [], [], 0 + for p, g in zip(params, grads, strict=False): + gb = g.numel() * g.element_size() + if cur_g and cur_b + gb > target_bytes: + chunks.append((cur_p, cur_g)) + cur_p, cur_g, cur_b = [], [], 0 + cur_p.append(p) + cur_g.append(g) + cur_b += gb + if cur_g: + chunks.append((cur_p, cur_g)) + return chunks + + def _allreduce_non_tensor_model_parallel_grads(model, config, tp_group=None): + # post-core_v0.15.0rc7 dev passes tp_group; core_v0.13.0 line omits it. + # Default-fill from parallel_state so the same body works for both call sites. + if tp_group is None: + tp_group = parallel_state.get_tensor_model_parallel_group() + if tp_group.size() <= 1: + return + + params_sum, grads_sum = [], [] + params_avg, grads_avg = [], [] + ddp_config = None + for model_chunk in model: + ddp_config = model_chunk.ddp_config + fsdp_on = _fsdp_flag(ddp_config) + for name, param in get_attr_wrapped_model(model_chunk, "named_parameters")(): + if not param.requires_grad: + continue + if getattr(param, "average_gradients_across_tp_domain", False): + target_params, target_grads = params_avg, grads_avg + elif (config.sequence_parallel and getattr(param, "sequence_parallel", False)) or ( + config.qk_layernorm and ("q_layernorm" in name or "k_layernorm" in name) + ): + target_params, target_grads = params_sum, grads_sum + else: + continue + + grad_attr = _grad_attr(param, fsdp_on) + grad = getattr(param, grad_attr) + if grad is None: + continue + target_params.append(param) + if fsdp_on and hasattr(grad, "_local_tensor"): + target_grads.append(grad._local_tensor.data) + else: + target_grads.append(_unshard_if_dtensor(grad).data) + + for params, grads, op in ( + (params_sum, grads_sum, torch.distributed.ReduceOp.SUM), + (params_avg, grads_avg, torch.distributed.ReduceOp.AVG), + ): + if not grads: + continue + fsdp_on = _fsdp_flag(ddp_config) + for p_chunk, g_chunk in _split_into_chunks(params, grads, _chunk_bytes): + coalesced = _flatten_dense_tensors(g_chunk) + torch.distributed.all_reduce(coalesced, op=op, group=tp_group) + for param, buf, synced in zip( + p_chunk, g_chunk, _unflatten_dense_tensors(coalesced, g_chunk), strict=False + ): + buf.copy_(synced) + grad_attr = _grad_attr(param, fsdp_on) + orig_grad = getattr(param, grad_attr) + if fsdp_on and hasattr(orig_grad, "_local_tensor"): + # buf already aliases orig_grad._local_tensor.data; + # restore original DTensor wrapper (post-rc7 dev semantics). + setattr(param, grad_attr, orig_grad) + else: + setattr(param, grad_attr, _reshard_if_dtensor(buf, orig_grad)) + del coalesced + + # The parent package re-exports a same-named function, shadowing the + # submodule attribute. Pull the real module out of sys.modules to setattr. + _fmg = sys.modules["megatron.core.distributed.finalize_model_grads"] + _fmg._allreduce_non_tensor_model_parallel_grads = _allreduce_non_tensor_model_parallel_grads + _fmg._allreduce_layernorm_grads = _allreduce_non_tensor_model_parallel_grads + + logger.info( + "slime grad coalesce patch applied to " + "megatron.core.distributed.finalize_model_grads." + "_allreduce_non_tensor_model_parallel_grads (chunk=%d MiB)", + _chunk_bytes // (1 << 20), + ) + +except ImportError as exc: + warnings.warn( + f"slime grad coalesce patch not applied — Megatron import failed ({exc!r}). " + "If this is a Megatron upgrade, the symbol layout may have changed; " + "without this patch, large-model TP grad sync may OOM.", + stacklevel=2, + ) From 41dc3b6d21d3c75b212965077a1cc4117932f06d Mon Sep 17 00:00:00 2001 From: Leo Fan <84952531+leofan-lab@users.noreply.github.com> Date: Mon, 11 May 2026 04:17:12 -0700 Subject: [PATCH 36/39] fix: harden retool rollout against multi-turn / retry desync (#1861) Co-authored-by: Claude Opus 4.7 (1M context) --- examples/retool/generate_with_retool.py | 65 ++++++++++++++++++++++--- 1 file changed, 57 insertions(+), 8 deletions(-) diff --git a/examples/retool/generate_with_retool.py b/examples/retool/generate_with_retool.py index a090af63d7..339b79dd45 100644 --- a/examples/retool/generate_with_retool.py +++ b/examples/retool/generate_with_retool.py @@ -216,6 +216,15 @@ async def generate(args, sample: Sample, sampling_params) -> Sample: """Custom generation function supporting tool calls""" assert not args.partial_rollout, "Partial rollout is not supported for " "this function at the moment." + # Retried samples (previously aborted / partial) arrive here with stale + # rollout state from the first attempt. Clear it so this generation starts + # clean; otherwise the concatenation below appends new tokens to old ones + # and downstream `slice_log_prob_with_cp` sees a length mismatch. + sample.rollout_log_probs = None + sample.response = "" + sample.response_length = 0 + sample.loss_mask = None + state = GenerateState(args) url = f"http://{args.sglang_router_ip}:{args.sglang_router_port}/generate" @@ -229,22 +238,37 @@ async def generate(args, sample: Sample, sampling_params) -> Sample: loss_masks = [] tool_call_count = 0 # Track actual tool call rounds + if args.rollout_max_context_len is not None: + max_context_length = args.rollout_max_context_len + else: + max_context_length = args.context_parallel_size * args.max_tokens_per_gpu + for turn in range(TOOL_CONFIGS["max_turns"]): # Check if total length exceeds max context length total_length = len(prompt_tokens_ids) + len(response_token_ids) - if args.rollout_max_context_len is not None: - max_context_length = args.rollout_max_context_len - else: - max_context_length = args.context_parallel_size * args.max_tokens_per_gpu if total_length >= max_context_length: sample.status = Sample.Status.TRUNCATED break + # Clamp per-turn max_new_tokens to the remaining context budget so a + # single turn cannot push total_length past max_context_length. Without + # this, a turn can append up to rollout_max_response_len tokens on top + # of a total that was just barely under the cap, producing samples + # that exceed the training-side max_tokens_per_gpu * cp_size budget + # and crash the partition/batch code (asserts or OOMs on an oversized + # partition). + remaining_budget = max_context_length - total_length + per_turn_sampling_params = dict(sampling_params) + per_turn_sampling_params["max_new_tokens"] = min( + sampling_params.get("max_new_tokens", remaining_budget), + remaining_budget, + ) + # Use token IDs instead of text current_token_ids = prompt_tokens_ids + response_token_ids payload = { "input_ids": current_token_ids, - "sampling_params": sampling_params, + "sampling_params": per_turn_sampling_params, "return_logprob": True, # Request log probabilities for training } @@ -285,9 +309,14 @@ async def generate(args, sample: Sample, sampling_params) -> Sample: sample.rollout_log_probs += cur_log_probs else: - cur_response = output["text"] - cur_response = postprocess_responses(cur_response) - cur_response_token_ids = state.tokenizer(cur_response, add_special_tokens=False)["input_ids"] + # sglang returned text but no output_token_logprobs — we cannot + # recover per-token logprobs for this turn, which would desync + # rollout_log_probs from response_token_ids and blow up + # `slice_log_prob_with_cp` downstream. Abort the sample so the + # fully_async rollout manager returns the whole group to the + # buffer for retry instead of poisoning the trainer. + sample.status = Sample.Status.ABORTED + return sample response += cur_response response_token_ids += cur_response_token_ids @@ -321,6 +350,26 @@ async def generate(args, sample: Sample, sampling_params) -> Sample: sample.rollout_log_probs ), f"Token/logp length mismatch at turn {turn}: {len(response_token_ids)} tokens vs {len(sample.rollout_log_probs)} logps" + # Tool output is appended verbatim and can push total_length past + # max_context_length (the per-turn generation was clamped to the + # remaining budget, but tool output is unconstrained). Trim tail + # tokens so the final sample fits the training budget exactly. + overflow = len(prompt_tokens_ids) + len(response_token_ids) - max_context_length + if overflow > 0: + response_token_ids = response_token_ids[:-overflow] + loss_masks = loss_masks[:-overflow] + if sample.rollout_log_probs is not None: + sample.rollout_log_probs = sample.rollout_log_probs[:-overflow] + # Resync the text field from the trimmed token list so + # reward_func's `sample.prompt + sample.response` matches what + # the model was actually trained on. decode(tokenize(text)) can + # be lossy on some tokenizers (whitespace / special-token + # collapse), but reward_func's regex is whitespace-robust and + # the trainer sees tokens, not text — so the drift is safe. + response = state.tokenizer.decode(response_token_ids) + sample.status = Sample.Status.TRUNCATED + break + if tool_call_count >= TOOL_CONFIGS["max_tool_calls"]: break From a7a3ee11139c8cbc846f726759dddc22ea9a90b2 Mon Sep 17 00:00:00 2001 From: knlnguyen1802 Date: Thu, 14 May 2026 13:35:05 +0800 Subject: [PATCH 37/39] Fix log file Signed-off-by: knlnguyen1802 --- slime/backends/vllm_utils/vllm_engine.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/slime/backends/vllm_utils/vllm_engine.py b/slime/backends/vllm_utils/vllm_engine.py index f431922545..c2902ea13f 100644 --- a/slime/backends/vllm_utils/vllm_engine.py +++ b/slime/backends/vllm_utils/vllm_engine.py @@ -97,11 +97,10 @@ def init(self, port=None, host=None, router_ip=None, router_port=None, **kwargs) env["NCCL_P2P_DISABLE"] = "1" env.setdefault("NCCL_IB_DISABLE", "1") - # self._log_file = tempfile.NamedTemporaryFile( - # prefix="vllm_engine_", suffix=".log", delete=False, mode="w" - # ) - log_path = "/data/n0090/SLIME_PJ/new_version/slime/vllm_server_log.log" - self._log_file = open(log_path,"w") + self._log_file = tempfile.NamedTemporaryFile( + prefix="vllm_engine_", suffix=".log", delete=False, mode="w" + ) + logger.info("Launching vLLM: cmd=%s, CUDA_VISIBLE_DEVICES=%s, log=%s", " ".join(cmd), dev_str, self._log_file.name) self.process = subprocess.Popen( From b015d727c52003efedcf923ee12ef870f8998bca Mon Sep 17 00:00:00 2001 From: knlnguyen1802 Date: Fri, 15 May 2026 11:39:32 +0800 Subject: [PATCH 38/39] Fix import engine group Signed-off-by: knlnguyen1802 --- slime/ray/rollout.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/slime/ray/rollout.py b/slime/ray/rollout.py index 42a3f04fd2..f6a361c626 100644 --- a/slime/ray/rollout.py +++ b/slime/ray/rollout.py @@ -1056,7 +1056,7 @@ def _start_vllm_rollout_servers(args, pg) -> dict[str, RolloutServer]: first_host = args.sglang_router_ip first_port = args.sglang_router_port - group = EngineGroup( + group = ServerGroup( args=args, pg=pg, all_engines=engines, @@ -1066,10 +1066,19 @@ def _start_vllm_rollout_servers(args, pg) -> dict[str, RolloutServer]: rank_offset=0, gpu_offset=0, sglang_overrides={}, + needs_offload=getattr(args, "offload_rollout", False), + model_path=getattr(args, "hf_checkpoint", None), router_ip=first_host, router_port=first_port, ) - return {"default": RolloutServer(engine_groups=[group], router_ip=first_host, router_port=first_port, model_name="default")} + return { + "default": RolloutServer( + server_groups=[group], + router_ip=first_host, + router_port=first_port, + model_name="default", + ) + } def _compute_rollout_offset(args) -> int: """Offset (in PG bundle slots) where rollout GPUs start.""" if args.debug_train_only or args.debug_rollout_only or args.colocate: From 26f979d370cfbd160db8b69db5a9d9551cb0ec8d Mon Sep 17 00:00:00 2001 From: knlnguyen1802 Date: Fri, 15 May 2026 14:31:53 +0800 Subject: [PATCH 39/39] Fix rebase code Signed-off-by: knlnguyen1802 --- run-qwen2.5-0.5B-vllm.sh | 3 ++- slime/ray/rollout.py | 46 +++++++++++++++++++++++----------------- 2 files changed, 28 insertions(+), 21 deletions(-) diff --git a/run-qwen2.5-0.5B-vllm.sh b/run-qwen2.5-0.5B-vllm.sh index ea010b7e47..397fac1c58 100644 --- a/run-qwen2.5-0.5B-vllm.sh +++ b/run-qwen2.5-0.5B-vllm.sh @@ -96,6 +96,7 @@ VLLM_ARGS=( --rollout-num-gpus-per-engine 1 --rollout-server-concurrency 512 --use-slime-router + --vllm-gpu-memory-utilization 0.2 --slime-router-middleware-paths slime.router.middleware_hub.radix_tree_middleware.RadixTreeMiddleware # Uncomment below to enable colocate mode (IPC weight transfer, same GPU) --colocate @@ -115,7 +116,7 @@ ray start --head --node-ip-address 127.0.0.1 --num-gpus 2 --disable-usage-stats ray job submit --address="http://127.0.0.1:8265" \ --runtime-env-json='{ "env_vars": { - "PYTHONPATH": "/root/Megatron-LM:/data/n0090/SLIME_PJ/new_version/local_package", + "PYTHONPATH": "/root/Megatron-LM", "CUDA_DEVICE_MAX_CONNECTIONS": "1", "NCCL_ALGO": "Ring", "NCCL_IB_DISABLE": "1", diff --git a/slime/ray/rollout.py b/slime/ray/rollout.py index f6a361c626..fbb6338f93 100644 --- a/slime/ray/rollout.py +++ b/slime/ray/rollout.py @@ -20,7 +20,6 @@ GPU_MEMORY_TYPE_CUDA_GRAPH = "cuda_graph" from slime.backends.sglang_utils.sglang_config import ModelConfig, ServerGroupConfig, SglangConfig -from slime.backends.sglang_utils.sglang_engine import SGLangEngine from slime.rollout.base_types import call_rollout_fn from slime.utils import logging_utils from slime.utils.health_monitor import RolloutHealthMonitor @@ -932,7 +931,7 @@ def addr(): def _start_router(args, *, has_pd_disaggregation: bool = False, force_new: bool = False) -> tuple[str, int]: - """Start sglang_router and return (router_ip, router_port). + """Start sgl router or slime router and return (router_ip, router_port). If ``args.sglang_router_ip`` is already set (e.g. by the user) and ``force_new`` is False, skip launching and return the existing values. @@ -949,28 +948,35 @@ def _start_router(args, *, has_pd_disaggregation: bool = False, force_new: bool if router_port is None: router_port = find_available_port(random.randint(3000, 4000)) - from sglang_router.launch_router import RouterArgs + if args.use_slime_router: + assert not has_pd_disaggregation, "slime router does not support PD disaggregation." + import copy - from slime.utils.http_utils import run_router + from slime.router.router import run_router - router_args = RouterArgs.from_cli_args(args, use_router_prefix=True) - router_args.host = router_ip - router_args.port = router_port - router_args.prometheus_port = find_available_port(random.randint(4000, 5000)) - router_args.log_level = "warn" - router_args.request_timeout_secs = args.sglang_router_request_timeout_secs + router_args = copy.copy(args) + router_args.sglang_router_ip = router_ip + router_args.sglang_router_port = router_port - if has_pd_disaggregation: - router_args.pd_disaggregation = True - # Disable circuit breaker to prevent RDMA transfer timeouts from - # marking decode workers as dead. Timeouts are transient (PCIe - # contention under high load) and do not indicate a dead server. - router_args.disable_circuit_breaker = True + else: + from sglang_router.launch_router import RouterArgs # noqa: delayed import — only needed for sglang router + + from slime.utils.http_utils import run_router + + router_args = RouterArgs.from_cli_args(args, use_router_prefix=True) + router_args.host = router_ip + router_args.port = router_port + router_args.prometheus_port = find_available_port(random.randint(4000, 5000)) + router_args.log_level = "warn" + router_args.request_timeout_secs = args.sglang_router_request_timeout_secs + + if hasattr(args, "sglang_router_policy") and args.sglang_router_policy: + router_args.policy = args.sglang_router_policy - # We will not use the health check from router. - router_args.disable_health_check = True + if has_pd_disaggregation: + router_args.pd_disaggregation = True - logger.info(f"Launch router with args: {router_args}") + logger.info(f"Launch router with args: {router_args}") process = multiprocessing.Process( target=run_router, @@ -981,7 +987,7 @@ def _start_router(args, *, has_pd_disaggregation: bool = False, force_new: bool # Wait 3 seconds time.sleep(3) assert process.is_alive() - logger.info(f"Router launched at {router_ip}:{router_port}, Prometheus port: {router_args.prometheus_port}") + logger.info(f"Router launched at {router_ip}:{router_port}") return router_ip, router_port