From 5350c7189dda91f0fcfe2d753542a38946c57acc Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Tue, 26 Mar 2024 03:40:55 +1100 Subject: [PATCH] tokenizer patching --- unsloth/chat_templates.py | 7 +- unsloth/models/_utils.py | 128 ----------------------------- unsloth/models/llama.py | 2 +- unsloth/tokenizer_utils.py | 159 +++++++++++++++++++++++++++++++++++++ 4 files changed, 165 insertions(+), 131 deletions(-) diff --git a/unsloth/chat_templates.py b/unsloth/chat_templates.py index 38e6bba958..9de15b98cb 100644 --- a/unsloth/chat_templates.py +++ b/unsloth/chat_templates.py @@ -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 = {} diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 6f7da0f32a..375702264b 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -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 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 diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index d3b48d9439..bc558c2ddb 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -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 diff --git a/unsloth/tokenizer_utils.py b/unsloth/tokenizer_utils.py index 75f99c1239..3add80c6ad 100644 --- a/unsloth/tokenizer_utils.py +++ b/unsloth/tokenizer_utils.py @@ -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 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