vLLM Stateful Weight Sync Sender

.claude/docs/architecture.md

@@ -11,7 +11,7 @@ skyrl/                  # Core library
 │   └── skyrl_train/    # FSDP/Megatron training backend
 │       ├── distributed/        # Dispatch, FSDP/Megatron strategies
 │       ├── inference_servers/  # HTTP inference path (RemoteInferenceClient, vLLM servers, router)
-│       ├── weight_sync/        # Weight extraction and transfer
+│       ├── weight_sync/        # WeightSources + trainer-side transfer engines
 │       └── workers/            # FSDP/Megatron workers
 ├── train/              # Training entrypoints, config, dataset, generators, trainer
 │   ├── config/         # Hydra YAML configs (ppo_base, megatron, skyrl_gym)

.claude/docs/backends/fsdp.md

@@ -6,7 +6,7 @@ Default backend (`trainer.strategy=fsdp`). Uses PyTorch FSDP2 for distributed tr
 
 - **FSDPConfig** in `skyrl/train/config.py`.
 - **FSDPStrategy** in `skyrl/backends/skyrl_train/distributed/fsdp_strategy.py`.
-- **FSDPWeightExtractor** for extracting weights from sharded parameters (in `skyrl/backends/skyrl_train/workers/fsdp/fsdp_worker.py`).
+- **FsdpWeightSource** (`skyrl/backends/skyrl_train/weight_sync/sources.py`) presents the sharded model to weight sync as vLLM's `WeightSource`; the worker builds it in `FSDPPolicyWorkerBase._build_weight_source`.
 
 ## CPU Offload
 

docs/content/docs/getting-started/inference_architecture.mdx

@@ -102,29 +102,31 @@ SkyRL uses the native weight syncing APIs in vLLM, with the following four-stage
 - `POST /update_weights` — Updates all or a subset of the weights. SkyRL uses chunked weight transfer for efficiency.
 - `POST /finish_weight_update` — Finishes the current weight update.
 
-<Callout type="info">
-For colocated training, SkyRL currently uses chunked transfers with CUDA IPC handles and currently implements a [custom `Worker` extension](https://github.com/NovaSky-AI/SkyRL/blob/fb87f35dfe9f74f71445da60c79549773b15ba5e/skyrl/backends/skyrl_train/inference_servers/new_inference_worker_wrap.py) as a transitional implementation, pending [vllm-project/vllm#39212](https://github.com/vllm-project/vllm/issues/39212). We plan to migrate to the native APIs in the next vLLM release.
-</Callout>
+The trainer side runs on vLLM's *trainer-send* abstraction: each training worker builds a `WeightSource` over its live model and hands it to a `TrainerWeightTransferEngine`, whose `send_weights()` drives the whole four-stage round trip above. See `skyrl/backends/skyrl_train/weight_sync/trainer_engines.py`.
+
+Four backends, selected by `generator.inference_engine.weight_sync_backend` plus `trainer.placement.colocate_all`:
 
-SkyRL implements two transfer strategies in `skyrl/backends/skyrl_train/weight_sync/`:
+- **`nccl`** — non-colocated. Tensor data is broadcast over NCCL from trainer rank 0 to all inference workers, concurrently with the `/update_weights` HTTP call that ships the metadata (both sides rendezvous inside the same NCCL calls). Combined with `/pause?mode=keep` and `/resume` so in-flight rollouts pause correctly during the sync.
+- **`ipc`** — colocated, where the trainer and inference engines share GPUs. Weights are packed into one reusable buffer and exchanged via CUDA IPC handles. Combined with `/sleep` and `/wake_up` for memory management — the inference engine sleeps to free VRAM during training, then wakes for rollouts.
+- **`delta`** — weights travel as compressed deltas against the base checkpoint through a shared filesystem or object store, for setups where the two sides are not NCCL-reachable.
+- **`sharded_rdt`** — the inference workers *pull* the slices they consume from the trainer ranks over NIXL/RDMA.
 
-- **`BroadcastTransferStrategy`** (`broadcast_strategy.py`) — used for non-colocated training. Tensor data is broadcast over NCCL from trainer rank 0 to all inference workers, concurrently with `/update_weights` HTTP calls that ship the metadata. This is used in combination with `/pause?mode=keep` and `/resume` so that in-flight rollouts are paused correctly during the sync.
-- **`CudaIpcTransferStrategy`** (`cuda_ipc_strategy.py`) — used for colocated training where the trainer and inference engines share GPUs. Weights are exchanged via CUDA IPC handles. Combined with `/sleep` and `/wake_up` for memory management — the inference engine sleeps to free VRAM during training, then wakes for rollouts.
+<Callout type="info">
+The receive side is a vLLM `WeightTransferEngine` in each inference worker. SkyRL registers its own under `skyrl_nccl` / `skyrl_ipc`, subclassing vLLM's to also reload the speculative-decoding drafter (a separate module the main model's `load_weights` never touches). A small [`Worker` extension](https://github.com/NovaSky-AI/SkyRL/blob/main/skyrl/backends/skyrl_train/inference_servers/new_inference_worker_wrap.py) remains for the two things an engine cannot reach: the checkpoint-delta `fetch_weights` call, and driving the allocator directly for KV offload.
+</Callout>
 
 Endpoint summary:
 
 | Endpoint | Plane | Purpose |
 |----------|-------|---------|
 | `/init_weight_transfer_engine` | fan-out | One-time communicator setup |
-| `/start_weight_update` * | fan-out | Begin a chunked update |
-| `/update_weights` * | fan-out | Send a tensor chunk |
-| `/finish_weight_update` * | fan-out | Commit the update |
+| `/start_weight_update` | fan-out | Begin a chunked update |
+| `/update_weights` | fan-out | Send a tensor chunk |
+| `/finish_weight_update` | fan-out | Commit the update |
 | `/pause`, `/resume` | fan-out | Generation control |
 | `/sleep`, `/wake_up` | fan-out | Colocated memory management |
 | `/v1/completions`, `/v1/chat/completions`, generate | routed | Generation |
 
-\* Currently we use custom `/collective_rpc` + worker methods that mimick the native APIs because SkyRL makes some vLLM fixes for Qwen 3.6 model loading. We will migrate to the native `/start_weight_update` and `/finish_weight_update` soon.
-
 ## End-to-end Weight Sync Flow
 
 **Non-colocated mode (NCCL broadcast):**

examples/train/rlm/openrouter_client.py

@@ -182,9 +182,6 @@ async def sleep(self, *args: Any, **kwargs: Any) -> None:
     async def wake_up(self, *args: Any, **kwargs: Any) -> None:
         pass
 
-    async def init_weight_update_communicator(self, *args: Any, **kwargs: Any) -> None:
-        pass
-
     async def update_named_weights(self, *args: Any, **kwargs: Any) -> None:
         pass
 

skyrl/backends/skyrl_train/inference_servers/base.py

@@ -4,10 +4,7 @@
 from skyrl.backends.skyrl_train.utils.routed_experts import RoutedExpertIndices
 
 if TYPE_CHECKING:
-    from skyrl.backends.skyrl_train.weight_sync import WeightUpdateRequest
-    from skyrl.backends.skyrl_train.weight_sync.transfer_strategy import (
-        WeightSyncInitInfo,
-    )
+    from skyrl.backends.skyrl_train.weight_sync import LoraLoadRequest
 
 MessageType = Dict[str, str]
 ConversationType = List[MessageType]
@@ -125,19 +122,16 @@ async def sleep(self, *args: Any, **kwargs: Any):
         raise NotImplementedError
 
     @abstractmethod
-    async def init_weight_update_communicator(self, init_info: "WeightSyncInitInfo"):
-        """Initialize weight update communicator from init info.
+    async def update_named_weights(self, request: "LoraLoadRequest | Dict[str, Any]"):
+        """Load weights the engine can reach itself, rather than transferring them.
 
-        Args:
-            init_info: WeightSyncInitInfo from the sender containing all info needed
-                to create the appropriate receiver.
+        The only caller is the LoRA path, which passes a
+        :class:`LoraLoadRequest` naming an adapter directory on disk. Tensor
+        transfer goes through the trainer-side engines (see
+        ``weight_sync/trainer_engines.py``).
         """
         raise NotImplementedError()
 
-    @abstractmethod
-    async def update_named_weights(self, request: "WeightUpdateRequest"):
-        raise NotImplementedError()
-
     @abstractmethod
     async def teardown(self):
         raise NotImplementedError