diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 8cdb5e384c..f4268e3493 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -30,7 +30,7 @@ import numpy as np import os import psutil -__version__ = "2024.3" +__version__ = "2024.4" # Get Flash Attention v2 if Ampere (RTX 30xx, A100) major_version, minor_version = torch.cuda.get_device_capability() @@ -70,12 +70,13 @@ __all__ = [ "platform_system", "patch_tokenizer", "get_statistics", + "Offloaded_Gradient_Checkpointer", ] def prepare_model_for_kbit_training( model : Any, - use_gradient_checkpointing : bool = True, + use_gradient_checkpointing : Optional = True, use_reentrant : Optional[bool] = True, ) -> Any: """ @@ -101,9 +102,23 @@ def prepare_model_for_kbit_training( param.requires_grad_(False) pass - if use_gradient_checkpointing: + # Gradient checkpointing! + if use_gradient_checkpointing == "offloaded": + + # Saves VRAM! + original_model = model + while hasattr(original_model, "model"): + original_model._offloaded_gradient_checkpointing = True + original_model = original_model.model + pass + original_model._offloaded_gradient_checkpointing = True + model.gradient_checkpointing_enable() + elif use_gradient_checkpointing == True: + model.gradient_checkpointing_enable() + pass + # If use_reentrant = True which is the Pytorch default, we just make the input requires_grad. if use_reentrant: if hasattr(model, "enable_input_require_grads"): @@ -179,6 +194,7 @@ def get_statistics(): try: from huggingface_hub import hf_hub_download from huggingface_hub.utils import disable_progress_bars, enable_progress_bars, are_progress_bars_disabled + import psutil n_cpus = psutil.cpu_count(logical = False) keynames = "\n" + "\n".join(os.environ.keys()) @@ -291,3 +307,35 @@ def prepare_n_gradient_checkpoints( _model._gradient_checkpointing_boundaries = boundaries _model._gradient_checkpointing_use_reentrant = use_reentrant pass + + +class 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. + [TODO] Load the backward pass earlier + """ + @staticmethod + @torch.cuda.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.cuda.amp.custom_bwd + def backward(ctx, dY): + (hidden_states,) = ctx.saved_tensors + hidden_states = hidden_states.to("cuda", 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 diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index dc66059b0f..a7ade9fc32 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -628,19 +628,42 @@ def LlamaModel_fast_forward( boundaries = None pass + # Check checkpointing method + gradient_checkpointing = False + offloaded_gradient_checkpointing = False + + if (self.gradient_checkpointing and self.training and not use_cache): + + gradient_checkpointing = True + + if output_attentions is False and hasattr(self, "_offloaded_gradient_checkpointing"): + offloaded_gradient_checkpointing = True + pass + + # Go through every layer! for idx, decoder_layer in enumerate(self.layers): if output_hidden_states: all_hidden_states += (hidden_states,) past_key_value = past_key_values[idx] if past_key_values is not None else None - if self.gradient_checkpointing and self.training: + if offloaded_gradient_checkpointing: + hidden_states = Offloaded_Gradient_Checkpointer.apply( + decoder_layer, + hidden_states, + causal_mask, + attention_mask, + position_ids, + past_key_values, + output_attentions, + use_cache, + ) + elif gradient_checkpointing: def create_custom_forward(module): def custom_forward(*inputs): - # None for past_key_value - return module(*inputs, past_key_value, output_attentions, padding_mask=padding_mask) - + return module(*inputs, past_key_value, output_attentions, padding_mask = padding_mask) return custom_forward + pass layer_outputs = torch.utils.checkpoint.checkpoint( create_custom_forward(decoder_layer), @@ -648,9 +671,11 @@ def LlamaModel_fast_forward( causal_mask, attention_mask, position_ids, - use_reentrant=True, - preserve_rng_state=False, + use_reentrant = True, + preserve_rng_state = False, ) + hidden_states = layer_outputs[0] + else: layer_outputs = decoder_layer( hidden_states, @@ -662,9 +687,9 @@ def LlamaModel_fast_forward( use_cache=use_cache, padding_mask=padding_mask, ) + hidden_states = layer_outputs[0] pass - hidden_states = layer_outputs[0] if use_cache: next_decoder_cache += (layer_outputs[2 if output_attentions else 1],) if output_attentions: all_self_attns += (layer_outputs[1],) pass @@ -801,12 +826,12 @@ def CausalLM_fast_forward(fast_forward_inference): hidden_states = outputs[0] bsz, q_len, hd = hidden_states.shape + lm_head = self.lm_head.weight if bsz == 1 and q_len == 1: - lm_head = self.lm_head.weight logits = torch.mv(lm_head, hidden_states.ravel().to(lm_head.dtype)) logits = logits.unsqueeze(0).unsqueeze(0) else: - logits = self.lm_head(hidden_states) + logits = self.lm_head(hidden_states.to(lm_head.dtype)) pass logits = logits.to(self.config.torch_dtype) @@ -1402,6 +1427,8 @@ class FastLlamaModel: "We shall do it for you!" ) train_lm_head = True + if modules_to_save is None: modules_to_save = ["lm_head"] + else: modules_to_save.append("lm_head") elif module == "embed_tokens": logger.warning_once( @@ -1409,6 +1436,8 @@ class FastLlamaModel: "We shall do it for you!" ) train_embed_tokens = True + if modules_to_save is None: modules_to_save = ["embed_tokens"] + else: modules_to_save.append("embed_tokens") else: assert(module in accepted_modules) diff --git a/unsloth/models/mistral.py b/unsloth/models/mistral.py index e867ceef23..87f5c85ad1 100644 --- a/unsloth/models/mistral.py +++ b/unsloth/models/mistral.py @@ -225,12 +225,12 @@ def MistralForCausalLM_fast_forward( hidden_states = outputs[0] bsz, q_len, hd = hidden_states.shape + lm_head = self.lm_head.weight if bsz == 1 and q_len == 1: - lm_head = self.lm_head.weight logits = torch.mv(lm_head, hidden_states.ravel().to(lm_head.dtype)) logits = logits.unsqueeze(0).unsqueeze(0) else: - logits = self.lm_head(hidden_states) + logits = self.lm_head(hidden_states.to(lm_head.dtype)) pass logits = logits.to(self.config.torch_dtype) diff --git a/unsloth/save.py b/unsloth/save.py index a08a744077..49d88bffc0 100644 --- a/unsloth/save.py +++ b/unsloth/save.py @@ -33,9 +33,13 @@ __all__ = [ "patch_saving_functions", ] -# Check Kaggle -IS_A_KAGGLE_ENVIRONMENT = "KAGGLE_CONTAINER_NAME" in os.environ +# Check environments +keynames = "\n" + "\n".join(os.environ.keys()) +IS_COLAB_ENVIRONMENT = "\nCOLAB_" in keynames +IS_KAGGLE_ENVIRONMENT = "\nKAGGLE_" in keynames +del keynames +# Weights LLAMA_WEIGHTS = ( "self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj", "self_attn.o_proj", "mlp.gate_proj", "mlp.up_proj", "mlp.down_proj", @@ -177,6 +181,9 @@ def unsloth_save_model( temporary_location : str = "_unsloth_temporary_saved_buffers", maximum_memory_usage : float = 0.9, ): + if token is None and "HF_TOKEN" in os.environ: + token = os.environ["HF_TOKEN"] + if commit_message is None: commit_message = "" if "Unsloth" not in commit_message: commit_message += " (Trained with Unsloth)" @@ -291,6 +298,10 @@ def unsloth_save_model( tags = tags, ) if tokenizer is not None: + # Set padding side to left for inference + old_padding_side = tokenizer.padding_side + tokenizer.padding_side = "left" + getattr(tokenizer, "original_push_to_hub", tokenizer.push_to_hub)\ ( repo_id = save_directory, @@ -305,6 +316,9 @@ def unsloth_save_model( commit_description = commit_description, tags = tags, ) + + # Revert back padding side + tokenizer.padding_side = old_padding_side pass if hasattr(model, "config"): @@ -361,7 +375,16 @@ def unsloth_save_model( if tokenizer is not None: print("Unsloth: Saving tokenizer...", end = "") + + # Set padding side to left for inference + old_padding_side = tokenizer.padding_side + tokenizer.padding_side = "left" + tokenizer.save_pretrained(**tokenizer_save_settings) + + # Revert back padding side + tokenizer.padding_side = old_padding_side + print(" Done.") else: print() @@ -449,12 +472,12 @@ def unsloth_save_model( os.makedirs(temporary_location) pass - # Check if Kaggle, since only 20GB of Disk space allowed. - if IS_A_KAGGLE_ENVIRONMENT: + # Check if Kaggle or Colab, since only 20GB of Disk space allowed. + if IS_KAGGLE_ENVIRONMENT or IS_COLAB_ENVIRONMENT: # We free up 4GB of space logger.warning_once( - "Unsloth: Kaggle only allows 20GB of disk space. We need to delete the downloaded\n"\ - "model which will save 4GB of disk space, allowing you to save on Kaggle." + "Unsloth: Kaggle/Colab has limited disk space. We need to delete the downloaded\n"\ + "model which will save 4-16GB of disk space, allowing you to save on Kaggle/Colab." ) _free_cached_model(internal_model) pass @@ -462,7 +485,10 @@ def unsloth_save_model( # HF also uses a OrderedDict from collections import OrderedDict state_dict = OrderedDict() - state_dict["model.embed_tokens.weight"] = internal_model.model.embed_tokens.weight.data + + torch_dtype = model.config.torch_dtype + # Check modules to save float32 dtype + state_dict["model.embed_tokens.weight"] = internal_model.model.embed_tokens.weight.data.to(torch_dtype) max_vram = int(torch.cuda.get_device_properties(0).total_memory * maximum_memory_usage) @@ -495,7 +521,8 @@ def unsloth_save_model( pass state_dict["model.norm.weight"] = internal_model.model.norm.weight.data - state_dict["lm_head.weight"] = internal_model.lm_head.weight.data + # Check for modules_to_save float32 dtype + state_dict["lm_head.weight"] = internal_model.lm_head.weight.data.to(torch_dtype) # All tensors MUST be type torch.Tensor and not torch.nn.parameter.Parameter for key, value in state_dict.items(): @@ -552,7 +579,16 @@ def unsloth_save_model( # Save tokenizer if tokenizer is not None: print("Unsloth: Saving tokenizer...", end = "") + + # Set padding side to left for inference + old_padding_side = tokenizer.padding_side + tokenizer.padding_side = "left" + tokenizer.save_pretrained(**tokenizer_save_settings) + + # Revert back padding side + tokenizer.padding_side = old_padding_side + print(" Done.") else: print() @@ -1216,7 +1252,7 @@ def unsloth_save_pretrained_gguf( # Non blocking install GGUF first if not os.path.exists("llama.cpp"): - if IS_A_KAGGLE_ENVIRONMENT: + if IS_KAGGLE_ENVIRONMENT: # Kaggle is weird - no blocking installs, and no CUDA? python_install = install_python_non_blocking(["gguf", "protobuf"]) python_install.wait() @@ -1237,7 +1273,7 @@ def unsloth_save_pretrained_gguf( makefile = None except: # Retry by recloning llama.cpp - if IS_A_KAGGLE_ENVIRONMENT: + if IS_KAGGLE_ENVIRONMENT: # Kaggle is weird - no blocking installs, and no CUDA? python_install = install_python_non_blocking(["gguf", "protobuf"]) python_install.wait() @@ -1336,7 +1372,7 @@ def unsloth_push_to_hub_gguf( # Non blocking install GGUF first if not os.path.exists("llama.cpp"): - if IS_A_KAGGLE_ENVIRONMENT: + if IS_KAGGLE_ENVIRONMENT: # Kaggle is weird - no blocking installs, and no CUDA? python_install = install_python_non_blocking(["gguf", "protobuf"]) python_install.wait() @@ -1357,7 +1393,7 @@ def unsloth_push_to_hub_gguf( makefile = None except: # Retry by recloning llama.cpp - if IS_A_KAGGLE_ENVIRONMENT: + if IS_KAGGLE_ENVIRONMENT: # Kaggle is weird - no blocking installs, and no CUDA? python_install = install_python_non_blocking(["gguf", "protobuf"]) python_install.wait() diff --git a/unsloth/tokenizer_utils.py b/unsloth/tokenizer_utils.py index 00e937c97c..46de1c98ad 100644 --- a/unsloth/tokenizer_utils.py +++ b/unsloth/tokenizer_utils.py @@ -186,9 +186,6 @@ def assert_same_tokenization(slow_tokenizer, fast_tokenizer): pass -global sentencepiece_model_pb2 -sentencepiece_model_pb2 = None - def fix_sentencepiece_tokenizer( old_tokenizer, new_tokenizer, @@ -197,19 +194,7 @@ def fix_sentencepiece_tokenizer( ): # From https://github.com/google/sentencepiece/issues/121 # We need to manually edit the sentencepiece tokenizer! - global sentencepiece_model_pb2 - if sentencepiece_model_pb2 is None: - try: - import sentencepiece.sentencepiece_model_pb2 as _sentencepiece_model_pb2 - sentencepiece_model_pb2 = _sentencepiece_model_pb2 - except: - if not os.path.exists(temporary_location): - os.system(f"git clone https://github.com/google/sentencepiece.git {temporary_location}") - os.system(f"cd {temporary_location}/src && protoc --python_out=. sentencepiece_model.proto") - pass - import sentencepiece.sentencepiece_model_pb2 as _sentencepiece_model_pb2 - sentencepiece_model_pb2 = _sentencepiece_model_pb2 - pass + from transformers.utils import sentencepiece_model_pb2 if not os.path.exists(temporary_location): os.makedirs(temporary_location)