[chore] Remove weight-sync hacks now handled natively by vLLM 0.23.0

.claude/docs/weight_sync.md

@@ -18,7 +18,7 @@ skyrl/backends/skyrl_train/weight_sync/
 
 vLLM worker-extension class (loaded via `--worker-extension-cls`):
 
-- `skyrl/backends/skyrl_train/inference_servers/new_inference_worker_wrap.py` — `NewInferenceWorkerWrap`. Three-phase chunked lifecycle.
+- `skyrl/backends/skyrl_train/inference_servers/new_inference_worker_wrap.py` — `NewInferenceWorkerWrap`. 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).
 
 The weight sync implementation relies on the native vLLM weight sync APIs - `WeightTransferEngine` abstractions as well as native RPC endpoints for weight updates.
 
@@ -29,10 +29,15 @@ The weight sync implementation relies on the native vLLM weight sync APIs - `Wei
 
 Strategy choice is decided by the sender (`get_transfer_strategy_cls`). The init info is expanded per server via `for_servers()` / `to_api_payload()` and pushed to the servers through the HTTP control plane (`init_weight_update_communicator` → vLLM's native `/init_weight_transfer_engine`); the receive side is vLLM's native weight-transfer engine, driven by `NewInferenceWorkerWrap`.
 
-## Lifecycle (`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).
+## Lifecycle — 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
 

skyrl/backends/skyrl_train/inference_servers/new_inference_worker_wrap.py

@@ -1,21 +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.
-
-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:
@@ -24,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).
-
-    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
+class NewInferenceWorkerWrap:
+    """SkyRL custom-behavior hook injected into the vLLM GPUWorker.
 
-    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

skyrl/backends/skyrl_train/inference_servers/remote_inference_client.py

@@ -1111,21 +1111,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
@@ -1135,81 +1130,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(

skyrl/backends/skyrl_train/weight_sync/broadcast_strategy.py

@@ -148,13 +148,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,
@@ -163,7 +161,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(