Clean up

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

Author: HollowMan6

docs/design-docs/checkpoint-engines.md

@@ -77,8 +77,8 @@ policy:
 ```
 
 The factory passes `bucket_size` in bytes plus the selected backend kwargs to
-the backend constructor. It also provides a backend-neutral default `device`
-unless the config already specifies one.
+the backend constructor. Backend-specific settings such as transfer device,
+cleanup behavior, and transport plugin name live in config.
 
 ## Backend Interface
 
@@ -100,7 +100,7 @@ from nemo_rl.utils.checkpoint_engines import (
 class MyCheckpointEngine(CheckpointEngine):
     cleanup_after_load = True
 
-    def __init__(self, bucket_size: int, device: str | torch.device = "cuda"):
+    def __init__(self, bucket_size: int, device: str | torch.device):
         self.bucket_size = bucket_size
         self.device = torch.device(device)
 
@@ -245,6 +245,8 @@ policy:
       backend: nixl
       engine_kwargs:
         nixl:
+          device: cuda
+          cleanup_after_load: false
           backend_name: UCX
           backend_init_params:
             ucx_error_handling_mode: peer

docs/guides/checkpoint-engine-refit.md

@@ -77,6 +77,8 @@ policy:
       backend: nixl
       engine_kwargs:
         nixl:
+          device: cpu
+          cleanup_after_load: false
           backend_name: UCX
           backend_init_params:
             ucx_error_handling_mode: peer

examples/configs/grpo_math_1B.yaml

@@ -330,7 +330,10 @@ policy:
       update_weights_bucket_megabytes: 2048
       engine_kwargs:
         # For plugin backends, key kwargs by the exact backend string.
-        nixl: {}
+        nixl:
+          device: cuda
+          cleanup_after_load: true
+          backend_name: UCX
     colocated:
       # true: generation shares training GPUs
       # false: uses dedicated generation resources

examples/configs/grpo_math_8B.yaml

@@ -65,7 +65,10 @@ policy:
       update_weights_bucket_megabytes: 2048
       engine_kwargs:
         # For plugin backends, key kwargs by the exact backend string.
-        nixl: {}
+        nixl:
+          device: cuda
+          cleanup_after_load: true
+          backend_name: UCX
 
 cluster:
   gpus_per_node: 8

nemo_rl/algorithms/grpo.py

@@ -72,6 +72,7 @@
     run_async_nemo_gym_rollout,
     run_multi_turn_rollout,
 )
+from nemo_rl.models.generation.constants import SGLANG_BACKEND, VLLM_BACKEND
 from nemo_rl.models.generation.interfaces import GenerationInterface
 from nemo_rl.models.generation.sglang import SGLangConfig, SGLangGeneration
 from nemo_rl.models.generation.vllm import VllmConfig, VllmGeneration
@@ -88,6 +89,7 @@
 from nemo_rl.utils.nsys import maybe_gpu_profile_step
 from nemo_rl.utils.timer import TimeoutChecker, Timer
 from nemo_rl.utils.venvs import create_local_venv_on_each_node
+from nemo_rl.weight_sync import create_weight_synchronizer
 
 # ===============================================================================
 # Configuration
@@ -1188,16 +1190,6 @@ def _create_advantage_estimator(master_config: MasterConfig):
     return adv_estimator
 
 
-def _flatten_checkpoint_engine_metadata(metadata_results: list[Any]) -> list[Any]:
-    metadata = []
-    for worker_metadata in metadata_results:
-        if isinstance(worker_metadata, list):
-            metadata.extend(worker_metadata)
-        else:
-            metadata.append(worker_metadata)
-    return metadata
-
-
 def _get_enabled_checkpoint_engine_config(
     checkpoint_engine_config: dict[str, Any] | None,
 ) -> dict[str, Any] | None:
@@ -1206,71 +1198,6 @@ def _get_enabled_checkpoint_engine_config(
     return checkpoint_engine_config
 
 
-def _refit_policy_generation_with_checkpoint_engine(
-    policy: ColocatablePolicyInterface,
-    policy_generation: GenerationInterface,
-    checkpoint_engine_config: dict[str, Any],
-    kv_scales: Optional[dict[str, float]] = None,
-) -> bool:
-    backend = checkpoint_engine_config["backend"]
-    bucket_size_bytes = (
-        checkpoint_engine_config["update_weights_bucket_megabytes"] * 1024 * 1024
-    )
-    engine_kwargs = checkpoint_engine_config["engine_kwargs"][backend]
-
-    try:
-        init_futures = policy.init_checkpoint_engine(
-            backend=backend,
-            bucket_size_bytes=bucket_size_bytes,
-            engine_kwargs=engine_kwargs,
-        ) + policy_generation.init_checkpoint_engine(
-            backend=backend,
-            bucket_size_bytes=bucket_size_bytes,
-            engine_kwargs=engine_kwargs,
-        )
-        ray.get(init_futures)
-
-        policy_metadata = _flatten_checkpoint_engine_metadata(
-            ray.get(policy.prepare_checkpoint_engine())
-        )
-        generation_metadata = _flatten_checkpoint_engine_metadata(
-            ray.get(policy_generation.prepare_checkpoint_engine())
-        )
-
-        train_world_size = len(policy_metadata)
-        rollout_world_size = len(generation_metadata)
-        metadata = policy_metadata + generation_metadata
-        ray.get(
-            policy.init_checkpoint_engine_process_group(
-                metadata=metadata,
-                train_world_size=train_world_size,
-                rollout_world_size=rollout_world_size,
-            )
-            + policy_generation.init_checkpoint_engine_process_group(
-                metadata=metadata,
-                train_world_size=train_world_size,
-                rollout_world_size=rollout_world_size,
-            )
-        )
-
-        futures_train = policy.send_weights_via_checkpoint_engine(kv_scales=kv_scales)
-        futures_inference = policy_generation.update_weights_from_checkpoint_engine()
-        ray.get(futures_train)
-        results = ray.get(futures_inference)
-        return all(result for result in results if result is not None)
-    finally:
-        try:
-            ray.get(
-                policy.finalize_checkpoint_engine()
-                + policy_generation.finalize_checkpoint_engine()
-            )
-        except Exception as finalize_error:
-            warnings.warn(
-                f"Failed to finalize checkpoint-engine refit state: {finalize_error}",
-                RuntimeWarning,
-            )
-
-
 def refit_policy_generation(
     policy: ColocatablePolicyInterface,
     policy_generation: GenerationInterface,
@@ -1300,101 +1227,44 @@ def refit_policy_generation(
             "policy.generation.checkpoint_engine.enabled."
         )
 
-    if colocated_inference:
-        policy.offload_before_refit()
-        policy_generation.prepare_for_generation(tags=["weights"])
+    if colocated_inference or checkpoint_engine_config is not None:
+        generation_backend = (
+            SGLANG_BACKEND
+            if isinstance(policy_generation, SGLangGeneration)
+            else VLLM_BACKEND
+        )
+        weight_sync = create_weight_synchronizer(
+            policy=policy,
+            generation=policy_generation,
+            generation_backend=generation_backend,
+            colocated=colocated_inference,
+            refit_buffer_size_gb=_refit_buffer_size_gb,
+        )
+        weight_sync.sync_weights(timer=timer, kv_scales=kv_scales)
+        return
 
-    # Create a context manager that does nothing when timer is None
     timer_context = (
         timer.time("prepare_for_generation/transfer_and_update_weights")
         if timer is not None
         else nullcontext()
     )
     with timer_context:
-        # update weights
-        update_success = False
-        if colocated_inference:
-            # get model param keys, which is grouped by size
-            if _refit_buffer_size_gb is not None:
-                buffer_size_bytes = _refit_buffer_size_gb * (1024**3)
-            else:
-                # Empirically sets ratio as 30% to maximize efficiency.
-                # The remaining 70% is a necessary buffer reserved for the parameter all-gathering across the expert-parallelism dimension.
-                memory_ratio = os.getenv("NRL_REFIT_BUFFER_MEMORY_RATIO", "0.3")
-                buffer_size_bytes = int(
-                    policy.get_free_memory_bytes() * float(memory_ratio)
-                )
-
-            if isinstance(policy_generation, SGLangGeneration):
-                sglang_url_to_gpu_uuids = (
-                    policy_generation.get_sglang_url_to_gpu_uuids()
-                )
-                # Stream weights via HTTP
-                flush_success = policy_generation.invalidate_kv_cache()
-                if not flush_success:
-                    print("SGLang KV cache invalidation failed before weight update. ")
-                futures_train = policy.stream_weights_via_http(
-                    sglang_url_to_gpu_uuids=sglang_url_to_gpu_uuids,
-                )
-                # Wait for all workers to complete
-                ray.get(futures_train)
-                update_success = True
-            else:
-                # Original ZMQ IPC path for vLLM
-                futures_train = policy.stream_weights_via_ipc_zmq(
-                    buffer_size_bytes=buffer_size_bytes
-                )
-                futures_inference = policy_generation.update_weights_via_ipc_zmq()
-                # wait for all futures to complete
-                ray.get(futures_train)
-                results = ray.get(futures_inference)
-                update_success = all(result for result in results if result is not None)
-        else:
-            if checkpoint_engine_config is not None:
-                if isinstance(policy_generation, SGLangGeneration):
-                    raise NotImplementedError(
-                        "SGLang does not support checkpoint-engine non-colocated refit."
-                    )
-                update_success = _refit_policy_generation_with_checkpoint_engine(
-                    policy,
-                    policy_generation,
-                    checkpoint_engine_config,
-                    kv_scales=kv_scales,
-                )
-            else:
-                # update weights through nccl
-                # SGLang haven't implemented non-colocated inference mode.
-                if isinstance(policy_generation, SGLangGeneration):
-                    raise NotImplementedError(
-                        "SGLang haven't implemented non-colocated inference mode. "
-                    )
-                futures_train = policy.broadcast_weights_for_collective(
-                    kv_scales=kv_scales
-                )
-                futures_inference = policy_generation.update_weights_from_collective()
-                # wait for all futures to complete
-                ray.get(futures_train)
-                results = ray.get(futures_inference)
-                update_success = all(result for result in results if result is not None)
-
-        # check if update is successful
-        if not update_success:
-            if colocated_inference:
-                error_tag = "cuda-ipc"
-            elif checkpoint_engine_config is not None:
-                error_tag = checkpoint_engine_config["backend"]
-            else:
-                error_tag = "nccl"
-            error_message = (
-                "❌ Error: Updating weights for the generation policy failed during refit.\n"
-                f"This often indicates an issue with {error_tag} or "
-                "a problem within the generation backend (e.g., vLLM worker).\n"
+        if isinstance(policy_generation, SGLangGeneration):
+            raise NotImplementedError(
+                "SGLang haven't implemented non-colocated inference mode. "
             )
-            raise RuntimeError(error_message)
+        futures_train = policy.broadcast_weights_for_collective(kv_scales=kv_scales)
+        futures_inference = policy_generation.update_weights_from_collective()
+        ray.get(futures_train)
+        results = ray.get(futures_inference)
+        update_success = all(result for result in results if result is not None)
 
-    if colocated_inference:
-        policy.offload_after_refit()
-        policy_generation.prepare_for_generation(tags=["kv_cache"])
+    if not update_success:
+        raise RuntimeError(
+            "❌ Error: Updating weights for the generation policy failed during refit.\n"
+            "This often indicates an issue with nccl or a problem within the "
+            "generation backend (e.g., vLLM worker).\n"
+        )
 
 
 def _log_mixed_rewards_and_advantages_information(

nemo_rl/models/generation/vllm/vllm_backend.py

@@ -49,7 +49,7 @@ def maybe_preinit_nixl_for_vllm_worker(
     if wrapper_state.get("_nrl_nixl_preinit_agent") is not None:
         return
 
-    backend_name = preinit_config.get("backend_name", "UCX")
+    backend_name = preinit_config["backend_name"]
     backend_init_params = preinit_config.get("backend_init_params")
 
     from nemo_rl.utils.checkpoint_engines.nixl import preinit_nixl_agent
@@ -122,7 +122,6 @@ def init_checkpoint_engine(
             backend,
             bucket_size_bytes=bucket_size_bytes,
             engine_kwargs=engine_kwargs,
-            default_device=self.device,
         )
 
     def prepare_checkpoint_engine(self) -> Any:

nemo_rl/models/generation/vllm/vllm_worker.py

@@ -284,7 +284,7 @@ def _patch_vllm_worker_nixl_preinit():
 
             nixl_kwargs = checkpoint_cfg["engine_kwargs"]["nixl"]
             preinit_config = {
-                "backend_name": nixl_kwargs.get("backend_name", "UCX"),
+                "backend_name": nixl_kwargs["backend_name"],
                 "backend_init_params": nixl_kwargs.get("backend_init_params"),
             }
             old_snippet = (

nemo_rl/models/policy/utils.py

@@ -239,7 +239,7 @@ def maybe_preinit_nixl_for_checkpoint_engine(config: Any) -> Any | None:
     from nemo_rl.utils.checkpoint_engines.nixl import preinit_nixl_agent
 
     return preinit_nixl_agent(
-        backend_name=nixl_kwargs.get("backend_name", "UCX"),
+        backend_name=nixl_kwargs["backend_name"],
         backend_init_params=nixl_kwargs.get("backend_init_params"),
     )
 

nemo_rl/models/policy/workers/base_policy_worker.py

@@ -58,16 +58,10 @@ def init_checkpoint_engine(
 
         from nemo_rl.utils.checkpoint_engine import create_checkpoint_engine
 
-        default_device = (
-            torch.device("cuda", torch.cuda.current_device())
-            if torch.cuda.is_available()
-            else torch.device("cpu")
-        )
         self.checkpoint_engine = create_checkpoint_engine(
             backend,
             bucket_size_bytes=bucket_size_bytes,
             engine_kwargs=engine_kwargs,
-            default_device=default_device,
         )
 
     def prepare_checkpoint_engine(self) -> Any:

nemo_rl/utils/checkpoint_engines/base.py

@@ -182,15 +182,12 @@ def create_checkpoint_engine(
     *,
     bucket_size_bytes: int,
     engine_kwargs: dict[str, Any],
-    default_device: Any,
 ) -> CheckpointEngine:
-    """Create a checkpoint engine with a backend-neutral device default."""
-    kwargs = dict(engine_kwargs)
-    kwargs.setdefault("device", default_device)
+    """Create a checkpoint engine from backend-specific configuration."""
     return CheckpointEngineRegistry.new(
         backend,
         bucket_size=bucket_size_bytes,
-        **kwargs,
+        **engine_kwargs,
     )