Optimize rl

unsloth/models/llama.py

@@ -2735,18 +2735,15 @@ def _for_training(m):
             m = m.model
         _for_training(m)
 
-        # Also re-enable training for embeddings for NEFTune
-        if hasattr(model, "get_input_embeddings"):
-            embeddings = model.get_input_embeddings()
-            if hasattr(embeddings, "training"): embeddings.training = True
-        pass
-        if hasattr(model, "get_output_embeddings"):
-            embeddings = model.get_output_embeddings()
-            if hasattr(embeddings, "training"): embeddings.training = True
-        pass
+        for get_embeddings_fn in (
+        model.get_input_embeddings, 
+        model.get_output_embeddings):
+
+            if hasattr(model, get_embeddings_fn.__name__):
+                embeddings = get_embeddings_fn()
+                if hasattr(embeddings, "training"):
+                    embeddings.training = True
         return model
-    pass
-pass
 
 from .rl import PatchFastRL
 PatchFastRL(FastLanguageModel = FastLlamaModel)

unsloth/models/rl.py

@@ -576,10 +576,14 @@ def patch_functions(RLTrainer, trainer_file, RLTrainer_name, all_imports, import
     old_init = init
 
     # Remove peft_config
-    init = init.replace("elif peft_config is None:", "elif False:")
-    init = init.replace("elif peft_config is not None:", "elif False:")
-    init = init.replace("if peft_config is None:", "if False:")
-    init = init.replace("if peft_config is not None:", "if False:")
+    replacements = {
+        "elif peft_config is None:": "elif False:",
+        "elif peft_config is not None:": "elif False:",
+        "if peft_config is None:": "if False:",
+        "if peft_config is not None:": "if False:",
+    }
+    for old, new in replacements.items():
+        init = init.replace(old, new)
     init = init.replace("get_peft_model(model, peft_config)", "model")
 
     # Set use_vllm if not set

unsloth/models/vision.py

@@ -50,7 +50,6 @@
 import os
 import gc
 import math
-import functools
 from typing import Optional, Tuple, List, Union
 import re, inspect, sys
 import types