From be098122f14aba8304f1a5b1963ebd16fe42deaa Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Sat, 24 Feb 2024 15:42:12 +1100 Subject: [PATCH] Update llama.py --- unsloth/models/llama.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 02681bd4cc..2c431f226a 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -510,10 +510,10 @@ def LlamaModel_fast_forward( # Mormalized from Gemma if self.config.model_type == "gemma": - # inputs_requires_grad = inputs_embeds.requires_grad - # if inputs_requires_grad: inputs_embeds.requires_grad_(False) - inputs_embeds = inputs_embeds * math_sqrt(self.config.hidden_size) - # if inputs_requires_grad: inputs_embeds.requires_grad_(True) + inputs_requires_grad = inputs_embeds.requires_grad + if inputs_requires_grad: inputs_embeds.requires_grad_(False) + inputs_embeds *= math_sqrt(self.config.hidden_size) + if inputs_requires_grad: inputs_embeds.requires_grad_(True) pass # Fix up attention mask by setting elements to 0 @@ -609,7 +609,7 @@ def LlamaModel_fast_forward( all_self_attns += (layer_outputs[1],) pass - hidden_states = self.norm(hidden_states) + hidden_states = fast_rms_layernorm(self.norm, hidden_states) # add hidden states from the last decoder layer if output_hidden_states: @@ -1400,7 +1400,7 @@ class FastLlamaModel: (down_proj.base_layer if hasattr(down_proj, "base_layer") else down_proj).bias is None: # https://stackoverflow.com/questions/50599045/python-replacing-a-function-within-a-class-of-a-module - # layer.mlp.forward = types.MethodType(apply_lora_mlp, layer.mlp) + layer.mlp.forward = types.MethodType(apply_lora_mlp, layer.mlp) n_mlp += 1 else: logger.warning_once( @@ -1420,7 +1420,7 @@ class FastLlamaModel: (k_proj.base_layer if hasattr(k_proj, "base_layer") else k_proj).bias is None and \ (v_proj.base_layer if hasattr(v_proj, "base_layer") else v_proj).bias is None: - # layer.self_attn.apply_qkv = apply_lora_qkv + layer.self_attn.apply_qkv = apply_lora_qkv n_qkv += 1 else: logger.warning_once( @@ -1434,7 +1434,7 @@ class FastLlamaModel: if hasattr(o_proj, "lora_A") and \ (o_proj.base_layer if hasattr(o_proj, "base_layer") else o_proj).bias is None: - # layer.self_attn.apply_o = apply_lora_o + layer.self_attn.apply_o = apply_lora_o n_o += 1 else: logger.warning_once(