From 9e932030ffb0d147fdce788389c5d76e1f93f70a Mon Sep 17 00:00:00 2001 From: DoubleMathew Date: Thu, 13 Nov 2025 23:26:49 -0600 Subject: [PATCH] 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> --- unsloth/models/llama.py | 12 +++++++++--- unsloth/models/loader.py | 14 ++++++++++++++ 2 files changed, 23 insertions(+), 3 deletions(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 306f95b41b..9537dd63c3 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -3214,9 +3214,15 @@ class FastLlamaModel: ) ): # 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( diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index 427e8ee75e..15f01cf15d 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -57,6 +57,7 @@ from ..device_type import ( # 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 @@ class FastLanguageModel(FastLlamaModel): ) # 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 @@ class FastModel(FastBaseModel): 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