ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Patch in tiled mlp (#3584)

Browse source 
* 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>
This commit is contained in:
DoubleMathew 2025-11-13 23:26:49 -06:00 • committed by GitHub
commit 9e932030ff
2 changed files with 23 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

12
unsloth/models/llama.py
View file

@ -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(

14
unsloth/models/loader.py
View file

@ -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