From 69f7dfdb901a72f99d9d019fab2ddba3135a3efa Mon Sep 17 00:00:00 2001 From: DoubleMathew Date: Sun, 30 Nov 2025 00:47:51 -0600 Subject: [PATCH] make unsloth_tiled_mlp a from_pretrained arg (#3655) * make unsloth_tiled_mlp a from_pretrained arg * adjust patching logic --- unsloth/models/loader.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index a8ddbfed2a..e1c13315f7 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -146,6 +146,7 @@ class FastLanguageModel(FastLlamaModel): disable_log_stats = True, qat_scheme = None, load_in_fp8 = False, # fp8 LoRA (True, False, 'block') + unsloth_tiled_mlp = False, *args, **kwargs, ): @@ -512,6 +513,7 @@ class FastLanguageModel(FastLlamaModel): disable_log_stats = disable_log_stats, qat_scheme = qat_scheme, load_in_fp8 = load_in_fp8, + unsloth_tiled_mlp = unsloth_tiled_mlp, *args, **kwargs, ) @@ -609,8 +611,10 @@ class FastLanguageModel(FastLlamaModel): # 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_choice = os.environ.get( + "UNSLOTH_TILED_MLP", "arctic" if unsloth_tiled_mlp else "0" + ) + if patch_tiled_mlp_choice != "0" or unsloth_tiled_mlp: patch_tiled_mlp(model, patch_options_str = patch_tiled_mlp_choice) return model, tokenizer @@ -674,6 +678,7 @@ class FastModel(FastBaseModel): disable_log_stats = True, qat_scheme = None, load_in_fp8 = False, # fp8 LoRA (True, False, 'block') + unsloth_tiled_mlp = False, *args, **kwargs, ): @@ -1240,8 +1245,10 @@ class FastModel(FastBaseModel): # 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_choice = os.environ.get( + "UNSLOTH_TILED_MLP", "arctic" if unsloth_tiled_mlp else "0" + ) + if patch_tiled_mlp_choice != "0" or unsloth_tiled_mlp: patch_tiled_mlp(model, patch_options_str = patch_tiled_mlp_choice) return model, tokenizer