ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Utils

Browse source
This commit is contained in:
Daniel Han 2024-11-02 19:08:57 -07:00
commit e097ba3e7b
2 changed files with 11 additions and 359 deletions
Download patch file Download diff file Expand all files Collapse all files

1
unsloth/kernels/cross_entropy_loss.py
View file

@ -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"],

369
unsloth/models/_utils.py
View file

@ -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 <unk>
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
"<pad>", # 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|>",
"<s>",
"</s>",
)
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():