Simplify GRPO resume vLLM handling

unsloth/models/rl.py

@@ -93,6 +93,62 @@ def vLLMSamplingParams(**kwargs):
     return sampling_params
 
 
+def _maybe_prepare_vllm_for_resume(trainer):
+    if not torch.cuda.is_available():
+        return
+
+    llm = getattr(trainer, "llm", None)
+    if llm is None:
+        llm = getattr(getattr(trainer, "model", None), "vllm_engine", None)
+
+    sleep_fn = getattr(llm, "sleep", None)
+    if callable(sleep_fn):
+        try:
+            sleep_mode = int(os.environ.get("VLLM_SLEEP_MODE", "1"))
+        except ValueError:
+            sleep_mode = 1
+
+        try:
+            signature = inspect.signature(sleep_fn)
+        except (TypeError, ValueError):
+            signature = None
+
+        try:
+            if signature is not None and len(signature.parameters) == 0:
+                sleep_fn()
+            else:
+                sleep_fn(sleep_mode)
+        except Exception as error:
+            logger.warning_once(
+                f"Unsloth: vLLM sleep() failed during resume cleanup: {error}"
+            )
+
+    import gc
+
+    for _ in range(3):
+        gc.collect()
+        torch.cuda.empty_cache()
+
+
+def _patch_resume_from_checkpoint_memory(trainer_class):
+    original_train = getattr(trainer_class, "train", None)
+    if original_train is None:
+        return
+    if getattr(original_train, "_unsloth_resume_guard", False):
+        return
+
+    def _unsloth_train_with_resume_guard(self, *args, **kwargs):
+        resume_from_checkpoint = kwargs.get("resume_from_checkpoint", None)
+        if resume_from_checkpoint is None and len(args) != 0:
+            resume_from_checkpoint = args[0]
+
+        if resume_from_checkpoint:
+            _maybe_prepare_vllm_for_resume(self)
+        return original_train(self, *args, **kwargs)
+
+    _unsloth_train_with_resume_guard._unsloth_resume_guard = True
+    trainer_class.train = _unsloth_train_with_resume_guard
+
 def PatchRL(FastLanguageModel):
     try:
         from trl.models.utils import unwrap_model_for_generation