feat(vllm): add delta-compressed collective refit

Adds optional delta-compressed weight transfer for
non-colocated vLLM collective refit.

This introduces a delta-aware packed weight transfer
protocol that can send either full weights or additive
deltas, with support for `dense`, `sparse_indices`, and
`sparse_bitmask` delta encodings. The trainer source rank
keeps a pinned CPU baseline of the last successfully
synced HF-format weights, computes deltas against that
baseline, and periodically sends full syncs
based on `full_sync_interval`.

The feature is disabled by default and only applies
to non-colocated vLLM refit. Colocated CUDA IPC,
vLLM FP8 weights, and ModelOpt quantized vLLM paths are rejected.

Author: HollowMan6

nemo_rl/data_plane/__init__.py

@@ -18,14 +18,11 @@
 detail of a specific adapter.
 """
 
-from nemo_rl.data_plane.codec import materialize
-from nemo_rl.data_plane.factory import build_data_plane_client
 from nemo_rl.data_plane.interfaces import (
     DataPlaneClient,
     DataPlaneConfig,
     KVBatchMeta,
 )
-from nemo_rl.data_plane.observability import MetricsDataPlaneClient, log_event
 
 __all__ = [
     "DataPlaneClient",
@@ -36,3 +33,25 @@
     "log_event",
     "materialize",
 ]
+
+
+def __getattr__(name: str):
+    if name == "materialize":
+        from nemo_rl.data_plane.codec import materialize
+
+        return materialize
+    if name == "build_data_plane_client":
+        from nemo_rl.data_plane.factory import build_data_plane_client
+
+        return build_data_plane_client
+    if name in {"MetricsDataPlaneClient", "log_event"}:
+        from nemo_rl.data_plane.observability import (
+            MetricsDataPlaneClient,
+            log_event,
+        )
+
+        return {
+            "MetricsDataPlaneClient": MetricsDataPlaneClient,
+            "log_event": log_event,
+        }[name]
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

nemo_rl/data_plane/interfaces.py

@@ -37,9 +37,18 @@
 
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
-from typing import Any, Callable, Literal, NotRequired, Sequence, TypedDict
-
-from tensordict import TensorDict
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Literal,
+    NotRequired,
+    Sequence,
+    TypedDict,
+)
+
+if TYPE_CHECKING:
+    from tensordict import TensorDict
 
 
 class DataPlaneConfig(TypedDict):

nemo_rl/models/generation/vllm/config.py

@@ -16,6 +16,10 @@
 
 from nemo_rl.models.generation.interfaces import GenerationConfig
 
+DeltaCompressionDType = Literal[
+    "fp16", "float16", "bf16", "bfloat16", "fp32", "float32"
+]
+
 
 class VllmSpecificArgs(TypedDict):
     tensor_parallel_size: int
@@ -41,9 +45,34 @@ class VllmSpecificArgs(TypedDict):
     tool_parser_plugin: NotRequired[str]
 
 
+class VllmDeltaCompressionConfig(TypedDict):
+    # Enables delta-compressed collective refit for non-colocated vLLM.
+    # Delta payloads use sparse indices to reduce collective payload bytes;
+    # vLLM still receives decoded dense tensors before load_weights.
+    # This keeps a pinned CPU baseline of the last synced HF-format weights
+    # on the source producer rank, roughly one extra model-weight copy.
+    # Recommended default: false.
+    enabled: bool
+    # Floating dtype used when computing and sending deltas.
+    # Recommended default: ${policy.precision}.
+    dtype: DeltaCompressionDType
+    # Number of successful refits between full baseline refreshes.
+    # Recommended default: 20.
+    full_sync_interval: int
+    # Maximum sparse-encoded payload bytes to bucket before broadcasting.
+    # Smaller values improve refit pipelining; larger values reduce broadcast
+    # call overhead. Recommended default: 5368709120 (5 GiB).
+    sparse_bucket_size_bytes: int
+    # Maximum decoded delta tensor bytes to batch before calling vLLM load_weights.
+    # Smaller values improve overlap with receives; larger values reduce loader
+    # call overhead. Recommended default: 536870912 (512 MiB).
+    delta_load_batch_size_bytes: int
+
+
 class VllmConfig(GenerationConfig):
     vllm_cfg: VllmSpecificArgs
     vllm_kwargs: NotRequired[dict[str, Any]]
+    delta_compression: NotRequired[VllmDeltaCompressionConfig | None]
 
     # quantization config
     quant_cfg: NotRequired[str | None]

nemo_rl/models/generation/vllm/vllm_backend.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 import gc
 import traceback
+from collections.abc import Iterator
 from typing import Any
 
 import torch
@@ -24,7 +25,10 @@
     rebuild_cuda_tensor_from_ipc,
 )
 from nemo_rl.utils.nsys import wrap_with_nvtx_name
-from nemo_rl.utils.packed_tensor import packed_broadcast_consumer
+from nemo_rl.utils.weight_transfer import (
+    additive_weight_load_context,
+    packed_weight_transfer_consumer,
+)
 
 try:
     import vllm  # noqa: F401
@@ -52,6 +56,9 @@ def fix_gpt_oss_export_transpose(key: str, weight: torch.Tensor) -> torch.Tensor
 
 
 class VllmInternalWorkerExtension:
+    state_dict_info: dict[str, Any] | None = None
+    delta_load_batch_size_bytes: int | None = None
+
     def init_collective(
         self,
         rank_prefix: int,
@@ -98,14 +105,24 @@ def maybe_init_zmq(self):
             self.zmq_socket.setsockopt(zmq.LINGER, 0)
             self.zmq_socket.connect(self.get_zmq_address())
 
-    def prepare_refit_info(self, state_dict_info: dict[str, Any]) -> None:
-        """Prepare state dict metadata for weight refitting and IPC streaming.
+    def prepare_refit_info(
+        self,
+        state_dict_info: dict[str, Any],
+        delta_load_batch_size_bytes: int | None = None,
+    ) -> None:
+        """Prepare state dict metadata for IPC/ZMQ weight refitting.
+
+        Collective refit receives tensor metadata from the transfer headers.
 
         Args:
             state_dict_info (dict): A dictionary containing the info for refit.
                 e.g. {tensor_name: (shape, dtype)}
+            delta_load_batch_size_bytes (int | None): Maximum decoded delta bytes
+                to batch before calling vLLM load_weights. None means delta
+                transfer is disabled.
         """
-        self.state_dict_info = state_dict_info  # pyrefly: ignore[implicitly-defined-attribute]  This class does not define __init__ so assignments like this should be ignored
+        self.state_dict_info = state_dict_info
+        self.delta_load_batch_size_bytes = delta_load_batch_size_bytes
 
     def _maybe_process_fp8_kv_cache(self) -> None:
         """Process weights after loading for FP8 KV cache (static scales)."""
@@ -213,7 +230,7 @@ def _load_draft_weights(
         draft_weights = self._trim_vocab_padding(draft_model, draft_weights)
         draft_model.load_weights(weights=draft_weights)
 
-    def _load_weights(self, weights):
+    def _load_weights(self, weights: list[tuple[str, torch.Tensor]]) -> None:
         """Load weights with GptOss transpose fix, FP8, and draft-weight support.
 
         Applies GPT-OSS down_proj transpose if needed, splits policy/draft
@@ -238,13 +255,35 @@ def _load_weights(self, weights):
 
         self._load_draft_weights(draft_weights)
 
+    def _load_weight_deltas(self, weights: list[tuple[str, torch.Tensor]]) -> None:
+        """Apply additive weight deltas through vLLM's regular loaders."""
+        with additive_weight_load_context(self._iter_delta_load_target_tensors()):
+            self._load_weights(weights)
+
+    def _iter_delta_load_target_tensors(self) -> Iterator[torch.Tensor]:
+        """Yield model-owned tensors that should receive additive delta loads."""
+        yield from self.model_runner.model.parameters()
+        yield from self.model_runner.model.buffers()
+
+        draft_owner = getattr(self.model_runner, "drafter", None)
+        draft_model = getattr(draft_owner, "model", None) if draft_owner else None
+        if draft_model is not None:
+            yield from draft_model.parameters()
+            yield from draft_model.buffers()
+
     @wrap_with_nvtx_name("vllm_internal_worker_extension/update_weights_via_ipc_zmq")
     def update_weights_via_ipc_zmq(self) -> bool:
         """Receive and update model weights via ZMQ IPC socket.
 
         Returns:
             bool: True if weights were successfully updated.
         """
+        state_dict_info = self.state_dict_info
+        assert state_dict_info is not None, (
+            "state_dict_info is not prepared. "
+            "Please call prepare_refit_info when initializing the worker."
+        )
+
         buffer = None
         weights = None
 
@@ -275,7 +314,7 @@ def update_weights_via_ipc_zmq(self) -> bool:
                 weights = []
                 offset = 0
                 for key in list_keys:
-                    shape, dtype = self.state_dict_info[key]  # pyrefly
+                    shape, dtype = state_dict_info[key]
                     if isinstance(shape, list):
                         shape = torch.Size(shape)
 
@@ -336,19 +375,14 @@ def update_weights_via_ipc_zmq(self) -> bool:
     )
     def update_weights_from_collective(self) -> bool:
         """Update the model weights from collective communication."""
-        assert self.state_dict_info is not None, (
-            "state_dict_info is not prepared. "
-            "Please call prepare_refit_info when initializing the worker."
-        )
-
-        load_model_weight_func = self._load_weights
-
         try:
-            packed_broadcast_consumer(
-                iterator=iter(self.state_dict_info.items()),
+            packed_weight_transfer_consumer(
                 group=self.model_update_group,
                 src=0,
-                post_unpack_func=load_model_weight_func,
+                load_full_weights_func=self._load_weights,
+                load_delta_weights_func=self._load_weight_deltas,
+                device=self.device,
+                delta_load_batch_size_bytes=self.delta_load_batch_size_bytes,
             )
 
             # Process weights after loading for FP8 KV cache

nemo_rl/models/generation/vllm/vllm_worker.py

@@ -664,6 +664,17 @@ def _get_raw_spec_counters(self) -> dict[str, float | list[float]]:
                     metrics[metric.name] = metric.value
         return metrics
 
+    def _get_delta_load_batch_size_bytes(self) -> int | None:
+        delta_config = self.cfg.get("delta_compression", None)
+        if delta_config is None or not delta_config["enabled"]:
+            return None
+        delta_load_batch_size_bytes = int(delta_config["delta_load_batch_size_bytes"])
+        if delta_load_batch_size_bytes < 1:
+            raise ValueError(
+                "delta_compression.delta_load_batch_size_bytes must be >= 1"
+            )
+        return delta_load_batch_size_bytes
+
 
 class VllmGenerationWorkerImpl(BaseVllmGenerationWorker):
     def _create_engine(self, llm_kwargs: dict[str, Any]) -> None:
@@ -908,7 +919,10 @@ def report_device_id(self) -> list[str]:
 
     def prepare_refit_info(self, state_dict_info: dict[str, Any]) -> None:
         """Prepare the info for refit."""
-        self.llm.collective_rpc("prepare_refit_info", args=(state_dict_info,))
+        self.llm.collective_rpc(
+            "prepare_refit_info",
+            args=(state_dict_info, self._get_delta_load_batch_size_bytes()),
+        )
 
     @wrap_with_nvtx_name("vllm_genertion_worker/update_weights_via_ipc_zmq")
     def update_weights_via_ipc_zmq(self) -> bool:

nemo_rl/models/generation/vllm/vllm_worker_async.py

@@ -1073,7 +1073,10 @@ async def report_device_id_async(self) -> list[str]:
 
     async def prepare_refit_info_async(self, state_dict_info: dict[str, Any]) -> None:
         """Async version of prepare_refit_info."""
-        await self.llm.collective_rpc("prepare_refit_info", args=(state_dict_info,))
+        await self.llm.collective_rpc(
+            "prepare_refit_info",
+            args=(state_dict_info, self._get_delta_load_batch_size_bytes()),
+        )
 
     async def update_weights_via_ipc_zmq_async(
         self,

nemo_rl/models/policy/workers/dtensor_policy_worker.py

@@ -92,7 +92,10 @@
     save_checkpoint,
 )
 from nemo_rl.utils.nsys import wrap_with_nvtx_name
-from nemo_rl.utils.packed_tensor import packed_broadcast_producer
+from nemo_rl.utils.weight_transfer import (
+    create_vllm_delta_transfer_tracker,
+    packed_weight_transfer_producer,
+)
 
 
 def _attach_context_parallel_hooks(model: nn.Module) -> None:
@@ -229,6 +232,9 @@ def __init__(
         configure_dynamo_cache()
 
         self.cfg = config
+        self.delta_weight_transfer_tracker = create_vllm_delta_transfer_tracker(
+            self.cfg.get("generation")
+        )
         # torch distributed init. Envars for rank, world_size, and master_addr and master_port are set from the ray remote call
         torch.distributed.init_process_group(backend="nccl")
         self.rank = torch.distributed.get_rank()
@@ -1795,6 +1801,11 @@ def prepare_refit_info(self) -> Optional[dict[str, Any]]:
             # all tensor will be casted to self.dtype in stream_weights_via_ipc_zmq/broadcast_weights_for_collective
             state_dict_info[name] = (tensor.shape, self.dtype)
 
+        if self.rank == 0 and self.delta_weight_transfer_tracker is not None:
+            self.delta_weight_transfer_tracker.prewarm_baseline_from_metadata(
+                state_dict_info
+            )
+
         return state_dict_info
 
     @torch.no_grad()
@@ -1879,14 +1890,15 @@ def _dtensor_post_iter_func(tensor, dtype):
             tensor = tensor.to(dtype, non_blocking=True)
             return tensor
 
-        # param_iterator will return (name, tensor), we only need tensor
-        dtensor_post_iter_func = lambda x: _dtensor_post_iter_func(x[1], self.dtype)
+        def _params_iterator():
+            for name, tensor in self.model.state_dict().items():
+                yield name, _dtensor_post_iter_func(tensor, self.dtype)
 
-        packed_broadcast_producer(
-            iterator=iter(self.model.state_dict().items()),
+        packed_weight_transfer_producer(
+            iterator=_params_iterator(),
             group=self.model_update_group,
             src=0,
-            post_iter_func=dtensor_post_iter_func,
+            delta_tracker=self.delta_weight_transfer_tracker,
         )
 
         # Manually move model to cpu for cpu offload case

nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py

@@ -73,7 +73,10 @@
 )
 from nemo_rl.utils.checkpoint import CheckpointingConfig
 from nemo_rl.utils.nsys import wrap_with_nvtx_name
-from nemo_rl.utils.packed_tensor import packed_broadcast_producer
+from nemo_rl.utils.weight_transfer import (
+    create_vllm_delta_transfer_tracker,
+    packed_weight_transfer_producer,
+)
 
 
 def dtensor_params_generator(
@@ -336,6 +339,9 @@ def __init__(
             self.sampling_params,
             _runtime_is_reward_model,  # Duplicate, already set as _is_reward_model
         ) = runtime_config
+        self.delta_weight_transfer_tracker = create_vllm_delta_transfer_tracker(
+            self.cfg.get("generation")
+        )
 
     @wrap_with_nvtx_name("dtensor_policy_worker_v2/train")
     def train(
@@ -875,6 +881,11 @@ def prepare_refit_info(self) -> Optional[dict[str, Any]]:
             for adapted_fqn, adapted_tensor in adapted_fqn_tensors:
                 state_dict_info[adapted_fqn] = (adapted_tensor.shape, self.dtype)
 
+        if self.rank == 0 and self.delta_weight_transfer_tracker is not None:
+            self.delta_weight_transfer_tracker.prewarm_baseline_from_metadata(
+                state_dict_info
+            )
+
         return state_dict_info
 
     @torch.no_grad()
@@ -985,14 +996,11 @@ def broadcast_weights_for_collective(
             )
             self.model = self.move_to_cuda(self.model)
 
-        # param_iterator will return (name, tensor), we only need tensor
-        dtensor_post_iter_func = lambda x: x[1]
-
-        packed_broadcast_producer(
+        packed_weight_transfer_producer(
             iterator=dtensor_params_generator(self.model, self.dtype),
             group=self.model_update_group,
             src=0,
-            post_iter_func=dtensor_post_iter_func,
+            delta_tracker=self.delta_weight_transfer_tracker,
         )
 
         # Manually move model to cpu for cpu offload case