ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

tokenizer patching

Browse source
This commit is contained in:
Daniel Han-Chen 2024-03-26 03:40:55 +11:00
commit 5350c7189d
4 changed files with 165 additions and 131 deletions
Download patch file Download diff file Expand all files Collapse all files

7
unsloth/chat_templates.py
View file

@ -20,11 +20,14 @@ __all__ = [
from transformers import StoppingCriteria, StoppingCriteriaList
from torch import LongTensor, FloatTensor
from transformers.models.llama.modeling_llama import logger
from .models._utils import patch_tokenizer
from .save import patch_saving_functions
import os
import shutil
from .tokenizer_utils import load_correct_tokenizer, fix_sentencepiece_tokenizer
from .tokenizer_utils import (
load_correct_tokenizer,
fix_sentencepiece_tokenizer,
patch_tokenizer,
)
CHAT_TEMPLATES = {}

128
unsloth/models/_utils.py
View file

@ -60,8 +60,6 @@ from xformers import __version__ as xformers_version
__all__ = [
"prepare_model_for_kbit_training",
"patch_tokenizer",
"check_tokenizer",
"xformers",
"xformers_attention",
"xformers_version",
@ -71,11 +69,6 @@ __all__ = [
]
IGNORED_TOKENIZER_CHECKING = frozenset((
"CodeLlamaTokenizerFast",
"CodeLlamaTokenizer",
))
def prepare_model_for_kbit_training(
model : Any,
use_gradient_checkpointing : bool = True,
@ -120,127 +113,6 @@ def prepare_model_for_kbit_training(
pass
def patch_tokenizer(model, tokenizer):
if model is not None:
model.config.update({"unsloth_version" : __version__})
if not hasattr(tokenizer, "pad_token") or tokenizer.pad_token is None:
# Fixes https://github.com/unslothai/unsloth/issues/5
if hasattr(tokenizer, "unk_token"):
tokenizer.add_special_tokens({"pad_token" : tokenizer.unk_token})
tokenizer.pad_token = tokenizer.unk_token
else:
name = model.config._name_or_path if model is not None else "Model"
logger.warning_one(
f"{name} does not have a padding or unknown token!\n"\
f"Will use the EOS token of id {tokenizer.eos_token_id} as padding."
)
assert(hasattr(tokenizer, "eos_token"))
tokenizer.add_special_tokens({"pad_token" : tokenizer.eos_token})
tokenizer.pad_token = tokenizer.eos_token
if model is not None:
config = model.config.update({"pad_token_id" : tokenizer.eos_token_id})
pass
return model, tokenizer
pass
def check_tokenizer(
model,
tokenizer,
model_name = "unsloth/llama-2-7b-bnb-4bit",
model_max_length = 4096,
padding_side = "right",
token = None,
_reload = True,
):
# Checks tokenizer for out of bounds ids.
# Mainly a fix for https://huggingface.co/berkeley-nest/Starling-LM-7B-alpha
# where <sep> had token id=32002.
# See https://huggingface.co/berkeley-nest/Starling-LM-7B-alpha/discussions/25
# Seems like the Fast tokenizer in Rust breaks things!
# We ignore some of them!
if tokenizer.__repr__().split("(", 1)[0] in IGNORED_TOKENIZER_CHECKING:
return tokenizer
pass
max_embedding_size = model.model.embed_tokens.weight.shape[0]
added_tokens_fast = tokenizer.added_tokens_decoder
added_tokens_fast = {index : str(value) for index, value in added_tokens_fast.items()}
sorted_keys = sorted(added_tokens_fast)
added_tokens_fast = {key : added_tokens_fast[key] for key in sorted_keys}
for j, index in enumerate(added_tokens_fast.keys()):
if index >= max_embedding_size:
bad_indices = list(added_tokens_fast.keys ())[j:]
bad_tokens = list(added_tokens_fast.values())[j:]
if not _reload:
# Try removing the token
added_tokens = [str(x) for x in tokenizer.added_tokens_decoder.values()]
special_tokens = tokenizer.special_tokens_map
import itertools
special_tokens = frozenset(
itertools.chain.from_iterable(
[x] if type(x) is str else x for x in special_tokens.values()
)
)
can_be_removed1 = [x for x in bad_tokens if x not in special_tokens]
can_be_removed2 = [x for x in can_be_removed1 if x in tokenizer._added_tokens_encoder.keys()]
# Check of extra tokens can in fact we removed!
if (len(can_be_removed1) == len(bad_tokens)) and \
(len(can_be_removed2) == len(bad_tokens)):
# Yes it can be fixed!
for bad_token in can_be_removed1:
remove_id = tokenizer._added_tokens_encoder[bad_token]
del tokenizer._added_tokens_decoder[remove_id]
del tokenizer._added_tokens_encoder[bad_token]
pass
# Confirm 1 more time!
if max(tokenizer.added_tokens_decoder.keys()) < max_embedding_size:
logger.warning_once(
f"Unsloth loaded a broken tokenizer `{model_name}`, but managed to repair it!\n"\
f"Tokens {bad_tokens} with ids {bad_indices} exceeds the max vocab size of {max_embedding_size}.\n"\
"We removed these bad tokens. If you think this is incorrect, fix your tokenizer first."
)
return tokenizer
pass
pass
# :( Failure
raise RuntimeError(
f"Unsloth tried to load `{model_name}`, but cannot succeed.\n"\
f"Tokens {bad_tokens} with ids {bad_indices} exceeds the max vocab size of {max_embedding_size}.\n"\
f"Fix your tokenizer since it'll perform out of bounds memory accesses."
)
pass
# Try slow tokenizer which can fix things!
tokenizer = AutoTokenizer.from_pretrained(
model_name,
model_max_length = model_max_length,
padding_side = padding_side,
token = token,
use_fast = False,
)
return check_tokenizer(
model = model,
tokenizer = tokenizer,
model_name = model_name,
model_max_length = model_max_length,
padding_side = padding_side,
token = token,
_reload = False,
)
break
pass
pass
return tokenizer
pass
# Weirdly LoraLayer.update_layer downcasts PEFT layers to float16??
# For mixed precision, we need it to be in float32 not float16.
from peft.tuners.lora.layer import LoraLayer

2
unsloth/models/llama.py
View file

@ -26,7 +26,7 @@ from transformers.modeling_attn_mask_utils import (
from ..kernels import *
from ._utils import *
from ._utils import __version__
from ..tokenizer_utils import load_correct_tokenizer
from ..tokenizer_utils import *
if HAS_FLASH_ATTENTION:
from flash_attn import flash_attn_func

159
unsloth/tokenizer_utils.py
View file

@ -21,9 +21,17 @@ import os
__all__ = [
"load_correct_tokenizer",
"fix_sentencepiece_tokenizer",
"check_tokenizer",
"patch_tokenizer",
]
IGNORED_TOKENIZER_CHECKING = frozenset((
"CodeLlamaTokenizerFast",
"CodeLlamaTokenizer",
))
def try_fix_tokenizer(tokenizer, prepend = True):
if hasattr(tokenizer, "_tokenizer"):
@ -277,3 +285,154 @@ def load_correct_tokenizer(
pass
pass
def patch_tokenizer(model, tokenizer):
if model is not None:
model.config.update({"unsloth_version" : __version__})
if not hasattr(tokenizer, "pad_token") or tokenizer.pad_token is None:
# Fixes https://github.com/unslothai/unsloth/issues/5
if hasattr(tokenizer, "unk_token"):
tokenizer.add_special_tokens({"pad_token" : tokenizer.unk_token})
tokenizer.pad_token = tokenizer.unk_token
else:
name = model.config._name_or_path if model is not None else "Model"
logger.warning_one(
f"{name} does not have a padding or unknown token!\n"\
f"Will use the EOS token of id {tokenizer.eos_token_id} as padding."
)
assert(hasattr(tokenizer, "eos_token"))
tokenizer.add_special_tokens({"pad_token" : tokenizer.eos_token})
tokenizer.pad_token = tokenizer.eos_token
if model is not None:
config = model.config.update({"pad_token_id" : tokenizer.eos_token_id})
pass
return model, tokenizer
pass
def check_tokenizer(
model,
tokenizer,
model_name = "unsloth/llama-2-7b-bnb-4bit",
model_max_length = 4096,
padding_side = "right",
token = None,
_reload = True,
):
# Checks tokenizer for out of bounds ids.
# Mainly a fix for https://huggingface.co/berkeley-nest/Starling-LM-7B-alpha
# where <sep> had token id=32002.
# See https://huggingface.co/berkeley-nest/Starling-LM-7B-alpha/discussions/25
# Seems like the Fast tokenizer in Rust breaks things!
# We ignore some of them!
if tokenizer.__repr__().split("(", 1)[0] in IGNORED_TOKENIZER_CHECKING:
return tokenizer
pass
max_embedding_size = model.model.embed_tokens.weight.shape[0]
added_tokens_fast = tokenizer.added_tokens_decoder
added_tokens_fast = {index : str(value) for index, value in added_tokens_fast.items()}
sorted_keys = sorted(added_tokens_fast)
added_tokens_fast = {key : added_tokens_fast[key] for key in sorted_keys}
for j, index in enumerate(added_tokens_fast.keys()):
if index >= max_embedding_size:
bad_indices = list(added_tokens_fast.keys ())[j:]
bad_tokens = list(added_tokens_fast.values())[j:]
if not _reload:
# Try removing the token
added_tokens = [str(x) for x in tokenizer.added_tokens_decoder.values()]
special_tokens = tokenizer.special_tokens_map
import itertools
special_tokens = frozenset(
itertools.chain.from_iterable(
[x] if type(x) is str else x for x in special_tokens.values()
)
)
can_be_removed1 = [x for x in bad_tokens if x not in special_tokens]
can_be_removed2 = [x for x in can_be_removed1 if x in tokenizer._added_tokens_encoder.keys()]
# Check of extra tokens can in fact we removed!
can_be_removed = \
(len(can_be_removed1) == len(bad_tokens)) and \
(len(can_be_removed2) == len(bad_tokens))
# Check if sep_token or other generic types
remove_generic = False
try_mapper = []
if not can_be_removed:
names = dir(tokenizer)
names = (x for x in names if x.endswith("_token") and x.count("_") == 1)
generic_tokens = [(x, getattr(tokenizer, x, None)) for x in names]
try_removal = []
for token in bad_tokens:
for (name_token, check_token) in generic_tokens:
if check_token == token:
try_removal.append(token)
try_mapper.append(name_token)
pass
pass
pass
# Recheck!
can_be_removed = len(try_removal) == len(bad_tokens)
if can_be_removed: remove_generic = True
can_be_removed1 = bad_tokens
pass
if can_be_removed:
# Yes it can be fixed!
for j, bad_token in enumerate(can_be_removed1):
remove_id = tokenizer._added_tokens_encoder[bad_token]
del tokenizer._added_tokens_decoder[remove_id]
del tokenizer._added_tokens_encoder[bad_token]
if remove_generic and try_removal[j] == bad_token:
# Remove sep token for example
setattr(tokenizer, try_mapper[j], None)
setattr(tokenizer, try_mapper[j] + "_id", None)
pass
pass
# Confirm 1 more time!
if max(tokenizer.added_tokens_decoder.keys()) < max_embedding_size:
logger.warning_once(
f"Unsloth loaded a broken tokenizer `{model_name}`, but managed to repair it!\n"\
f"Tokens {bad_tokens} with ids {bad_indices} exceeds the max vocab size of {max_embedding_size}.\n"\
"We removed these bad tokens. If you think this is incorrect, fix your tokenizer first."
)
return convert_to_fast_tokenizer(tokenizer)
pass
pass
# :( Failure
raise RuntimeError(
f"Unsloth tried to load `{model_name}`, but cannot succeed.\n"\
f"Tokens {bad_tokens} with ids {bad_indices} exceeds the max vocab size of {max_embedding_size}.\n"\
f"Fix your tokenizer since it'll perform out of bounds memory accesses."
)
pass
# Try slow tokenizer which can fix things!
tokenizer = AutoTokenizer.from_pretrained(
model_name,
model_max_length = model_max_length,
padding_side = padding_side,
token = token,
use_fast = False,
)
return check_tokenizer(
model = model,
tokenizer = tokenizer,
model_name = model_name,
model_max_length = model_max_length,
padding_side = padding_side,
token = token,
_reload = False,
)
break
pass
pass
return convert_to_fast_tokenizer(tokenizer)
pass