Refactor and Optimize

Signed-off-by: Hollow Man <hollowman@opensuse.org>

Author: HollowMan6

examples/configs/distillation_math.yaml

@@ -206,6 +206,8 @@ policy: &POLICY_BASE
             dtype: ${...precision}
             transport: "sparse_indices"  # dense, sparse_indices, or sparse_bitmask
             full_sync_interval: 20
+            sparse_bucket_size_bytes: 5368709120  # 5 GiB
+            delta_load_batch_size_bytes: 536870912  # 512 MiB
 
         colocated:
             # true: generation shares training GPUs

examples/configs/grpo_math_1B.yaml

@@ -314,6 +314,8 @@ policy:
       dtype: ${policy.precision}
       transport: "sparse_indices"  # dense, sparse_indices, or sparse_bitmask
       full_sync_interval: 20
+      sparse_bucket_size_bytes: 5368709120  # 5 GiB
+      delta_load_batch_size_bytes: 536870912  # 512 MiB
     colocated:
       # true: generation shares training GPUs
       # false: uses dedicated generation resources

nemo_rl/models/generation/vllm/config.py

@@ -62,6 +62,14 @@ class VllmDeltaCompressionConfig(TypedDict):
     # 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):

nemo_rl/models/generation/vllm/vllm_backend.py

@@ -25,7 +25,6 @@
     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,
@@ -58,7 +57,7 @@ def fix_gpt_oss_export_transpose(key: str, weight: torch.Tensor) -> torch.Tensor
 
 class VllmInternalWorkerExtension:
     state_dict_info: dict[str, Any] | None = None
-    use_delta_weight_transfer: bool = False
+    delta_load_batch_size_bytes: int | None = None
 
     def init_collective(
         self,
@@ -109,18 +108,21 @@ def maybe_init_zmq(self):
     def prepare_refit_info(
         self,
         state_dict_info: dict[str, Any],
-        use_delta_weight_transfer: bool,
+        delta_load_batch_size_bytes: int | None = None,
     ) -> None:
-        """Prepare state dict metadata for weight refitting and IPC streaming.
+        """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)}
-            use_delta_weight_transfer (bool): Whether collective refit receives
-                full weights only or the delta-aware full/delta protocol.
+            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
-        self.use_delta_weight_transfer = use_delta_weight_transfer
+        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)."""
@@ -332,28 +334,15 @@ def update_weights_via_ipc_zmq(self) -> bool:
     )
     def update_weights_from_collective(self) -> bool:
         """Update the model weights from collective communication."""
-        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."
-        )
-
         try:
-            if not self.use_delta_weight_transfer:
-                packed_broadcast_consumer(
-                    iterator=iter(state_dict_info.items()),
-                    group=self.model_update_group,
-                    src=0,
-                    post_unpack_func=self._load_weights,
-                )
-            else:
-                packed_weight_transfer_consumer(
-                    group=self.model_update_group,
-                    src=0,
-                    load_full_weights_func=self._load_weights,
-                    load_delta_weights_func=self._load_weight_deltas,
-                    device=self.device,
-                )
+            packed_weight_transfer_consumer(
+                group=self.model_update_group,
+                src=0,
+                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
             self._maybe_process_fp8_kv_cache()

nemo_rl/models/generation/vllm/vllm_worker.py

@@ -35,7 +35,6 @@
 from nemo_rl.models.huggingface.common import ModelFlag
 from nemo_rl.models.policy.utils import is_vllm_v1_engine_enabled
 from nemo_rl.utils.nsys import wrap_with_nvtx_name
-from nemo_rl.utils.weight_transfer import get_vllm_delta_transfer_config
 
 
 # Use a base class to share some functions to avoid code duplication.
@@ -136,9 +135,6 @@ def __init__(
                           the vLLM worker subprocess (e.g. for quantization configs).
         """
         self.cfg = config
-        self.use_delta_weight_transfer = (
-            get_vllm_delta_transfer_config(self.cfg) is not None
-        )
         self.model_name = self.cfg["model_name"]
         self.tensor_parallel_size = self.cfg["vllm_cfg"]["tensor_parallel_size"]
         self.pipeline_parallel_size = self.cfg["vllm_cfg"]["pipeline_parallel_size"]
@@ -666,6 +662,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:
@@ -912,7 +919,7 @@ 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.use_delta_weight_transfer),
+            args=(state_dict_info, self._get_delta_load_batch_size_bytes()),
         )
 
     @wrap_with_nvtx_name("vllm_genertion_worker/update_weights_via_ipc_zmq")

nemo_rl/models/generation/vllm/vllm_worker_async.py

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

nemo_rl/models/policy/workers/dtensor_policy_worker.py

@@ -93,7 +93,7 @@
 from nemo_rl.utils.nsys import wrap_with_nvtx_name
 from nemo_rl.utils.weight_transfer import (
     create_vllm_delta_transfer_tracker,
-    dispatch_packed_weight_transfer,
+    packed_weight_transfer_producer,
 )
 
 
@@ -1877,7 +1877,7 @@ def _params_iterator():
                 yield name, _dtensor_post_iter_func(tensor, self.dtype)
 
         params_iterator = _params_iterator()
-        dispatch_packed_weight_transfer(
+        packed_weight_transfer_producer(
             iterator=params_iterator,
             group=self.model_update_group,
             src=0,

nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py

@@ -74,7 +74,7 @@
 from nemo_rl.utils.nsys import wrap_with_nvtx_name
 from nemo_rl.utils.weight_transfer import (
     create_vllm_delta_transfer_tracker,
-    dispatch_packed_weight_transfer,
+    packed_weight_transfer_producer,
 )
 
 
@@ -979,7 +979,7 @@ def broadcast_weights_for_collective(
             self.model = self.move_to_cuda(self.model)
 
         params_iterator = dtensor_params_generator(self.model, self.dtype)
-        dispatch_packed_weight_transfer(
+        packed_weight_transfer_producer(
             iterator=params_iterator,
             group=self.model_update_group,
             src=0,

nemo_rl/models/policy/workers/megatron_policy_worker.py

@@ -92,7 +92,7 @@
 from nemo_rl.utils.nsys import wrap_with_nvtx_name
 from nemo_rl.utils.weight_transfer import (
     create_vllm_delta_transfer_tracker,
-    dispatch_packed_weight_transfer,
+    packed_weight_transfer_producer,
 )
 
 TokenizerType = TypeVar("TokenizerType", bound=PreTrainedTokenizerBase)
@@ -1125,7 +1125,7 @@ def broadcast_weights_for_collective(
     ) -> None:
         """Broadcast the weights for collective communication."""
         params_iterator = self._iter_params_with_optional_kv_scales(kv_scales=kv_scales)
-        dispatch_packed_weight_transfer(
+        packed_weight_transfer_producer(
             iterator=params_iterator,
             group=self.model_update_group,
             src=0,

nemo_rl/utils/torch_dtypes.py

@@ -37,3 +37,14 @@
     **G_CANONICAL_FLOAT_DTYPE_MAP,
     "float64": torch.float64,
 }
+
+for _float8_dtype_name in (
+    "float8_e4m3fn",
+    "float8_e5m2",
+    "float8_e4m3fnuz",
+    "float8_e5m2fnuz",
+):
+    _float8_dtype = getattr(torch, _float8_dtype_name, None)
+    if _float8_dtype is not None:
+        G_TENSOR_DTYPE_MAP[_float8_dtype_name] = _float8_dtype
+del _float8_dtype, _float8_dtype_name