Bug fixes (#306)
* Approx gelu * Update llama.py * Update __init__.py * Update __init__.py * Update _utils.py * Update geglu.py * Update gemma.py * Update rms_layernorm.py * Update rms_layernorm.py * Update rms_layernorm.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Fix Gemma merging * Update rms_layernorm.py * Update gemma.py * Update pyproject.toml * Layernorms * Gemma precision * Update gemma.py * sqrt * Update gemma.py * Update save.py * RoPE and Gemma precision * Update rms_layernorm.py * Fix warning * Update chat_templates.py * Update chat_templates.py * Update save.py * Update save.py * Update save.py * Update chat_templates.py * Update llama.py * model_name * Update loader.py * Tokenizer overwritten * Update llama.py * Update llama.py * Update llama.py * Update save.py * Accuracy * Revert * Update save.py * Update fast_lora.py * Update fast_lora.py * Update fast_lora.py * Update fast_lora.py * Update fast_lora.py * Update chat_templates.py * Update save.py * Update save.py * Update llama.py * Update llama.py * Account for DoRA * Update llama.py * Update save.py * GGUF incorrect * Update save.py * Update pyproject.toml * kaggle new * Update pyproject.toml * Update pyproject.toml * upcasting * Fix Colab * Update pyproject.toml * Update pyproject.toml * Update pyproject.toml * Update pyproject.toml * Update pyproject.toml * Update pyproject.toml * Update pyproject.toml * Update pyproject.toml * Update chat_templates.py * Update chat_templates.py * Update chat_templates.py * Update chat_templates.py * Update chat_templates.py * Update pyproject.toml * Update pyproject.toml * Update pyproject.toml * Update rope_embedding.py * Update rope_embedding.py * Fix bugs * Update fast_lora.py * Update fast_lora.py * Update README.md * Update README.md * GGUF * Update save.py * Update save.py * Update save.py * Update save.py * Update README.md * Update README.md * Bugs * Update fast_lora.py * Update pyproject.toml * Update fast_lora.py * Update __init__.py * Update fast_lora.py * dtype * Update llama.py * Update llama.py * Update llama.py * dtype * Update mistral.py * trust_remote_code * lm_head * Update llama.py * save_pretrained_settings * Update save.py * Update save.py * Update save.py * Update save.py * Update save.py * Update save.py * Update save.py * Update save.py * Update save.py * Update save.py * Update save.py * Update save.py * state_dict * Update save.py * whoami * Update llama.py * Update save.py * Update llama.py * Patch tokenizer * Update chat_templates.py * Heal tokenizers * Update chat_templates.py * Update mapper.py * Update tokenizer_utils.py * Update tokenizer_utils.py * Update tokenizer_utils.py * Update tokenizer_utils.py * Update tokenizer_utils.py * Update chat_templates.py * tokenizer patching * patch_tokenizer * Update chat_templates.py * Update tokenizer_utils.py * Update chat_templates.py * Update chat_templates.py * Update chat_templates.py * Update tokenizer_utils.py * Edit * Update mistral.py * Update mistral.py * Stats * Update mistral.py * attention_mask * Update llama.py * Update llama.py * batch * Temp fix batch inference * Update llama.py * Update gemma.py * Fix inference * swiglu * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update mistral.py * Update llama.py * fast inference * model * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update utils.py * Update llama.py * Update utils.py * inference * Update llama.py * Update llama.py * Update llama.py * overhead * Update llama.py * Update llama.py * compile * Update gemma.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update utils.py * Update utils.py * lora mamtul * Update llama.py * Update llama.py * Update llama.py * offloaded checkpointing * Update llama.py * Update llama.py * Update _utils.py * Update _utils.py * Update _utils.py * Update llama.py * Update llama.py * Update gemma.py * Revert "Update gemma.py" This reverts commit e3c3c5f3fa3d04a87f854056f6b547ced610d712. * Update _utils.py * Update _utils.py * Update _utils.py * Saving * sentencepiece_model_pb2 * Update llama.py * Update save.py * Update llama.py * padding side
This commit is contained in:
Daniel Han
GitHub
parent
b5200a33f1
commit
4290c8a278
5 changed files with 140 additions and 42 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue