Post patch
This commit is contained in:
Daniel Han
parent
0a18ffd66a
commit
e692ccf727
4 changed files with 6 additions and 228 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue