From e097ba3e7bbae5b31b6eb1790a3818d28593c0b6 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 2 Nov 2024 19:08:57 -0700 Subject: [PATCH] Utils --- unsloth/kernels/cross_entropy_loss.py | 1 + unsloth/models/_utils.py | 369 +------------------------- 2 files changed, 11 insertions(+), 359 deletions(-) diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py index f8b9b4245e..f0dbb66f49 100644 --- a/unsloth/kernels/cross_entropy_loss.py +++ b/unsloth/kernels/cross_entropy_loss.py @@ -19,6 +19,7 @@ from .utils import calculate_settings, MAX_FUSED_SIZE, triton_tanh from transformers.models.llama.modeling_llama import logger from packaging.version import Version + @triton.heuristics({ "DO_SOFTCAPPING": lambda args: args["DO_SOFTCAPPING" ], "DO_LOGIT_SCALING": lambda args: args["DO_LOGIT_SCALING"], diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index e1ed649560..4829e4d88c 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -56,6 +56,16 @@ import numpy as np import warnings, subprocess, re, inspect, psutil, os, math from packaging.version import Version +from unsloth_zoo.tokenizer_utils import ( + patch_tokenizer, +) +from unsloth_zoo.gradient_checkpointing import ( + Unsloth_Offloaded_Gradient_Checkpointer, + unsloth_offloaded_gradient_checkpoint, + patch_gradient_checkpointing, + unpatch_gradient_checkpointing, +) + # ============================================= # Disable some warnings which can get annoying warnings.filterwarnings(action = "ignore", category = UserWarning, module = "torch") @@ -131,7 +141,6 @@ pass # ============================================= # torch.cuda.amp.custom_fwd is deprecated >= 2.4 -import torch torch_version = torch.__version__ if Version(torch_version) < Version("2.4.0"): torch_amp_custom_fwd = torch.cuda.amp.custom_fwd @@ -457,228 +466,6 @@ def prepare_model_for_kbit_training( return model pass - -def patch_tokenizer(model, tokenizer): - """ - Phi3's pad_token isn't set. We set it to <|placeholder... - Llama-3 is <|reserved... - Llama-2 is - Check if pad_token is not the same as eos_token otherwise the loss will ignore it!! - Fixes https://github.com/unslothai/unsloth/issues/5 - """ - possible_reserved_tokens = ( - "<|finetune_right_pad_id|>", # Llama-3.1 - "", # Mistral Nemo - "<|reserved", # Llama-3 - "<|placeholder", # Phi-3 - "[control", # Mistral type models - ) - joiner = "\1\0=+=\0\1" - number_repetitions = 3 - 1 # Number of reserved tokens needed - - if model is not None: - model.config.update({"unsloth_version" : __version__}) - - # First remove pad and unk tokens if they are known to be BOS / EOS - possible_bad_tokens = ( - "<|endoftext|>", - "<|im_start|>", - "<|im_end|>", - "<|begin_of_text|>", - "<|end_of_text|>", - "", - "", - ) - input_ids = tokenizer(list(possible_bad_tokens), add_special_tokens = False).input_ids - possible_bad_tokens = frozenset(token for token, input_id in zip(possible_bad_tokens, input_ids) if len(input_id) == 1) - - if hasattr(tokenizer, "pad_token") and tokenizer.pad_token in possible_bad_tokens: - print(f"Unsloth: Pad token was {tokenizer.pad_token} which is not a good idea. We shall fix this.") - tokenizer.pad_token = None - pass - - has_bad_unk_token = False - if hasattr(tokenizer, "unk_token") and tokenizer.unk_token in possible_bad_tokens: - print(f"Unsloth: Unk token was {tokenizer.unk_token} which is not a good idea. We shall fix this.") - tokenizer.unk_token = None - has_bad_unk_token = True - pass - - # Now check pad token again - bad_pad_token = False - if hasattr(tokenizer, "pad_token") and tokenizer.pad_token is not None: - # Check if pad_token is not the same as eos_token otherwise the loss will ignore it!! - bad_pad_token = tokenizer.eos_token == tokenizer.pad_token - elif hasattr(tokenizer, "pad_token") and tokenizer.pad_token is None: - bad_pad_token = True - else: - bad_pad_token = False - pass - - # Check if unknown token is broken - fixed_unk_token = False - - if (hasattr(tokenizer, "unk_token") and tokenizer.unk_token is not None) or has_bad_unk_token: - - eos_token = getattr(tokenizer, "eos_token", None) - bos_token = getattr(tokenizer, "bos_token", None) - - old_unk_token = tokenizer.unk_token - if (old_unk_token == eos_token) or (old_unk_token == bos_token) or has_bad_unk_token: - has_broken_unk = True - # Use the unicode replacement characters - possible_replacements = [ - "\uFFFD", # Original replacement char - "\uFFFC", # Another option - "\u2753", # Red Question mark emoji - "\u2754", # White Question mark emoji - "\u00BF", # Inverted question mark - ] - for replacement_char in possible_replacements: - char = tokenizer(replacement_char, add_special_tokens = False).input_ids - if len(char) == 1: - # Get actual token representation - try: char = tokenizer.convert_ids_to_tokens(char[0]) - except: continue - tokenizer.unk_token = char - fixed_unk_token = True - break - pass - pass - - if not fixed_unk_token: # Still broken! - raise RuntimeError( - f"Unsloth: Tried fixing the unk_token = {old_unk_token}, but couldn't!" - ) - pass - - logger.warning_once( - f"Unsloth: unk_token = {old_unk_token} is the same as the EOS or BOS tokens. "\ - f"We fixed it by changing it to {tokenizer.unk_token}." - ) - pass - pass - - if bad_pad_token: - # Find a better pad token - added_tokens = [str(x) for x in tokenizer.added_tokens_decoder.values()] - all_added_tokens = joiner.join(added_tokens[::-1]) - all_added_tokens += joiner - - final_pad_token = None - final_good_match = False - - for possible_reserved_token in possible_reserved_tokens: - possible_reserved_token = re.escape(possible_reserved_token) - found = re.finditer(f"{possible_reserved_token}", all_added_tokens) - first_match = None - good_match = False - for j, x in enumerate(found): - if j == 0: first_match = x - if j >= number_repetitions: - good_match = True - break - pass - pass - - if first_match is None: continue - - # If it ends with |> or > etc, then set it as a good pad token! - start = first_match.span(0)[0] - possible_pad_token = first_match.group(0) - end = all_added_tokens.find(joiner, start) - first_match = all_added_tokens[start:end] - - if first_match is not None: - good_match = possible_pad_token.endswith((">", "|>", "]", ")")) - pass - possible_pad_token = first_match - - # Replace current pad token if another exact match is found - if not final_good_match and good_match: - final_good_match = True - final_pad_token = possible_pad_token - break - else: - final_good_match = False - final_pad_token = possible_pad_token - pass - pass - possible_pad_token = final_pad_token - - # Try unk_token if it wasn't fixed - if possible_pad_token is None and not fixed_unk_token and hasattr(tokenizer, "unk_token"): - possible_pad_token = tokenizer.unk_token - pass - - # Check pad token's id must be less than vocab size - if possible_pad_token is not None: - check_pad_token = tokenizer(possible_pad_token, add_special_tokens = False).input_ids - if len(check_pad_token) != 1: - possible_pad_token = None - if model is not None and check_pad_token[0] >= model.config.vocab_size: - possible_pad_token = None - pass - - if possible_pad_token is None: - # Failure to find a good replacement!! We shall manually add one! - new_pad_token = "<|PAD_TOKEN|>" - while new_pad_token in tokenizer.get_vocab(): - new_pad_token = f"<{new_pad_token}>" - pass - possible_pad_token = new_pad_token - pass - - name = model.config._name_or_path if model is not None else "Model" - logger.warning_once( - f"{name} does not have a padding token! Will use pad_token = {possible_pad_token}." - ) - - # Edit pad_token - tokenizer.add_special_tokens({"pad_token" : possible_pad_token}) - tokenizer.pad_token = possible_pad_token - if model is not None: - - # Edit all config with new pad token - current_model = model - while hasattr(current_model, "model") and hasattr(current_model, "config"): - current_model.config.update({"pad_token_id" : tokenizer.pad_token_id}) - current_model = current_model.model - if hasattr(current_model, "model") and hasattr(current_model, "config"): - current_model.config.update({"pad_token_id" : tokenizer.pad_token_id}) - pass - - # Generation edit pad token - if getattr(model, "generation_config") is not None: - model.generation_config.update(pad_token_id = tokenizer.pad_token_id) - else: - if model is not None: - - if model.config.pad_token_id is None: - - # Edit all config with new pad token - current_model = model - while hasattr(current_model, "model") and hasattr(current_model, "config"): - current_model.config.update({"pad_token_id" : tokenizer.pad_token_id}) - current_model = model - if hasattr(current_model, "model") and hasattr(current_model, "config"): - current_model.config.update({"pad_token_id" : tokenizer.pad_token_id}) - pass - - # Generation edit pad token - if getattr(model, "generation_config") is not None: - model.generation_config.update(pad_token_id = tokenizer.pad_token_id) - pass - pass - - if model is not None: - if getattr(model, "generation_config") is not None: - model.generation_config.update(max_length = model.config.max_position_embeddings) - - return model, tokenizer -pass - - # ============================================= # Weirdly LoraLayer.update_layer downcasts PEFT layers to float16?? # For mixed precision, we need it to be in float32 not float16. @@ -820,142 +607,6 @@ def get_statistics(): pass -def _calculate_n_gradient_checkpoints( - n_layers : int, - method : Optional[Union[str, int]] = "sqrt", -) -> List[int]: - assert(type(n_layers) is int and n_layers > 0) - - if method is None: method = "sqrt" - - if method == "sqrt": - n_checkpoints = int(n_layers**0.5) - elif type(method) is int and method > 0: - n_checkpoints = int(np.ceil(n_layers / method)) - else: - raise ValueError("method must be 'sqrt' or an int >0 and <= n_layers.") - - size = n_layers // n_checkpoints - sizes = np.full(n_checkpoints, size, dtype = int) - leftovers = n_layers % n_checkpoints - # We append leftovers from the right - for k in range(leftovers): - sizes[n_checkpoints-1-k] += 1 - boundaries = np.hstack((0, np.cumsum(sizes))) - boundaries = boundaries.tolist() - return boundaries -pass - - -def calculate_n_gradient_checkpoints( - n_layers : int, - layers_per_checkpoint : Optional[Union[str, int]] = "sqrt", -) -> List[int]: - assert(type(n_layers) is int and n_layers > 0) - - if layers_per_checkpoint is None or layers_per_checkpoint == 1: - return None - - boundaries = _calculate_n_gradient_checkpoints(n_layers, layers_per_checkpoint) - - assert(boundaries[0] == 0 and boundaries[-1] == n_layers) - assert(min(boundaries) == 0 and max(boundaries) == n_layers) - assert(np.diff(boundaries).min() >= 0) - return boundaries -pass - - -def prepare_n_gradient_checkpoints( - model : Any, - layers_per_checkpoint : Optional[Union[str, int]] = "sqrt", - use_reentrant : Optional[bool] = True, -) -> None: - """ - Calculates where to place the gradient checkpoints given n_layers. - - Args: - model: Any LlamaModel with layers. - layers_per_checkpoint (`Union[str, int]`, *optional*): - Can either be `sqrt` or an integer for how many layers per checkpoint you want. - The more, the less memory usage, but can be slower. Default is `sqrt`. - Choose 1 for Pytorch gradient checkpointing. 2 to wrap 2 layers in 1 module etc. - use_reentrant (`bool`, *optional*): - https://github.com/pytorch/pytorch/blob/main/torch/utils/checkpoint.py#L354 - Optimal gradient checkpointing algorithm `use_reentrant=False` which will - be the default in future Pytorch versions doesn't seem to work?? - """ - _model = None - if hasattr(model, "layers"): - _model = model - elif hasattr(model, "model"): - if hasattr(model.model, "layers"): - _model = model.model - if _model is None: - raise TypeError("`model` or `model.model` does not have attribute `layers`. Are you sure this is a model?") - pass - - if use_reentrant is False: - use_reentrant = True - pass - - n_layers = len(_model.layers) - boundaries = calculate_n_gradient_checkpoints(n_layers, layers_per_checkpoint) - _model._gradient_checkpointing_boundaries = boundaries - _model._gradient_checkpointing_use_reentrant = use_reentrant -pass - - -class Unsloth_Offloaded_Gradient_Checkpointer(torch.autograd.Function): - """ - Saves VRAM by smartly offloading to RAM. - Tiny hit to performance, since we mask the movement via non blocking calls. - """ - @staticmethod - @torch_amp_custom_fwd - def forward(ctx, forward_function, hidden_states, *args): - saved_hidden_states = hidden_states.to("cpu", non_blocking = True) - with torch.no_grad(): - output = forward_function(hidden_states, *args) - ctx.save_for_backward(saved_hidden_states) - ctx.forward_function = forward_function - ctx.args = args - return output - pass - - @staticmethod - @torch_amp_custom_bwd - def backward(ctx, dY): - (hidden_states,) = ctx.saved_tensors - hidden_states = hidden_states.to("cuda:0", non_blocking = True).detach() - hidden_states.requires_grad_(True) - with torch.enable_grad(): - (output,) = ctx.forward_function(hidden_states, *ctx.args) - torch.autograd.backward(output, dY) - return (None, hidden_states.grad,) + (None,)*len(ctx.args) - pass -pass - - -# @torch._disable_dynamo -def unsloth_offloaded_gradient_checkpoint(function, *args, use_reentrant = None, **kwargs): - return Unsloth_Offloaded_Gradient_Checkpointer.apply(function, *args) -pass - -import torch.utils -def patch_gradient_checkpointing(): - if torch.utils.checkpoint.checkpoint.__name__ == "unsloth_offloaded_gradient_checkpoint": return - torch.utils.checkpoint._old_checkpoint = torch.utils.checkpoint.checkpoint - torch.utils.checkpoint.checkpoint = unsloth_offloaded_gradient_checkpoint -pass - -def unpatch_gradient_checkpointing(): - if hasattr(torch.utils.checkpoint, "_old_checkpoint"): - torch.utils.checkpoint.checkpoint = torch.utils.checkpoint._old_checkpoint - del torch.utils.checkpoint._old_checkpoint - pass -pass - - # ============================================= # Regional torch 2.5 Recompilation - weirdly very slow?? def patch_regional_compilation():