Patch in tiled mlp (#3584)

* Patch in tiled mlp

* Update unsloth/models/llama.py

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: 3 people

unsloth/models/llama.py

@@ -3214,9 +3214,15 @@ def patch_peft_model(
                         )
                     ):
                         # https://stackoverflow.com/questions/50599045/python-replacing-a-function-within-a-class-of-a-module
-                        mlp_module.forward = types.MethodType(
-                            _apply_lora_mlp, mlp_module
-                        )
+                        if hasattr(mlp_module, "_unsloth_forward"):
+                            # then we've patched the mlp to use TiledMLP
+                            mlp_module._unsloth_forward = types.MethodType(
+                                _apply_lora_mlp, mlp_module
+                            )
+                        else:
+                            mlp_module.forward = types.MethodType(
+                                _apply_lora_mlp, mlp_module
+                            )
                         n_mlp += 1
                     else:
                         logger.warning_once(

unsloth/models/loader.py

@@ -57,6 +57,7 @@
 # https://github.com/huggingface/transformers/pull/26037 allows 4 bit loading!
 from unsloth_zoo.utils import Version, _get_dtype
 from unsloth_zoo.hf_utils import dtype_from_config
+from unsloth_zoo.tiled_mlp import patch_tiled_mlp
 
 transformers_version = Version(transformers_version)
 SUPPORTS_FOURBIT = transformers_version >= Version("4.37")
@@ -566,6 +567,13 @@ def from_pretrained(
             )
             # Patch it as well!
             model = dispatch_model.patch_peft_model(model, use_gradient_checkpointing)
+
+        # Patch Tiled MLP
+        # to turn on set UNSLOTH_TILED_MLP to "arctic", "target", or "target:{GB}""
+        patch_tiled_mlp_choice = os.environ.get("UNSLOTH_TILED_MLP", "0")
+        if patch_tiled_mlp_choice != "0":
+            patch_tiled_mlp(model, patch_options_str = patch_tiled_mlp_choice)
+
         return model, tokenizer
 
 
@@ -1138,6 +1146,12 @@ def from_pretrained(
             print("Unsloth: Applying QAT to mitigate quantization degradation")
             model = _prepare_model_for_qat(model, qat_scheme)
 
+        # Patch Tiled MLP
+        # to turn on set UNSLOTH_TILED_MLP to "arctic", "target", or "target:{GB}""
+        patch_tiled_mlp_choice = os.environ.get("UNSLOTH_TILED_MLP", "0")
+        if patch_tiled_mlp_choice != "0":
+            patch_tiled_mlp(model, patch_options_str = patch_tiled_mlp_choice)
+
         return model, tokenizer