fix: skip_reference_policy_logprobs_calculation=true crashes training

nemo_rl/algorithms/grpo.py

@@ -1769,6 +1769,10 @@ def grpo_train(
                                 timer=timer,
                             )["reference_logprobs"]
                         )
+                    else:
+                        train_data["reference_policy_logprobs"] = torch.zeros_like(
+                            train_data["prev_logprobs"]
+                        )
 
                     del logprob_data
                     del extra_multimodal_data
@@ -2799,12 +2803,18 @@ def async_grpo_train(
                         train_data,
                         timer=timer,
                     )["logprobs"]
-                    reference_logprobs = policy.get_reference_policy_logprobs(
-                        train_data,
-                        timer=timer,
-                    )["reference_logprobs"]
                     train_data["prev_logprobs"] = fprop_logprobs
-                    train_data["reference_policy_logprobs"] = reference_logprobs
+
+                    if not master_config["grpo"].get(
+                        "skip_reference_policy_logprobs_calculation"
+                    ):
+                        reference_logprobs = policy.get_reference_policy_logprobs(
+                            train_data,
+                            timer=timer,
+                        )["reference_logprobs"]
+                        train_data["reference_policy_logprobs"] = reference_logprobs
+                    else:
+                        train_data["reference_policy_logprobs"] = torch.zeros_like(fprop_logprobs)
 
                     (
                         max_seq_mult_prob_error,

nemo_rl/models/policy/workers/base_policy_worker.py

@@ -139,6 +139,15 @@ def get_reference_policy_logprobs(
           We use the convention that the logprob of the first token is 0 so that the sequence length is maintained.
           The logprob of input token i is specified at position i in the output logprobs tensor.
         """
+        # When reference model was never initialized (e.g.,
+        # skip_reference_policy_logprobs_calculation=true), return zeros
+        # matching the expected shape to avoid crashes downstream.
+        if not self._has_reference_model():
+            logprobs = self.get_logprobs(data=data, micro_batch_size=micro_batch_size)
+            return_data = BatchedDataDict[ReferenceLogprobOutputSpec]()
+            return_data["reference_logprobs"] = torch.zeros_like(logprobs["logprobs"]).cpu()
+            return return_data
+
         with self.use_reference_model():
             reference_logprobs = self.get_logprobs(
                 data=data, micro_batch_size=micro_batch_size
@@ -148,6 +157,15 @@ def get_reference_policy_logprobs(
         return_data["reference_logprobs"] = reference_logprobs["logprobs"].cpu()
         return return_data
 
+    def _has_reference_model(self) -> bool:
+        """Check if a reference model has been initialized."""
+        # DTensor v2 uses reference_model_state_dict, others use reference_state_dict
+        for attr in ("reference_model_state_dict", "reference_state_dict"):
+            val = getattr(self, attr, None)
+            if val is not None:
+                return True
+        return False
+
     def finish_training(self, *args: Any, **kwargs: Any) -> None:
         # Placeholder implementation
         pass

nemo_rl/models/policy/workers/dtensor_policy_worker.py

@@ -1667,6 +1667,11 @@ def use_reference_model(self) -> Generator[None, None, None]:
                   is different from the current policy, making filtered logprobs incompatible.
         On exit: Restores original references and re-flips cuda/cpu, restores sampling_params.
         """
+        # If reference model was never initialized, yield without swapping
+        if not hasattr(self, "reference_model_state_dict") or self.reference_model_state_dict is None:
+            yield
+            return
+
         with torch.no_grad():
             # Save train model state_dict
             curr_state_dict = get_cpu_state_dict(

nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py

@@ -785,6 +785,11 @@ def use_reference_model(self) -> Generator[None, None, None]:
                   is different from the current policy, making filtered logprobs incompatible.
         On exit: Restores original references and re-flips cuda/cpu, restores sampling_params.
         """
+        # If reference model was never initialized, yield without swapping
+        if not hasattr(self, "reference_model_state_dict") or self.reference_model_state_dict is None:
+            yield
+            return
+
         with torch.no_grad():
             # Save train model state_dict
             curr_state_dict = get_cpu_state_dict(

nemo_rl/models/policy/workers/megatron_policy_worker.py

@@ -534,6 +534,11 @@ def use_reference_model(self):
                   is different from the current policy, making filtered logprobs incompatible.
         On exit: Restores original references and re-flips cuda/cpu, restores sampling_params.
         """
+        # If reference model was never initialized, yield without swapping
+        if not hasattr(self, "reference_state_dict") or self.reference_state_dict is None:
+            yield
+            return
+
         ## disable overlap param gather when swapping weights
         if self.should_disable_forward_pre_hook:
             self.disable_forward_pre_hook()