diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index 5b835129f7..7dc62010f0 100755 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -156,7 +156,6 @@ def _patch_resume_from_checkpoint_memory(trainer_class): trainer_class.train = _unsloth_train_with_resume_guard - def _maybe_prepare_vllm_for_resume(trainer): if not torch.cuda.is_available(): return @@ -185,10 +184,12 @@ def _maybe_prepare_vllm_for_resume(trainer): trainer._unsloth_resume_wake_vllm = True import gc + for _ in range(3): gc.collect() torch.cuda.empty_cache() -pass + + def _patch_resume_from_checkpoint_memory(trainer_class): @@ -206,11 +207,13 @@ def _patch_resume_from_checkpoint_memory(trainer_class): if resume_from_checkpoint: _maybe_prepare_vllm_for_resume(self) return original_train(self, *args, **kwargs) - pass + _unsloth_train_with_resume_guard._unsloth_resume_guard = True trainer_class.train = _unsloth_train_with_resume_guard -pass + + + def PatchRL(FastLanguageModel): try: diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index 476286e0ba..221a42a385 100755 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -1311,14 +1311,18 @@ RL_CONFIG_CHANGES["grpo_trainer"].append(grpo_trainer_fix_batch_size) def grpo_trainer_fix_generation_batch_size(RLTrainer_source, RLConfig_source): - if "generation_batch_size" not in RLConfig_source: return "" - if "steps_per_generation" not in RLConfig_source: return "" + if "generation_batch_size" not in RLConfig_source: + return "" + if "steps_per_generation" not in RLConfig_source: + return "" - check_generation_batch_size = \ - "if generation_batch_size is not None and steps_per_generation is not None:\n"\ - " generation_batch_size = None\n" + check_generation_batch_size = ( + "if generation_batch_size is not None and steps_per_generation is not None:\n" + " generation_batch_size = None\n" + ) return check_generation_batch_size -pass + + RL_CONFIG_CHANGES["grpo_trainer"].append(grpo_trainer_fix_generation_batch_size)