diff --git a/requirements.txt b/requirements.txt index 427680d16..8b016fb03 100644 --- a/requirements.txt +++ b/requirements.txt @@ -13,6 +13,7 @@ qwen_vl_utils # for VLM ray[default] ring_flash_attn sglang-router>=0.2.3 +vllm-router>=0.1.14 tensorboard transformers wandb 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 822b80177..bb27e4a6d 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,11 +1,18 @@ +from __future__ import annotations + +import logging +import os import socket import time +import traceback from argparse import Namespace from collections.abc import Callable, Mapping, Sequence +from typing import Any import ray import torch import torch.distributed as dist +import torch.multiprocessing as mp from megatron.core import mpu from ray import ObjectRef from ray.actor import ActorHandle @@ -16,6 +23,187 @@ 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: str, + master_port: int, + world_size: int, + device: int, + cvd: str, + env_snapshot: dict[str, str], +) -> None: + """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": [gpu_tensor, ...]} + {"op": "send_packed", "named_tensors": [(name, gpu_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 as _torch # noqa: PLC0415 — subprocess needs fresh import + import torch.multiprocessing # noqa: F401, PLC0415 — register CUDA IPC reducers + + _torch.cuda.set_device(device) + + from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator # noqa: PLC0415 + from vllm.distributed.utils import StatelessProcessGroup # noqa: PLC0415 + + 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 t in cmd["tensors"]: + comm.broadcast(t, src=0, stream=_torch.cuda.current_stream()) + _torch.cuda.synchronize() + conn.send("ok") + + elif op == "send_packed": + # Prefer NCCLWeightTransferEngine.trainer_send_weights (newer vLLM). Some pip builds omit + # NCCLTrainerSendWeightsArgs but still ship packed_broadcast_producer. + try: + from vllm.distributed.weight_transfer.nccl_engine import ( # noqa: PLC0415 + NCCLTrainerSendWeightsArgs, + NCCLWeightTransferEngine, + ) + + trainer_args = NCCLTrainerSendWeightsArgs( + group=comm, + packed=True, + ) + NCCLWeightTransferEngine.trainer_send_weights( + iterator=iter(cmd["named_tensors"]), + trainer_args=trainer_args, + ) + except ImportError: + from vllm.distributed.weight_transfer.packed_tensor import ( # noqa: PLC0415 + DEFAULT_PACKED_BUFFER_SIZE_BYTES, + DEFAULT_PACKED_NUM_BUFFERS, + packed_broadcast_producer, + ) + + packed_broadcast_producer( + iterator=iter(cmd["named_tensors"]), + group=comm, + src=0, + post_iter_func=lambda x: x[1], + buffer_size_bytes=DEFAULT_PACKED_BUFFER_SIZE_BYTES, + num_buffers=DEFAULT_PACKED_NUM_BUFFERS, + ) + _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. 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): + ctx = mp.get_context("spawn") + self._parent_conn, child_conn = ctx.Pipe() + + 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.""" + 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.""" + gpu_pairs = [] + for name, t in named_tensors: + data = t.data if hasattr(t, "data") else t + 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: + 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() + + +def _is_vllm_backend(args: Namespace) -> bool: + return getattr(args, "rollout_backend", "sglang") == "vllm" + class UpdateWeightFromDistributed: """ @@ -106,34 +294,65 @@ 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 + use_vllm_packed = self._use_vllm_packed() + if use_vllm_packed and self._is_pp_src_rank: + logger.info( + "Using vLLM packed weight sync (bucketed; metadata + trainer_send_weights per bucket)" + ) - 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 use_vllm_packed: + buffer_size = 0 + converted_named_tensors: list[tuple[str, torch.Tensor]] = [] + pbar = ( + tqdm(desc=f"[{self._group_name}] Update weights (vLLM packed)", 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, + flush_packed=True, + ) + if converted_named_tensors and self._is_pp_src_rank: + self._update_weights_vllm_packed(converted_named_tensors) + if pbar is not None: + pbar.update(1) + 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: @@ -147,6 +366,37 @@ def update_weights(self) -> None: ray.get([engine.continue_generation.remote() for engine in self.rollout_engines]) dist.barrier(group=get_gloo_group()) + 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 + if any(".experts." in name for name, _ in named_params_and_buffers(self.args, self.model)): + return False + 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.""" + 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_weight_from_distributed( self, name: str, @@ -154,6 +404,8 @@ def _update_weight_from_distributed( converted_named_tensors: list[tuple[str, torch.Tensor]], buffer_size: int, pbar: tqdm | None = None, + *, + flush_packed: bool = False, ) -> int | None: """ Non-expert: gather TP → rm pad → HF → buffer (flush if full). All gather, PP source buffers. @@ -165,7 +417,14 @@ def _update_weight_from_distributed( param_size = param.numel() * param.element_size() if buffer_size + param_size > self.args.update_weight_buffer_size: - self._update_bucket_weights_from_distributed(converted_named_tensors, pbar=pbar) + if converted_named_tensors: + if flush_packed: + self._update_weights_vllm_packed(converted_named_tensors) + converted_named_tensors.clear() + if pbar is not None: + pbar.update(1) + else: + self._update_bucket_weights_from_distributed(converted_named_tensors, pbar=pbar) buffer_size = 0 converted_named_tensors += convert_to_hf(self.args, self.model_name, name, param, self.quantization_config) buffer_size += param_size @@ -239,7 +498,6 @@ def _update_bucket_weights_from_distributed( """ Lock → broadcast → clear → unlock → pbar++. Lock prevents NCCL deadlock. """ - # lock the rollout engines to prevent dead lock on broadcast. while not ray.get(self.rollout_engine_lock.acquire.remote()): time.sleep(0.1) @@ -249,12 +507,15 @@ def _update_bucket_weights_from_distributed( self.weight_version, self.rollout_engines, converted_named_tensors, + use_vllm=_is_vllm_backend(self.args), + packed=False, ) ray.get(refs) converted_named_tensors.clear() ray.get(self.rollout_engine_lock.release.remote()) - pbar.update(1) + if pbar is not None: + pbar.update(1) def connect_rollout_engines_from_distributed( @@ -262,13 +523,18 @@ def connect_rollout_engines_from_distributed( group_name: str, rollout_engines: Sequence[ActorHandle], engine_gpu_counts: Sequence[int] | None = None, -) -> dist.ProcessGroup: +) -> Any: """ Create NCCL group: training rank 0 + all engine GPUs. Blocks until joined. ``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, 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) @@ -279,7 +545,6 @@ def connect_rollout_engines_from_distributed( master_port = sock.getsockname()[1] world_size = sum(engine_gpu_counts) + 1 # +1 for training rank 0 - # Compute cumulative rank offsets: engine i starts at cumulative[i] + 1. cumulative = [0] for c in engine_gpu_counts: cumulative.append(cumulative[-1] + c) @@ -295,52 +560,96 @@ 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, - ) + + torch.cuda.synchronize() + torch.cuda.empty_cache() + + if _is_vllm_backend(args): + 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", + 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 -def disconnect_rollout_engines_from_distributed(args, group_name, model_update_groups, rollout_engines): +def disconnect_rollout_engines_from_distributed( + args: Namespace, + group_name: str, + model_update_groups: Any, + rollout_engines: Sequence[ActorHandle], +) -> None: """ 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): + if isinstance(model_update_groups, _NcclBridge): + model_update_groups.shutdown() + elif model_update_groups is not None: + dist.destroy_process_group(model_update_groups) ray.get(refs) def update_weights_from_distributed( group_name: str, - group: dist.ProcessGroup, + group: Any, weight_version: int, 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). - """ - 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), - ) - 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() + 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: dict[str, Any] = { + "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(**kwargs) for engine in rollout_engines] + + if use_vllm and packed: + group.send_weights_packed(list(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: + 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/megatron_utils/update_weight/update_weight_from_tensor.py b/slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py index dbba6aeb5..48167b953 100644 --- a/slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py +++ b/slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py @@ -14,6 +14,7 @@ from ..sglang import FlattenedTensorBucket, MultiprocessingSerializer from .hf_weight_iterator_base import HfWeightIteratorBase from .update_weight_from_distributed import ( + _is_vllm_backend, connect_rollout_engines_from_distributed, disconnect_rollout_engines_from_distributed, post_process_weights, @@ -199,6 +200,8 @@ def _send_hf_params(self, hf_named_tensors) -> tuple[list[ObjectRef], Any]: self.weight_version, self.distributed_rollout_engines, hf_named_tensors, + use_vllm=_is_vllm_backend(self.args), + packed=False, ) if refs_distributed: all_refs.extend(refs_distributed) diff --git a/slime/backends/sglang_utils/sglang_engine.py b/slime/backends/sglang_utils/sglang_engine.py index c28e13d5f..c4465b3c9 100644 --- a/slime/backends/sglang_utils/sglang_engine.py +++ b/slime/backends/sglang_utils/sglang_engine.py @@ -404,8 +404,16 @@ def destroy_weights_update_group(self, group_name): pass def update_weights_from_distributed( - self, names, dtypes, shapes, group_name, flush_cache=False, weight_version: str | None = None + self, + names, + dtypes, + shapes, + group_name, + flush_cache=False, + weight_version: str | None = None, + packed: bool = False, ): + del packed payload = { "names": names, "dtypes": [str(dtype).replace("torch.", "") for dtype in dtypes], diff --git a/slime/backends/vllm_utils/__init__.py b/slime/backends/vllm_utils/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/slime/backends/vllm_utils/vllm_engine.py b/slime/backends/vllm_utils/vllm_engine.py new file mode 100644 index 000000000..ce86ab2df --- /dev/null +++ b/slime/backends/vllm_utils/vllm_engine.py @@ -0,0 +1,707 @@ +from __future__ import annotations + +import ipaddress +import logging +import multiprocessing +import os +import time + +import requests +from urllib.parse import quote + +from slime.ray.ray_actor import RayActor +from slime.utils.http_utils import get_host_info + +logger = logging.getLogger(__name__) + +_spawn_ctx = multiprocessing.get_context("spawn") + +# vLLM sleep/wake only supports these tags (SGLang also uses ``cuda_graph``, which must be dropped). +_VLLM_WAKE_TAGS = frozenset({"weights", "kv_cache"}) + + +def _normalize_vllm_wake_tags(tags: list[str] | None) -> list[str] | None: + if not tags: + return tags + normalized = [t for t in tags if t in _VLLM_WAKE_TAGS] + dropped = set(tags) - set(normalized) + if dropped: + logger.debug("vLLM wake_up: dropped tags not supported by vLLM: %s", sorted(dropped)) + return normalized or None + + +def get_base_gpu_id(args, rank): + """First local GPU index on this node for rollout engine *rank* (colocate vs actor[/critic]-offset layout).""" + num_gpus = min(args.num_gpus_per_node, args.rollout_num_gpus_per_engine) + if args.colocate: + start_index = (rank * num_gpus) % args.num_gpus_per_node + else: + num_actor_gpus = 0 if args.debug_rollout_only else args.actor_num_gpus_per_node * args.actor_num_nodes + start_index = (num_actor_gpus + rank * num_gpus) % args.num_gpus_per_node + if args.use_critic: + num_critic_gpus = args.critic_num_gpus_per_node * args.critic_num_nodes + start_index = (num_actor_gpus + num_critic_gpus + rank * num_gpus) % args.num_gpus_per_node + return start_index + + +def _to_local_gpu_id(physical_gpu_id: int) -> int: + cvd = os.environ.get("CUDA_VISIBLE_DEVICES") + if not cvd: + return physical_gpu_id + visible = [int(x) for x in cvd.split(",") if x.strip() != ""] + if physical_gpu_id in visible: + return visible.index(physical_gpu_id) + if 0 <= physical_gpu_id < len(visible): + return physical_gpu_id + raise RuntimeError( + f"GPU id {physical_gpu_id} is not valid under CUDA_VISIBLE_DEVICES={cvd}. " + f"Expected one of {visible} (physical) or 0..{len(visible)-1} (local)." + ) + + +def _format_v6_uri(addr: str | None) -> str | None: + if not addr or addr.startswith("["): + return addr + try: + if ipaddress.ip_address(addr).version == 6: + return f"[{addr}]" + except ValueError: + pass + return addr + + +def _exec_vllm_cmd(cmd: list[str], env: dict[str, str]) -> None: + """Entry point for multiprocessing child process.""" + os.execvpe(cmd[0], cmd, env) + + +def launch_server_process( + *, + bind_host: str, + server_port: int, + args, + rank: int, + visible_devices: str, + model_path: str, +) -> multiprocessing.Process: + """Spawn ``vllm serve`` (OpenAI API server) in a subprocess. + + Contrasts with SGLang's launcher, which starts the HTTP server in-process from ``ServerArgs``. + """ + env = os.environ.copy() + env.pop("PYTORCH_CUDA_ALLOC_CONF", None) + env.setdefault("NCCL_CUMEM_ENABLE", "0") + env["CUDA_VISIBLE_DEVICES"] = visible_devices + env.setdefault("VLLM_SERVER_DEV_MODE", "1") + + host_for_subprocess = bind_host.strip("[]") + model = getattr(args, "vllm_model", None) or model_path + tp = args.rollout_num_gpus_per_engine + seed = getattr(args, "seed", 1234) + rank + + cmd = [ + "vllm", + "serve", + str(model), + "--tensor-parallel-size", + str(tp), + "--port", + str(server_port), + "--host", + host_for_subprocess, + "--seed", + str(seed), + "--trust-remote-code", + "--gpu-memory-utilization", + str(getattr(args, "vllm_gpu_memory_utilization", 0.4)), + ] + if getattr(args, "vllm_weight_sync_mode", "auto") == "native": + cmd += ["--weight-transfer-config", '{"backend":"nccl"}'] + if getattr(args, "offload_rollout", False) or getattr(args, "vllm_enable_sleep_mode", False): + cmd += ["--enable-sleep-mode"] + if getattr(args, "vllm_enforce_eager", False): + cmd += ["--enforce-eager"] + if getattr(args, "fp16", False): + cmd += ["--dtype", "float16"] + if getattr(args, "vllm_kv_cache_memory_bytes", None) is not None: + cmd += ["--kv-cache-memory-bytes", str(args.vllm_kv_cache_memory_bytes)] + if args.rollout_max_context_len is not None: + cmd += ["--max-model-len", str(args.rollout_max_context_len)] + + logger.info("Launching vLLM server: %s", " ".join(cmd)) + + p = _spawn_ctx.Process(target=_exec_vllm_cmd, args=(cmd, env)) + p.start() + return p + + +def _wait_server_healthy(base_url: str, process: multiprocessing.Process | None, timeout_s: float = 300.0) -> None: + """Wait until the vLLM server responds on ``GET /health`` (SGLang stacks typically use ``GET /health_generate``).""" + start = time.time() + while True: + try: + response = requests.get(f"{base_url}/health", timeout=3) + if response.status_code == 200: + return + except requests.RequestException: + pass + + if process is not None and not process.is_alive(): + raise RuntimeError(f"vLLM server exited unexpectedly with code {process.exitcode}") + if time.time() - start > timeout_s: + raise TimeoutError(f"Timeout waiting for vLLM server healthy: {base_url}") + time.sleep(2) + + +class VLLMEngine(RayActor): + """Ray actor for vLLM OpenAI HTTP rollout (connect or spawn local ``vllm serve``).""" + + def __init__( + self, + args, + rank: int, + worker_type: str = "regular", + base_gpu_id: int | None = None, + model_path: str | None = None, + sglang_overrides: dict | None = None, + num_gpus_per_engine: int | None = None, + ): + self.args = args + self.rank = rank + self.worker_type = worker_type + self.base_gpu_id = base_gpu_id + self.model_path = model_path or args.hf_checkpoint + # Uniform Ray ``start_engines`` kwargs; unused when launching vLLM over HTTP. + self.sglang_overrides = sglang_overrides or {} + self.num_gpus_per_engine = num_gpus_per_engine + self.process: multiprocessing.Process | None = None + self._weight_version: str | None = None + self._sync_mode = getattr(args, "vllm_weight_sync_mode", "auto") + self._warned_sync_fallback = False + self._pending_reload_version: str | None = None + self._is_local_server = not args.rollout_external + self._native_weight_update_ready = False + # Slime runs one vLLM HTTP process per logical engine; multi-node worker rank is not used. + self.node_rank = 0 + + def _http_base(self) -> str: + return f"http://{self.server_host}:{self.server_port}" + + def init( + self, + dist_init_addr, + port, + nccl_port, + host=None, + disaggregation_bootstrap_port=None, + router_ip=None, + router_port=None, + ): + del dist_init_addr, nccl_port, disaggregation_bootstrap_port + + self.router_ip = router_ip if router_ip is not None else self.args.sglang_router_ip + self.router_port = router_port if router_port is not None else self.args.sglang_router_port + + host = host or get_host_info()[1] + self.server_host = _format_v6_uri(host) + self.server_port = port + + if self.worker_type != "regular": + logger.warning( + "vLLMEngine: worker_type=%s is not used by current vLLM deployment (treated as regular).", + self.worker_type, + ) + + if self.args.rollout_external: + self._init_external() + else: + self._init_normal() + + if self.node_rank == 0 and self.router_ip and self.router_port: + self._register_worker_with_router() + + def _register_worker_with_router(self) -> None: + worker_url = self._http_base() + payload = {"url": worker_url, "worker_type": self.worker_type} + response = requests.post( + f"http://{self.router_ip}:{self.router_port}/workers", + json=payload, + timeout=30, + ) + response.raise_for_status() + + def _deregister_worker_from_router(self) -> None: + if self.node_rank != 0 or not self.router_ip or not self.router_port: + return + worker_url = self._http_base() + try: + all_workers = requests.get( + f"http://{self.router_ip}:{self.router_port}/workers", timeout=30 + ).json()["workers"] + for worker in all_workers: + if worker["url"] == worker_url: + response = requests.delete( + f"http://{self.router_ip}:{self.router_port}/workers/{quote(worker_url, safe='')}", + timeout=30, + ) + response.raise_for_status() + return + logger.warning("Worker %s not found in vllm-router during shutdown.", worker_url) + except Exception as e: + logger.warning("Failed to list/remove worker on vllm-router: %s", e) + + def _init_external(self) -> None: + logger.info("Use external vLLM engine (rank=%s) at %s:%s", self.rank, self.server_host, self.server_port) + base = self._http_base() + _wait_server_healthy(base, process=None) + self._wait_external_config_ready() + + def _wait_external_config_ready(self) -> None: + """External engine: best-effort ``GET /server_info`` TP check (non-fatal).""" + try: + # SGLang external mode uses ``/get_server_info``; vLLM exposes ``/server_info``. + actual = requests.get(f"{self._http_base()}/server_info", params={"config_format": "json"}, timeout=30) + actual.raise_for_status() + body = actual.json() + except requests.RequestException as e: + logger.warning("External vLLM: could not GET /server_info (non-fatal): %s", e) + return + + expect_tp = self.args.rollout_num_gpus_per_engine + parallel_cfg = body.get("vllm_config", {}).get("parallel_config", {}) + actual_tp = parallel_cfg.get("tensor_parallel_size") + if actual_tp is not None and actual_tp != expect_tp: + logger.warning( + "External vLLM server_info TP mismatch: expect=%s actual=%s (weak check)", + expect_tp, + actual_tp, + ) + + def _init_normal(self) -> None: + logger.info("Launch vLLM OpenAI api_server at: %s:%s", self.server_host, self.server_port) + num_gpus = min(self.args.num_gpus_per_node, self.args.rollout_num_gpus_per_engine) + base = self.base_gpu_id if self.base_gpu_id is not None else get_base_gpu_id(self.args, self.rank) + base = _to_local_gpu_id(base) + visible_devices = ",".join(str(base + i) for i in range(num_gpus)) + + bind_host = self.server_host + self.process = launch_server_process( + bind_host=bind_host, + server_port=self.server_port, + args=self.args, + rank=self.rank, + visible_devices=visible_devices, + model_path=self.model_path, + ) + _wait_server_healthy(self._http_base(), process=self.process) + + def _restart_local_server(self) -> None: + if not self._is_local_server: + logger.warning("Skip vLLM reload for external server mode.") + return + if self.process and self.process.is_alive(): + self.process.terminate() + try: + self.process.join(timeout=15) + except Exception: + pass + if self.process.is_alive(): + self.process.kill() + self.process.join(timeout=10) + self._init_normal() + + def _post_json(self, endpoint: str, payload: dict, timeout: float) -> requests.Response: + url = f"{self._http_base()}/{endpoint.lstrip('/')}" + return requests.post(url, json=payload, timeout=timeout) + + def _post_vllm_update_weights_http(self, update_info: dict) -> dict: + """POST ``/update_weights`` with ``{"update_info": ...}`` (vLLM RLHF control plane). + + Same contract as upstream ``examples/online_serving/new_weight_syncing/rlhf_http_nccl.py``: + no ``start_weight_update`` / ``finish_weight_update`` wrapper. + """ + timeout_s = float( + os.environ.get( + "SLIME_VLLM_WEIGHT_TRANSFER_UPDATE_TIMEOUT_SEC", + os.environ.get("SLIME_VLLM_WEIGHT_TRANSFER_HTTP_TIMEOUT_SEC", "900"), + ) + ) + response = self._post_json("update_weights", {"update_info": update_info}, timeout=timeout_s) + response.raise_for_status() + try: + return response.json() + except Exception: + return {"ok": True, "raw": response.text} + + def _run_vllm_weight_update(self, update_info: dict, *, is_checkpoint_format: bool = False): + """Backward-compatible alias for non-NCCL ``update_info`` shapes (e.g. tensor/IPC path).""" + del is_checkpoint_format + return self._post_vllm_update_weights_http(update_info) + + def health_generate(self, timeout: float = 5.0) -> bool: + """Return True if ``GET /health`` succeeds (SGLang uses ``GET /health_generate`` for the same role).""" + if self.node_rank != 0: + return True + response = requests.get(f"{self._http_base()}/health", timeout=timeout) + response.raise_for_status() + return True + + def update_weights_from_tensor( + self, + serialized_named_tensors: list[str], + load_format: str | None = None, + flush_cache: bool = False, + weight_version: str | None = None, + ): + """ + Post tensor metadata via ``/update_weights`` when native weight transfer is ready; otherwise record + ``weight_version`` and return a reload placeholder (no HTTP update on the fallback path). + Contrasts with SGLang, which posts to ``update_weights_from_tensor`` with a different payload shape. + """ + del load_format + if self.node_rank != 0: + return + + if weight_version is not None: + self._weight_version = str(weight_version) + self._pending_reload_version = self._weight_version + if flush_cache: + self.flush_cache() + + update_info = { + "serialized_named_tensors": serialized_named_tensors, + "format": "serialized_named_tensors", + "weight_version": self._weight_version, + } + if self._native_weight_update_ready: + try: + return self._run_vllm_weight_update(update_info, is_checkpoint_format=False) + except Exception as e: + if self._sync_mode == "native": + raise RuntimeError(f"Native vLLM tensor weight update failed: {e}") from e + logger.warning("Native vLLM tensor weight update failed, fallback: %s", e) + self._native_weight_update_ready = False + + self._pending_reload_version = self._weight_version + if not self._warned_sync_fallback: + logger.warning( + "vLLM tensor weight update fallback to reload-on-continue " + "(init_weight_transfer_engine not ready or update failed)." + ) + self._warned_sync_fallback = True + if self._sync_mode == "native": + raise RuntimeError("Native mode requested but weight transfer is not ready.") + return {"ok": True, "mode": "reload", "weight_version": self._weight_version} + + def flush_cache(self): + """Clear prefix cache via ``POST /reset_prefix_cache`` (SGLang uses ``GET /flush_cache``).""" + if self.node_rank != 0: + return + reset_running = bool(getattr(self.args, "vllm_reset_prefix_cache_reset_running", False)) + reset_external = bool(getattr(self.args, "vllm_reset_prefix_cache_reset_external", False)) + params = {"reset_running_requests": reset_running, "reset_external": reset_external} + for _ in range(60): + try: + response = requests.post(f"{self._http_base()}/reset_prefix_cache", params=params, timeout=60) + if response.status_code == 200: + return + except requests.ConnectionError: + raise + except Exception as e: + logger.info("Error resetting vLLM prefix cache: %s", e) + time.sleep(1) + continue + raise TimeoutError("Timeout while resetting vLLM prefix cache (reset_prefix_cache).") + + def get_url(self): + """Worker HTTP base URL, or ``None`` when ``node_rank != 0``.""" + if self.node_rank != 0: + return None + return self._http_base() + + def shutdown(self): + logger.info("Shutdown vLLM engine %s:%s...", self.server_host, self.server_port) + self._deregister_worker_from_router() + if self.args.rollout_external: + return + + if self.process is None or not self.process.is_alive(): + return + pid = self.process.pid + try: + from vllm.utils.system_utils import kill_process_tree + + kill_process_tree(pid) + except Exception as e: + logger.warning("vLLM kill_process_tree failed (%s); terminate root only.", e) + if self.process.is_alive(): + self.process.terminate() + try: + self.process.join(timeout=15) + except Exception: + pass + if self.process.is_alive(): + self.process.kill() + try: + self.process.join(timeout=30) + except Exception: + pass + self.process = None + + def get_weight_version(self): + """ + Prefer ``_weight_version`` if weight sync already set it; else try ``GET /v1/models`` for a stable id string. + + SGLang exposes ``GET /get_weight_version``; vLLM has no name-equivalent route, so semantics differ from that endpoint. + """ + if self.node_rank != 0: + return + if self._weight_version is not None: + return self._weight_version + try: + r = requests.get(f"{self._http_base()}/v1/models", timeout=10) + r.raise_for_status() + data = r.json().get("data") or [] + if data and isinstance(data[0], dict) and "id" in data[0]: + return str(data[0]["id"]) + except requests.RequestException as e: + logger.info("get_weight_version: /v1/models failed (%s)", e) + return None + + def release_memory_occupation(self): + """ + ``POST /sleep`` when sleep mode is enabled (SGLang: ``POST /release_memory_occupation``); otherwise a no-op placeholder dict. + """ + self.flush_cache() + if not getattr(self.args, "vllm_enable_sleep_mode", False): + return {"ok": True, "sleep_mode": False, "note": "vLLM sleep mode disabled; no /sleep call."} + # vLLM ``POST /sleep`` reads ``level`` from query params, not JSON body + # (``vllm.entrypoints.serve.sleep.api_router.sleep``). + level = int(getattr(self.args, "vllm_sleep_level", 1)) + response = requests.post( + f"{self._http_base()}/sleep", + params={"level": level}, + timeout=30, + ) + response.raise_for_status() + try: + return response.json() + except Exception: + return {"ok": True, "raw": response.text} + + def resume_memory_occupation(self, tags: list[str] | None = None): + """``POST /wake_up`` when sleep mode is on (SGLang: ``POST /resume_memory_occupation``); else a small placeholder dict.""" + if not getattr(self.args, "vllm_enable_sleep_mode", False): + return {"ok": True, "sleep_mode": False} + tags = _normalize_vllm_wake_tags(tags) + # vLLM ``POST /wake_up`` uses ``query_params.getlist("tags")``, not JSON. + # Omit params when ``tags`` is empty so the server wakes all tags (see api_router.wake_up). + wake_params: list[tuple[str, str]] | None = ( + [("tags", t) for t in tags] if tags else None + ) + response = requests.post( + f"{self._http_base()}/wake_up", + params=wake_params, + timeout=30, + ) + response.raise_for_status() + try: + return response.json() + except Exception: + return {"ok": True, "raw": response.text} + + def check_weights(self, action: str): + """No vLLM ``weights_checker`` route; return a placeholder (SGLang posts to ``/weights_checker``).""" + del action + return {"ok": True, "supported": False, "note": "vLLM has no weights_checker endpoint."} + + def init_weights_update_group(self, master_address, master_port, rank_offset, world_size, group_name, backend): + """ + Call ``POST /init_weight_transfer_engine`` with an ``init_info`` block (SGLang: ``/init_weights_update_group``). + + ``group_name`` / ``backend`` are accepted for a uniform caller signature but are not sent to vLLM. + If ``vllm_weight_sync_mode`` is not ``native``, the HTTP call is skipped (SGLang still posts to its endpoint). + """ + del group_name, backend + if self._sync_mode != "native": + return {"ok": True, "mode": self._sync_mode, "skipped": True} + + payload = { + "init_info": { + "master_address": master_address, + "master_port": master_port, + "rank_offset": rank_offset, + "world_size": world_size, + } + } + init_timeout_s = float(os.environ.get("SLIME_VLLM_WEIGHT_TRANSFER_HTTP_TIMEOUT_SEC", "900")) + last_error = None + for attempt in range(1, 4): + try: + response = self._post_json("init_weight_transfer_engine", payload, timeout=init_timeout_s) + response.raise_for_status() + self._native_weight_update_ready = True + try: + return response.json() + except Exception: + return {"ok": True, "raw": response.text} + except Exception as e: + last_error = e + self._native_weight_update_ready = False + if attempt < 3: + logger.warning("init_weight_transfer_engine attempt %s/3 failed: %s", attempt, e) + time.sleep(2 * attempt) + if self._sync_mode == "native": + raise RuntimeError(f"vLLM init_weight_transfer_engine failed: {last_error}") from last_error + logger.warning("vLLM native weight transfer init failed, fallback: %s", last_error) + return {"ok": False, "error": str(last_error)} + + def destroy_weights_update_group(self, group_name): + """No vLLM destroy call; return ``None`` (SGLang may ``POST /destroy_weights_update_group`` and swallow errors).""" + del group_name + return None + + def update_weights_from_distributed( + self, + names, + dtypes, + shapes, + group_name, + flush_cache=False, + weight_version: str | None = None, + packed: bool = True, + ): + """NCCL/native path posts ``/update_weights`` (SGLang: ``POST /update_weights_from_distributed``).""" + del group_name + if weight_version is not None: + self._weight_version = str(weight_version) + if flush_cache: + self.flush_cache() + dtype_names = [str(d).replace("torch.", "") for d in dtypes] + if self._native_weight_update_ready: + # Payload matches vLLM NCCL weight transfer (see upstream rlhf_http_nccl example). + update_info = { + "names": names, + "dtype_names": dtype_names, + "shapes": [list(s) for s in shapes], + "packed": bool(packed), + } + try: + return self._post_vllm_update_weights_http(update_info) + except Exception as e: + if self._sync_mode == "native": + raise RuntimeError(f"Native vLLM weight update failed: {e}") from e + logger.warning("Native vLLM weight update failed, fallback: %s", e) + self._native_weight_update_ready = False + + self._pending_reload_version = self._weight_version + if self._sync_mode == "native": + raise RuntimeError("Native mode requested but weight transfer is not ready.") + if not self._warned_sync_fallback and self._sync_mode in ("auto", "reload"): + logger.warning("vLLM weight sync mode=%s: reload-on-continue path may apply.", self._sync_mode) + self._warned_sync_fallback = True + return {"ok": True, "mode": self._sync_mode, "weight_version": self._weight_version} + + def update_weights_from_disk(self, model_path: str, load_format: str | None = None): + """``POST /collective_rpc`` with ``reload_weights`` and ``weights_path`` (SGLang uses a dedicated disk API).""" + if self.node_rank != 0: + return + del load_format + response = requests.post( + f"{self._http_base()}/collective_rpc", + json={ + "method": "reload_weights", + "kwargs": {"weights_path": model_path, "is_checkpoint_format": True}, + }, + timeout=600, + ) + response.raise_for_status() + try: + return response.json() + except Exception: + return {"ok": True, "raw": response.text} + + def pause_generation(self): + """``POST /pause`` with mode query (SGLang: ``POST /pause_generation``); returns the ``requests.Response``.""" + if self.node_rank != 0: + return None + mode = getattr(self.args, "vllm_pause_mode", "keep") + response = requests.post( + f"{self._http_base()}/pause", + params={"mode": mode, "clear_cache": "false"}, + json={}, + timeout=120, + ) + response.raise_for_status() + return response + + def continue_generation(self): + """ + ``POST /resume`` (SGLang: ``POST /continue_generation``); may restart the local child process after a pending reload. + """ + if self.node_rank != 0: + return None + response = requests.post(f"{self._http_base()}/resume", json={}, timeout=120) + response.raise_for_status() + if self._pending_reload_version is not None: + logger.info("Reload vLLM server after weight update, version=%s", self._pending_reload_version) + self._restart_local_server() + self._pending_reload_version = None + return response + + def post_process_weights( + self, + restore_weights_before_load: bool = False, + post_process_quantization: bool = False, + ): + """No vLLM HTTP hook (SGLang: ``POST /post_process_weights``); return a noop placeholder dict.""" + del restore_weights_before_load, post_process_quantization + return {"ok": True, "noop": True, "note": "vLLM post_process is internal to load; no HTTP API."} + + def start_profile( + self, + output_dir: str | None = None, + start_step: int | None = None, + num_steps: int | None = None, + activities: list[str] | None = None, + profile_by_stage: bool = False, + with_stack: bool | None = None, + record_shapes: bool | None = None, + ): + """``POST /start_profile`` with an empty JSON body; kwargs are not forwarded and may be ignored by the server.""" + if self.node_rank != 0: + return None + if any( + x is not None and x is not False + for x in ( + output_dir, + start_step, + num_steps, + activities, + profile_by_stage, + with_stack, + record_shapes, + ) + ): + logger.warning("vLLM start_profile: extra kwargs may be ignored by server; see vLLM profiling docs.") + response = requests.post(f"{self._http_base()}/start_profile", json={}, timeout=30) + response.raise_for_status() + return response + + def stop_profile(self): + """POST ``/stop_profile`` to stop an active server-side profile.""" + if self.node_rank != 0: + return None + response = requests.post(f"{self._http_base()}/stop_profile", json={}, timeout=30) + response.raise_for_status() + return response + + def simulate_crash(self): + if self.args.rollout_external or not getattr(self, "process", None): + logger.info( + "simulate_crash called but no local engine process exists (rollout_external=%s); skip kill", + self.args.rollout_external, + ) + return + logger.info("Simulating crash on vLLM engine %s:%s...", self.server_host, self.server_port) + self.shutdown() diff --git a/slime/ray/rollout.py b/slime/ray/rollout.py index cde7a9e06..90a4fc718 100644 --- a/slime/ray/rollout.py +++ b/slime/ray/rollout.py @@ -5,6 +5,7 @@ import os import random import time +from argparse import Namespace from pathlib import Path from typing import Any @@ -15,7 +16,6 @@ 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_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 @@ -35,6 +35,31 @@ logger = logging.getLogger(__name__) +def _sanitize_vllm_router_args(ra: Any) -> Any: + """Replace negative int fields with dataclass defaults (sglang CLI may use -1; vllm-router rejects it).""" + from vllm_router.router_args import RouterArgs as VR + + fixes: dict[str, Any] = {} + for f in dataclasses.fields(VR): + val = getattr(ra, f.name, None) + if not isinstance(val, int) or val >= 0: + continue + if f.default is not dataclasses.MISSING: + fixes[f.name] = f.default + elif f.default_factory is not dataclasses.MISSING: # type: ignore[attr-defined] + fixes[f.name] = f.default_factory() # type: ignore[misc] + else: + logger.warning("vllm-router: field %r is negative (%s); leaving as-is", f.name, val) + return dataclasses.replace(ra, **fixes) if fixes else ra + + +def _vllm_router_args_from_cli(args: Namespace) -> Any: + from vllm_router.router_args import RouterArgs + + ra = RouterArgs.from_cli_args(args, use_router_prefix=True) + return _sanitize_vllm_router_args(ra) + + @dataclasses.dataclass class ServerGroup: """A group of homogeneous SGLang engines with the same configuration. @@ -88,7 +113,14 @@ def start_engines(self, port_cursors: dict[int, int] | None = None) -> tuple[lis pg, reordered_bundle_indices, reordered_gpu_ids = self.pg - RolloutRayActor = ray.remote(SGLangEngine) + if getattr(self.args, "rollout_backend", "sglang") == "vllm": + from slime.backends.vllm_utils.vllm_engine import VLLMEngine + + 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)): @@ -135,6 +167,7 @@ def start_engines(self, port_cursors: dict[int, int] | None = None) -> tuple[lis rank=global_rank, worker_type=self.worker_type, base_gpu_id=base_gpu_id, + model_path=self.model_path, sglang_overrides=self.sglang_overrides, num_gpus_per_engine=self.num_gpus_per_engine, ) @@ -432,6 +465,18 @@ def _try_ci_fault_injection(self): def dispose(self): for monitor in self._health_monitors: monitor.stop() + # Release inference workers (vLLM / SGLang). debug_rollout_only still hits this path at train.py end. + shutdown_refs = [] + for srv in self.servers.values(): + for group in srv.server_groups: + for eng in group.all_engines: + if eng is not None: + shutdown_refs.append(eng.shutdown.remote()) + if shutdown_refs: + try: + ray.get(shutdown_refs) + except Exception as e: + logger.warning("Engine shutdown during dispose failed (non-fatal): %s", e) logging_utils.finish_tracking(self.args) @property @@ -909,12 +954,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). - - 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. - When ``force_new`` is True (multi-model), always allocate a fresh port. - """ + """Start the rollout HTTP gateway: vllm-router when ``rollout_backend=vllm``, else sglang-router.""" if not force_new and args.sglang_router_ip is not None: return args.sglang_router_ip, args.sglang_router_port @@ -926,38 +966,61 @@ 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 - + use_vllm_router = getattr(args, "rollout_backend", "vllm") == "vllm" from slime.utils.http_utils import run_router - router_args = RouterArgs.from_cli_args(args, use_router_prefix=True) + if use_vllm_router: + router_args = _vllm_router_args_from_cli(args) + impl = "vllm" + else: + from sglang_router.launch_router import RouterArgs + + router_args = RouterArgs.from_cli_args(args, use_router_prefix=True) + impl = "sglang" + 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.log_level = "warning" if use_vllm_router else "warn" router_args.request_timeout_secs = args.sglang_router_request_timeout_secs 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 - - # We will not use the health check from router. - router_args.disable_health_check = True + if use_vllm_router and hasattr(router_args, "vllm_pd_disaggregation"): + router_args.vllm_pd_disaggregation = True + elif hasattr(router_args, "pd_disaggregation"): + router_args.pd_disaggregation = True + if hasattr(router_args, "disable_circuit_breaker"): + router_args.disable_circuit_breaker = True + + # MiniLB is PD-only in vllm-router; non-PD rollout needs the Rust Router (omit SLIME_VLLM_ROUTER_USE_RUST). + if ( + use_vllm_router + and has_pd_disaggregation + and os.environ.get("SLIME_VLLM_ROUTER_USE_RUST", "") != "1" + ): + router_args.mini_lb = True + + if use_vllm_router: + if any(f.name == "disable_health_check" for f in dataclasses.fields(type(router_args))): + router_args.disable_health_check = True + else: + router_args.disable_health_check = True - logger.info(f"Launch router with args: {router_args}") + logger.info("Launch HTTP router (impl=%s) with args: %s", impl, router_args) - process = multiprocessing.Process( + process = multiprocessing.get_context("spawn").Process( target=run_router, - args=(router_args,), + args=((impl, router_args),), ) - process.daemon = True # Set the process as a daemon + process.daemon = True process.start() - # Wait 3 seconds time.sleep(3) - assert process.is_alive() + if not process.is_alive(): + raise RuntimeError( + f"Router subprocess exited (exitcode={process.exitcode}), impl={impl!r}. " + "For vllm-router non-PD mode, install the Rust router (pip wheel with Router); " + "MiniLB is only valid with PD disaggregation. See slime.utils.http_utils run_router logs." + ) logger.info(f"Router launched at {router_ip}:{router_port}, Prometheus port: {router_args.prometheus_port}") return router_ip, router_port diff --git a/slime/rollout/vllm_rollout.py b/slime/rollout/vllm_rollout.py new file mode 100644 index 000000000..adefad96d --- /dev/null +++ b/slime/rollout/vllm_rollout.py @@ -0,0 +1,992 @@ +import asyncio +import copy +import inspect +import json +import logging +import uuid +from argparse import Namespace +from collections.abc import Callable +from contextlib import contextmanager +from typing import Any + +import numpy as np +import vllm_router # noqa: F401 — same side-effect as ``import sglang_router`` in sglang rollout +from tqdm import tqdm + +from slime.rollout.base_types import 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.misc import SingletonMeta, load_function +from slime.utils.processing_utils import ( + build_processor_kwargs, + encode_image_for_rollout_engine, + load_processor, + load_tokenizer, +) +from slime.utils.trace_utils import trace_function, trace_span +from slime.utils.types import Sample + +from .rm_hub import async_rm, batched_async_rm + +__all__ = ["generate_rollout", "get_model_url"] + +logger = logging.getLogger(__name__) + +_PROCESSOR_PROMPT_KEYS = {"input_ids", "attention_mask"} + + +def _coerce_flat_int_token_ids(ids: Any) -> list[int]: + """Turn tokenizer / processor output into a flat ``list[int]`` for vLLM ``/v1/completions`` JSON. + + vLLM deserializes ``prompt`` as ``StringOrArray`` (Rust): a JSON string or a JSON array of integers. + Nested lists, numpy/torch scalars, or non-int elements cause ``422`` deserialization errors. + """ + if ids is None: + return [] + if isinstance(ids, str): + raise TypeError("token ids must not be a str; use the string ``prompt`` field for text prompts") + x = ids + if hasattr(x, "tolist") and not isinstance(x, (list, tuple, str, bytes)): + x = x.tolist() + if isinstance(x, (list, tuple)): + out: list[int] = [] + for item in x: + out.extend(_coerce_flat_int_token_ids(item)) + return out + return [int(x)] + + +def _prepare_prompt_ids(sample: Sample, tokenizer, processor: Any) -> list[int]: + raw_multimodal_inputs = sample.multimodal_inputs or {} + has_multimodal_inputs = any(value is not None for value in raw_multimodal_inputs.values()) + reuse_existing_input_ids = bool(sample.tokens) and ( + sample.multimodal_train_inputs is not None or not has_multimodal_inputs + ) + + if processor and has_multimodal_inputs and not reuse_existing_input_ids: + processor_output = processor(text=sample.prompt, **build_processor_kwargs(raw_multimodal_inputs)) + prompt_ids = processor_output["input_ids"][0] + if sample.multimodal_train_inputs is None: + sample.multimodal_train_inputs = { + k: v for k, v in processor_output.items() if k not in _PROCESSOR_PROMPT_KEYS + } or None + return _coerce_flat_int_token_ids(prompt_ids) + + if reuse_existing_input_ids: + return _coerce_flat_int_token_ids(sample.tokens) + + return _coerce_flat_int_token_ids(tokenizer.encode(sample.prompt, add_special_tokens=False)) + + +def _base_dataset_prompt_ids(sample: Sample, tokenizer, processor: Any) -> list[int]: + """Token ids for the dataset prompt only (never reuse ``sample.tokens``). + + Used for partial-continuation budgeting to match ``dev_vllm`` ``sglang_rollout``: + ``max_new_tokens -= len(sample.tokens) - len(base_prompt_ids)`` when ``sample.response`` is non-empty. + """ + raw_multimodal_inputs = sample.multimodal_inputs or {} + has_multimodal_inputs = any(value is not None for value in raw_multimodal_inputs.values()) + if processor and has_multimodal_inputs: + processor_output = processor(text=sample.prompt, **build_processor_kwargs(raw_multimodal_inputs)) + prompt_ids = processor_output["input_ids"][0] + return _coerce_flat_int_token_ids(prompt_ids) + return _coerce_flat_int_token_ids(tokenizer.encode(sample.prompt, add_special_tokens=False)) + + +def get_model_url(args: Namespace, model_name: str, endpoint: str = "/v1/completions") -> str: + """Return the router URL for a named model. + + Use this in custom rollout functions to route requests to a specific + model when multiple models are deployed via ``--sglang-config``:: + + url = get_model_url(args, "ref", "/v1/completions") + resp = await post(url, json=payload) + + Falls back to the default router if *model_name* is not found or + ``sglang_model_routers`` is not set. + """ + routers = getattr(args, "sglang_model_routers", None) + if routers and model_name in routers: + ip, port = routers[model_name] + return f"http://{ip}:{port}{endpoint}" + return f"http://{args.sglang_router_ip}:{args.sglang_router_port}{endpoint}" + + +async def _router_worker_urls(args: Namespace) -> list[str]: + """Resolve worker base URLs from the vLLM router (same HTTP shape as SGLang router).""" + base = f"http://{args.sglang_router_ip}:{args.sglang_router_port}" + try: + response = await get(f"{base}/workers") + return [worker["url"] for worker in response["workers"]] + except Exception: + response = await get(f"{base}/list_workers") + return list(response["urls"]) + + +async def _resume_vllm_workers(urls: list[str]) -> None: + """Call ``POST /resume`` on each worker after ``pause?mode=abort`` so engines accept traffic again.""" + if not urls: + return + logger.info("vLLM rollout: resuming workers after abort drain: %s", urls) + resume_tasks = [post(f"{url.rstrip('/')}/resume", {}, max_retries=3) for url in urls] + resume_results = await asyncio.gather(*resume_tasks, return_exceptions=True) + for url, result in zip(urls, resume_results, strict=False): + if isinstance(result, Exception): + logger.warning("Failed to resume vLLM worker at %s: %s", url, result) + + +def _openai_meta_from_completion_choice(args: Namespace, choice: dict, usage: dict | None) -> dict[str, Any]: + """Build a minimal ``meta_info``-like dict for :meth:`Sample.update_from_meta_info`.""" + fr = choice.get("finish_reason") or "stop" + if isinstance(fr, dict): + finish = fr + else: + if fr == "length": + typ = "length" + elif fr in ("abort", "cancelled"): + typ = "abort" + else: + typ = "stop" + finish = {"type": typ} + meta: dict[str, Any] = {"finish_reason": finish} + if usage: + meta["prompt_tokens"] = usage.get("prompt_tokens", 0) + meta["completion_tokens"] = usage.get("completion_tokens", 0) + return meta + + +def _apply_vllm_routed_experts( + args: Namespace, + sample: Sample, + _output: dict, + choice: dict, +) -> None: + """Populate ``sample.rollout_routed_experts`` from vLLM ``choices[].routed_experts`` when enabled. + + vLLM exposes MoE routing replay on the **per-completion** ``CompletionOutput`` as a single + ``routed_experts`` ndarray (shape ``[num_positions, num_layers, topk]``); the V1 scheduler + builds it once per finished request from KV slot indices for ``request.num_tokens - 1`` + positions — there is **no** separate ``prompt_routed_experts`` key on the HTTP completion + payload (confirmed absent in upstream ``vllm``; see ``CompletionOutput`` in + ``vllm/outputs.py`` and ``_get_routed_experts`` in ``vllm/v1/core/sched/scheduler.py``). + When the OpenAI layer forwards it, it appears as an extra field on the choice object + (Pydantic ``extra="allow"`` on ``CompletionResponseChoice``). + """ + if not getattr(args, "use_rollout_routing_replay", False): + return + gen_re = choice.get("routed_experts") + if gen_re is None: + return + arr = np.asarray(gen_re, dtype=np.int32) + n_tok = len(sample.tokens) + expected_rows = max(0, n_tok - 1) + if arr.ndim != 3: + logger.warning(f"Unexpected routed_experts ndim={arr.ndim} shape={arr.shape}") + return + if arr.shape[0] == n_tok: + arr = arr[:-1] + elif arr.shape[0] != expected_rows: + logger.warning( + f"routed_experts row count {arr.shape[0]} not in {{{expected_rows}, {n_tok}}}; " + "skipping rollout_routed_experts assign", + ) + return + nl = getattr(args, "num_layers", None) + mtk = getattr(args, "moe_router_topk", None) + if nl is not None and mtk is not None and (arr.shape[1] != nl or arr.shape[2] != mtk): + logger.warning( + f"routed_experts shape {arr.shape} does not match args (num_layers={nl}, moe_router_topk={mtk})", + ) + return + sample.rollout_routed_experts = arr + + +def _fallback_tokens_from_text_and_completion_logprobs( + tokenizer, choice: dict, completion_text: str +) -> tuple[list[int], list[float]]: + """Fallback when vLLM did not return ``token_ids``: approximate from string logprobs or re-tokenize.""" + lp = choice.get("logprobs") + if not lp or not isinstance(lp, dict): + toks = tokenizer.encode(completion_text, add_special_tokens=False) + return toks, [0.0] * len(toks) + + token_logprobs = lp.get("token_logprobs") + tokens_field = lp.get("tokens") + if token_logprobs and tokens_field and len(token_logprobs) == len(tokens_field): + new_toks: list[int] = [] + new_lps: list[float] = [] + for piece, logp in zip(tokens_field, token_logprobs, strict=False): + if logp is None: + continue + ids = tokenizer.encode(piece, add_special_tokens=False) + if not ids: + continue + per = float(logp) / max(len(ids), 1) + new_toks.extend(ids) + new_lps.extend([per] * len(ids)) + if new_toks: + return new_toks, new_lps + + toks = tokenizer.encode(completion_text, add_special_tokens=False) + return toks, [0.0] * len(toks) + + +def _vllm_engine_tokens_and_logprobs( + tokenizer, + choice: dict[str, Any], + completion_text: str, + *, + is_chat: bool, +) -> tuple[list[int], list[float]]: + """Parse engine ``token_ids`` and per-token logprobs from an OpenAI choice (requires ``return_token_ids``).""" + tids_raw = choice.get("token_ids") + if isinstance(tids_raw, list) and tids_raw and all(isinstance(x, int) for x in tids_raw): + tids = [int(x) for x in tids_raw] + lps: list[float] = [] + lp = choice.get("logprobs") + if not isinstance(lp, dict): + return tids, [0.0] * len(tids) + + if is_chat: + content = lp.get("content") + if isinstance(content, list) and len(content) == len(tids): + for item in content: + if isinstance(item, dict): + lps.append(float(item.get("logprob", 0.0))) + else: + lps.append(0.0) + return tids, lps + if isinstance(content, list) and content: + # Partial content: pad / truncate to token_ids length if possible + for i in range(len(tids)): + if i < len(content) and isinstance(content[i], dict): + lps.append(float(content[i].get("logprob", 0.0))) + else: + lps.append(0.0) + return tids, lps + return tids, [0.0] * len(tids) + + token_logprobs = lp.get("token_logprobs") + if isinstance(token_logprobs, list) and len(token_logprobs) == len(tids): + for p in token_logprobs: + lps.append(0.0 if p is None else float(p)) + return tids, lps + + return tids, [0.0] * len(tids) + + if is_chat: + lp = choice.get("logprobs") + content = lp.get("content") if isinstance(lp, dict) else None + if isinstance(content, list) and content: + tids_ch: list[int] = [] + lps_ch: list[float] = [] + for item in content: + if not isinstance(item, dict): + continue + tok = item.get("token") + if not isinstance(tok, str): + continue + ids = tokenizer.encode(tok, add_special_tokens=False) + if not ids: + continue + lv = float(item.get("logprob", 0.0)) + per = lv / max(len(ids), 1) + tids_ch.extend(ids) + lps_ch.extend([per] * len(ids)) + if tids_ch: + return tids_ch, lps_ch + + return _fallback_tokens_from_text_and_completion_logprobs(tokenizer, choice, completion_text) + + +def _align_engine_tokens_and_logprobs( + new_response_tokens: list[int], new_response_log_probs: list[float] +) -> tuple[list[int], list[float]]: + """Pad or truncate logprobs so ``len(.) == len(new_response_tokens)`` (SGLang always has matched OTL pairs).""" + n = len(new_response_tokens) + if n == 0: + return [], [] + m = len(new_response_log_probs) + if m == n: + return new_response_tokens, [float(x) for x in new_response_log_probs] + if m > n: + return new_response_tokens, [float(x) for x in new_response_log_probs[:n]] + return new_response_tokens, [float(x) for x in new_response_log_probs] + [0.0] * (n - m) + + +class GenerateState(metaclass=SingletonMeta): + """ + The global state for the generation process. + """ + + def __init__(self, args: Namespace) -> None: + # persistent state for the generation process + self.args = args + self.tokenizer = load_tokenizer(args.hf_checkpoint, trust_remote_code=True) + 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 + ) + self.sampling_params: dict[str, Any] = dict( + temperature=args.rollout_temperature, + top_p=args.rollout_top_p, + top_k=args.rollout_top_k, + max_new_tokens=args.rollout_max_response_len, + stop=args.rollout_stop, + stop_token_ids=args.rollout_stop_token_ids, + skip_special_tokens=args.rollout_skip_special_tokens, + no_stop_trim=True, + spaces_between_special_tokens=False, + ) + + if getattr(args, "sglang_enable_deterministic_inference", False): + sampling_seed_base = args.rollout_seed + self.group_sampling_seeds = [sampling_seed_base + i for i in range(args.n_samples_per_prompt)] + + # dp rank balancing + self.dp_counts = [0] * (args.sglang_dp_size or 1) + self.dp_rank = 0 + + self.reset() + + @contextmanager + def dp_rank_context(self): + candidates = [i for i, count in enumerate(self.dp_counts) if count == min(self.dp_counts)] + dp_rank = int(np.random.choice(candidates)) + self.dp_counts[dp_rank] += 1 + self.dp_rank = dp_rank + try: + yield dp_rank + finally: + self.dp_counts[dp_rank] -= 1 + assert self.dp_counts[dp_rank] >= 0 + + def reset(self) -> None: + self.remaining_batch_size = 0 + self.pendings = set() + self.aborted = False + + def submit_generate_tasks(self, samples: list[list[Sample]]) -> None: + for group in samples: + self.pendings.add( + asyncio.create_task( + # submit a group of samples as a single task. + generate_and_rm_group( + self.args, + group, + sampling_params=self.sampling_params.copy(), + evaluation=False, + ) + ) + ) + self.remaining_batch_size += len(samples) + + +def _build_inference_sampling_params(sampling_params: dict[str, Any]) -> dict[str, Any]: + """Map rollout ``sampling_params`` to vLLM ``/inference/v1/generate`` ``sampling_params`` body.""" + sp: dict[str, Any] = { + "max_tokens": sampling_params["max_new_tokens"], + "temperature": sampling_params["temperature"], + "top_p": sampling_params["top_p"], + "logprobs": 1, + } + tk = sampling_params.get("top_k") + if tk is not None and tk > 0: + sp["top_k"] = tk + if sampling_params.get("stop"): + sp["stop"] = sampling_params["stop"] + if sampling_params.get("stop_token_ids"): + sp["stop_token_ids"] = sampling_params["stop_token_ids"] + if sampling_params.get("seed") is not None: + sp["seed"] = sampling_params["seed"] + if sampling_params.get("skip_special_tokens") is not None: + sp["skip_special_tokens"] = bool(sampling_params["skip_special_tokens"]) + return sp + + +def _mm_render_response_to_generate_body(render_data: Any, model: str) -> dict[str, Any]: + """Turn ``/v1/chat/completions/render`` JSON into a ``/inference/v1/generate`` request body (minus ``sampling_params``). + + vLLM stable docs use a flat dict with ``token_ids`` and optional ``features``; some builds return + ``[conversation, engine_prompts]`` from the render route — normalize both. + """ + if isinstance(render_data, dict) and isinstance(render_data.get("token_ids"), list): + body = copy.deepcopy(render_data) + body.setdefault("model", model) + return body + + if isinstance(render_data, list) and len(render_data) >= 2: + engine_prompts = render_data[1] + if not isinstance(engine_prompts, list) or not engine_prompts: + raise ValueError("chat/render: expected non-empty engine_prompts list") + p = engine_prompts[0] + if not isinstance(p, dict): + raise ValueError("chat/render: engine_prompts[0] must be a dict") + token_ids = p.get("prompt_token_ids") or p.get("token_ids") + if not isinstance(token_ids, list) or not token_ids: + raise ValueError("chat/render: missing prompt_token_ids / token_ids on engine prompt") + body: dict[str, Any] = {"token_ids": [int(x) for x in token_ids], "model": model} + if p.get("features") is not None: + body["features"] = p["features"] + elif isinstance(p.get("multi_modal_data"), dict): + try: + body["features"] = json.dumps(p["multi_modal_data"], default=str) + except TypeError: + pass + if p.get("cache_salt") is not None: + body["cache_salt"] = p["cache_salt"] + return body + + raise ValueError( + "chat/render: unexpected JSON shape; expected a dict with token_ids or " + "[conversation, engine_prompts] list" + ) + + +def _build_completion_payload( + args: Namespace, sampling_params: dict[str, Any], prompt_field: str | list[int] +) -> dict[str, Any]: + """Map shared ``sampling_params`` to vLLM OpenAI ``/v1/completions`` body.""" + body: dict[str, Any] = { + "model": args.hf_checkpoint, + "prompt": prompt_field, + "max_tokens": sampling_params["max_new_tokens"], + "temperature": sampling_params["temperature"], + "top_p": sampling_params["top_p"], + "logprobs": 1, + "return_token_ids": True, + } + tk = sampling_params.get("top_k") + if tk is not None and tk > 0: + body["top_k"] = tk + if sampling_params.get("stop"): + body["stop"] = sampling_params["stop"] + if sampling_params.get("stop_token_ids"): + body["stop_token_ids"] = sampling_params["stop_token_ids"] + if sampling_params.get("skip_special_tokens"): + body["skip_special_tokens"] = True + if sampling_params.get("seed") is not None: + body["seed"] = sampling_params["seed"] + return body + + +async def generate(args: Namespace, sample: Sample, sampling_params: dict[str, Any]) -> Sample: + """Generate using vLLM OpenAI-compatible HTTP API on the router host/port.""" + if args.ci_test: + assert isinstance(sample.prompt, str) + + state = GenerateState(args) + base = f"http://{args.sglang_router_ip}:{args.sglang_router_port}" + + assert ( + sample.status == Sample.Status.PENDING or sample.status == Sample.Status.ABORTED + ), f"Sample status is {sample.status}" + + prompt_ids = _prepare_prompt_ids(sample, state.tokenizer, state.processor) + base_prompt_ids = _base_dataset_prompt_ids(sample, state.tokenizer, state.processor) + + params = dict(sampling_params) + if len(sample.response) > 0: + params["max_new_tokens"] -= len(sample.tokens) - len(base_prompt_ids) + + assert params["max_new_tokens"] >= 0, ( + f"max_new_tokens: {params['max_new_tokens']} should not be less than 0 " + f"(after partial continuation adjustment; tokens={len(sample.tokens)}, base_prompt={len(base_prompt_ids)})" + ) + if params["max_new_tokens"] == 0: + sample.status = Sample.Status.TRUNCATED + return sample + + images = sample.multimodal_inputs.get("images") if sample.multimodal_inputs else None + + if not sample.tokens: + sample.tokens = prompt_ids + + # Use session_id for consistent hashing routing (SGLang Model Gateway) + headers = None + if sample.session_id: + if getattr(args, "router_policy", None) == "consistent_hashing": + headers = {"X-SMG-Routing-Key": sample.session_id} + + if images: + # Disaggregated MM flow: render (preprocess) then tokens-only generate — see vLLM docs + # ``examples/online_serving/disaggregated_serving`` (``/v1/chat/completions/render`` + + # ``/inference/v1/generate``). + content: list[dict[str, Any]] = [{"type": "text", "text": sample.prompt}] + for image in images: + data_url = encode_image_for_rollout_engine(image) + content.append({"type": "image_url", "image_url": {"url": data_url}}) + render_payload = { + "model": args.hf_checkpoint, + "messages": [{"role": "user", "content": content}], + } + render_url = f"{base}/v1/chat/completions/render" + with trace_span(sample, "vllm_mm_render", attrs={"model": args.hf_checkpoint}): + render_data = await post(render_url, render_payload, headers=headers) + generate_body = _mm_render_response_to_generate_body(render_data, args.hf_checkpoint) + generate_body["sampling_params"] = _build_inference_sampling_params(params) + gen_url = f"{base}/inference/v1/generate" + with trace_span(sample, "vllm_mm_generate", attrs={"max_tokens": params["max_new_tokens"]}): + output = await post(gen_url, generate_body, headers=headers) + choice = output["choices"][0] + skip_sp = params.get("skip_special_tokens") + skip_decode = True if skip_sp is None else bool(skip_sp) + out_ids = choice.get("token_ids") or [] + text = ( + state.tokenizer.decode(out_ids, skip_special_tokens=skip_decode) + if isinstance(out_ids, list) and out_ids + else "" + ) + usage = output.get("usage") + meta = _openai_meta_from_completion_choice(args, choice, usage) + new_response_tokens, new_response_log_probs = _vllm_engine_tokens_and_logprobs( + state.tokenizer, choice, text, is_chat=True + ) + new_response_tokens, new_response_log_probs = _align_engine_tokens_and_logprobs( + new_response_tokens, new_response_log_probs + ) + else: + url = f"{base}/v1/completions" + # vLLM OpenAI ``prompt``: JSON string or flat array of integers. On partial continuation, send full + # ``sample.tokens`` as integer ids (aligned with dev_vllm ``sglang_rollout`` + vLLM backend). + if len(sample.response) > 0: + prompt_field = _coerce_flat_int_token_ids(sample.tokens) + elif isinstance(sample.prompt, str): + prompt_field = sample.prompt + else: + prompt_field = prompt_ids + payload = _build_completion_payload(args, params, prompt_field) + with trace_span(sample, "vllm_completions", attrs={"max_new_tokens": params["max_new_tokens"]}): + output = await post(url, payload, headers=headers) + choice = output["choices"][0] + text = choice.get("text") or "" + usage = output.get("usage") + meta = _openai_meta_from_completion_choice(args, choice, usage) + new_response_tokens, new_response_log_probs = _vllm_engine_tokens_and_logprobs( + state.tokenizer, choice, text, is_chat=False + ) + new_response_tokens, new_response_log_probs = _align_engine_tokens_and_logprobs( + new_response_tokens, new_response_log_probs + ) + + if new_response_tokens: + meta["output_token_logprobs"] = [ + [float(lp), int(tid)] for lp, tid in zip(new_response_log_probs, new_response_tokens, strict=True) + ] + + # Update sample with tokens directly - avoiding re-tokenization + sample.tokens = sample.tokens + new_response_tokens + sample.response_length += len(new_response_tokens) + sample.response += 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 + + _apply_vllm_routed_experts(args, sample, output, choice) + + sample.update_from_meta_info(args, meta) + return sample + + +@trace_function("generate_and_rm", target="sample") +async def generate_and_rm( + args: Namespace, + sample: Sample | list[Sample], + sampling_params: dict[str, Any], + evaluation: bool = False, +) -> Sample | list[Sample]: + if isinstance(sample, list): + return await asyncio.gather( + *[generate_and_rm(args, s, sampling_params, evaluation=evaluation) for s in sample] + ) + + # mask previous off-policy generation for partial rollout + if args.partial_rollout and args.mask_offpolicy_in_partial_rollout and sample.response_length > 0: + sample.loss_mask = [0] * sample.response_length + + # For samples with existing response, check if they're complete + if sample.status == Sample.Status.COMPLETED or sample.status == Sample.Status.TRUNCATED: + assert sample.response is not None + if not args.group_rm: + assert sample.reward is not None + return sample + + state = GenerateState(args) + + # generate + async with state.semaphore: + if state.aborted: + sample.status = Sample.Status.ABORTED + return sample + + with state.dp_rank_context() as _: + # Check sample.generate_function_path for per-sample custom_generate_function_path (e.g., from eval dataset config) + custom_func_path = getattr(sample, "generate_function_path", None) or args.custom_generate_function_path + + if custom_func_path is not None: + custom_generate_func = load_function(custom_func_path) + # if signature has evaluation, pass evaluation + if "evaluation" in inspect.signature(custom_generate_func).parameters: + sample = await custom_generate_func(args, sample, sampling_params, evaluation=evaluation) + else: + sample = await custom_generate_func(args, sample, sampling_params) + else: + sample = await generate(args, sample, sampling_params) + + # for the rm that need the whole group, we will not do the rm here + if args.group_rm: + return sample + + if isinstance(sample, list): + samples = sample + if any(sample.status == Sample.Status.ABORTED for sample in samples): + return samples + + samples_need_reward = [sample for sample in samples if sample.reward is None] + with trace_span(samples_need_reward, "reward_model"): + rewards = await batched_async_rm(args, samples_need_reward) + for sample, reward in zip(samples_need_reward, rewards, strict=False): + sample.reward = reward + return samples + else: + if sample.status == Sample.Status.ABORTED: + return sample + # Some custom generate paths may have already filled the reward. + if sample.reward is None: + with trace_span(sample, "reward_model"): + sample.reward = await async_rm(args, sample) + + return sample + + +@trace_function( + "generate_and_rm_group", + target="group", + attrs_getter=lambda args, group, sampling_params, evaluation=False: {"group_size": len(group)}, +) +async def generate_and_rm_group( + args: Namespace, group: list[Sample], sampling_params: dict[str, Any], evaluation: bool = False +) -> list[Sample]: + state = GenerateState(args) + + if state.aborted: + return group + + # Generate a unique session_id for each sample in the group + for sample in group: + if sample.session_id is None: + sample.session_id = str(uuid.uuid4()) + + tasks = [] + for idx, sample in enumerate(group): + current_sampling_params = sampling_params.copy() + if getattr(args, "sglang_enable_deterministic_inference", False): + seed = state.group_sampling_seeds[idx] + current_sampling_params["seed"] = seed + tasks.append( + asyncio.create_task(generate_and_rm(args, sample, current_sampling_params, evaluation=evaluation)) + ) + + group = await asyncio.gather(*tasks) + + # for the rm that need the whole group, we will do the rm here + if not state.aborted and args.group_rm: + with trace_span(group, "group_reward_model"): + rewards = await batched_async_rm(args, group) + for sample, reward in zip(group, rewards, strict=False): + sample.reward = reward + + return group + + +async def abort(args: Namespace, rollout_id: int) -> list[list[Sample]]: + aborted_samples: list[list[Sample]] = [] + + state = GenerateState(args) + assert not state.aborted + state.aborted = True + + urls: list[str] = [] + paused_workers = False + if state.pendings: + urls = await _router_worker_urls(args) + logger.info("vLLM rollout abort (pause) for workers: %s", urls) + pause_tasks = [post(f"{url.rstrip('/')}/pause?mode=abort", {}, max_retries=3) for url in urls] + pause_results = await asyncio.gather(*pause_tasks, return_exceptions=True) + for url, result in zip(urls, pause_results, strict=False): + if isinstance(result, Exception): + logger.warning("Failed to pause/abort worker at %s: %s", url, result) + paused_workers = True + + # make sure all the pending tasks are finished + count = 0 + while state.pendings: + done, state.pendings = await asyncio.wait(state.pendings, return_when=asyncio.FIRST_COMPLETED) + + if not args.partial_rollout: + continue + + # for partial rollout, collect the partial samples into the data buffer + for task in done: + group = task.result() + for sample in group: + if sample.response and "start_rollout_id" not in sample.metadata: + sample.metadata["start_rollout_id"] = rollout_id + aborted_samples.append(group) + count += len(group) + + if args.partial_rollout: + logger.info(f"Collected {count} partial samples into the data buffer") + + state.pendings = set() + if paused_workers: + await _resume_vllm_workers(urls) + return aborted_samples + + +async def generate_rollout_async( + args: Namespace, rollout_id: int, data_source: Callable[[int], list[list[Sample]]] +) -> tuple[RolloutFnTrainOutput, list[list[Sample]]]: + """An example to implement the generate_rollout function for an rule based rm rollout generation. + + Args: + args: the whole args + rollout_id: int, the id of the rollout, used for deterministic data generation + data_source: the data source to fetch + + Returns: + tuple[RolloutFnTrainOutput, list[list[Sample]]]: + - data: a list of groups of samples generated by the rollout, length equals `rollout_batch_size` + - aborted_samples: any partial groups collected during abort when partial_rollout is enabled + """ + assert args.rollout_global_dataset + + state = GenerateState(args) + + # instantiate data filters + dynamic_filter = ( + load_function(args.dynamic_sampling_filter_path) if args.dynamic_sampling_filter_path is not None else None + ) + + metric_gatherer = MetricGatherer() + + # target_data_size is the total number of valid samples to get + target_data_size = args.rollout_batch_size + + data = [] + all_data = [] + do_print = True + pbar = tqdm(total=target_data_size * args.n_samples_per_prompt, desc="Rollout generation") + while len(data) < target_data_size: + while state.remaining_batch_size < target_data_size: + # get samples from the buffer and submit the generation requests. + samples = data_source(args.over_sampling_batch_size) + state.submit_generate_tasks(samples) + + # wait for the generation to finish + done, state.pendings = await asyncio.wait(state.pendings, return_when=asyncio.FIRST_COMPLETED) + for task in done: + group: list[Sample] = task.result() + + if do_print: + sample = group[0][0] if isinstance(group[0], list) else group[0] + logger.info( + f"First rollout sample: {[str(sample.prompt) + sample.response]}, label: {str(sample.label)[:100]}, reward: {sample.reward}", + ) + do_print = False + + assert len(group) == args.n_samples_per_prompt + all_data.append(group) + dynamic_filter_output = call_dynamic_filter(dynamic_filter, args, group) + if not dynamic_filter_output.keep: + metric_gatherer.on_dynamic_filter_drop(reason=dynamic_filter_output.reason) + state.remaining_batch_size -= 1 + continue + + # add the samples to the data + # NOTE: here we have not stored all the unused samples back to the data buffer. + if len(data) < target_data_size: + data.append(group) + pbar.update(args.n_samples_per_prompt) + + pbar.close() + sample = data[-1][0][0] if isinstance(data[-1][0], list) else data[-1][0] + logger.info( + f"Finish rollout: {[str(sample.prompt) + sample.response]}, label: {str(sample.label)[:100]}, reward: {sample.reward}", + ) + + # there are still some unfinished requests, abort them + aborted_samples = await abort(args, rollout_id) + + assert len(data) == args.rollout_batch_size, f"Got {len(data)} samples, expected {args.rollout_batch_size}" + data = sorted(data, key=lambda group: group[0][0].index if isinstance(group[0], list) else group[0].index) + all_samples = sorted( + all_data, key=lambda group: group[0][0].index if isinstance(group[0], list) else group[0].index + ) + + # reset the global state to prevent effects on the next rollout or eval. + state.reset() + if args.rollout_sample_filter_path is not None: + filter_func = load_function(args.rollout_sample_filter_path) + filter_func(args, data) + + # There can be circumstances where users want to process all samples including filtered ones. + if args.rollout_all_samples_process_path is not None: + process_func = load_function(args.rollout_all_samples_process_path) + process_func(args, all_samples, data_source) + + return RolloutFnTrainOutput(samples=data, metrics=metric_gatherer.collect()), aborted_samples + + +EVAL_PROMPT_DATASET = {} + + +async def eval_rollout(args: Namespace, rollout_id: int) -> tuple[dict[str, dict[str, list[Any]]], list[list[Sample]]]: + assert not args.group_rm, "Group RM is not supported for eval rollout" + + coros = [] + for dataset_cfg in getattr(args, "eval_datasets", []) or []: + coros.append(eval_rollout_single_dataset(args, rollout_id, dataset_cfg)) + results_list = await asyncio.gather(*coros) + results = {} + for r in results_list: + results.update(r) + return RolloutFnEvalOutput(data=results), [] + + +async def eval_rollout_single_dataset( + args: Namespace, rollout_id: int, dataset_cfg: EvalDatasetConfig +) -> dict[str, dict[str, list[Any]]]: + """An example to implement the eval_rollout function for an rule based rm rollout generation. + + Args: + args: the whole args + rollout_id: int, the id of the rollout, used for deterministic data generation + dataset_cfg: configuration of the dataset + """ + assert not args.group_rm, "Group RM is not supported for eval rollout" + + global EVAL_PROMPT_DATASET + + cache_key = dataset_cfg.cache_key + (args.hf_checkpoint, args.apply_chat_template) + if cache_key not in EVAL_PROMPT_DATASET: + tokenizer = load_tokenizer(args.hf_checkpoint, trust_remote_code=True) + processor = load_processor(args.hf_checkpoint, trust_remote_code=True) + EVAL_PROMPT_DATASET[cache_key] = Dataset( + path=dataset_cfg.path, + tokenizer=tokenizer, + processor=processor, + max_length=args.eval_max_prompt_len, + prompt_key=dataset_cfg.input_key, + label_key=dataset_cfg.label_key, + multimodal_keys=args.multimodal_keys, + metadata_key=dataset_cfg.metadata_key, + tool_key=dataset_cfg.tool_key, + apply_chat_template=args.apply_chat_template, + apply_chat_template_kwargs=args.apply_chat_template_kwargs, + ) + dataset = EVAL_PROMPT_DATASET[cache_key] + + base_sampling_params = dict( + temperature=dataset_cfg.temperature, + top_p=dataset_cfg.top_p, + top_k=dataset_cfg.top_k, + max_new_tokens=dataset_cfg.max_response_len, + stop=args.rollout_stop, + stop_token_ids=args.rollout_stop_token_ids, + skip_special_tokens=args.rollout_skip_special_tokens, + no_stop_trim=True, + spaces_between_special_tokens=False, + ) + + tasks = [] + # do multiple samples for eval prompts + sample_index = 0 + for _i, prompt_sample in enumerate(dataset.samples): + for j in range(dataset_cfg.n_samples_per_eval_prompt): + # use the same prompt for multiple samples + sample = copy.deepcopy(prompt_sample) + sample.index = sample_index + sample_index += 1 + sample.metadata = dataset_cfg.inject_metadata(getattr(sample, "metadata", None)) + sample.generate_function_path = getattr(dataset_cfg, "custom_generate_function_path", None) + sampling_params = base_sampling_params + if getattr(args, "sglang_enable_deterministic_inference", False): + sampling_params = base_sampling_params.copy() + sampling_params["seed"] = args.rollout_seed + j + tasks.append( + asyncio.create_task( + generate_and_rm( + args, + sample, + sampling_params=sampling_params, + evaluation=True, + ) + ) + ) + + data = [] + do_print = True + pbar = tqdm(total=len(tasks), desc=f"Eval {dataset_cfg.name}", disable=not do_print) + for coro in asyncio.as_completed(tasks): + sample = await coro + if do_print: + logged_sample = sample[0] if isinstance(sample, list) else sample + logger.info( + "eval_rollout_single_dataset example data: " + f"{[str(logged_sample.prompt) + logged_sample.response]} " + f"reward={logged_sample.reward}" + ) + do_print = False + if isinstance(sample, list): + data.extend(sample) + else: + data.append(sample) + pbar.update(1) + pbar.close() + + data.sort(key=lambda sample: sample.index) + + reward_key = args.eval_reward_key or args.reward_key + return { + dataset_cfg.name: { + "rewards": [sample.reward if not reward_key else sample.reward[reward_key] for sample in data], + "truncated": [sample.status == Sample.Status.TRUNCATED for sample in data], + "samples": data, + } + } + + +def generate_rollout( + args: Namespace, rollout_id: int, data_source: Any, evaluation: bool = False +) -> RolloutFnTrainOutput | RolloutFnEvalOutput: + """An example to implement the generate_rollout function for an rule based rm rollout generation. + + Args: + args: the whole args + rollout_id: int, the id of the rollout, used for deterministic data generation + data_source: the data source to get and store samples + evaluation: bool, whether the rollout is for evaluation or not + + Returns: + RolloutFnTrainOutput | RolloutFnEvalOutput: the output of the rollout + """ + assert args.rollout_global_dataset + if evaluation: + output, _ = run(eval_rollout(args, rollout_id)) + return output + + output, aborted_samples = run(generate_rollout_async(args, rollout_id, data_source.get_samples)) + if aborted_samples: + data_source.add_samples(aborted_samples) + return output \ No newline at end of file diff --git a/slime/utils/arguments.py b/slime/utils/arguments.py index e8a173078..316444afd 100644 --- a/slime/utils/arguments.py +++ b/slime/utils/arguments.py @@ -205,6 +205,37 @@ def add_train_arguments(parser): # rollout def add_rollout_arguments(parser): + parser.add_argument( + "--rollout-backend", + type=str, + choices=["sglang", "vllm"], + default="vllm", + help="Backend for rollout inference service.", + ) + parser.add_argument( + "--vllm-gpu-memory-utilization", + type=float, + default=0.55, + help="GPU memory utilization target for vLLM server.", + ) + parser.add_argument( + "--vllm-enforce-eager", + action="store_true", + default=False, + help="Enable --enforce-eager when launching vLLM server.", + ) + parser.add_argument( + "--vllm-weight-sync-mode", + type=str, + choices=["auto", "native", "reload"], + default="native", + help=( + "vLLM weight sync policy: 'native' launches vLLM with --weight-transfer-config and " + "init_weight_transfer_engine (Megatron NCCL / NcclBridge sync). " + "'reload' / 'auto' skip native init here (reload-on-continue or other fallbacks may apply)." + ), + ) + parser.add_argument( "--hf-checkpoint", type=str, @@ -429,6 +460,23 @@ def add_rollout_arguments(parser): default=1, help="Interval for updating the weights", ) + _vllm_packed = parser.add_mutually_exclusive_group() + _vllm_packed.add_argument( + "--vllm-weight-sync-packed", + dest="vllm_weight_sync_packed", + action="store_true", + help=( + "When rollout-backend is vllm: use one-shot packed weight transfer for dense models (no MoE experts). " + "Automatically disabled for MoE or compressed-tensors quantization." + ), + ) + _vllm_packed.add_argument( + "--no-vllm-weight-sync-packed", + dest="vllm_weight_sync_packed", + action="store_false", + help="Disable vLLM packed weight sync; use per-bucket NCCL via NcclBridge instead.", + ) + parser.set_defaults(vllm_weight_sync_packed=True) parser.add_argument( "--keep-old-actor", action="store_true", @@ -1767,6 +1815,12 @@ def slime_validate_args(args): if args.eval_function_path is None: args.eval_function_path = args.rollout_function_path + + if args.rollout_backend == "vllm": + if args.rollout_function_path == "slime.rollout.sglang_rollout.generate_rollout": + args.rollout_function_path = "slime.rollout.vllm_rollout.generate_rollout" + if args.eval_function_path == "slime.rollout.sglang_rollout.generate_rollout": + args.eval_function_path = "slime.rollout.vllm_rollout.generate_rollout" if args.num_steps_per_rollout is not None: global_batch_size = args.rollout_batch_size * args.n_samples_per_prompt // args.num_steps_per_rollout diff --git a/slime/utils/http_utils.py b/slime/utils/http_utils.py index 7ce395c4d..d78a08cab 100644 --- a/slime/utils/http_utils.py +++ b/slime/utils/http_utils.py @@ -114,16 +114,32 @@ def _wrap_ipv6(host): return host -def run_router(args): +def run_router(payload): + """Start the HTTP router gateway in a child process. + + ``payload`` is either ``(impl, router_args)`` with ``impl`` in ``{"sglang","vllm"}``, + or a legacy single ``router_args`` object (treated as ``sglang``). + """ try: - from sglang_router.launch_router import launch_router + if isinstance(payload, tuple) and len(payload) == 2: + impl, router_args = payload + else: + impl, router_args = "sglang", payload - router = launch_router(args) + if impl == "vllm": + from vllm_router.launch_router import launch_router + else: + from sglang_router.launch_router import launch_router + + router = launch_router(router_args) if router is None: return 1 return 0 - except Exception as e: - logger.info(e) + except Exception: + logger.exception( + "run_router failed (impl=%s). For vllm-router, ensure the package is installed and RouterArgs are valid.", + payload[0] if isinstance(payload, tuple) and len(payload) == 2 else "sglang", + ) return 1