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

Train on responses only (#770)

Browse source 
* Update gemma2.py

* Update llama.py

* Update llama.py

* Update gemma2.py

* init

* Update gemma2.py

* Update gemma2.py

* Update _utils.py

* Update _utils.py

* Update _utils.py

* Update _utils.py

* Update _utils.py

* Update gemma2.py

* Update gemma2.py

* Update gemma2.py

* All RoPE Scaling support

* cleanup

* Update llama.py

* Update llama.py

* Update _utils.py

* Update _utils.py

* exec

* exec

* Attention_Module

* attention_module

* imports

* exec

* Update llama.py

* Update llama.py

* boolean mask

* revert masking

* Update llama.py

* Update save.py

* Update llama.py

* Update gemma2.py

* Update gemma2.py

* Update gemma2.py

* Update utils.py

* retry

* Update gemma2.py

* Update gemma2.py

* Update gemma2.py

* Update _utils.py

* Update _utils.py

* Update gemma2.py

* Update chat_templates.py

* Gemma 2 Ollama support

* Update llama.py

* Update llama.py

* error handling

* Update _utils.py

* Update _utils.py

* Stats for debugging

* Update _utils.py

* Update _utils.py

* Debugging

* Update tokenizer_utils.py

* Update _utils.py

* Update cross_entropy_loss.py

* Update cross_entropy_loss.py

* Update cross_entropy_loss.py

* Update rms_layernorm.py

* Update rms_layernorm.py

* Update rms_layernorm.py

* Update rms_layernorm.py

* Update rms_layernorm.py

* Update rms_layernorm.py

* Update rms_layernorm.py

* Check exec, eval

* Update _utils.py

* Update _utils.py

* Images

* Bug fixes

* Update pyproject.toml

* Bug fixes

* Update _utils.py

* Update _utils.py

* Deprecation fix

* Update chat_templates.py

* Now permitting use of pre-installed llama.cpp (#763)

* Now permitting use of pre-installed llama.cpp

* Update save.py

---------

Co-authored-by: Giuseppe Strafforello <giuseppe.strafforello@titantechnologies.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>

* Update save.py

* Deprecation & compile

* typo

* Update chat_templates.py

* Update chat_templates.py

* train_on_responses_only

* Update llama.py

* Update llama.py

* Update save.py

* Update gemma2.py

---------

Co-authored-by: pepistrafforello <pepi.strafforello@gmail.com>
Co-authored-by: Giuseppe Strafforello <giuseppe.strafforello@titantechnologies.com>
This commit is contained in:
Daniel Han 2024-07-14 22:41:04 -07:00 • committed by GitHub
commit 7b3b216fd3
6 changed files with 280 additions and 81 deletions
Download patch file Download diff file Expand all files Collapse all files

166
unsloth/chat_templates.py
View file

@ -21,6 +21,7 @@ __all__ = [
"to_sharegpt",
"standardize_sharegpt",
"apply_chat_template",
"train_on_responses_only",
"test_construct_chat_template",
]
@ -1063,7 +1064,6 @@ default_system_message = \
"Below are some instructions that describe some tasks. Write responses that appropriately complete each request.",
extra_eos_tokens = None,
):
"""
Creates a Ollama modelfile and a HF Jinja template from a custom
@ -1072,6 +1072,9 @@ extra_eos_tokens = None,
You must use {INPUT}, {OUTPUT} twice, and {SYSTEM} is optional.
"""
# Strip only the left
chat_template = chat_template.lstrip()
assert(tokenizer is not None)
if extra_eos_tokens is None: extra_eos_tokens = []
@ -1128,19 +1131,47 @@ extra_eos_tokens = None,
chat_template = re.sub(r"{OUTPUT}", r"{OUTPUT}" + eos, chat_template)
pass
# O(N^2) search finding 2 repeatted pieces of text
j = len(chat_template)-1
at_least_one = False
while j > 0:
found = chat_template.rfind(chat_template[j:], 0, j)
if found == -1: break
j -= 1
at_least_one = True
pass
if j > 0: j += 1
else: raise RuntimeError(error_msg)
# This forces you to provide 2 input and outputs
final_combined_check = False
if not at_least_one: raise RuntimeError(error_msg)
try:
# O(N^2) search finding 2 repeatted pieces of text
j = len(chat_template)-1
at_least_one = False
while j > 0:
found = chat_template.rfind(chat_template[j:], 0, j)
if found == -1: break
j -= 1
at_least_one = True
pass
if j > 0: j += 1
else: raise RuntimeError(error_msg)
if not at_least_one: raise RuntimeError(error_msg)
# Must be equivalent to left
final_combined_check = True
except:
# Simple 1 singular input and output
system_count = chat_template.count("{SYSTEM}")
input_count = chat_template.count("{INPUT}")
output_count = chat_template.count("{OUTPUT}")
if system_count > 1:
raise RuntimeError("You must only provide 1 {SYSTEM} in the chat template")
if input_count > 1:
raise RuntimeError("You must only provide 1 {INPUT} in the chat template")
if output_count > 1:
raise RuntimeError("You must only provide 1 {OUTPUT} in the chat template")
if system_count != 0:
j = next(re.finditer(r"\{SYSTEM\}[\s]{0,}", chat_template)).span(0)[1]
else:
j = 0
pass
# Must be equivalent to the original text
final_combined_check = False
pass
# Repeatted text
instruction_response = chat_template[j:]
@ -1153,6 +1184,8 @@ extra_eos_tokens = None,
# 2nd Instruction, Output pair
right = chat_template[j:]
final_combined_check = left if final_combined_check else chat_template
# Isolate input
extra_eos_tokens_regex = "|".join(f"(?:{re.escape(x)})" for x in extra_eos_tokens)
if len(extra_eos_tokens_regex) != 0:
@ -1170,13 +1203,14 @@ extra_eos_tokens = None,
output_part = right[input_end:]
# Isolate system
system_part = left[:left.find(input_part)]
where_system = left.find(input_part)
system_part = left[:where_system if where_system != -1 else len(left)]
# Check if the user provided a correct prompt
combined = system_part + input_part + output_part
if combined != left:
combined_changed = combined.replace('\n', '\\n')
left_changed = left .replace('\n', '\\n')
if combined != final_combined_check:
combined_changed = combined .replace('\n', '\\n')
left_changed = final_combined_check.replace('\n', '\\n')
raise RuntimeError(
"Unsloth: The prompt template you provided isn't correct. You gave:\n"\
f"{combined_changed}\n\n"\
@ -1285,6 +1319,15 @@ extra_eos_tokens = None,
jinja_template = "{{ bos_token }}" + jinja_template
pass
# Fix missing loop_messages
if "{% set loop_messages = messages %}" not in jinja_template:
jinja_template = jinja_template.replace(
"{% for message in loop_messages %}",
"{% for message in messages %}",
1, # Only replace the first one
)
pass
# Check if system part is the same!
jinja_template = re.sub(
r"\{\% if messages\[0\]\['role'\] \=\= 'system' \%\}\{\{ '(.+?)' \}\}"\
@ -1300,8 +1343,11 @@ extra_eos_tokens = None,
if not jinja_template.startswith("{{ bos_token }}"):
jinja_template = "{{ bos_token }}" + jinja_template
pass
return modelfile, jinja_template
# Get instruction and output parts for train_on_inputs = False
input_part = input_part [:input_part .find("{INPUT}")]
output_part = output_part[:output_part.find("{OUTPUT}")]
return modelfile, jinja_template, input_part, output_part
pass
@ -1327,7 +1373,7 @@ def test_construct_chat_template():
extra_eos_tokens = None
modelfile, jinja_template = construct_chat_template(
modelfile, jinja_template, _, _ = construct_chat_template(
tokenizer = tokenizer,
chat_template = chat_template,
extra_eos_tokens = extra_eos_tokens,
@ -1380,7 +1426,7 @@ extra_eos_tokens = None,
You must use {INPUT}, {OUTPUT} twice, and {SYSTEM} is optional.
"""
modelfile, jinja_template = construct_chat_template(
modelfile, jinja_template, input_part, output_part = construct_chat_template(
tokenizer = tokenizer,
chat_template = chat_template,
default_system_message = default_system_message,
@ -1391,12 +1437,90 @@ extra_eos_tokens = None,
texts = [tokenizer.apply_chat_template(convo, tokenize = False, add_generation_prompt = False) for convo in convos]
return { "text" : texts, }
pass
tokenizer.chat_template = jinja_template
tokenizer._ollama_modelfile = modelfile
tokenizer._unsloth_input_part = input_part
tokenizer._unsloth_output_part = output_part
return dataset.map(formatting_prompts_func, batched = True,)
pass
def train_on_responses_only(
trainer,
instruction_part = None,
response_part = None,
):
"""
Trains only on responses and not on the instruction by masking out
the labels with -100 for the instruction part.
"""
tokenizer = trainer.tokenizer
if not hasattr(tokenizer, "_unsloth_input_part") or \
not hasattr(tokenizer, "_unsloth_output_part"):
if instruction_part is None or response_part is None:
raise ValueError("Unsloth: instruction_part and response_part must be given!")
pass
elif (instruction_part is not None or response_part is not None) and \
(hasattr(tokenizer, "_unsloth_input_part") or hasattr(tokenizer, "_unsloth_output_part")):
raise ValueError("Unsloth: Your tokenizer already has instruction and response parts set - do not give custom ones!")
else:
instruction_part = tokenizer._unsloth_input_part
response_part = tokenizer._unsloth_output_part
pass
instruction_ids = tokenizer(instruction_part, add_special_tokens = False).input_ids
response_ids = tokenizer(response_part, add_special_tokens = False).input_ids
instruction_length = len(instruction_ids)
response_length = len(response_ids)
max_length = max(instruction_length, response_length)
def _train_on_responses_only(examples):
input_ids_ = examples["input_ids"]
all_labels = []
for input_ids in input_ids_:
labels = [-100] * len(input_ids)
m = len(input_ids) - max_length
first_response = response_ids[0]
first_instruction = instruction_ids[0]
j = 0
while j < m:
if input_ids[j] == first_response:
if input_ids[j : j+response_length] == response_ids:
j = j + response_length
start = j
while j < m:
if input_ids[j] == first_instruction and input_ids[j : j+instruction_length] == instruction_ids:
j = j + instruction_length
labels[start : j] = input_ids[start : j]
break
elif j == (m-1):
j = m
labels[start:] = input_ids[start:]
break
pass
j += 1
pass
pass
pass
j += 1
pass
all_labels.append(labels)
pass
return { "labels" : all_labels }
pass
trainer.train_dataset = trainer.train_dataset.map(_train_on_responses_only, batched = True)
return trainer
pass
def create_stopping_criteria(tokenizer, stop_word = "eos_token"):
class StoppingCriteriaSub(StoppingCriteria):
__slots__ = "stop_token", "single_match", "length",

14
unsloth/kernels/fast_lora.py
View file

@ -19,6 +19,8 @@ from .utils import (
get_lora_parameters,
get_lora_parameters_bias,
matmul_lora,
torch_amp_custom_fwd,
torch_amp_custom_bwd,
)
@ -61,7 +63,7 @@ class LoRA_MLP(torch.autograd.Function):
Don't forget to see our blog post for more details!
"""
@staticmethod
@torch.cuda.amp.custom_fwd
@torch_amp_custom_fwd
def forward(ctx, X : torch.Tensor,
gateW, gateW_quant, gateA, gateB, gateS,
upW, upW_quant, upA, upB, upS,
@ -87,7 +89,7 @@ class LoRA_MLP(torch.autograd.Function):
@staticmethod
@torch.cuda.amp.custom_bwd
@torch_amp_custom_bwd
def backward(ctx, dY : torch.Tensor):
gateW, gateW_quant, gateS, upW, upW_quant, upS, downW, downW_quant, downS, \
_backward_function = ctx.custom_saved_tensors
@ -223,7 +225,7 @@ class LoRA_QKV(torch.autograd.Function):
dC/dBv = A.T @ X.T @ D(Wv)
"""
@staticmethod
@torch.cuda.amp.custom_fwd
@torch_amp_custom_fwd
def forward(ctx, X : torch.Tensor,
QW, QW_quant, QA, QB, QS,
KW, KW_quant, KA, KB, KS,
@ -244,7 +246,7 @@ class LoRA_QKV(torch.autograd.Function):
pass
@staticmethod
@torch.cuda.amp.custom_bwd
@torch_amp_custom_bwd
def backward(ctx, dQ, dK, dV):
QW, QW_quant, QS, KW, KW_quant, KS, VW, VW_quant, VS = \
ctx.custom_saved_tensors
@ -352,7 +354,7 @@ class LoRA_W(torch.autograd.Function):
dC/dBv = A.T @ X.T @ D(Wv)
"""
@staticmethod
@torch.cuda.amp.custom_fwd
@torch_amp_custom_fwd
def forward(ctx, X : torch.Tensor,
W, W_quant, A, B, S):
dtype = X.dtype
@ -363,7 +365,7 @@ class LoRA_W(torch.autograd.Function):
pass
@staticmethod
@torch.cuda.amp.custom_bwd
@torch_amp_custom_bwd
def backward(ctx, dY : torch.Tensor):
W, W_quant, S = ctx.custom_saved_tensors
A, B, X = ctx.saved_tensors

13
unsloth/kernels/utils.py
View file

@ -16,6 +16,18 @@ import triton
MAX_FUSED_SIZE = 65536
next_power_of_2 = triton.next_power_of_2
# torch.cuda.amp.custom_fwd is deprecated >= 2.4
import torch
from packaging.version import Version
if Version(torch.__version__) < Version("2.4.0"):
torch_amp_custom_fwd = torch.cuda.amp.custom_fwd
torch_amp_custom_bwd = torch.cuda.amp.custom_bwd
else:
torch_amp_custom_fwd = torch.amp.custom_fwd(device_type = "cuda")
torch_amp_custom_bwd = torch.amp.custom_bwd(device_type = "cuda")
pass
def calculate_settings(n):
BLOCK_SIZE = next_power_of_2(n)
if BLOCK_SIZE > MAX_FUSED_SIZE:
@ -32,7 +44,6 @@ pass
import bitsandbytes as bnb
get_ptr = bnb.functional.get_ptr
import ctypes
import torch
cdequantize_blockwise_fp32 = bnb.functional.lib.cdequantize_blockwise_fp32
cdequantize_blockwise_fp16_nf4 = bnb.functional.lib.cdequantize_blockwise_fp16_nf4
cdequantize_blockwise_bf16_nf4 = bnb.functional.lib.cdequantize_blockwise_bf16_nf4

60
unsloth/models/_utils.py
View file

@ -35,6 +35,8 @@ __all__ = [
"patch_linear_scaling",
"check_nvidia",
"create_boolean_mask",
"torch_amp_custom_fwd",
"torch_amp_custom_bwd",
]
import torch
@ -92,6 +94,19 @@ for model_name in model_architectures:
pass
# =============================================
# =============================================
# torch.cuda.amp.custom_fwd is deprecated >= 2.4
import torch
from packaging.version import Version
if Version(torch.__version__) < Version("2.4.0"):
torch_amp_custom_fwd = torch.cuda.amp.custom_fwd
torch_amp_custom_bwd = torch.cuda.amp.custom_bwd
else:
torch_amp_custom_fwd = torch.amp.custom_fwd(device_type = "cuda")
torch_amp_custom_bwd = torch.amp.custom_bwd(device_type = "cuda")
pass
# =============================================
# =============================================
# Get Flash Attention v2 if Ampere (RTX 30xx, A100)
import bitsandbytes as bnb
@ -176,11 +191,22 @@ torch_compile_arguments = [
"config.cuda.use_fast_math = True",
"config.cuda.compile_opt_level = '-O2'",
]
# Torch dynamo arguments
torch_dynamo_arguments = [
"config.accumulated_cache_size_limit = 512", # Bump up a bit from 256
"config.suppress_errors = True", # Supress errors for now
"config.do_not_emit_runtime_asserts = True",
]
import torch._inductor.config as config
for _try_compile_argument in torch_compile_arguments:
try: exec(_try_compile_argument)
except: pass
pass
import torch._dynamo.config as config
for _try_dynamo_argument in torch_dynamo_arguments:
try: exec(_try_dynamo_argument)
except: pass
pass
torch_compile_options = {
"epilogue_fusion" : True,
"max_autotune" : True,
@ -358,15 +384,13 @@ except:
pass
# =============================================
def _get_statistics(statistics = None):
import psutil
def _get_statistics(statistics = None, force_download = True):
# We log some basic stats about which environment is being used.
# We simply download a README.md file from HF - all data is made public.
# This is simply so we can check if some envs are broken or not.
# You can disable this by commenting the below out
try:
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())
@ -382,21 +406,12 @@ def _get_statistics(statistics = None):
else: statistics = "other"
if statistics is not None:
disabled = False
if not are_progress_bars_disabled():
disable_progress_bars()
disabled = True
pass
from transformers import AutoModelForCausalLM
stats_model = AutoModelForCausalLM.from_pretrained(
f"unslothai/{statistics}",
force_download = True,
force_download = force_download,
)
del stats_model
if disabled:
enable_progress_bars()
pass
pass
except:
pass
@ -408,7 +423,14 @@ def get_statistics():
# We simply download a README.md file from HF - all data is made public.
# This is simply so we can check if some envs are broken or not.
# You can disable this by commenting the below out
from huggingface_hub.utils import disable_progress_bars, enable_progress_bars, are_progress_bars_disabled
disabled = False
if not are_progress_bars_disabled():
disable_progress_bars()
disabled = True
pass
_get_statistics(None)
_get_statistics("repeat", force_download = False)
try:
vram = torch.cuda.get_device_properties(0).total_memory / 1024 / 1024 / 1024
if vram <= 8 : vram = 8
@ -423,6 +445,12 @@ def get_statistics():
except:
pass
pass
try:
devices = torch.cuda.device_count()
_get_statistics(f"{devices if devices <= 8 else 9}")
except:
pass
if disabled: enable_progress_bars()
pass
@ -517,7 +545,7 @@ class Unsloth_Offloaded_Gradient_Checkpointer(torch.autograd.Function):
Tiny hit to performance, since we mask the movement via non blocking calls.
"""
@staticmethod
@torch.cuda.amp.custom_fwd
@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():
@ -529,7 +557,7 @@ class Unsloth_Offloaded_Gradient_Checkpointer(torch.autograd.Function):
pass
@staticmethod
@torch.cuda.amp.custom_bwd
@torch_amp_custom_bwd
def backward(ctx, dY):
(hidden_states,) = ctx.saved_tensors
hidden_states = hidden_states.to("cuda:0", non_blocking = True).detach()

11
unsloth/models/gemma2.py
View file

@ -70,6 +70,17 @@ def fast_rms_layernorm_gemma2_compiled(layernorm, X, gemma = True):
pass
# Flex Attention in torch 2.5 and higher
# try:
# from torch.nn.attention._flex_attention import _flex_attention
# from functools import lru_cache
# @lru_cache
# def create_block_mask_from_score_mod(score_mod, B, H, M, N):
# SPARSE_BLOCK = 128
# block_mask = _create_block_mask(score_mod, B, H, M, N, device = "cuda:0")
# return block_mask
# Logit softcapping
@torch.compile(fullgraph = True, dynamic = True, options = torch_compile_options)
def gemma2_attention(Q, K, V, causal_mask, self, bsz, q_len):

97
unsloth/save.py
View file

@ -840,6 +840,21 @@ def install_llama_cpp_blocking(use_cuda = False):
pass
def get_executable(executables):
# Get system locations (System Path).split(system separator)
system_directories = os.environ.get("PATH").split(os.pathsep)
for directory in system_directories:
for executable in executables:
path = os.path.join(directory, executable)
# Check if the executable exists and is executable
if os.path.exists(path) and os.access(path, os.X_OK): return path
pass
pass
return None
pass
def save_to_gguf(
model_type : str,
model_dtype : str,
@ -932,48 +947,56 @@ def save_to_gguf(
)
pass
print("Unsloth: [0] Installing llama.cpp. This will take 3 minutes...")
if _run_installer is not None:
error = _run_installer.wait()
# Determine whether the system already has llama.cpp installed and the scripts are executable
quantize_location = get_executable(["llama-quantize", "quantize"])
convert_location = get_executable(["convert-hf-to-gguf.py", "convert_hf_to_gguf.py"])
if quantize_location is not None and convert_location is not None:
print("Unsloth: llama.cpp found in the system. We shall skip installation.")
else:
error = 0
install_llama_cpp_blocking()
pass
print("Unsloth: [0] Installing llama.cpp. This will take 3 minutes...")
if _run_installer is not None:
error = _run_installer.wait()
else:
error = 0
install_llama_cpp_blocking()
pass
# Check if successful. If not install 10th latest release
# Check if successful. If not install 10th latest release
# Careful llama.cpp/quantize changed to llama.cpp/llama-quantize
# and llama.cpp/main changed to llama.cpp/llama-cli
# See https://github.com/ggerganov/llama.cpp/pull/7809
quantize_location = None
if os.path.exists("llama.cpp/quantize"):
quantize_location = "llama.cpp/quantize"
elif os.path.exists("llama.cpp/llama-quantize"):
quantize_location = "llama.cpp/llama-quantize"
else:
raise RuntimeError(
"Unsloth: The file 'llama.cpp/llama-quantize' or 'llama.cpp/quantize' does not exist.\n"\
"But we expect this file to exist! Maybe the llama.cpp developers changed the name?"
)
pass
# Careful llama.cpp/quantize changed to llama.cpp/llama-quantize
# and llama.cpp/main changed to llama.cpp/llama-cli
# See https://github.com/ggerganov/llama.cpp/pull/7809
quantize_location = None
if os.path.exists("llama.cpp/quantize"):
quantize_location = "llama.cpp/quantize"
elif os.path.exists("llama.cpp/llama-quantize"):
quantize_location = "llama.cpp/llama-quantize"
else:
raise RuntimeError(
"Unsloth: The file 'llama.cpp/llama-quantize' or 'llama.cpp/quantize' does not exist.\n"\
"But we expect this file to exist! Maybe the llama.cpp developers changed the name?"
)
pass
# See https://github.com/unslothai/unsloth/pull/730
# Filenames changed again!
convert_location = None
if os.path.exists("llama.cpp/convert-hf-to-gguf.py"):
convert_location = "llama.cpp/convert-hf-to-gguf.py"
elif os.path.exists("llama.cpp/convert_hf_to_gguf.py"):
convert_location = "llama.cpp/convert_hf_to_gguf.py"
else:
raise RuntimeError(
"Unsloth: The file 'llama.cpp/convert-hf-to-gguf.py' or 'llama.cpp/convert_hf_to_gguf.py' does not exist.\n"\
"But we expect this file to exist! Maybe the llama.cpp developers changed the name?"
)
pass
# See https://github.com/unslothai/unsloth/pull/730
# Filenames changed again!
convert_location = None
if os.path.exists("llama.cpp/convert-hf-to-gguf.py"):
convert_location = "llama.cpp/convert-hf-to-gguf.py"
elif os.path.exists("llama.cpp/convert_hf_to_gguf.py"):
convert_location = "llama.cpp/convert_hf_to_gguf.py"
else:
raise RuntimeError(
"Unsloth: The file 'llama.cpp/convert-hf-to-gguf.py' or 'llama.cpp/convert_hf_to_gguf.py' does not exist.\n"\
"But we expect this file to exist! Maybe the llama.cpp developers changed the name?"
)
pass
if error != 0 or quantize_location is None or convert_location is None:
print(f"Unsloth: llama.cpp error code = {error}.")
install_llama_cpp_old(-10)
if error != 0 or quantize_location is None or convert_location is None:
print(f"Unsloth: llama.cpp error code = {error}.")
install_llama_cpp_old(-10)
pass
pass
# Determine maximum first_conversion state