diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 32d20aee5f..0ccbe9117f 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -19,6 +19,14 @@ import warnings import gc warnings.filterwarnings(action = "ignore", category = UserWarning, module = "torch") import bitsandbytes as bnb +from transformers.models.llama.modeling_llama import logger + +__version__ = "2023.12" +__all__ = [ + "prepare_model_for_kbit_training", + "patch_tokenizer", + "print_unsloth_message", +] def prepare_model_for_kbit_training( @@ -59,3 +67,38 @@ def prepare_model_for_kbit_training( return model pass + + +def patch_tokenizer(model, tokenizer): + 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: + logger.warning_one( + f"{model.config._name_or_path} 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 + config = model.config.update({"pad_token_id" : tokenizer.eos_token_id}) + pass + return model, tokenizer +pass + + +def print_unsloth_message(name): + SUPPORTS_BFLOAT16 = torch.cuda.is_bf16_supported() + gpu_stats = torch.cuda.get_device_properties(0) + max_memory = round(gpu_stats.total_memory / 1024 / 1024 / 1024, 3) + + statistics = \ + f"==((====))== Unsloth: Fast {name} patching release {__version__}\n"\ + f" \\\ /| GPU: {gpu_stats.name}. Max memory: {max_memory} GB\n"\ + f"O^O/ \_/ \\ CUDA compute capability = {gpu_stats.major}.{gpu_stats.minor}\n"\ + f"\ / Pytorch version: {torch.__version__}. CUDA Toolkit = {torch.version.cuda}\n"\ + f' "-____-" bfloat16 support = {str(SUPPORTS_BFLOAT16).upper()}\n' + print(statistics) +pass diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 9e9c3cca5d..a3a91ab33d 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -22,9 +22,7 @@ from transformers.models.llama.modeling_llama import ( CausalLMOutputWithPast, ) from ..kernels import * -from ._utils import ( - prepare_model_for_kbit_training, -) +from ._utils import * # Get Flash Attention v2 if Ampere (RTX 30xx, A100) major_version, minor_version = torch.cuda.get_device_capability() @@ -615,18 +613,8 @@ class FastLlamaModel: device_map = "sequential", rope_scaling = None, ): - gpu_stats = torch.cuda.get_device_properties(0) - max_memory = round(gpu_stats.total_memory / 1024 / 1024 / 1024, 3) SUPPORTS_BFLOAT16 = torch.cuda.is_bf16_supported() - - statistics = \ - "==((====))== Unsloth: Fast Llama patching release 2023.12\n"\ - f" \\\ /| GPU: {gpu_stats.name}. Max memory: {max_memory} GB\n"\ - f"O^O/ \_/ \\ CUDA compute capability = {gpu_stats.major}.{gpu_stats.minor}\n"\ - f"\ / Pytorch version: {torch.__version__}. CUDA Toolkit = {torch.version.cuda}\n"\ - f' "-____-" bfloat16 support = {str(SUPPORTS_BFLOAT16).upper()}\n' - print(statistics) - + print_unsloth_message("Mistral") FastLlamaModel.pre_patch() if dtype is None: @@ -676,22 +664,7 @@ class FastLlamaModel: token = token, ) - if not hasattr(tokenizer, "pad_token"): - # 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: - logger.warning_one( - f"{model_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 - config = model.config.update({"pad_token_id" : tokenizer.eos_token_id}) - pass - + model, tokenizer = patch_tokenizer(model, tokenizer) model = FastLlamaModel.post_patch(model) # Patch up QKV / O and MLP diff --git a/unsloth/models/mistral.py b/unsloth/models/mistral.py index 0c13f998f7..8b1fd0fd84 100644 --- a/unsloth/models/mistral.py +++ b/unsloth/models/mistral.py @@ -231,18 +231,8 @@ class FastMistralModel(FastLlamaModel): device_map = "sequential", # rope_scaling = None, Mistral does not support RoPE scaling ): - gpu_stats = torch.cuda.get_device_properties(0) - max_memory = round(gpu_stats.total_memory / 1024 / 1024 / 1024, 3) SUPPORTS_BFLOAT16 = torch.cuda.is_bf16_supported() - - statistics = \ - "==((====))== Unsloth: Fast Mistral patching release 2023.12\n"\ - f" \\\ /| GPU: {gpu_stats.name}. Max memory: {max_memory} GB\n"\ - f"O^O/ \_/ \\ CUDA compute capability = {gpu_stats.major}.{gpu_stats.minor}\n"\ - f"\ / Pytorch version: {torch.__version__}. CUDA Toolkit = {torch.version.cuda}\n"\ - f' "-____-" bfloat16 support = {str(SUPPORTS_BFLOAT16).upper()}\n' - print(statistics) - + print_unsloth_message("Mistral") FastMistralModel.pre_patch() if dtype is None: @@ -277,22 +267,7 @@ class FastMistralModel(FastLlamaModel): token = token, ) - if not hasattr(tokenizer, "pad_token"): - # 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: - logger.warning_one( - f"{model_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 - config = model.config.update({"pad_token_id" : tokenizer.eos_token_id}) - pass - + model, tokenizer = patch_tokenizer(model, tokenizer) model = FastMistralModel.post_patch(model) # Patch up QKV / O and MLP