[test] attempt to fix CI test for PT 2.0

tests/test_ppo_trainer.py

@@ -246,7 +246,6 @@ def test_ppo_step_with_no_ref_sgd(self):
         for stat in EXPECTED_STATS:
             assert stat in train_stats.keys()
 
-    @unittest.skip("TODO: fix this test")
     def test_ppo_step_with_no_ref_sgd_lr_scheduler(self):
         # initialize dataset
         dummy_dataset = self._init_dummy_dataset()

trl/import_utils.py

@@ -12,7 +12,26 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import importlib
+import sys
+
+
+if sys.version_info[0] < 3.8:
+    _is_python_greater_3_8 = False
+else:
+    _is_python_greater_3_8 = True
 
 
 def is_peft_available():
     return importlib.util.find_spec("peft") is not None
+
+
+def is_torch_greater_2_0():
+    if _is_python_greater_3_8:
+        from importlib.metadata import version
+
+        torch_version = version("torch")
+    else:
+        import pkg_resources
+
+        torch_version = pkg_resources.get_distribution("torch").version
+    return torch_version >= "2.0"

trl/trainer/ppo_trainer.py

@@ -41,6 +41,7 @@
     stack_dicts,
     stats_to_np,
 )
+from ..import_utils import is_torch_greater_2_0
 from ..models import SUPPORTED_ARCHITECTURES, PreTrainedModelWrapper, create_reference_model
 from . import AdaptiveKLController, BaseTrainer, FixedKLController, PPOConfig
 
@@ -248,8 +249,16 @@ def __init__(
 
         self.lr_scheduler = lr_scheduler
         if self.lr_scheduler is not None:
-            if not isinstance(self.lr_scheduler, torch.optim.lr_scheduler._LRScheduler):
-                raise ValueError("lr_scheduler must be a torch.optim.lr_scheduler._LRScheduler")
+            lr_scheduler_class = (
+                torch.optim.lr_scheduler._LRScheduler
+                if not is_torch_greater_2_0()
+                else torch.optim.lr_scheduler.LRScheduler
+            )
+
+            if not isinstance(self.lr_scheduler, lr_scheduler_class):
+                raise ValueError(
+                    "lr_scheduler must be a torch.optim.lr_scheduler._LRScheduler or torch.optim.lr_scheduler.LRScheduler (for torch >= 2.0)"
+                )
 
         if self.config.adap_kl_ctrl:
             self.kl_ctl = AdaptiveKLController(self.config.init_kl_coef, self.config.target, self.config.horizon)