diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 3e3b8ff22c..95a032b698 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -95,7 +95,7 @@ def prepare_model_for_kbit_training( # Freeze all parameters except LoRA for name, param in model.named_parameters(): - if ".lora_A." in name or ".lora_B." in name: + if ".lora_A." in name or ".lora_B." in name or ".lora_magnitude_vector" in name: param.requires_grad_(True) else: param.requires_grad_(False) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 9e78c656b0..861a12a7c9 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -511,6 +511,7 @@ def LlamaModel_fast_forward( # Mormalized from Gemma IS_GEMMA = self.config.model_type == "gemma" + train_embed_tokens = self.embed_tokens.weight.requires_grad if IS_GEMMA: # Match Gemma exactly by casting to bfloat16 / float16 @@ -519,7 +520,7 @@ def LlamaModel_fast_forward( # & 2048**0.5 = 45.2500 in bfloat16, whilst 45.2548 in float32 normalizer = torch.tensor(math_sqrt(self.config.hidden_size), dtype = inputs_embeds.dtype) - if self.embed_tokens.weight.requires_grad: + if train_embed_tokens: # Careful we must not do an inplace op! inputs_embeds = inputs_embeds * normalizer else: @@ -538,7 +539,8 @@ def LlamaModel_fast_forward( # Fix up attention mask by setting elements to 0 # Specifically for DPO - if self._has_no_labels and (attention_mask is not None) and (past_key_values is None): + if self._has_no_labels and (attention_mask is not None) and (past_key_values is None) and \ + (not train_embed_tokens): # Careful for inference the attention_mask is size (1, kv_seq_len) # Whilst the input_embeds is size (1, 1, 4096) inputs_requires_grad = inputs_embeds.requires_grad @@ -1234,6 +1236,7 @@ class FastLlamaModel: random_state = 3407, max_seq_length = 2048, # not used anymore use_rslora = False, + modules_to_save = None, init_lora_weights = True, loftq_config = {}, **kwargs, @@ -1320,15 +1323,45 @@ class FastLlamaModel: accepted_modules = frozenset(("q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj",),) model.config.update({"unsloth_version" : __version__}) + + train_lm_head = False + train_embed_tokens = False + final_modules = [] for module in target_modules: - assert(module in accepted_modules) + if module == "lm_head": + logger.warning_once( + "Unsloth: `lm_head` should be placed in `modules_to_save` and not `target_modules`."\ + "We shall do it for you!" + ) + train_lm_head = True + + elif module == "embed_tokens": + logger.warning_once( + "Unsloth: `embed_tokens` should be placed in `modules_to_save` and not `target_modules`."\ + "We shall do it for you!" + ) + train_embed_tokens = True + + else: + assert(module in accepted_modules) + final_modules.append(module) + pass + + # Check modules_to_save + if modules_to_save is not None: + for module in modules_to_save: + if module == "lm_head": + train_lm_head = True + elif module == "embed_tokens": + train_embed_tokens = True + pass pass # Get LoRA arguments = dict( r = r, lora_alpha = lora_alpha, - target_modules = target_modules, + target_modules = final_modules, lora_dropout = lora_dropout, bias = bias, task_type = TaskType.CAUSAL_LM, @@ -1336,6 +1369,7 @@ class FastLlamaModel: init_lora_weights = init_lora_weights, loftq_config = loftq_config, use_rslora = use_rslora, + modules_to_save = modules_to_save, **kwargs, ) if not SUPPORTS_LOFTQ: del arguments["loftq_config"] @@ -1345,6 +1379,14 @@ class FastLlamaModel: model = _get_peft_model(model, lora_config) model = FastLlamaModel.patch_peft_model(model, use_gradient_checkpointing) + + # Now patch lm_head and embed_tokens + if train_embed_tokens: + model.model.model.embed_tokens.requires_grad_(True) + if train_lm_head: + model.model.lm_head.requires_grad_(True) + pass + return model pass @@ -1435,9 +1477,12 @@ class FastLlamaModel: if hasattr(gate_proj, "lora_A") and \ hasattr( up_proj, "lora_A") and \ hasattr(down_proj, "lora_A") and \ - (gate_proj.base_layer if hasattr(gate_proj, "base_layer") else gate_proj).bias is None and \ - ( up_proj.base_layer if hasattr( up_proj, "base_layer") else up_proj).bias is None and \ - (down_proj.base_layer if hasattr(down_proj, "base_layer") else down_proj).bias is None: + ((gate_proj.base_layer if hasattr(gate_proj, "base_layer") else gate_proj).bias is None) and \ + (( up_proj.base_layer if hasattr( up_proj, "base_layer") else up_proj).bias is None) and \ + ((down_proj.base_layer if hasattr(down_proj, "base_layer") else down_proj).bias is None) and \ + (not hasattr(gate_proj, "lora_magnitude_vector")) and \ + (not hasattr( up_proj, "lora_magnitude_vector")) and \ + (not hasattr(down_proj, "lora_magnitude_vector")): # 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) @@ -1456,9 +1501,12 @@ class FastLlamaModel: if hasattr(q_proj, "lora_A") and \ hasattr(k_proj, "lora_A") and \ hasattr(v_proj, "lora_A") and \ - (q_proj.base_layer if hasattr(q_proj, "base_layer") else q_proj).bias is None and \ - (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: + ((q_proj.base_layer if hasattr(q_proj, "base_layer") else q_proj).bias is None) and \ + ((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) and \ + (not hasattr(q_proj, "lora_magnitude_vector")) and \ + (not hasattr(k_proj, "lora_magnitude_vector")) and \ + (not hasattr(v_proj, "lora_magnitude_vector")): layer.self_attn.apply_qkv = apply_lora_qkv n_qkv += 1 @@ -1472,7 +1520,8 @@ class FastLlamaModel: # O attention patching o_proj = layer.self_attn.o_proj if hasattr(o_proj, "lora_A") and \ - (o_proj.base_layer if hasattr(o_proj, "base_layer") else o_proj).bias is None: + ((o_proj.base_layer if hasattr(o_proj, "base_layer") else o_proj).bias is None) and \ + (not hasattr(o_proj, "lora_magnitude_vector")): layer.self_attn.apply_o = apply_lora_o n_o += 1