model_type

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):

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,