From fa2fbdbc19ca10f7429fc8f6f41bc961f7b29d74 Mon Sep 17 00:00:00 2001 From: haoaaron Date: Wed, 1 Jul 2026 20:43:13 +0000 Subject: [PATCH 1/2] [chore] Remove weight-sync hacks now handled natively by vLLM 0.23.0 vLLM 0.23.0 performs RL weight sync natively, so the SkyRL worker-extension methods that routed weight updates through /collective_rpc are no longer needed: - Drop skyrl_start_weight_update / update_weights_ipc / update_weights_nccl / skyrl_finish_weight_update. RemoteInferenceClient now drives vLLM's native /start_weight_update, /update_weights, and /finish_weight_update endpoints. - CUDA IPC and broadcast (NCCL) senders delegate to vLLM's own IPCWeightTransferEngine / NCCLWeightTransferEngine.trainer_send_weights and call update_named_weights instead of the SkyRL wrap. The set_current_vllm_config wrap is gone: vLLM 0.23.0 (vllm-project/vllm#44613) dropped the get_current_vllm_config() read from the FlashInfer MoE kernels. - NewInferenceWorkerWrap is now a near-empty hook that only applies the #44814 layerwise-reload numel patch at import time (remove once on vLLM >= 0.23.1). - Token-in-token-out: request return_token_ids=True (vllm-project/vllm#22587) and read choice["token_ids"] instead of re-tokenizing the output text. - Bump the minimum vLLM assertion to 0.23.0 and update weight_sync.md. Signed-off-by: haoaaron Co-Authored-By: Claude Opus 4.8 (1M context) --- .claude/docs/weight_sync.md | 15 +- .../ray_wrapped_inference_engine.py | 2 +- .../remote_inference_engine.py | 16 +- .../new_inference_worker_wrap.py | 185 ++++-------------- .../remote_inference_client.py | 82 +------- .../weight_sync/broadcast_strategy.py | 14 +- .../weight_sync/cuda_ipc_strategy.py | 134 +++++-------- .../inference_servers/test_weight_sync.py | 77 ++++---- .../test_pause_and_continue_generation.py | 11 +- 9 files changed, 177 insertions(+), 359 deletions(-) diff --git a/.claude/docs/weight_sync.md b/.claude/docs/weight_sync.md index c25eb43ec7..d2dd074e6d 100644 --- a/.claude/docs/weight_sync.md +++ b/.claude/docs/weight_sync.md @@ -19,8 +19,8 @@ skyrl/backends/skyrl_train/weight_sync/ vLLM worker-extension classes (loaded via `--worker-extension-cls`): -- `skyrl/backends/skyrl_train/inference_servers/vllm_worker.py` — `WorkerWrap`. **Legacy path.** One or more calls to `load_weights(request)`. -- `skyrl/backends/skyrl_train/inference_servers/new_inference_worker_wrap.py` — `NewInferenceWorkerWrap`. **New path** (`_SKYRL_USE_NEW_INFERENCE=1`, default). Three-phase chunked lifecycle. +- `skyrl/backends/skyrl_train/inference_servers/vllm_worker.py` — `WorkerWrap`. **Legacy path** (`_SKYRL_USE_NEW_INFERENCE=0`, local Ray-actor engines). One or more calls to `load_weights(request)`. Not yet migrated to the native API. +- `skyrl/backends/skyrl_train/inference_servers/new_inference_worker_wrap.py` — `NewInferenceWorkerWrap`. **New path** (`_SKYRL_USE_NEW_INFERENCE=1`, default, remote inference servers). As of vLLM 0.23.0 this is a near-empty **hook** class: weight sync runs entirely through vLLM's native endpoints (below), and the class only exists as the place to add SkyRL-specific worker overrides when upstream lacks something we need. It currently applies the `#44814` layerwise-reload numel patch at import time (remove once vLLM ≥ 0.23.1). ## Transfer Strategies @@ -36,10 +36,13 @@ vLLM worker-extension classes (loaded via `--worker-extension-cls`): 2. `load_weights(request)` — per sync. Receives weights using strategy-specific weight receiver and loads weights into vLLM. 3. `teardown_weight_receiver()` — on shutdown. -**New path** (`NewInferenceWorkerWrap`): -1. `start_weight_update(is_checkpoint_format=True)` — initializes layerwise reload (moves layers to meta device, wraps loaders). -2. `update_weights_chunk(update_info)` — called repeatedly. Unpacks the SkyRL packed CUDA-IPC payload, slices the contiguous buffer per param, calls `model.load_weights(weights=...)` under `set_current_vllm_config`. -3. `finish_weight_update()` — runs `finalize_layerwise_reload` (quantization repacking, attention weight postprocessing). +**New path** — vLLM 0.23.0 native endpoints (driven by `RemoteInferenceClient`; the GPUWorker's `weight_transfer_engine` is selected via `weight_transfer_config`, already passed at `inference_servers/utils.py`): +1. `/init_weight_transfer_engine` — once per session. +2. `/start_weight_update` (`is_checkpoint_format=True`) — initializes layerwise reload (moves layers to meta device, wraps loaders). +3. `/update_weights` — called per chunk. For CUDA IPC the sender delegates packing to vLLM's `IPCWeightTransferEngine.trainer_send_weights` (packed mode) and the worker's `weight_transfer_engine` unpacks via `packed_ipc_consumer`; for NCCL the sender uses `NCCLWeightTransferEngine.trainer_send_weights`. +4. `/finish_weight_update` — runs `finalize_layerwise_reload`. + +The `set_current_vllm_config` wrap the worker used to apply is no longer needed: vLLM 0.23.0 (vllm-project/vllm#44613) dropped the `get_current_vllm_config()` read from the FlashInfer MoE kernels. ## Convention: vLLM imports diff --git a/skyrl/backends/skyrl_train/inference_engines/ray_wrapped_inference_engine.py b/skyrl/backends/skyrl_train/inference_engines/ray_wrapped_inference_engine.py index 268174c3f9..2de2317a90 100644 --- a/skyrl/backends/skyrl_train/inference_engines/ray_wrapped_inference_engine.py +++ b/skyrl/backends/skyrl_train/inference_engines/ray_wrapped_inference_engine.py @@ -158,7 +158,7 @@ def create_ray_wrapped_inference_engines( ) if "dev" not in vllm.__version__: - assert version.parse(vllm.__version__) >= version.parse("0.18.0"), "SkyRL-Train requires vLLM >= 0.18.0" + assert version.parse(vllm.__version__) >= version.parse("0.23.0"), "SkyRL-Train requires vLLM >= 0.23.0" else: raise ValueError(f"Unsupported backend: {backend}") diff --git a/skyrl/backends/skyrl_train/inference_engines/remote_inference_engine.py b/skyrl/backends/skyrl_train/inference_engines/remote_inference_engine.py index 2faddb110d..efbdeae552 100644 --- a/skyrl/backends/skyrl_train/inference_engines/remote_inference_engine.py +++ b/skyrl/backends/skyrl_train/inference_engines/remote_inference_engine.py @@ -180,6 +180,10 @@ async def generate( payload = sampling_params.copy() payload["model"] = self.model_name payload["prompt"] = prompt_token_ids + # Ask vLLM to return the sampled token IDs directly (added in + # vLLM 0.23.0 via https://github.com/vllm-project/vllm/pull/22587) + # so generation is token-in-token-out rather than re-tokenized. + payload["return_token_ids"] = True request_url = f"{self.url}/v1/completions" else: raise ValueError(f"Invalid engine backend: {self.engine_backend}") @@ -198,10 +202,14 @@ async def generate( text = choice["text"] outputs.append(text) finish_reasons.append(choice["finish_reason"]) - # TODO(Charlie): this is not token-in-token-out because vLLM does not support - # returning token IDs via HTTP requests. Fix after this vLLM PR is merged: - # https://github.com/vllm-project/vllm/pull/22587 - output_ids.append(self.tokenizer.encode(text, add_special_tokens=False)) + # Token-in-token-out: vLLM returns the sampled token IDs when + # `return_token_ids=True` is set on the request (see payload + # above). Fall back to re-tokenizing only if the field is + # absent (older server), which is lossy. + token_ids = choice.get("token_ids") + if token_ids is None: + token_ids = self.tokenizer.encode(text, add_special_tokens=False) + output_ids.append(token_ids) else: raise ValueError(f"Invalid engine backend: {self.engine_backend}") diff --git a/skyrl/backends/skyrl_train/inference_servers/new_inference_worker_wrap.py b/skyrl/backends/skyrl_train/inference_servers/new_inference_worker_wrap.py index b2d154504a..403caa64e7 100644 --- a/skyrl/backends/skyrl_train/inference_servers/new_inference_worker_wrap.py +++ b/skyrl/backends/skyrl_train/inference_servers/new_inference_worker_wrap.py @@ -1,23 +1,28 @@ """ -vLLM Worker Extension for native weight sync with chunked transfer support. - -This module provides NewInferenceWorkerWrap, a vLLM worker extension that -enables chunked weight updates from training to inference using the -start/update/finish lifecycle: - - skyrl_start_weight_update -> one or more update_weights_ipc -> skyrl_finish_weight_update - -This separates the layerwise reload initialization/finalization from individual -chunk transfers, allowing weights to be sent in bounded-memory chunks rather -than all at once. - -Used only with the new inference path (_SKYRL_USE_NEW_INFERENCE=1). - -TODO: Once https://github.com/vllm-project/vllm/pull/39212 lands, vLLM will -natively support start_weight_update / update_weights / finish_weight_update -on GPUWorker with dedicated HTTP endpoints. At that point this worker extension -can be removed and SkyRL can call the native endpoints directly instead of -routing through /collective_rpc. +vLLM worker-extension hook for SkyRL (new inference path). + +vLLM 0.23.0 handles RL weight sync natively: the trainer drives +``/init_weight_transfer_engine`` -> ``/start_weight_update`` -> +``/update_weights`` (packed CUDA IPC or NCCL) -> ``/finish_weight_update`` +against the inference server, and the GPUWorker's ``weight_transfer_engine`` +(selected via ``weight_transfer_config``) receives and loads the weights. SkyRL +therefore no longer needs to inject worker methods to receive/load weights, and +the previous ``skyrl_start_weight_update`` / ``update_weights_ipc`` / +``update_weights_nccl`` / ``skyrl_finish_weight_update`` methods are gone. + +This class is kept deliberately as the **injection point for SkyRL-specific +worker behavior**: it is passed to vLLM via ``--worker-extension-cls`` and mixed +into the GPUWorker. Add a targeted method/override here only when upstream vLLM +lacks a feature/fix we need, or when our pinned vLLM version doesn't have it yet. + +The one piece of custom behavior currently required is applying the ``#44814`` +layerwise-reload numel patch (see ``layerwise_reload.patch_numel_loaded``): the +fix is not in vLLM 0.23.0 (it lands in 0.23.1), and the native +``finish_weight_update`` -> ``finalize_layerwise_reload`` -> ``get_numel_loaded`` +over-counts elements for composed weight loaders, silently dropping trailing +params (e.g. Mamba ``mixer.D``). We patch the vLLM function globally as soon as +this module is imported in the worker process (vLLM imports it there when +resolving ``--worker-extension-cls``). Remove once we bump to vLLM >= 0.23.1. Usage: Pass as --worker-extension-cls to vLLM: @@ -26,137 +31,33 @@ skyrl.backends.skyrl_train.inference_servers.new_inference_worker_wrap.NewInferenceWorkerWrap """ -import torch - from skyrl.backends.skyrl_train.inference_servers.layerwise_reload import ( - LayerwiseReloadWorkerMixin, + patch_numel_loaded, ) VLLM_NEW_INFERENCE_WORKER_EXTENSION_CLS = f"{__name__}.NewInferenceWorkerWrap" -class NewInferenceWorkerWrap(LayerwiseReloadWorkerMixin): - """ - vLLM worker extension for chunked weight sync (new inference path). +class NewInferenceWorkerWrap: + """SkyRL custom-behavior hook injected into the vLLM GPUWorker. - Provides a three-phase weight update protocol via collective_rpc: - 1. skyrl_start_weight_update: Prepare model for receiving weights - 2. update_weights_ipc: Receive and load one chunk of weights - 3. skyrl_finish_weight_update: Finalize the model after all chunks - - Attributes accessed from the host GPUWorker (via mixin inheritance): - self.weight_transfer_engine - self.model_runner - self.model_config - self.device + Intentionally near-empty: vLLM 0.23.0 performs weight sync natively, so + there are no SkyRL receive/load methods here anymore. Keep this class as + the place to add worker-side overrides when upstream is missing something + we need (or our pinned version is). See the module docstring. """ - def update_weights_ipc(self, update_info: dict) -> None: - """ - Receive and load a single chunk of weights. - - SkyRL packs each chunk's tensors into a single contiguous CUDA buffer and sends - one IPC handle per rank plus per-param `sizes` metadata. We rebuild - the packed tensor here, slice it per param, and hand the list to - model.load_weights (checkpoint format) or copy per-param directly - (kernel format). - - Args: - update_info: Dict with keys: - - names: list[str] - - dtype_names: list[str] - - shapes: list[list[int]] - - sizes: list[int] (element count per param; used for slicing) - - ipc_handles_pickled: b64(pickle({gpu_uuid: (func, args)})) - """ - if not getattr(self, "_skyrl_weight_update_active", False): - raise RuntimeError("skyrl_start_weight_update must be called before update_weights_ipc.") - - if self.weight_transfer_engine is None: - raise RuntimeError( - "Weight transfer not configured. " "Please set weight_transfer_config to enable weight transfer." - ) - - # --- unpack SkyRL packed CUDA IPC format --- - import base64 - import pickle - - names = update_info["names"] - shapes = update_info["shapes"] - sizes = update_info["sizes"] - pickled = update_info["ipc_handles_pickled"] - handles = pickle.loads(base64.b64decode(pickled)) - - device_index = torch.cuda.current_device() - physical_gpu_id = str(torch.cuda.get_device_properties(device_index).uuid) - if physical_gpu_id not in handles: - raise ValueError(f"IPC handle not found for GPU UUID {physical_gpu_id}. " f"Available: {list(handles)}") - func, args = handles[physical_gpu_id] - # Remap device index to the LOCAL current-device. - list_args = list(args) - list_args[6] = device_index - packed_tensor = func(*list_args) - - weights: list[tuple[str, torch.Tensor]] = [] - offset = 0 - for name, shape, size in zip(names, shapes, sizes): - weights.append((name, packed_tensor[offset : offset + size].view(*shape))) - offset += size - - # process_weights_after_loading reads get_current_vllm_config() (e.g. - # flashinfer_cutlass_moe needs the compilation config to build kernels), - # and vllm only sets that context around init_device / load_model. - from vllm.config import set_current_vllm_config - - model = self.model_runner.model - with set_current_vllm_config(self.vllm_config), torch.device(self.device): - if self._skyrl_is_checkpoint_format: - model.load_weights(weights=weights) - else: - for name, weight in weights: - param = model.get_parameter(name) - param.copy_(weight) - - # Ensure consumption of packed_tensor finishes before we return (and - # before the sender drops its reference on the next barrier). - torch.accelerator.synchronize() - - def update_weights_nccl(self, update_info: dict) -> None: - """ - Receive a batched weight update via vLLM's NCCL weight transfer engine. - - Alternative to update_weights_ipc for the broadcast (non-IPC) sender: - the trainer initiates an NCCL broadcast via - NCCLWeightTransferEngine.trainer_send_weights, and each inference - worker calls weight_transfer_engine.receive_weights here. - - Routed through this skyrl wrap (rather than vLLM's native - /update_weights endpoint) so the load is wrapped with - set_current_vllm_config — process_weights_after_loading on MoE - models can otherwise instantiate kernels (e.g. FlashInfer CUTLASS) - whose __init__ reads get_current_vllm_config(). - - TODO: remove once the upstream vLLM patch lands (vllm-project/vllm - weight-sync-fix), then route via the native /update_weights endpoint. - https://github.com/vllm-project/vllm/pull/42577 - """ - if not getattr(self, "_skyrl_weight_update_active", False): - raise RuntimeError("skyrl_start_weight_update must be called before update_weights_nccl.") - - if self.weight_transfer_engine is None: - raise RuntimeError( - "Weight transfer not configured. Please set weight_transfer_config to enable weight transfer." - ) - - from vllm.config import set_current_vllm_config - - typed_update_info = self.weight_transfer_engine.parse_update_info(update_info) - model = self.model_runner.model + pass - with set_current_vllm_config(self.vllm_config), torch.device(self.device): - self.weight_transfer_engine.receive_weights( - typed_update_info, - load_weights=model.load_weights, - ) - torch.accelerator.synchronize() +# Apply the #44814 layerwise-reload numel patch at import time. vLLM loads this +# module in each worker process when resolving ``--worker-extension-cls``; the +# native ``/finish_weight_update`` path (now used instead of SkyRL's old +# ``skyrl_finish_weight_update``, which used to apply the patch lazily) would +# otherwise hit the un-capped ``get_numel_loaded``. Guarded so the module stays +# importable where vllm isn't available (non-Linux / CPU CI). Remove once we +# bump to vLLM >= 0.23.1. +try: + patch_numel_loaded() +except Exception: # pragma: no cover - vllm not importable in this process + pass diff --git a/skyrl/backends/skyrl_train/inference_servers/remote_inference_client.py b/skyrl/backends/skyrl_train/inference_servers/remote_inference_client.py index 82e08bee2b..d8e9709f4f 100644 --- a/skyrl/backends/skyrl_train/inference_servers/remote_inference_client.py +++ b/skyrl/backends/skyrl_train/inference_servers/remote_inference_client.py @@ -1106,21 +1106,16 @@ async def update_named_weights( {"update_info": update_info}, ) - # TODO: Once https://github.com/vllm-project/vllm/pull/39212 lands, switch - # these three methods from /collective_rpc to the native vLLM endpoints - # (/start_weight_update, /update_weights, /finish_weight_update) and remove - # the NewInferenceWorkerWrap worker extension. - async def start_weight_update( self, is_checkpoint_format: bool = True, ) -> Dict[str, Any]: """ - Start a new chunked weight update via /collective_rpc. + Start a new chunked weight update via vLLM's native /start_weight_update. - Calls the NewInferenceWorkerWrap.skyrl_start_weight_update method on all - workers. For checkpoint-format weights this initializes layerwise - reload. Must be called before any update_weights_ipc calls. + For checkpoint-format weights this initializes layerwise reload on each + worker. Must be called before any update_named_weights calls. Pair with + finish_weight_update. Args: is_checkpoint_format: True if weights are in checkpoint format @@ -1130,81 +1125,22 @@ async def start_weight_update( Dict mapping server_url to response. """ return await self._call_all_servers( - "/collective_rpc", - { - "method": "skyrl_start_weight_update", - "kwargs": {"is_checkpoint_format": is_checkpoint_format}, - }, - ) - - async def update_weights_ipc( - self, - update_info: Dict[str, Any], - ) -> Dict[str, Any]: - """ - Send a single weight chunk via /collective_rpc. - - Calls NewInferenceWorkerWrap.update_weights_ipc on all workers. - Can be called multiple times between skyrl_start_weight_update and - skyrl_finish_weight_update. - - Args: - update_info: Dict with backend-specific update info (names, - dtype_names, shapes, ipc_handles_pickled or packed flag). - - Returns: - Dict mapping server_url to response. - """ - return await self._call_all_servers( - "/collective_rpc", - { - "method": "update_weights_ipc", - "kwargs": {"update_info": update_info}, - }, - ) - - async def update_weights_nccl( - self, - update_info: Dict[str, Any], - ) -> Dict[str, Any]: - """ - Send batched weight update via /collective_rpc to the broadcast receiver. - - Calls NewInferenceWorkerWrap.update_weights_nccl on all workers, - which routes weight_transfer_engine.receive_weights through the - set_current_vllm_config wrap. Used by the broadcast (NCCL) sender as - a temporary substitute for vLLM's native /update_weights endpoint - until the upstream patch (vllm-project/vllm weight-sync-fix) lands. - - Args: - update_info: Dict with backend-specific update info (names, - dtype_names, shapes, packed flag, etc.) — same shape vLLM's - native /update_weights expects. - - Returns: - Dict mapping server_url to response. - """ - return await self._call_all_servers( - "/collective_rpc", - { - "method": "update_weights_nccl", - "kwargs": {"update_info": update_info}, - }, + "/start_weight_update", + {"is_checkpoint_format": is_checkpoint_format}, ) async def finish_weight_update(self) -> Dict[str, Any]: """ - Finish the current chunked weight update via /collective_rpc. + Finish the current chunked weight update via vLLM's native /finish_weight_update. - Calls NewInferenceWorkerWrap.skyrl_finish_weight_update on all workers. For checkpoint-format weights, runs layerwise postprocessing. Returns: Dict mapping server_url to response. """ return await self._call_all_servers( - "/collective_rpc", - {"method": "skyrl_finish_weight_update"}, + "/finish_weight_update", + {}, ) async def load_lora_adapter( diff --git a/skyrl/backends/skyrl_train/weight_sync/broadcast_strategy.py b/skyrl/backends/skyrl_train/weight_sync/broadcast_strategy.py index 9ee5d439b0..339182bd5d 100644 --- a/skyrl/backends/skyrl_train/weight_sync/broadcast_strategy.py +++ b/skyrl/backends/skyrl_train/weight_sync/broadcast_strategy.py @@ -185,13 +185,11 @@ def weight_iterator() -> Iterator[Tuple[str, torch.Tensor]]: for chunk in chunks: yield from zip(chunk.names, chunk.tensors) - # Route via the skyrl wrap (start_weight_update + update_weights_nccl - # + finish_weight_update) rather than vLLM's native /update_weights so - # the receive is wrapped with set_current_vllm_config. Matches how - # CUDA IPC already routes through skyrl's wrap. - # TODO: switch back to update_named_weights once the upstream vLLM - # patch lands (vllm-project/vllm weight-sync-fix). - # https://github.com/vllm-project/vllm/pull/42577 + # Drive vLLM's native weight-sync lifecycle (/start_weight_update + + # /update_weights + /finish_weight_update). The MoE set_current_vllm_config + # wrap that previously forced routing through SkyRL's worker extension is + # no longer needed: vLLM 0.23.0 (vllm-project/vllm#44613) dropped the + # get_current_vllm_config() read from the FlashInfer MoE kernels. if torch.distributed.get_rank() == 0: from vllm.distributed.weight_transfer.nccl_engine import ( NCCLWeightTransferEngine, @@ -200,7 +198,7 @@ def weight_iterator() -> Iterator[Tuple[str, torch.Tensor]]: await self._inference_client.start_weight_update(is_checkpoint_format=True) update_info = {**weight_metadata, "packed": True} - update_task = asyncio.create_task(self._inference_client.update_weights_nccl(update_info)) + update_task = asyncio.create_task(self._inference_client.update_named_weights(update_info)) # Run in thread so the HTTP update_task can progress concurrently await asyncio.to_thread( diff --git a/skyrl/backends/skyrl_train/weight_sync/cuda_ipc_strategy.py b/skyrl/backends/skyrl_train/weight_sync/cuda_ipc_strategy.py index bdf2c7408e..6757e65417 100644 --- a/skyrl/backends/skyrl_train/weight_sync/cuda_ipc_strategy.py +++ b/skyrl/backends/skyrl_train/weight_sync/cuda_ipc_strategy.py @@ -4,6 +4,7 @@ from training workers to inference engines using CUDA IPC handles. """ +import asyncio import base64 import copy import pickle @@ -175,93 +176,66 @@ async def _send_chunks_vllm_native( chunks: Iterable[WeightChunk], weight_metadata: Optional[Dict[str, list]] = None, ) -> None: - """Send weights chunk-by-chunk via vLLM native IPC (new inference path). - - Uses the start/update/finish lifecycle to enable chunked transfers. - Per chunk, all tensors are packed into a single contiguous CUDA buffer - (one dtype per chunk, guaranteed by the weight extractor) and one IPC - handle is created for the packed buffer per rank. - - All ranks iterate chunks (weight extraction may use collective ops). - Per chunk, each rank packs + creates one IPC handle, handles are - all_gather_object'd into a single {gpu_uuid: handle} dict, and rank 0 - sends the dict (plus per-param `sizes` metadata) via - update_weights_ipc. The receiver rebuilds the packed tensor, slices - it per param, and loads into vLLM. - - TODO: Once https://github.com/vllm-project/vllm/pull/39212 lands, - replace update_weights_ipc with the native /update_weights endpoint - and start/finish with /start_weight_update and /finish_weight_update. + """Send weights via vLLM's native CUDA IPC weight transfer. + + vLLM 0.23.0 receives weights natively: the inference worker's + ``weight_transfer_engine`` (``weight_transfer_config.backend="ipc"``) + unpacks the packed CUDA-IPC payload via ``packed_ipc_consumer`` and + loads it. We delegate the trainer side to vLLM's own + ``IPCWeightTransferEngine.trainer_send_weights`` (packed mode), which + handles uint8 packing, the cross-rank IPC-handle all-gather, bounded + per-chunk buffers, and barriers. We drive the lifecycle over the + inference client: ``/start_weight_update`` -> ``/update_weights`` + (invoked by the ``send_mode`` callback, once per packed chunk) -> + ``/finish_weight_update``. + + ``trainer_send_weights`` is synchronous; we run it in a worker thread + and bridge its rank-0 ``send_mode`` callback back to this event loop so + the native ``IPCWeightTransferUpdateInfo`` payload is fanned out to all + inference servers via the (async) client. """ + from vllm.distributed.weight_transfer.ipc_engine import ( + IPCTrainerSendWeightsArgs, + IPCWeightTransferEngine, + ) + + loop = asyncio.get_running_loop() rank = torch.distributed.get_rank() - world_size = torch.distributed.get_world_size() - device = torch.cuda.current_device() - gpu_uuid = str(torch.cuda.get_device_properties(device).uuid) - dtype = str_to_torch_dtype(self._init_info.model_dtype_str) - dtype_name = self._init_info.model_dtype_str.split(".")[-1] + + def weight_iterator() -> Iterator[Tuple[str, torch.Tensor]]: + for chunk in chunks: + yield from zip(chunk.names, chunk.tensors) + + def send_update(update_info: Any) -> None: + # Called by trainer_send_weights on rank 0 (inside the worker + # thread). Build the HTTP-friendly native update payload (IPC + # handles base64-pickled, mirroring vLLM's own 'http' send_mode) + # and fan out to all servers via the async client, blocking until + # done. We read fields off update_info directly rather than + # asdict() to avoid deep-copying the IPC handle args. + fields = { + "update_kind": update_info.update_kind, + "names": update_info.names, + "dtype_names": update_info.dtype_names, + "shapes": update_info.shapes, + "tensor_sizes": update_info.tensor_sizes, + "packed": update_info.packed, + "ipc_handles_pickled": base64.b64encode(pickle.dumps(update_info.ipc_handles)).decode("utf-8"), + } + future = asyncio.run_coroutine_threadsafe(self._inference_client.update_named_weights(fields), loop) + future.result() if rank == 0: await self._inference_client.start_weight_update(is_checkpoint_format=True) torch.distributed.barrier() - for chunk in chunks: - # --- pack all tensors in this chunk into one contiguous buffer --- - # Chunk tensors share a single dtype by construction (see - # weight_extractor_utils.py), so offsets in element units are safe. - names: List[str] = [] - dtype_names: List[str] = [] - shapes: List[List[int]] = [] - sizes: List[int] = [] - - total_numel = sum(t.numel() for t in chunk.tensors) - packed_tensor = torch.empty( - total_numel, - device=device, - dtype=dtype, - requires_grad=False, - ) - - offset = 0 - for name, tensor, shape in zip(chunk.names, chunk.tensors, chunk.shapes): - size = tensor.numel() - packed_tensor[offset : offset + size].copy_(tensor.detach().reshape(-1)) - offset += size - names.append(name) - dtype_names.append(dtype_name) - shapes.append(list(shape) if not isinstance(shape, list) else shape) - sizes.append(size) - - # --- one IPC handle per rank for the packed buffer --- - ipc_handle: IpcHandle = reduce_tensor(packed_tensor) - local_handle_dict: Dict[str, IpcHandle] = {gpu_uuid: ipc_handle} - gathered: List[Optional[Dict[str, IpcHandle]]] = [None] * world_size - torch.distributed.all_gather_object(gathered, local_handle_dict) - - torch.distributed.barrier() - torch.cuda.synchronize() - - if rank == 0: - merged_handles: Dict[str, IpcHandle] = {} - for d in gathered: - if d is not None: - merged_handles.update(d) - - pickled = base64.b64encode(pickle.dumps(merged_handles)).decode("utf-8") - chunk_update_info: Dict[str, Any] = { - "names": names, - "dtype_names": dtype_names, - "shapes": shapes, - "sizes": sizes, - "ipc_handles_pickled": pickled, - } - await self._inference_client.update_weights_ipc(chunk_update_info) - - # Keep packed_tensor alive past the barrier so the receiver's - # rebuilt view has valid backing storage while it copies into - # the model. Post-barrier drops the local ref safely. - torch.distributed.barrier() - torch.cuda.ipc_collect() - torch.cuda.synchronize() + # Blocking: packs each chunk into a uint8 buffer, all-gathers IPC + # handles across ranks, and (rank 0) invokes send_update per chunk. + await asyncio.to_thread( + IPCWeightTransferEngine.trainer_send_weights, + iterator=weight_iterator(), + trainer_args=IPCTrainerSendWeightsArgs(send_mode=send_update, packed=True), + ) if rank == 0: await self._inference_client.finish_weight_update() diff --git a/tests/backends/skyrl_train/gpu/gpu_ci/inference_servers/test_weight_sync.py b/tests/backends/skyrl_train/gpu/gpu_ci/inference_servers/test_weight_sync.py index e6585846c1..a4b124d624 100644 --- a/tests/backends/skyrl_train/gpu/gpu_ci/inference_servers/test_weight_sync.py +++ b/tests/backends/skyrl_train/gpu/gpu_ci/inference_servers/test_weight_sync.py @@ -287,18 +287,15 @@ async def test_update_weights_flow(self, weight_update_env): "packed": True, } print( - f"[Step 3] Calling update_weights_nccl with {len(update_info['names'])} names, packed={update_info['packed']}" + f"[Step 3] Calling update_named_weights with {len(update_info['names'])} names, packed={update_info['packed']}" ) - # Use SkyRL's chunked weight-sync API (skyrl_start_weight_update -> - # update_weights_nccl -> skyrl_finish_weight_update) rather than vLLM's - # native /update_weights endpoint, which in vLLM 0.22.0+ requires - # vLLM's own native start_weight_update to be called first. - # skyrl_start_weight_update is local (layerwise-reload init), so it is - # safe to call while the trainer is blocked on the NCCL send; the - # actual receive happens in update_weights_nccl. + # Drive vLLM's native weight-sync lifecycle. start_weight_update + # (layerwise-reload init) is local, so it is safe to call while the + # trainer is blocked on the NCCL send; the actual receive happens in + # the native /update_weights call below. await client.start_weight_update() - result = await client.update_weights_nccl(update_info) - print(f"[Step 3] update_weights_nccl returned: {list(result.keys())}") + result = await client.update_named_weights(update_info) + print(f"[Step 3] update_named_weights returned: {list(result.keys())}") for server_url, resp in result.items(): assert resp["status"] == 200, f"Server {server_url} update weights failed: {resp}" await client.finish_weight_update() @@ -348,45 +345,38 @@ def ready(self): def create_ipc_update_info(self) -> dict: """Create a single packed CUDA-IPC buffer for all model parameters. - Matches SkyRL's ``update_weights_ipc`` contract (the packed format - produced by ``CudaIpcTransferStrategy``): all parameters are copied into - one contiguous CUDA buffer, a single IPC handle is created for that - buffer, and per-parameter ``sizes`` let the receiver slice it back out. - This differs from vLLM's native ``/update_weights`` (one handle per - parameter), which we no longer use. + Produces the native ``IPCWeightTransferUpdateInfo`` packed payload that + vLLM's ``packed_ipc_consumer`` expects: all parameters packed into one + contiguous uint8 buffer (via vLLM's ``pack_tensors``), a single IPC + handle for that buffer, per-parameter byte ``tensor_sizes``, and + ``packed=True``. The handle stores only the ``rebuild_cuda_tensor`` args + tuple (the consumer uses the global rebuild func). """ from torch.multiprocessing.reductions import reduce_tensor + from vllm.distributed.weight_transfer.packed_tensor import pack_tensors gpu_uuid = str(torch.cuda.get_device_properties(torch.cuda.current_device()).uuid) - params = list(self.model.named_parameters()) - # The model is loaded in a single dtype (bfloat16), so element offsets - # into one packed buffer are well-defined across all parameters. - dtype = params[0][1].dtype - total_numel = sum(p.numel() for _, p in params) - packed_tensor = torch.empty(total_numel, device=self.device, dtype=dtype) - - names, dtype_names, shapes, sizes = [], [], [], [] - offset = 0 - for name, param in params: - size = param.numel() - packed_tensor[offset : offset + size].copy_(param.detach().reshape(-1)) - offset += size - names.append(name) - dtype_names.append(str(param.dtype).split(".")[-1]) - shapes.append(list(param.shape)) - sizes.append(size) + items = [(name, param.detach()) for name, param in self.model.named_parameters()] + total_bytes = sum(t.numel() * t.element_size() for _, t in items) + packed = pack_tensors( + iter(items), + post_iter_func=lambda item: item[1], + buffer_size_bytes=max(total_bytes, 1), + ) + assert packed is not None # Keep the packed buffer alive so the IPC handle stays valid on the receiver. - self._tensor_refs = [packed_tensor] + self._tensor_refs = [packed.packed_tensor] - ipc_handle = reduce_tensor(packed_tensor) - pickled = base64.b64encode(pickle.dumps({gpu_uuid: ipc_handle})).decode("utf-8") + _, ipc_args = reduce_tensor(packed.packed_tensor) + pickled = base64.b64encode(pickle.dumps({gpu_uuid: ipc_args})).decode("utf-8") return { - "names": names, - "dtype_names": dtype_names, - "shapes": shapes, - "sizes": sizes, + "names": packed.names, + "dtype_names": [str(dt).split(".")[-1] for dt in packed.dtypes], + "shapes": packed.shapes, + "tensor_sizes": packed.tensor_sizes, + "packed": True, "ipc_handles_pickled": pickled, } @@ -480,11 +470,10 @@ async def test_update_weights_ipc(self, ipc_weight_update_env): update_info = ray.get(trainer.create_ipc_update_info.remote()) print(f"[Step 3] Created handles for {len(update_info['names'])} parameters") - # Use SkyRL's chunked weight-sync API (skyrl_start_weight_update -> - # update_weights_ipc -> skyrl_finish_weight_update) rather than vLLM's - # native /update_weights endpoint. + # Drive vLLM's native weight-sync lifecycle (start_weight_update -> + # native /update_weights -> finish_weight_update). await client.start_weight_update() - result = await client.update_weights_ipc(update_info) + result = await client.update_named_weights(update_info) for server_url, resp_data in result.items(): assert resp_data["status"] == 200, f"Server {server_url} IPC update failed: {resp_data}" await client.finish_weight_update() diff --git a/tests/backends/skyrl_train/gpu/gpu_ci/test_pause_and_continue_generation.py b/tests/backends/skyrl_train/gpu/gpu_ci/test_pause_and_continue_generation.py index 26c3514036..959555e059 100644 --- a/tests/backends/skyrl_train/gpu/gpu_ci/test_pause_and_continue_generation.py +++ b/tests/backends/skyrl_train/gpu/gpu_ci/test_pause_and_continue_generation.py @@ -59,6 +59,9 @@ async def test_continue_generation_vllm_engine_chat_completion(ray_init_fixture) # Ensure logprobs and token ids are returned for accumulation checks "logprobs": True, "top_logprobs": 1, + # vLLM 0.23.0+ returns sampled token IDs on the choice when set + # (https://github.com/vllm-project/vllm/pull/22587). + "return_token_ids": True, } async with InferenceEngineState.create( cfg=cfg, @@ -165,7 +168,13 @@ def send_request(i: int): assert ( len(prompt_tokens) == out["usage"]["prompt_tokens"] ), f"Request {i} expected {len(prompt_tokens)} tokens from prompt, got {out['usage']['prompt_tokens']}" - # TODO(Charlie): after we bump vllm such that it supports returnining tokens, check `choice["token_ids"]` + # vLLM 0.23.0+ returns sampled token IDs on the choice (we set + # `return_token_ids=True` above, https://github.com/vllm-project/vllm/pull/22587). + token_ids = choice.get("token_ids") + assert token_ids is not None, f"Request {i} missing `token_ids` on choice: {choice.keys()}" + assert ( + len(token_ids) == sampling_params["max_tokens"] + ), f"Request {i} expected {sampling_params['max_tokens']} token_ids, got {len(token_ids)}" # TODO(Charlie): after we add model version to the output, check that as well finally: shutdown_server(host=SERVER_HOST, port=SERVER_PORT, max_wait_seconds=5) From d77c14d1ebd710072bfe55e58352b321706cd81b Mon Sep 17 00:00:00 2001 From: haoaaron Date: Thu, 2 Jul 2026 00:09:54 +0000 Subject: [PATCH 2/2] [fix] Import Iterator in cuda_ipc_strategy (NameError in CUDA IPC weight sync) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `_send_chunks_vllm_native` annotates the inner `weight_iterator()` with `Iterator[...]`, but `Iterator` was never imported — so the annotation raised `NameError: name 'Iterator' is not defined` the moment the CUDA IPC (colocated) sender ran a weight sync. Not caught earlier because: py_compile only checks syntax; the GPU test_weight_sync.py IPC case drives the client endpoints directly and bypasses this method; and both gsm8k smoke runs used NCCL broadcast (broadcast_strategy.py, which does import Iterator). The colocated megatron logprobs roundtrip test is the first to exercise this path. Also drops two dead imports (reduce_tensor, str_to_torch_dtype) left over from the native-weight-sync refactor (ruff --fix). Co-Authored-By: Claude Opus 4.8 (1M context) --- skyrl/backends/skyrl_train/weight_sync/cuda_ipc_strategy.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/skyrl/backends/skyrl_train/weight_sync/cuda_ipc_strategy.py b/skyrl/backends/skyrl_train/weight_sync/cuda_ipc_strategy.py index 4ccc089ced..cbeaf45e7e 100644 --- a/skyrl/backends/skyrl_train/weight_sync/cuda_ipc_strategy.py +++ b/skyrl/backends/skyrl_train/weight_sync/cuda_ipc_strategy.py @@ -15,6 +15,7 @@ Callable, Dict, Iterable, + Iterator, List, Optional, Tuple, @@ -27,7 +28,6 @@ from skyrl.train.config import InferenceEngineConfig import torch -from torch.multiprocessing.reductions import reduce_tensor from skyrl.backends.skyrl_train.weight_sync.base import WeightChunk, WeightUpdateRequest from skyrl.backends.skyrl_train.weight_sync.transfer_strategy import ( @@ -35,7 +35,6 @@ WeightTransferSender, WeightTransferStrategy, ) -from skyrl.train.utils.utils import str_to_torch_dtype # IPC handle type: (rebuild_func, args) returned by reduce_tensor IpcHandle = Tuple[Callable[..., torch.Tensor], Tuple[Any, ...]]