feat: Unify policy phase transtions and offload methods

Signed-off-by: Saurabh Mishra <sauramishra@nvidia.com>

Author: saumishr

examples/nemo_gym/run_grpo_nemo_gym.py

@@ -26,7 +26,6 @@
 from wandb import Table
 
 from nemo_rl.algorithms.grpo import (
-    ColocatablePolicyInterface,
     EnvironmentInterface,
     GenerationInterface,
     Logger,
@@ -71,7 +70,7 @@ def parse_args() -> tuple[argparse.Namespace, list[str]]:
 
 # These types are directly imported from grpo_train since if something about the architecture changes we want to immediately fail.
 def collect_trajectories(
-    policy: ColocatablePolicyInterface,
+    policy,
     policy_generation: GenerationInterface,
     val_dataloader: StatefulDataLoader,
     tokenizer: TokenizerType,

nemo_rl/algorithms/distillation.py

@@ -54,7 +54,7 @@
 )
 from nemo_rl.models.generation.vllm import VllmConfig, VllmGeneration
 from nemo_rl.models.policy import PolicyConfig
-from nemo_rl.models.policy.interfaces import ColocatablePolicyInterface
+from nemo_rl.models.policy.interfaces import OffloadMode, PolicyTrainerInterface
 from nemo_rl.models.policy.lm_policy import Policy
 from nemo_rl.utils.checkpoint import CheckpointingConfig, CheckpointManager
 from nemo_rl.utils.logger import (
@@ -162,8 +162,8 @@ def setup(
     train_dataset: AllTaskProcessedDataset,
     val_dataset: Optional[AllTaskProcessedDataset],
 ) -> tuple[
-    ColocatablePolicyInterface,  # student_policy
-    ColocatablePolicyInterface,  # teacher_policy
+    PolicyTrainerInterface,  # student_policy
+    PolicyTrainerInterface,  # teacher_policy
     Optional[GenerationInterface],  # student_generation
     Optional[WeightSynchronizer],  # weight_sync
     StatefulDataLoader,
@@ -401,7 +401,7 @@ def setup(
         init_optimizer=False,
         init_reference_model=False,
     )
-    teacher_policy.offload_after_refit()
+    teacher_policy.finish_training()
 
     # ==========================
     #    Student Generation Interface
@@ -499,8 +499,8 @@ def setup(
 
 
 def distillation_train(
-    student_policy: ColocatablePolicyInterface,
-    teacher_policy: ColocatablePolicyInterface,
+    student_policy: PolicyTrainerInterface,
+    teacher_policy: PolicyTrainerInterface,
     student_generation: Optional[GenerationInterface],
     dataloader: StatefulDataLoader,
     val_dataloader: Optional[StatefulDataLoader],
@@ -682,7 +682,7 @@ def distillation_train(
 
                 print("▶ Preparing for teacher logprob inference...", flush=True)
                 with timer.time("teacher_logprob_inference_prep"):
-                    teacher_policy.prepare_for_lp_inference()
+                    teacher_policy.finish_training(offload_mode=OffloadMode.EVAL_ONLY)
 
                 print("▶ Computing teacher logprobs...", flush=True)
                 with timer.time("teacher_logprob_inference"):
@@ -696,7 +696,7 @@ def distillation_train(
 
                 print("▶ Preparing for training...", flush=True)
                 with timer.time("training_prep"):
-                    teacher_policy.offload_after_refit()
+                    teacher_policy.finish_training()
                     student_policy.prepare_for_training()  # set model train and reload optim to GPU
                     if weight_sync is not None:
                         weight_sync.mark_stale()

nemo_rl/algorithms/grpo.py

@@ -70,7 +70,7 @@
 from nemo_rl.models.generation.sglang import SGLangConfig, SGLangGeneration
 from nemo_rl.models.generation.vllm import VllmConfig, VllmGeneration
 from nemo_rl.models.policy import PolicyConfig
-from nemo_rl.models.policy.interfaces import ColocatablePolicyInterface
+from nemo_rl.models.policy.interfaces import OffloadMode, PolicyTrainerInterface
 from nemo_rl.models.policy.lm_policy import Policy
 from nemo_rl.utils.checkpoint import CheckpointingConfig, CheckpointManager
 from nemo_rl.utils.logger import (
@@ -220,7 +220,7 @@ def setup(
     val_dataset: Optional[AllTaskProcessedDataset],
     processor: Optional[AutoProcessor] = None,
 ) -> tuple[
-    ColocatablePolicyInterface,
+    PolicyTrainerInterface,
     Optional[GenerationInterface],
     tuple[RayVirtualCluster, RayVirtualCluster],
     Optional[WeightSynchronizer],
@@ -1089,7 +1089,7 @@ def _extract_prompt_only_messages(message_logs: list) -> list:
 
 
 def refit_policy_generation(
-    policy: ColocatablePolicyInterface,
+    policy: PolicyTrainerInterface,
     policy_generation: GenerationInterface,
     colocated_inference: bool,
     _refit_buffer_size_gb: Optional[int] = None,
@@ -1117,26 +1117,23 @@ def refit_policy_generation(
         DeprecationWarning,
         stacklevel=2,
     )
+    from nemo_rl.models.policy.interfaces import OffloadMode
+
     if colocated_inference:
-        policy.offload_before_refit()
+        policy.finish_training(offload_mode=OffloadMode.OPTIMIZER_ONLY)
         policy_generation.prepare_for_generation(tags=["weights"])
 
-    # 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)
@@ -1146,41 +1143,33 @@ def refit_policy_generation(
                 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:
-            # 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:
             error_tag = "cuda-ipc" if colocated_inference else "nccl"
             error_message = (
@@ -1191,7 +1180,7 @@ def refit_policy_generation(
             raise RuntimeError(error_message)
 
     if colocated_inference:
-        policy.offload_after_refit()
+        policy.finish_training(offload_mode=OffloadMode.FULL)
         policy_generation.prepare_for_generation(tags=["kv_cache"])
 
 
@@ -1308,7 +1297,7 @@ def compute_and_apply_seq_logprob_error_masking(
 
 
 def grpo_train(
-    policy: ColocatablePolicyInterface,
+    policy: PolicyTrainerInterface,
     policy_generation: Optional[GenerationInterface],
     wrapped_dataloader: StatefulDataLoader | MultipleDataloaderWrapper,
     val_dataloader: Optional[StatefulDataLoader],
@@ -1460,7 +1449,7 @@ def grpo_train(
                         # Compute KV scales if needed for FP8 quantization
                         if sync_kv_scales and kv_scales_cache is None:
                             print("▶ Computing KV cache scales...", flush=True)
-                            policy.prepare_for_lp_inference()
+                            policy.finish_training(offload_mode=OffloadMode.EVAL_ONLY)
                             # Align with training data processing to ensure parallel training compatibility
                             calib_flat, calib_input_lengths = (
                                 batched_message_log_to_flat_message(
@@ -1494,7 +1483,7 @@ def grpo_train(
                         )
                     else:
                         if colocated_inference:
-                            policy.offload_after_refit()  # unload optimizer to make space for generation
+                            policy.finish_training()
                         policy_generation.prepare_for_generation()
 
                 dynamic_sampling_num_gen_batches += 1
@@ -1727,7 +1716,7 @@ def grpo_train(
                 memory_tracker.snapshot_start_of_stage("Computing logprobs", dir())
                 print("▶ Preparing for logprob inference...", flush=True)
                 with timer.time("logprob_inference_prep"):
-                    policy.prepare_for_lp_inference()
+                    policy.finish_training(offload_mode=OffloadMode.EVAL_ONLY)
 
                 print("▶ Computing logprobs...", flush=True)
                 with timer.time("policy_and_reference_logprobs"):
@@ -1841,7 +1830,7 @@ def grpo_train(
                         )
                     else:
                         if colocated_inference:
-                            policy.offload_after_refit()  # unload optimizer to make space for generation
+                            policy.finish_training()
                         policy_generation.prepare_for_generation()
                     val_metrics, validation_timings = validate(
                         policy_generation,
@@ -2350,7 +2339,7 @@ def validate(
 
 
 def async_grpo_train(
-    policy: ColocatablePolicyInterface,
+    policy: PolicyTrainerInterface,
     policy_generation: Optional[GenerationInterface],
     dataloader: StatefulDataLoader,
     val_dataloader: Optional[StatefulDataLoader],
@@ -2768,7 +2757,7 @@ def async_grpo_train(
                 # Training phase (same as sync version)
                 print("▶ Preparing for logprob inference...")
                 with timer.time("logprob_inference_prep"):
-                    policy.prepare_for_lp_inference()
+                    policy.finish_training(offload_mode=OffloadMode.EVAL_ONLY)
 
                 print("▶ Computing logprobs...")
                 with timer.time("policy_and_reference_logprobs"):
@@ -3048,7 +3037,7 @@ def async_grpo_train(
                             os.path.join(checkpoint_path, "train_dataloader.pt"),
                         )
                         checkpointer.finalize_checkpoint(checkpoint_path)
-                    policy.offload_after_refit()
+                    policy.finish_training()
 
             # Logging
             # Log training data (match sync GRPO logging payload for parity)

nemo_rl/models/generation/interfaces.py

@@ -220,7 +220,7 @@ class GenerationOutputSpec(TypedDict):
 
 
 class GenerationInterface(ABC):
-    """Abstract base class defining the interface for RL policies."""
+    """Abstract base class defining the interface for generation backends."""
 
     @abstractmethod
     def init_collective(
@@ -237,10 +237,12 @@ def generate(
 
     @abstractmethod
     def prepare_for_generation(self, *args: Any, **kwargs: Any) -> bool:
+        """Transition to generation phase. Load weights, allocate KV cache, etc."""
         pass
 
     @abstractmethod
     def finish_generation(self, *args: Any, **kwargs: Any) -> bool:
+        """Transition out of generation phase. Free GPU resources for training."""
         pass
 
     @property

nemo_rl/models/policy/interfaces.py

@@ -12,9 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from abc import ABC, abstractmethod
+from enum import Enum
 from typing import Any, Optional, TypedDict
 
-import ray
 import torch
 
 from nemo_rl.algorithms.loss.interfaces import LossFunction
@@ -23,6 +23,26 @@
 from nemo_rl.utils.timer import Timer
 
 
+class OffloadMode(Enum):
+    """Controls how aggressively to offload during finish_training().
+
+    EVAL_ONLY: keep model on GPU in eval mode, optionally offload optimizer
+        based on worker config (offload_optimizer_for_logprob). Used by
+        algorithm code before logprob / KV-scale inference.
+    OPTIMIZER_ONLY: offload optimizer, keep model on GPU in its current mode.
+        Used by colocated weight synchronizers before weight transfer (model
+        must stay on GPU for CUDA IPC / HTTP streaming).
+    FULL: offload everything appropriate for the deployment topology. In
+        colocated mode this offloads model + optimizer; in non-colocated mode
+        this sets eval mode, keeps model on GPU, and optionally offloads
+        optimizer.
+    """
+
+    EVAL_ONLY = "eval_only"
+    OPTIMIZER_ONLY = "optimizer_only"
+    FULL = "full"
+
+
 class LogprobOutputSpec(TypedDict):
     """logprobs: Tensor of log probabilities."""
 
@@ -48,8 +68,8 @@ class TopkLogitsOutputSpec(TypedDict):
     topk_indices: torch.Tensor
 
 
-class PolicyInterface(ABC):
-    """Abstract base class defining the interface for RL policies."""
+class PolicyTrainerInterface(ABC):
+    """Abstract base class defining the interface for RL policy training."""
 
     @abstractmethod
     def get_logprobs(
@@ -148,10 +168,18 @@ def calibrate_qkv_fp8_scales(
 
     @abstractmethod
     def prepare_for_training(self, *args: Any, **kwargs: Any) -> None:
+        """Transition to training phase. Load model and optimizer to GPU, set train mode."""
         pass
 
     @abstractmethod
     def finish_training(self, *args: Any, **kwargs: Any) -> None:
+        """Transition out of training phase. Free GPU resources.
+
+        Accepts an optional ``offload_mode`` kwarg (OffloadMode enum):
+          - EVAL_ONLY: model on GPU in eval mode (for logprob inference).
+          - OPTIMIZER_ONLY: offload optimizer only (for weight staging).
+          - FULL (default): full offload based on deployment topology.
+        """
         pass
 
     @abstractmethod
@@ -163,49 +191,5 @@ def shutdown(self) -> bool:
         pass
 
 
-class ColocatablePolicyInterface(PolicyInterface):
-    @abstractmethod
-    def init_collective(
-        self, ip: str, port: int, world_size: int, *, train_world_size: int
-    ) -> list[ray.ObjectRef]:
-        pass
-
-    @abstractmethod
-    def offload_before_refit(self) -> None:
-        pass
-
-    @abstractmethod
-    def offload_after_refit(self) -> None:
-        pass
-
-    @abstractmethod
-    def prepare_refit_info(self) -> Optional[dict[str, Any]]:
-        pass
-
-    @abstractmethod
-    def stream_weights_via_ipc_zmq(
-        self, *args: Any, **kwargs: Any
-    ) -> list[ray.ObjectRef]:
-        pass
-
-    def stream_weights_via_http(
-        self, sglang_url_to_gpu_uuids: dict[str, list[str]]
-    ) -> list[ray.ObjectRef]:
-        """Stream model weights to SGLang servers via HTTP API.
-
-        Args:
-            sglang_url_to_gpu_uuids: Dict mapping SGLang server URL to list of GPU UUIDs it uses
-        """
-        raise NotImplementedError(
-            "stream_weights_via_http is not implemented for this policy worker"
-        )
-
-    @abstractmethod
-    def broadcast_weights_for_collective(
-        self, kv_scales: Optional[dict[str, float]] = None
-    ) -> list[ray.ObjectRef]:
-        pass
-
-    @abstractmethod
-    def prepare_for_lp_inference(self) -> None:
-        pass
+# Backward compatibility
+PolicyInterface = PolicyTrainerInterface