diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 8f0fc30568..ba6392eab8 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -66,6 +66,7 @@ from unsloth_zoo.patching_utils import ( patch_layernorm, patch_torch_compile, patch_regional_compilation, + patch_model_and_tokenizer, ) from unsloth_zoo.gradient_checkpointing import ( Unsloth_Offloaded_Gradient_Checkpointer, diff --git a/unsloth/models/gemma.py b/unsloth/models/gemma.py index 1ec116b2ea..1d9a0c1334 100644 --- a/unsloth/models/gemma.py +++ b/unsloth/models/gemma.py @@ -340,56 +340,8 @@ class FastGemmaModel(FastLlamaModel): @staticmethod def post_patch(model, tokenizer): - # Torch.compile fails on embedding matrix?? - # Workaround randomnly fixes it for torch versions < 2.2 - model.model.embed_tokens = torch.nn.Embedding.from_pretrained(model.model.embed_tokens.weight) - model.config.update({"unsloth_version" : __version__}) - - # We also do this for the lm_head - lm_head = torch.nn.Linear(1, 1, bias = None) - del lm_head.weight - lm_head.weight = model.lm_head.weight - lm_head.in_features = lm_head.weight.shape[1] - lm_head.out_features = lm_head.weight.shape[0] - model.lm_head = lm_head - - # Gemma has tied weights! This means lm_head == embed_tokens - if model.model.embed_tokens.weight.data_ptr() != model.lm_head.weight.data_ptr(): - lm_head = torch.nn.Linear(1, 1, bias = None) - del lm_head.weight - lm_head.weight = model.model.embed_tokens.weight - lm_head.in_features = lm_head.weight.shape[1] - lm_head.out_features = lm_head.weight.shape[0] - model.lm_head = lm_head - pass - - # Also patch all dtypes - BnB seems to not allocate the correct type? - # BnB default dtype seems to be float16! - correct_dtype = lm_head.weight.dtype - - for name, module in model.named_modules(): - if isinstance(module, (Bnb_Linear4bit, Peft_Linear4bit)): - weight = module.weight - quant_state = weight.quant_state - - if type(quant_state) is list: - # BnB seems to have float16 as default! - module.weight.quant_state[2] = correct_dtype # Cast to correct dtype - else: - # https://github.com/TimDettmers/bitsandbytes/pull/763/files - quant_state.dtype = correct_dtype - pass - pass - # Downcast RoPE embedding to correct data type - # RoPE must be done in float32 for Gemma - # if (name.endswith("rotary_emb") or hasattr(module, "cos_cached")) \ - # and (module.cos_cached.dtype != correct_dtype): - - # module.cos_cached = module.cos_cached.to(correct_dtype) - # module.sin_cached = module.sin_cached.to(correct_dtype) - # pass - # pass - pass + # Gemma does not downcast RoPE + model, tokenizer = patch_model_and_tokenizer(model, tokenizer, downcast_rope = False) # Add 1 to weight # return output * (1 + self.weight) diff --git a/unsloth/models/gemma2.py b/unsloth/models/gemma2.py index 54d8f628cb..4eb9d64313 100644 --- a/unsloth/models/gemma2.py +++ b/unsloth/models/gemma2.py @@ -491,56 +491,8 @@ class FastGemma2Model(FastLlamaModel): @staticmethod def post_patch(model, tokenizer): - # Torch.compile fails on embedding matrix?? - # Workaround randomnly fixes it for torch versions < 2.2 - model.model.embed_tokens = torch.nn.Embedding.from_pretrained(model.model.embed_tokens.weight) - model.config.update({"unsloth_version" : __version__}) - - # We also do this for the lm_head - lm_head = torch.nn.Linear(1, 1, bias = None) - del lm_head.weight - lm_head.weight = model.lm_head.weight - lm_head.in_features = lm_head.weight.shape[1] - lm_head.out_features = lm_head.weight.shape[0] - model.lm_head = lm_head - - # Gemma has tied weights! This means lm_head == embed_tokens - if model.model.embed_tokens.weight.data_ptr() != model.lm_head.weight.data_ptr(): - lm_head = torch.nn.Linear(1, 1, bias = None) - del lm_head.weight - lm_head.weight = model.model.embed_tokens.weight - lm_head.in_features = lm_head.weight.shape[1] - lm_head.out_features = lm_head.weight.shape[0] - model.lm_head = lm_head - pass - - # Also patch all dtypes - BnB seems to not allocate the correct type? - # BnB default dtype seems to be float16! - correct_dtype = lm_head.weight.dtype - - for name, module in model.named_modules(): - if isinstance(module, (Bnb_Linear4bit, Peft_Linear4bit)): - weight = module.weight - quant_state = weight.quant_state - - if type(quant_state) is list: - # BnB seems to have float16 as default! - module.weight.quant_state[2] = correct_dtype # Cast to correct dtype - else: - # https://github.com/TimDettmers/bitsandbytes/pull/763/files - quant_state.dtype = correct_dtype - pass - pass - # Downcast RoPE embedding to correct data type - # RoPE must be done in float32 for Gemma - # if (name.endswith("rotary_emb") or hasattr(module, "cos_cached")) \ - # and (module.cos_cached.dtype != correct_dtype): - - # module.cos_cached = module.cos_cached.to(correct_dtype) - # module.sin_cached = module.sin_cached.to(correct_dtype) - # pass - # pass - pass + # Gemma does not downcast RoPE + model, tokenizer = patch_model_and_tokenizer(model, tokenizer, downcast_rope = False) # Add 1 to weight # return output * (1 + self.weight) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 8c4b7fd253..4e83e69d60 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -57,8 +57,6 @@ from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING from transformers import set_seed as transformers_set_seed from peft import LoraConfig, TaskType, get_peft_model as _get_peft_model from peft import PeftModelForCausalLM -from bitsandbytes.nn import Linear4bit as Bnb_Linear4bit -from peft.tuners.lora import Linear4bit as Peft_Linear4bit from ..save import patch_saving_functions import re, os, inspect, math, sys try: @@ -1798,30 +1796,6 @@ class FastLlamaModel: internal_model = internal_model.model pass internal_model._saved_temp_tokenizer = tokenizer - - # Also fix torch_dtype - internal_model = model - while hasattr(internal_model, "model"): - if hasattr(internal_model, "config"): - if internal_model.config.torch_dtype == "float32": - internal_model.config.torch_dtype = torch.float32 - elif internal_model.config.torch_dtype == "bfloat16": - internal_model.config.torch_dtype = torch.bfloat16 - elif internal_model.config.torch_dtype == "float16": - internal_model.config.torch_dtype = torch.float16 - pass - pass - internal_model = internal_model.model - pass - if hasattr(internal_model, "config"): - if internal_model.config.torch_dtype == "float32": - internal_model.config.torch_dtype = torch.float32 - elif internal_model.config.torch_dtype == "bfloat16": - internal_model.config.torch_dtype = torch.bfloat16 - elif internal_model.config.torch_dtype == "float16": - internal_model.config.torch_dtype = torch.float16 - pass - pass return model, tokenizer pass @@ -1829,108 +1803,7 @@ class FastLlamaModel: @staticmethod def post_patch(model, tokenizer): - # Torch.compile fails on embedding matrix?? - try: old_input_embedding = model.get_input_embeddings ().weight - except: return model, tokenizer - - # Maybe not all models have a lm_head? - try: old_output_embedding = model.get_output_embeddings().weight - except: old_output_embedding = torch.zeros(0) - - # Check for tied weights as well - is_tied = (old_input_embedding.data_ptr() == old_output_embedding.data_ptr()) \ - or (model.config.tie_word_embeddings) - - # Check pad token's id -> we need to expand the embedding - if len(tokenizer) > old_input_embedding.shape[0]: - # Workaround randomnly fixes it for torch versions < 2. - requires_grad = old_input_embedding.requires_grad - old_input_embedding.requires_grad_(False) - old_input_embedding.resize_(len(tokenizer), old_input_embedding.shape[1]) - old_input_embedding.requires_grad_(requires_grad) - - # Fix up all vocab sizes - current_model = model - while hasattr(current_model, "model") and hasattr(current_model, "config"): - if hasattr(current_model.config, "vocab_size"): - current_model.config.update({"vocab_size" : len(tokenizer)}) - current_model = current_model.model - if hasattr(current_model, "model") and hasattr(current_model, "config"): - if hasattr(current_model.config, "vocab_size"): - current_model.config.update({"vocab_size" : len(tokenizer)}) - pass - pass - - model.set_input_embeddings( - torch.nn.Embedding.from_pretrained( - old_input_embedding, - padding_idx = getattr(model.config, "pad_token_id", None), - ) - ) - model.config.update({"unsloth_version" : __version__}) - - # We also do this for the lm_head - if old_output_embedding.numel() != 0: - - requires_grad = old_output_embedding.requires_grad - lm_head = torch.nn.Linear(1, 1, bias = None) - del lm_head.weight - - lm_head.weight = old_output_embedding if not is_tied else old_input_embedding - lm_head.in_features = lm_head.weight.shape[1] - lm_head.out_features = lm_head.weight.shape[0] - - lm_head.weight.requires_grad_(requires_grad) - model.set_output_embeddings(lm_head) - if hasattr(model, "lm_head"): model.lm_head = lm_head - - correct_dtype = lm_head.weight.dtype - else: - correct_dtype = old_input_embedding.dtype - pass - - # Must tie lm_head and embed_tokens if they are tied! - # Otherwise error will occur on saving models ie use save_model - if is_tied: model.tie_weights() - - # Also patch all dtypes - BnB seems to not allocate the correct type? - # BnB default dtype seems to be float16! - for name, module in model.named_modules(): - if isinstance(module, (Bnb_Linear4bit, Peft_Linear4bit)): - weight = module.weight - quant_state = weight.quant_state - - if type(quant_state) is list: - # BnB seems to have float16 as default! - module.weight.quant_state[2] = correct_dtype # Cast to correct dtype - else: - # https://github.com/TimDettmers/bitsandbytes/pull/763/files - quant_state.dtype = correct_dtype - pass - pass - # Downcast RoPE embedding to correct data type - if (name.endswith("rotary_emb") or hasattr(module, "cos_cached")): - - if hasattr(module, "cos_cached") and \ - (module.cos_cached.dtype != correct_dtype): - - module.cos_cached = module.cos_cached.to(correct_dtype) - module.sin_cached = module.sin_cached.to(correct_dtype) - - elif hasattr(module, "short_cos_cached") and \ - (module.short_cos_cached.dtype != correct_dtype): - - module.short_cos_cached = module.short_cos_cached.to(correct_dtype) - module.short_sin_cached = module.short_sin_cached.to(correct_dtype) - pass - pass - pass - - # Clear deleted GPU items - for _ in range(3): - gc.collect() - torch.cuda.empty_cache() - return model, tokenizer + model, tokenizer = patch_model_and_tokenizer(model, tokenizer, downcast_rope = True) pass