Weight sync refactor + move sparse nccl engine

docs/training/layerwise.md

@@ -58,7 +58,7 @@ class Fp8PerTensorOnlineLinearMethod(LinearMethodBase):
 
 ### High Level Weight Transfer API
 
-The layerwise reloading system is integrated with the post-training weight transfer system. To use layerwise reloading in conjunction to the weight transfer system, follow the examples found [here](../../examples/rl/). Layerwise reloading is controlled by the `WeightTransferUpdateInfo.is_checkpoint_format` flag and is set to `True` by default.
+The layerwise reloading system is integrated with the post-training weight transfer system. To use layerwise reloading in conjunction to the weight transfer system, follow the examples found [here](../../examples/rl/). Checkpoint-format weight transfer engines (e.g. the NCCL and IPC backends) run layerwise reloading automatically inside their `start_weight_update`/`finish_weight_update` lifecycle.
 
 ### Mid Level `reload_weights` API
 

docs/training/weight_transfer/README.md

@@ -17,6 +17,7 @@ The weight transfer system follows a **four-phase protocol** with a pluggable ba
 | ------- | --------- | -------- |
 | [NCCL](nccl.md) | NCCL broadcast | Separate GPUs for training and inference |
 | [IPC](ipc.md) | CUDA IPC handles | Colocated training and inference on same GPU |
+| [sparse_nccl](nccl.md#sparse-nccl) | NCCL broadcast | Sparse flat-index weight patches (TP=1/PP=1) |
 
 ## Configuration
 
@@ -41,7 +42,7 @@ vllm serve my-model \
     --weight-transfer-config '{"backend": "nccl"}'
 ```
 
-The `backend` field accepts `"nccl"` (default) or `"ipc"`.
+The `backend` field accepts `"nccl"` (default), `"ipc"`, or `"sparse_nccl"`.
 
 ## API Endpoints
 
@@ -69,7 +70,7 @@ Both backends provide static methods that the trainer calls to send weights. The
 EngineClass.trainer_init(init_info)
 
 # 2. Start weight update on inference side
-llm.start_weight_update(is_checkpoint_format=True)
+llm.start_weight_update()
 
 # 3. Send weights to inference workers
 EngineClass.trainer_send_weights(

docs/training/weight_transfer/base.md

@@ -11,15 +11,23 @@ The `WeightTransferEngine` is a generic abstract class parameterized by two data
 
 ### Abstract Methods
 
-Subclasses must implement these four methods:
+Subclasses must implement these methods:
 
 | Method | Side | Description |
 | ------ | ---- | ----------- |
 | `init_transfer_engine(init_info)` | Inference | Initialize the communication channel on each inference worker |
-| `receive_weights(update_info, load_weights)` | Inference | Receive weights and call `load_weights` incrementally |
+| `start_weight_update()` | Inference | Prepare for an update (e.g. begin layerwise reload); no-op for in-place engines |
+| `finish_weight_update()` | Inference | Finalize the update (e.g. finalize layerwise reload); no-op for in-place engines |
+| `receive_weights(update_info)` | Inference | Receive weights and load them into `self.model` |
 | `shutdown()` | Inference | Clean up resources |
 | `trainer_send_weights(iterator, trainer_args)` | Trainer | Static method to send weights from the trainer process |
 
+The base class provides two methods:
+
+1. `__init__` : Engines receive `config` (`WeightTransferConfig`),  `vllm_config` (`VllmConfig`), `device` (`torch.device`) and  `model` (`nn.Module`)  
+2. `update_weights(update_info_dict)`:  Thin wrapper for `receive_weights`: parses
+the dict into user-specified data type, calls `receive_weights`, and synchronizes the device. Subclasses implement `receive_weights`.
+
 ### Request Classes
 
 The API-level request classes provide backend-agnostic serialization using plain dictionaries. The engine's `parse_init_info` and `parse_update_info` methods convert these dictionaries into typed dataclasses.
@@ -81,7 +89,7 @@ class MyUpdateInfo(WeightTransferUpdateInfo):
 ### 2. Implement the Engine
 
 ```python
-from collections.abc import Callable, Iterator
+from collections.abc import Iterator
 from typing import Any
 import torch
 
@@ -93,18 +101,25 @@ class MyWeightTransferEngine(WeightTransferEngine[MyInitInfo, MyUpdateInfo]):
         # Set up connection to trainer using init_info.endpoint, etc.
         ...
 
-    def receive_weights(
-        self,
-        update_info: MyUpdateInfo,
-        load_weights: Callable[[list[tuple[str, torch.Tensor]]], None],
-    ) -> None:
-        # Receive each weight and call load_weights incrementally
+    def start_weight_update(self) -> None:
+        # Checkpoint-format engines: run initialize_layerwise_reload(self.model).
+        # In-place engines: no-op
+        ...
+
+    def finish_weight_update(self) -> None:
+        # Checkpoint-format engines: run finalize_layerwise_reload(...).
+        # In-place engines: no-op
+        ...
+
+    def receive_weights(self, update_info: MyUpdateInfo) -> None:
+        weights = []
         for name, dtype_name, shape in zip(
             update_info.names, update_info.dtype_names, update_info.shapes
         ):
             dtype = getattr(torch, dtype_name)
             weight = self._fetch_weight(name, shape, dtype)
-            load_weights([(name, weight)])
+            weights.append((name, weight))
+        self.model.load_weights(weights)
 
     def shutdown(self) -> None:
         # Clean up resources
@@ -121,9 +136,6 @@ class MyWeightTransferEngine(WeightTransferEngine[MyInitInfo, MyUpdateInfo]):
             ...
 ```
 
-!!! important
-    The `load_weights` callable passed to `receive_weights` should be called **incrementally** (one or a few weights at a time) rather than accumulating all weights first. This avoids GPU out-of-memory errors with large models.
-
 ### 3. Register with the Factory
 
 ```python
@@ -147,15 +159,16 @@ Once registered, users can select your backend via `WeightTransferConfig(backend
 
 ## WeightTransferEngineFactory
 
-The factory uses a registry pattern with lazy loading. Built-in engines (`nccl` and `ipc`) are registered at import time but their modules are only loaded when the backend is actually requested. This avoids importing heavy dependencies (like NCCL communicators) when they aren't needed.
+The factory uses a registry pattern with lazy loading. Built-in engines (`nccl`, `ipc`, and `sparse_nccl`) are registered at import time but their modules are only loaded when the backend is actually requested. This avoids importing heavy dependencies (like NCCL communicators) when they aren't needed.
 
 ```python
 from vllm.distributed.weight_transfer.factory import WeightTransferEngineFactory
 
 # Create an engine from config
 engine = WeightTransferEngineFactory.create_engine(
     config=weight_transfer_config,
-    parallel_config=parallel_config,
+    vllm_config=vllm_config,
+    device=device,
     model=model,
 )
 ```

docs/training/weight_transfer/ipc.md

@@ -55,7 +55,7 @@ trainer_args = IPCTrainerSendWeightsArgs(
     llm_handle=llm_actor_handle,
 )
 # start
-ray.get(llm_actor_handle.start_weight_update.remote(is_checkpoint_format=True))
+ray.get(llm_actor_handle.start_weight_update.remote())
 # send weights
 IPCWeightTransferEngine.trainer_send_weights(
     iterator=model.named_parameters(),
@@ -80,7 +80,7 @@ trainer_args = IPCTrainerSendWeightsArgs(
 # start
 base_url = "http://localhost:8000"
 url = f"{base_url}/start_weight_update"
-response = requests.post(url, json={"is_checkpoint_format": True}, timeout=60)
+response = requests.post(url, json={}, timeout=60)
 response.raise_for_status()
 # send weights
 IPCWeightTransferEngine.trainer_send_weights(

docs/training/weight_transfer/nccl.md

@@ -11,7 +11,7 @@ The NCCL weight transfer engine uses [NCCL](https://developer.nvidia.com/nccl) b
 ## How It Works
 
 1. The trainer and all inference workers join a shared NCCL process group using `StatelessProcessGroup` (vLLM's torch.distributed-independent group abstraction).
-2. The trainer broadcasts weights to all workers simultaneously. Each worker receives and loads weights incrementally.
+2. The trainer broadcasts weights to all workers simultaneously. Each worker receives and loads the weights.
 3. Optionally, **packed tensor broadcasting** batches multiple small tensors into larger buffers with double/triple buffering and CUDA stream overlap for higher throughput. This implementation is based on [NeMo-RL's packed tensor](https://github.com/NVIDIA-NeMo/RL/blob/main/nemo_rl/utils/packed_tensor.py).
 
 ## Initialization
@@ -93,7 +93,7 @@ remaining three steps are:
 from vllm.distributed.weight_transfer.base import WeightTransferUpdateRequest
 
 # 1. Start the weight update
-llm.start_weight_update(is_checkpoint_format=True)
+llm.start_weight_update()
 
 # 2. Receive weights (can be called multiple times for chunked transfers)
 llm.update_weights(
@@ -116,19 +116,22 @@ must match the order in which the trainer iterates over its parameters.
 
 `start_weight_update` must be called before `update_weights`, and
 `finish_weight_update` must be called after all weight chunks have been
-transferred. The `is_checkpoint_format` flag controls whether layerwise reload
-processing is applied (`True` for checkpoint-format weights, `False` for
-pre-processed kernel-format weights).
+transferred. The NCCL engine receives checkpoint-format weights and applies
+layerwise reload processing automatically inside `start_weight_update` /
+`finish_weight_update`.
 
-Sparse NCCL patches still use `update_kind="sparse_flat"` inside
-`update_info`, but they should be wrapped in
-`start_weight_update(is_checkpoint_format=False)` because sparse patches apply
-directly to runtime/kernel-format parameters. The current sparse MVP requires
-`TP=1` and `PP=1`.
+## Sparse NCCL
+
+Sparse, flat-index weight patches use a separate backend,
+`WeightTransferConfig(backend="sparse_nccl")`, implemented by
+`SparseNCCLWeightTransferEngine`. It shares only NCCL process-group
+initialization with the dense engine; patches are applied directly in place to
+existing parameters (no layerwise reload). The current sparse MVP requires
+`TP=1` and `PP=1`. See the example below.
 
 ## Examples
 
 - [RLHF with NCCL weight syncing (offline, Ray)](../../../examples/rl/rlhf_nccl.py) - Trainer on one GPU, 2x tensor-parallel vLLM engine on two others, with packed NCCL weight broadcast
-- [RLHF with sparse NCCL weight syncing (offline, Ray)](../../../examples/rl/rlhf_sparse_nccl.py) - Dense-vs-sparse equivalence demo with a real model on a 2-GPU trainer/inference setup; sparse patches use `start_weight_update(is_checkpoint_format=False)` and currently require `TP=1` and `PP=1`
+- [RLHF with sparse NCCL weight syncing (offline, Ray)](../../../examples/rl/rlhf_sparse_nccl.py) - Dense-vs-sparse equivalence demo with a real model on a 2-GPU trainer/inference setup; sparse patches use `backend="sparse_nccl"` and currently require `TP=1` and `PP=1`
 - [RLHF with async weight syncing (offline, Ray)](../../../examples/rl/rlhf_async_new_apis.py) - Async generation with mid-flight pause, weight sync, resume, and validation against a fresh model
 - [RLHF with NCCL weight syncing (online serving, HTTP)](../../../examples/rl/rlhf_http_nccl.py) - Weight transfer with a running vLLM HTTP server using HTTP control plane and NCCL data plane

examples/rl/rlhf_async_new_apis.py

@@ -306,7 +306,7 @@ def generate(self, token_ids: list[int], max_new_tokens: int) -> list[int]:
 
 ray.get(llm.pause_after_n_tokens.remote())
 
-ray.get(llm.start_weight_update.remote(is_checkpoint_format=True))
+ray.get(llm.start_weight_update.remote())
 
 inference_handle = llm.update_weights.remote(
     WeightTransferUpdateRequest(

examples/rl/rlhf_http_ipc.py

@@ -80,14 +80,10 @@ def init_weight_transfer_engine(base_url: str) -> None:
     response.raise_for_status()
 
 
-def start_weight_update(
-    base_url: str,
-    is_checkpoint_format: bool = True,
-) -> None:
+def start_weight_update(base_url: str) -> None:
     """Start a weight update via HTTP endpoint."""
     url = f"{base_url}/start_weight_update"
-    payload = {"is_checkpoint_format": is_checkpoint_format}
-    response = requests.post(url, json=payload, timeout=60)
+    response = requests.post(url, json={}, timeout=60)
     response.raise_for_status()
 
 
@@ -170,7 +166,7 @@ def main():
     pause_generation(BASE_URL)
 
     # Start weight update, broadcast via IPC, then finish
-    start_weight_update(BASE_URL, is_checkpoint_format=False)
+    start_weight_update(BASE_URL)
 
     print("Broadcasting weights via CUDA IPC (HTTP)...")
     trainer_args = IPCTrainerSendWeightsArgs(send_mode="http", url=BASE_URL)

examples/rl/rlhf_http_nccl.py

@@ -83,14 +83,10 @@ def init_weight_transfer_engine(
     response.raise_for_status()
 
 
-def start_weight_update(
-    base_url: str,
-    is_checkpoint_format: bool = True,
-) -> None:
+def start_weight_update(base_url: str) -> None:
     """Start a weight update via HTTP endpoint."""
     url = f"{base_url}/start_weight_update"
-    payload = {"is_checkpoint_format": is_checkpoint_format}
-    response = requests.post(url, json=payload, timeout=60)
+    response = requests.post(url, json={}, timeout=60)
     response.raise_for_status()
 
 
@@ -223,7 +219,7 @@ def main():
         shapes.append(list(p.shape))
 
     # Start weight update
-    start_weight_update(BASE_URL, is_checkpoint_format=True)
+    start_weight_update(BASE_URL)
 
     # Start the update_weights call in a separate thread since it will block
     # waiting for NCCL broadcasts

examples/rl/rlhf_ipc.py

@@ -139,7 +139,7 @@ def broadcast_weights(
 
 ray.get(train_model.init_weight_transfer.remote())
 # Start weight update, sync weights, then finish
-ray.get(llm.start_weight_update.remote(is_checkpoint_format=True))
+ray.get(llm.start_weight_update.remote())
 ray.get(train_model.broadcast_weights.remote(llm))
 ray.get(llm.finish_weight_update.remote())
 

examples/rl/rlhf_ipc_fsdp_ep.py

@@ -277,15 +277,8 @@ def init_weight_transfer(self):
             ]
         )
 
-    def start_weight_update(self, is_checkpoint_format: bool = True):
-        ray.get(
-            [
-                actor.start_weight_update.remote(
-                    is_checkpoint_format=is_checkpoint_format
-                )
-                for actor in self.llm_actors
-            ]
-        )
+    def start_weight_update(self):
+        ray.get([actor.start_weight_update.remote() for actor in self.llm_actors])
 
     def finish_weight_update(self):
         ray.get([actor.finish_weight_update.remote() for actor in self.llm_actors])
@@ -392,7 +385,7 @@ def main():
     ray.get(inference_engine.wake_up.remote(tags=["weights"]))
 
     print("[sync] Starting weight update...")
-    ray.get(inference_engine.start_weight_update.remote(is_checkpoint_format=True))
+    ray.get(inference_engine.start_weight_update.remote())
 
     print("[sync] Packed IPC transfer FSDP → vLLM...")
     ray.get(

examples/rl/rlhf_nccl.py

@@ -202,7 +202,7 @@ def broadcast_weights(self, packed: bool = True):
 names, dtype_names, shapes = ray.get(train_model.get_weight_metadata.remote())
 
 # Start weight update
-ray.get(llm.start_weight_update.remote(is_checkpoint_format=True))
+ray.get(llm.start_weight_update.remote())
 
 # Issue update_weights call with NCCL-specific update info
 # packed=True enables efficient batched tensor broadcasting

examples/rl/rlhf_nccl_fsdp_ep.py

@@ -299,7 +299,7 @@ async def main():
     print(f"[sync] Got metadata for {len(names)} parameters.")
 
     print("[sync] Starting weight update...")
-    await engine.start_weight_update(is_checkpoint_format=True)
+    await engine.start_weight_update()
 
     print("[sync] Broadcasting weights from FSDP → vLLM...")
     broadcast_handles = [