Fix Llama-3 (#366)
* Fix prompt * Update chat_templates.py * fix_untrained_tokens * Update llama.py * add tokens * Update _utils.py * Update tokenizer_utils.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * pad_token * Update chat_templates.py * Update chat_templates.py * tokenizer * Update save.py * Update chat_templates.py * Update chat_templates.py
This commit is contained in:
Daniel Han
GitHub
parent
dab0082a3d
commit
ac812404ef
5 changed files with 203 additions and 10 deletions
|
|
@ -23,10 +23,7 @@ from transformers.models.llama.modeling_llama import logger
|
|||
from .save import patch_saving_functions
|
||||
import os
|
||||
import shutil
|
||||
from .tokenizer_utils import (
|
||||
load_correct_tokenizer,
|
||||
fix_sentencepiece_tokenizer,
|
||||
)
|
||||
from .tokenizer_utils import *
|
||||
from .models._utils import patch_tokenizer
|
||||
|
||||
CHAT_TEMPLATES = {}
|
||||
|
|
@ -266,7 +263,7 @@ llama3_template = \
|
|||
"{{ '<|start_header_id|>assistant<|end_header_id|>\n\n' }}"\
|
||||
"{% endif %}"
|
||||
llama3_template_eos_token = "eos_token"
|
||||
CHAT_TEMPLATES["llama-3"] = (llama3_template, gemma_chatml_eos_token,)
|
||||
CHAT_TEMPLATES["llama-3"] = (llama3_template, llama3_template_eos_token,)
|
||||
|
||||
|
||||
def get_chat_template(
|
||||
|
|
@ -288,6 +285,8 @@ def get_chat_template(
|
|||
is_fast_tokenizer = getattr(tokenizer, "is_fast", False)
|
||||
old_padding_side = tokenizer.padding_side
|
||||
|
||||
same_padding_token = False
|
||||
|
||||
if type(chat_template) in (list, tuple,):
|
||||
chat_template, stop_word = chat_template
|
||||
assert(type(chat_template) is str)
|
||||
|
|
@ -342,10 +341,24 @@ def get_chat_template(
|
|||
if skipped != len(token_mapping):
|
||||
new_tokenizer = tokenizer._tokenizer.from_str(string_vocab)
|
||||
|
||||
# Careful on pad_token
|
||||
old_pad_token = tokenizer.pad_token
|
||||
if old_pad_token == tokenizer.eos_token:
|
||||
old_pad_token = stop_word
|
||||
same_padding_token = True
|
||||
pass
|
||||
|
||||
if map_eos_token:
|
||||
new_tokenizer = tokenizer.__class__(tokenizer_object = new_tokenizer, eos_token = stop_word)
|
||||
new_tokenizer = tokenizer.__class__(
|
||||
tokenizer_object = new_tokenizer,
|
||||
eos_token = stop_word,
|
||||
pad_token = old_pad_token,
|
||||
)
|
||||
else:
|
||||
new_tokenizer = tokenizer.__class__(tokenizer_object = new_tokenizer)
|
||||
new_tokenizer = tokenizer.__class__(
|
||||
tokenizer_object = new_tokenizer,
|
||||
pad_token = old_pad_token,
|
||||
)
|
||||
pass
|
||||
|
||||
# Must fix the sentence piece tokenizer since there's no tokenizer.model file!
|
||||
|
|
@ -380,6 +393,13 @@ def get_chat_template(
|
|||
string_vocab = string_vocab.replace(old_eos_token, stop_word)
|
||||
pass
|
||||
new_tokenizer = tokenizer._tokenizer.from_str(string_vocab)
|
||||
|
||||
# Careful on pad_token
|
||||
if old_pad_token == old_eos_token:
|
||||
old_pad_token = stop_word
|
||||
same_padding_token = True
|
||||
pass
|
||||
|
||||
new_tokenizer = tokenizer.__class__(
|
||||
tokenizer_object = new_tokenizer,
|
||||
bos_token = old_bos_token,
|
||||
|
|
@ -424,9 +444,11 @@ def get_chat_template(
|
|||
new_pad_token = getattr(tokenizer, "pad_token", None)
|
||||
new_bos_token = getattr(tokenizer, "bos_token", None)
|
||||
new_unk_token = getattr(tokenizer, "unk_token", None)
|
||||
if old_pad_token != new_pad_token: tokenizer.pad_token = old_pad_token
|
||||
if old_bos_token != new_bos_token: tokenizer.bos_token = old_bos_token
|
||||
if old_unk_token != new_unk_token: tokenizer.unk_token = old_unk_token
|
||||
if not same_padding_token:
|
||||
if old_pad_token != new_pad_token: tokenizer.pad_token = old_pad_token
|
||||
pass
|
||||
|
||||
# stopping_criteria = create_stopping_criteria(tokenizer, stop_word)
|
||||
|
||||
|
|
|
|||
|
|
@ -349,3 +349,4 @@ class Unsloth_Offloaded_Gradient_Checkpointer(torch.autograd.Function):
|
|||
return (None, hidden_states.grad,) + (None,)*len(ctx.args)
|
||||
pass
|
||||
pass
|
||||
|
||||
|
|
|
|||
|
|
@ -1445,6 +1445,10 @@ class FastLlamaModel:
|
|||
"gate_proj", "up_proj", "down_proj",),)
|
||||
model.config.update({"unsloth_version" : __version__})
|
||||
|
||||
if type(modules_to_save) is tuple:
|
||||
modules_to_save = list(modules_to_save)
|
||||
pass
|
||||
|
||||
train_lm_head = False
|
||||
train_embed_tokens = False
|
||||
final_modules = []
|
||||
|
|
@ -1472,6 +1476,29 @@ class FastLlamaModel:
|
|||
final_modules.append(module)
|
||||
pass
|
||||
|
||||
# Check if we added new tokens!
|
||||
if hasattr(model, "_need_to_train_embeddings"):
|
||||
if not train_lm_head or not train_embed_tokens:
|
||||
print(
|
||||
"Unsloth: You added new tokens but did not specify if you wanted to "\
|
||||
"train the lm_head and embed_tokens.\nWe must turn it on for you."
|
||||
)
|
||||
train_lm_head = True
|
||||
train_embed_tokens = True
|
||||
|
||||
if modules_to_save is None: modules_to_save = ["embed_tokens"]
|
||||
else: modules_to_save.append("embed_tokens")
|
||||
|
||||
if modules_to_save is None: modules_to_save = ["lm_head"]
|
||||
else: modules_to_save.append("lm_head")
|
||||
pass
|
||||
pass
|
||||
|
||||
# First fix untrained tokens
|
||||
if train_embed_tokens or train_lm_head:
|
||||
fix_untrained_tokens(model, eps = 1e-16)
|
||||
pass
|
||||
|
||||
# Check modules_to_save
|
||||
if modules_to_save is not None:
|
||||
for module in modules_to_save:
|
||||
|
|
@ -1479,8 +1506,15 @@ class FastLlamaModel:
|
|||
train_lm_head = True
|
||||
elif module == "embed_tokens":
|
||||
train_embed_tokens = True
|
||||
else:
|
||||
raise TypeError(
|
||||
f"Unsloth: Module = {module} is not allowed. Only 'lm_head' and 'embed_tokens' is allowed."
|
||||
)
|
||||
pass
|
||||
pass
|
||||
if isinstance(modules_to_save, (tuple, list)):
|
||||
modules_to_save = list(set(modules_to_save))
|
||||
pass
|
||||
|
||||
# Get LoRA
|
||||
arguments = dict(
|
||||
|
|
|
|||
|
|
@ -922,9 +922,16 @@ def save_to_gguf(
|
|||
f"The output location will be {final_location}\n"\
|
||||
"This will take 3 minutes...")
|
||||
|
||||
# We first check if tokenizer.model exists in the model_directory
|
||||
if os.path.exists(f"{model_directory}/tokenizer.model"):
|
||||
vocab_type = "hfft"
|
||||
else:
|
||||
vocab_type = "bpe"
|
||||
pass
|
||||
|
||||
if use_fast_convert:
|
||||
command = f"python llama.cpp/convert.py {model_directory} "\
|
||||
f"--outfile {final_location} --vocab-type hfft "\
|
||||
f"--outfile {final_location} --vocab-type {vocab_type} "\
|
||||
f"--outtype {first_conversion} --concurrency {n_cpus}"
|
||||
else:
|
||||
# Need to fix convert-hf-to-gguf.py for some models!
|
||||
|
|
|
|||
|
|
@ -18,11 +18,15 @@ from transformers import PreTrainedTokenizerFast
|
|||
import re
|
||||
import os
|
||||
from transformers.models.llama.modeling_llama import logger
|
||||
from peft import PeftModelForCausalLM
|
||||
import torch
|
||||
|
||||
__all__ = [
|
||||
"load_correct_tokenizer",
|
||||
"fix_sentencepiece_tokenizer",
|
||||
"check_tokenizer",
|
||||
"fix_untrained_tokens",
|
||||
"add_new_tokens",
|
||||
]
|
||||
|
||||
|
||||
|
|
@ -255,7 +259,11 @@ def fix_sentencepiece_tokenizer(
|
|||
|
||||
# And load it!
|
||||
from transformers import AutoTokenizer
|
||||
tokenizer = AutoTokenizer.from_pretrained(temporary_location, eos_token = new_tokenizer.eos_token)
|
||||
tokenizer = AutoTokenizer.from_pretrained(
|
||||
temporary_location,
|
||||
eos_token = new_tokenizer.eos_token,
|
||||
pad_token = new_tokenizer.pad_token,
|
||||
)
|
||||
return tokenizer
|
||||
pass
|
||||
|
||||
|
|
@ -466,3 +474,124 @@ def check_tokenizer(
|
|||
pass
|
||||
return convert_to_fast_tokenizer(tokenizer)
|
||||
pass
|
||||
|
||||
|
||||
@torch.inference_mode
|
||||
def fix_untrained_tokens(model, eps = 1e-16):
|
||||
"""
|
||||
Llama-3 for eg has untrained vectors in the base model.
|
||||
These include <|eot_id|>, <|start_header_id|>, <|end_header_id|>
|
||||
We reset them to the mean of the rest of the tokens
|
||||
"""
|
||||
embedding_matrix = model.get_input_embeddings ().weight.data
|
||||
lm_head_matrix = model.get_output_embeddings().weight.data
|
||||
|
||||
# Get untrained tokens
|
||||
indicator_untrained = torch.amax(embedding_matrix, axis = 1) <= eps
|
||||
where_untrained = torch.where(indicator_untrained)[0]
|
||||
n_untrained = where_untrained.shape[0]
|
||||
n_trained = embedding_matrix.shape[0] - n_untrained
|
||||
if n_untrained != 0:
|
||||
print(
|
||||
f"Unsloth: Not an error, but your model has {n_untrained} untrained tokens.\n"\
|
||||
"We shall set them to the mean of the other trained tokens."
|
||||
)
|
||||
pass
|
||||
|
||||
# First set untrained to all 0s - sometimes it's not! 1e-23 for bfloat16
|
||||
embedding_matrix[where_untrained] = 0
|
||||
lm_head_matrix [where_untrained] = 0
|
||||
|
||||
# Find sum
|
||||
sum_embedding = torch.sum(embedding_matrix, dtype = torch.float32, axis = 0)
|
||||
sum_lm_head = torch.sum(lm_head_matrix, dtype = torch.float32, axis = 0)
|
||||
|
||||
# Find correct average by dividing by sum of trained tokens
|
||||
mean_embedding = (sum_embedding / n_trained).to(embedding_matrix.dtype)
|
||||
mean_lm_head = (sum_lm_head / n_trained).to(lm_head_matrix .dtype)
|
||||
|
||||
# Set them to the mean
|
||||
embedding_matrix[where_untrained] = mean_embedding
|
||||
lm_head_matrix [where_untrained] = mean_lm_head
|
||||
|
||||
return mean_embedding, mean_lm_head
|
||||
pass
|
||||
|
||||
|
||||
@torch.inference_mode
|
||||
def add_new_tokens(
|
||||
model,
|
||||
tokenizer,
|
||||
new_tokens = [],
|
||||
method = "mean",
|
||||
interpolation = 0.05,
|
||||
):
|
||||
"""
|
||||
Smartly resizes the tokenizer and adds new tokens to the model.
|
||||
We also disregard untrained tokens by removing them from the mean calculation.
|
||||
"""
|
||||
assert(isinstance(new_tokens, (list, tuple)))
|
||||
assert(len(new_tokens) > 0)
|
||||
assert(method == "mean" or method == "interpolation")
|
||||
assert(interpolation >= 0 and interpolation <= 1)
|
||||
|
||||
# Check if tokens already exist
|
||||
overlapping_tokens = set(new_tokens) & set(tokenizer.vocab.keys())
|
||||
if len(overlapping_tokens) != 0:
|
||||
print(
|
||||
f"Unsloth: You're adding new_tokens = {new_tokens}\n"\
|
||||
f"There are tokens which are overlapping = {list(overlapping_tokens)}\n"\
|
||||
f"We shall safely ignore these overlapping tokens."
|
||||
)
|
||||
new_tokens = [x for x in new_tokens if x not in overlapping_tokens]
|
||||
pass
|
||||
|
||||
# Get mean of trained tokens
|
||||
mean_embedding, mean_lm_head = fix_untrained_tokens(model)
|
||||
mean_embedding = mean_embedding.to(torch.float32)
|
||||
mean_lm_head = mean_lm_head .to(torch.float32)
|
||||
|
||||
# Add tokens!
|
||||
old_length = len(tokenizer)
|
||||
tokenizer.add_tokens(new_tokens)
|
||||
model.resize_token_embeddings(len(tokenizer))
|
||||
|
||||
# If we use interpolation, we interpolate between the mean embeddings and
|
||||
# the Word2Vec sum of the other vectors
|
||||
embedding_matrix = model.get_input_embeddings ().weight.data
|
||||
lm_head_matrix = model.get_output_embeddings().weight.data
|
||||
|
||||
if method == "interpolation":
|
||||
print(
|
||||
"Unsloth: You are using interpolation to add new tokens.\n"\
|
||||
f"We shall set new tokens = mean(embeddings)*{1-interpolation} + mean(new_tokens)*{interpolation}"
|
||||
)
|
||||
for j, token in enumerate(new_tokens):
|
||||
input_ids = tokenizer(token, add_special_tokens = False).input_ids
|
||||
mean_embedding_token = embedding_matrix[input_ids].mean(axis = 0, dtype = torch.float32)
|
||||
mean_lm_head_token = lm_head_matrix [input_ids].mean(axis = 0, dtype = torch.float32)
|
||||
|
||||
# Interpolate
|
||||
mean_embedding_token = mean_embedding*(1-interpolation) + mean_embedding_token*interpolation
|
||||
mean_lm_head_token = mean_lm_head *(1-interpolation) + mean_lm_head_token *interpolation
|
||||
|
||||
# Set the new vector
|
||||
embedding_matrix[old_length+j] = mean_embedding_token
|
||||
lm_head_matrix [old_length+j] = mean_lm_head_token
|
||||
pass
|
||||
else:
|
||||
# Now set the new tokens to the mean!
|
||||
embedding_matrix[old_length:] = mean_embedding
|
||||
lm_head_matrix [old_length:] = mean_lm_head
|
||||
pass
|
||||
|
||||
# We set a flag to say we need to train embeddings
|
||||
internal_model = model
|
||||
while hasattr(internal_model, "model"):
|
||||
internal_model._need_to_train_embeddings = True
|
||||
internal_model = internal_model.model
|
||||
pass
|
||||
internal_model._need_to_train_embeddings = True
|
||||
|
||||
return
|
||||
pass
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue