From 11233cb3d40077a4f92a0f57b0dc4ac9a9dcdcd3 Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Fri, 23 Feb 2024 02:10:11 +1100 Subject: [PATCH] model_type --- unsloth/models/gemma.py | 5 ++++- unsloth/models/llama.py | 15 ++++++++------- 2 files changed, 12 insertions(+), 8 deletions(-) diff --git a/unsloth/models/gemma.py b/unsloth/models/gemma.py index 439a7aecbb..3b6564a7f3 100644 --- a/unsloth/models/gemma.py +++ b/unsloth/models/gemma.py @@ -319,17 +319,20 @@ class FastGemmaModel(FastLlamaModel): from transformers.models.gemma.modeling_gemma import GemmaRMSNorm # Freeze all parameters except LoRA + # We do this first since += 1 seems to not be liked by requires_grad = True for name, param in model.named_parameters(): if ".lora_A." in name or ".lora_B." in name: param.requires_grad_(True) else: param.requires_grad_(False) pass + + print("Unsloth: Patching Gemma RMS Layernorm + 1") for name, module in model.named_modules(): if isinstance(module, GemmaRMSNorm): module.weight += 1.0 # return output * (1 + self.weight) pass - + # Clear deleted GPU items import gc for _ in range(3): diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 074cb0fad1..32d159510c 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -1317,14 +1317,15 @@ class FastLlamaModel: pass # Get activation function - if model.config.mod == "swiglu": - apply_lora_mlp = apply_lora_mlp_swiglu - elif activation_function == "geglu": - apply_lora_mlp = apply_lora_mlp_geglu - else: - raise NotImplementedError(f"Unsloth: {activation_function} is not yet implemented!") - pass + model_type = model_config.model_type + if model_type == "llama": apply_lora_mlp = apply_lora_mlp_swiglu + elif model_type == "mistral": apply_lora_mlp = apply_lora_mlp_swiglu + elif model_type == "gemma": apply_lora_mlp = apply_lora_mlp_geglu + else: + raise NotImplementedError(f"Unsloth: {model_type} is not yet implemented!") + pass + model = prepare_model_for_kbit_training( model, use_gradient_checkpointing = use_gradient_checkpointing,