Bug fixes (#1245)
* Fix TRL * Update mistral.py * Patch processing_class * Update tokenizer_utils.py * Update tokenizer_utils.py * Update tokenizer_utils.py * Update tokenizer_utils.py * Update tokenizer_utils.py * Update tokenizer_utils.py * Installation guide (#1165) * chore: update chat_templates.py (#1166) orginal -> original * Disable Flex Attention * Update tokenizer_utils.py * Update _utils.py * n_items * Update cross_entropy_loss.py * Fix DPO, ORPO * Update _utils.py * Update _utils.py * fix/transformers-unpack (#1180) * Fix DPO, ORPO (#1177) * Fix TRL * Update mistral.py * Patch processing_class * Update tokenizer_utils.py * Update tokenizer_utils.py * Update tokenizer_utils.py * Update tokenizer_utils.py * Update tokenizer_utils.py * Update tokenizer_utils.py * Installation guide (#1165) * chore: update chat_templates.py (#1166) orginal -> original * Disable Flex Attention * Update tokenizer_utils.py * Update _utils.py * n_items * Update cross_entropy_loss.py * Fix DPO, ORPO * Update _utils.py --------- Co-authored-by: timothelaborie <97834767+timothelaborie@users.noreply.github.com> Co-authored-by: Ikko Eltociear Ashimine <eltociear@gmail.com> * Add warning for missing Unpack and KwargsForCausalLM in older Transformers versions --------- Co-authored-by: Daniel Han <danielhanchen@gmail.com> Co-authored-by: timothelaborie <97834767+timothelaborie@users.noreply.github.com> Co-authored-by: Ikko Eltociear Ashimine <eltociear@gmail.com> * Update cross_entropy_loss.py * Update _utils.py * Update _utils.py * donot upcast lm_head and embeddings to float32 (#1186) * Cleanup upcast logs (#1188) * Fix/phi-longrope (#1193) * Enhance rotary embedding handling in LlamaAttention and LongRopeRotaryEmbedding * Typo * Improve rotary embedding handling in LlamaAttention to prevent errors with short KV cache * Update llama.py * Update llama.py --------- Co-authored-by: Daniel Han <danielhanchen@gmail.com> * Update transformers * Unk token issues * Update _utils.py * Fix pad token * Update llama.py * Typo * ignored labels * Revert "ignored labels" This reverts commit4b25138ac7. * More patching * Update _utils.py * Update _utils.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Feat/all tmp (#1219) * Update save.py Check whether path is in /tmp dir for Kaggle environment * Update save.py Move temporary_location to /tmp in Kaggle * Enhance Kaggle environment support in save and tokenizer utilities --------- Co-authored-by: dendarrion <37800703+dendarrion@users.noreply.github.com> Co-authored-by: Erland366 <erland.pg366@gmail.com> * Bug fixes * Update pyproject.toml * Update _utils.py * Update __init__.py * Update __init__.py * Update _utils.py * Update _utils.py * Update _utils.py * Update _utils.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Tied weights * Revert "Tied weights" This reverts commit820cd4efef. * Tied weights * Utils * CE Loss patching * Update __init__.py * Update __init__.py * Patching * Update cross_entropy_loss.py * CE Loss * Update _utils.py * Update _utils.py * CE Loss * Update _utils.py * Update _utils.py * Layernorm * Update _utils.py * Update _utils.py * Post patch * Update _utils.py * Update llama.py * Update _utils.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * typing * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * int64 * Update _utils.py * Update cross_entropy_loss.py * constexpr * constexpr * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update _utils.py * Update _utils.py * Update _utils.py * CE * Update cross_entropy_loss.py * Update _utils.py * Update llama.py * Update _utils.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 utils.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 * Update rms_layernorm.py * Update rms_layernorm.py * Update rms_layernorm.py * Update rms_layernorm.py * Update rms_layernorm.py * typing * Update rope_embedding.py * types * Disable compiling * Update _utils.py * Update _utils.py * Forward hook * Update _utils.py * Update llama.py * Update _utils.py * Update llama.py * Update llama.py * Update _utils.py * Update pyproject.toml --------- Co-authored-by: timothelaborie <97834767+timothelaborie@users.noreply.github.com> Co-authored-by: Ikko Eltociear Ashimine <eltociear@gmail.com> Co-authored-by: Edd <68678137+Erland366@users.noreply.github.com> Co-authored-by: Datta Nimmaturi <datta.nimmaturi@nutanix.com> Co-authored-by: dendarrion <37800703+dendarrion@users.noreply.github.com> Co-authored-by: Erland366 <erland.pg366@gmail.com>
This commit is contained in:
Daniel Han
GitHub
parent
a9afc2f291
commit
4806f58aab
14 changed files with 287 additions and 727 deletions
|
|
@ -33,10 +33,10 @@ exclude = ["images*"]
|
|||
|
||||
[project.optional-dependencies]
|
||||
huggingface = [
|
||||
"unsloth_zoo",
|
||||
"unsloth_zoo>=2024.11.1",
|
||||
"packaging",
|
||||
"tyro",
|
||||
"transformers>=4.44.2",
|
||||
"transformers>=4.46.1",
|
||||
"datasets>=2.16.0",
|
||||
"sentencepiece>=0.2.0",
|
||||
"tqdm",
|
||||
|
|
@ -244,10 +244,10 @@ colab-ampere-torch220 = [
|
|||
"flash-attn>=2.6.3",
|
||||
]
|
||||
colab-new = [
|
||||
"unsloth_zoo",
|
||||
"unsloth_zoo>=2024.11.1",
|
||||
"packaging",
|
||||
"tyro",
|
||||
"transformers>=4.44.2",
|
||||
"transformers>=4.46.1",
|
||||
"datasets>=2.16.0",
|
||||
"sentencepiece>=0.2.0",
|
||||
"tqdm",
|
||||
|
|
|
|||
|
|
@ -27,13 +27,6 @@ import numpy as np
|
|||
# pass
|
||||
# pass
|
||||
|
||||
# Check for unsloth_zoo
|
||||
try:
|
||||
import unsloth_zoo
|
||||
except:
|
||||
raise ImportError("Unsloth: Please install unsloth_zoo via `pip install unsloth-zoo`")
|
||||
pass
|
||||
|
||||
# Unsloth currently does not work on multi GPU setups - sadly we are a 2 brother team so
|
||||
# enabling it will require much more work, so we have to prioritize. Please understand!
|
||||
# We do have a beta version, which you can contact us about!
|
||||
|
|
@ -60,6 +53,14 @@ pass
|
|||
# Reduce VRAM usage by reducing fragmentation
|
||||
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
|
||||
|
||||
# Hugging Face Hub faster downloads
|
||||
if "HF_HUB_ENABLE_HF_TRANSFER" not in os.environ:
|
||||
os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
|
||||
pass
|
||||
|
||||
# Log Unsloth is being used
|
||||
os.environ["UNSLOTH_IS_PRESENT"] = "1"
|
||||
|
||||
try:
|
||||
import torch
|
||||
except ModuleNotFoundError:
|
||||
|
|
@ -71,12 +72,6 @@ except Exception as exception:
|
|||
raise exception
|
||||
pass
|
||||
|
||||
# Hugging Face Hub faster downloads (only enable during Colab and Kaggle sessions)
|
||||
keynames = "\n" + "\n".join(os.environ.keys())
|
||||
if "\nCOLAB_" in keynames or "\nKAGGLE_" in keynames:
|
||||
os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
|
||||
pass
|
||||
|
||||
# We support Pytorch 2
|
||||
# Fixes https://github.com/unslothai/unsloth/issues/38
|
||||
torch_version = torch.__version__.split(".")
|
||||
|
|
@ -165,6 +160,13 @@ if "SPACE_AUTHOR_NAME" not in os.environ and "SPACE_REPO_NAME" not in os.environ
|
|||
pass
|
||||
pass
|
||||
|
||||
# Check for unsloth_zoo
|
||||
try:
|
||||
import unsloth_zoo
|
||||
except:
|
||||
raise ImportError("Unsloth: Please install unsloth_zoo via `pip install unsloth-zoo`")
|
||||
pass
|
||||
|
||||
from .models import *
|
||||
from .save import *
|
||||
from .chat_templates import *
|
||||
|
|
|
|||
|
|
@ -14,8 +14,8 @@
|
|||
|
||||
from .cross_entropy_loss import (
|
||||
fast_cross_entropy_loss,
|
||||
patch_llama_for_causal_lm,
|
||||
unpatch_llama_for_causal_lm,
|
||||
post_patch_loss_function,
|
||||
patch_loss_functions,
|
||||
)
|
||||
from .rms_layernorm import (
|
||||
fast_rms_layernorm,
|
||||
|
|
@ -25,7 +25,6 @@ from .rms_layernorm import (
|
|||
from .layernorm import (
|
||||
fast_layernorm,
|
||||
patch_layernorm,
|
||||
unpatch_layernorm,
|
||||
)
|
||||
from .rope_embedding import fast_rope_embedding, inplace_rope_embedding
|
||||
from .swiglu import swiglu_fg_kernel, swiglu_DWf_DW_dfg_kernel
|
||||
|
|
@ -54,8 +53,12 @@ from .flex_attention import (
|
|||
create_flex_attention_sliding_window_mask,
|
||||
)
|
||||
|
||||
try:
|
||||
print("🦥 Unsloth: Will patch your computer to enable 2x faster free finetuning.")
|
||||
except:
|
||||
print("Unsloth: Will patch your computer to enable 2x faster free finetuning.")
|
||||
import os
|
||||
if "UNSLOTH_ZOO_IS_PRESENT" not in os.environ:
|
||||
try:
|
||||
print("🦥 Unsloth: Will patch your computer to enable 2x faster free finetuning.")
|
||||
except:
|
||||
print("Unsloth: Will patch your computer to enable 2x faster free finetuning.")
|
||||
pass
|
||||
pass
|
||||
del os
|
||||
|
|
|
|||
|
|
@ -17,24 +17,31 @@ import triton.language as tl
|
|||
import torch
|
||||
from .utils import calculate_settings, MAX_FUSED_SIZE, triton_tanh
|
||||
from transformers.models.llama.modeling_llama import logger
|
||||
from packaging.version import Version
|
||||
|
||||
from unsloth_zoo.loss_utils import (
|
||||
patch_loss_functions as _patch_loss_functions,
|
||||
post_patch_loss_function,
|
||||
)
|
||||
|
||||
|
||||
@triton.heuristics({
|
||||
"DO_SOFTCAPPING": lambda args: args["DO_SOFTCAPPING" ],
|
||||
"DO_LOGIT_SCALING": lambda args: args["DO_LOGIT_SCALING"],
|
||||
"DO_SOFTCAPPING": lambda args: bool(args["DO_SOFTCAPPING" ]),
|
||||
"DO_LOGIT_SCALING": lambda args: bool(args["DO_LOGIT_SCALING"]),
|
||||
})
|
||||
@triton.jit
|
||||
def _cross_entropy_forward(
|
||||
logits_ptr, logits_row_stride,
|
||||
loss_ptr,
|
||||
logsumexp_ptr,
|
||||
labels_ptr,
|
||||
VOCAB_SIZE : tl.constexpr,
|
||||
BLOCK_SIZE : tl.constexpr,
|
||||
DO_SOFTCAPPING : tl.constexpr,
|
||||
SOFTCAP : tl.constexpr,
|
||||
DO_LOGIT_SCALING: tl.constexpr,
|
||||
LOGIT_SCALE : tl.constexpr,
|
||||
logits_ptr ,
|
||||
logits_row_stride ,
|
||||
loss_ptr ,
|
||||
logsumexp_ptr ,
|
||||
labels_ptr ,
|
||||
VOCAB_SIZE ,
|
||||
BLOCK_SIZE : tl.constexpr,
|
||||
DO_SOFTCAPPING ,
|
||||
SOFTCAP ,
|
||||
DO_LOGIT_SCALING ,
|
||||
LOGIT_SCALE ,
|
||||
):
|
||||
"""
|
||||
Cross Entropy Loss = 1/n sum [ -yi log(Pi) ]
|
||||
|
|
@ -57,7 +64,7 @@ def _cross_entropy_forward(
|
|||
This ensures exp(x - max(x))'s maximum is 1 as exp(0) = 1.
|
||||
"""
|
||||
row_idx = tl.program_id(0)
|
||||
logits_ptr += row_idx * logits_row_stride.to(tl.int64)
|
||||
logits_ptr += row_idx * tl.cast(logits_row_stride, tl.int64)
|
||||
loss_ptr += row_idx
|
||||
logsumexp_ptr += row_idx
|
||||
labels_ptr += row_idx
|
||||
|
|
@ -71,7 +78,7 @@ def _cross_entropy_forward(
|
|||
# Go logit scaling for Cohere: t * x
|
||||
if DO_LOGIT_SCALING: logits = LOGIT_SCALE * logits
|
||||
# Do logit softcapping for Gemma 2: t * tanh(1/t * x)
|
||||
if DO_SOFTCAPPING: logits = SOFTCAP * triton_tanh(logits / SOFTCAP)
|
||||
if DO_SOFTCAPPING: logits = SOFTCAP * triton_tanh(logits.to(tl.float32) / SOFTCAP).to(logits.dtype)
|
||||
|
||||
logits = logits.to(tl.float32)
|
||||
c = tl.max(logits, 0)
|
||||
|
|
@ -92,22 +99,23 @@ pass
|
|||
|
||||
|
||||
@triton.heuristics({
|
||||
"DO_SOFTCAPPING": lambda args: args["DO_SOFTCAPPING" ],
|
||||
"DO_LOGIT_SCALING": lambda args: args["DO_LOGIT_SCALING"],
|
||||
"DO_SOFTCAPPING": lambda args: bool(args["DO_SOFTCAPPING" ]),
|
||||
"DO_LOGIT_SCALING": lambda args: bool(args["DO_LOGIT_SCALING"]),
|
||||
})
|
||||
@triton.jit
|
||||
def _chunked_cross_entropy_forward(
|
||||
logits_ptr, logits_row_stride,
|
||||
loss_ptr,
|
||||
logsumexp_ptr,
|
||||
labels_ptr,
|
||||
VOCAB_SIZE : tl.constexpr,
|
||||
N_CHUNKS : tl.constexpr,
|
||||
BLOCK_SIZE : tl.constexpr,
|
||||
DO_SOFTCAPPING : tl.constexpr,
|
||||
SOFTCAP : tl.constexpr,
|
||||
DO_LOGIT_SCALING: tl.constexpr,
|
||||
LOGIT_SCALE : tl.constexpr,
|
||||
logits_ptr ,
|
||||
logits_row_stride ,
|
||||
loss_ptr ,
|
||||
logsumexp_ptr ,
|
||||
labels_ptr ,
|
||||
VOCAB_SIZE ,
|
||||
N_CHUNKS ,
|
||||
BLOCK_SIZE : tl.constexpr,
|
||||
DO_SOFTCAPPING ,
|
||||
SOFTCAP ,
|
||||
DO_LOGIT_SCALING ,
|
||||
LOGIT_SCALE ,
|
||||
):
|
||||
"""
|
||||
256K vocab divided in 4 chunks
|
||||
|
|
@ -135,7 +143,7 @@ def _chunked_cross_entropy_forward(
|
|||
"""
|
||||
row_idx = tl.program_id(0)
|
||||
chunk_idx = tl.program_id(1)
|
||||
logits_ptr += row_idx * logits_row_stride.to(tl.int64)
|
||||
logits_ptr += row_idx * tl.cast(logits_row_stride, tl.int64)
|
||||
loss_ptr += row_idx
|
||||
logsumexp_ptr += row_idx * N_CHUNKS + chunk_idx
|
||||
labels_ptr += row_idx
|
||||
|
|
@ -149,7 +157,7 @@ def _chunked_cross_entropy_forward(
|
|||
# Go logit scaling for Cohere: t * x
|
||||
if DO_LOGIT_SCALING: logits = LOGIT_SCALE * logits
|
||||
# Do logit softcapping for Gemma 2: t * tanh(1/t * x)
|
||||
if DO_SOFTCAPPING: logits = SOFTCAP * triton_tanh(logits / SOFTCAP)
|
||||
if DO_SOFTCAPPING: logits = SOFTCAP * triton_tanh(logits.to(tl.float32) / SOFTCAP).to(logits.dtype)
|
||||
|
||||
logits = logits.to(tl.float32)
|
||||
c = tl.max(logits, 0)
|
||||
|
|
@ -174,21 +182,23 @@ pass
|
|||
|
||||
|
||||
@triton.heuristics({
|
||||
"DO_SOFTCAPPING": lambda args: args["DO_SOFTCAPPING" ],
|
||||
"DO_LOGIT_SCALING": lambda args: args["DO_LOGIT_SCALING"],
|
||||
"DO_SOFTCAPPING": lambda args: bool(args["DO_SOFTCAPPING" ]),
|
||||
"DO_LOGIT_SCALING": lambda args: bool(args["DO_LOGIT_SCALING"]),
|
||||
})
|
||||
@triton.jit
|
||||
def _cross_entropy_backward(
|
||||
logits_ptr, logits_row_stride,
|
||||
dloss_ptr, dloss_row_stride,
|
||||
logsumexp_ptr,
|
||||
labels_ptr,
|
||||
VOCAB_SIZE : tl.constexpr,
|
||||
BLOCK_SIZE : tl.constexpr,
|
||||
DO_SOFTCAPPING : tl.constexpr,
|
||||
SOFTCAP : tl.constexpr,
|
||||
DO_LOGIT_SCALING: tl.constexpr,
|
||||
LOGIT_SCALE : tl.constexpr,
|
||||
logits_ptr ,
|
||||
logits_row_stride ,
|
||||
dloss_ptr ,
|
||||
dloss_row_stride ,
|
||||
logsumexp_ptr ,
|
||||
labels_ptr ,
|
||||
VOCAB_SIZE ,
|
||||
BLOCK_SIZE : tl.constexpr,
|
||||
DO_SOFTCAPPING ,
|
||||
SOFTCAP ,
|
||||
DO_LOGIT_SCALING ,
|
||||
LOGIT_SCALE ,
|
||||
):
|
||||
"""
|
||||
CE_i = -y log(P) = y * (log[sum(exp(x))] - x)
|
||||
|
|
@ -208,7 +218,7 @@ def _cross_entropy_backward(
|
|||
row_idx = tl.program_id(0)
|
||||
block_idx = tl.program_id(1)
|
||||
|
||||
logits_ptr += row_idx * logits_row_stride.to(tl.int64)
|
||||
logits_ptr += row_idx * tl.cast(logits_row_stride, tl.int64)
|
||||
dloss_ptr += row_idx * dloss_row_stride
|
||||
col_offsets = block_idx*BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
|
||||
mask = col_offsets < VOCAB_SIZE
|
||||
|
|
@ -228,9 +238,10 @@ def _cross_entropy_backward(
|
|||
pass
|
||||
|
||||
# Do logit softcapping for Gemma 2: t * tanh(1/t * x)
|
||||
partial = x
|
||||
if DO_SOFTCAPPING:
|
||||
# d/dx [t * tanh(1/t * x)] = 1 - tanh^2(1/t * x)
|
||||
partial = triton_tanh(x / SOFTCAP)
|
||||
partial = triton_tanh(x.to(tl.float32) / SOFTCAP).to(x.dtype)
|
||||
x = SOFTCAP * partial
|
||||
pass
|
||||
|
||||
|
|
@ -261,16 +272,20 @@ MAX_FUSED_SIZE = 65536 # 2**16
|
|||
|
||||
class Fast_CrossEntropyLoss(torch.autograd.Function):
|
||||
@staticmethod
|
||||
def forward(ctx, logits, labels, logit_softcapping = 0, logit_scaling = 0):
|
||||
def forward(ctx, logits, labels, logit_softcapping : float = 0, logit_scaling : float = 0):
|
||||
n_rows : int
|
||||
vocab_size : int
|
||||
n_rows, vocab_size = logits.shape
|
||||
|
||||
div, mod = divmod(vocab_size, MAX_FUSED_SIZE)
|
||||
n_chunks = div + (mod != 0)
|
||||
n_chunks : int = div + (mod != 0)
|
||||
losses = torch.empty(n_rows, dtype = torch.float32, device = "cuda:0")
|
||||
|
||||
DO_SOFTCAPPING = (logit_softcapping != 0)
|
||||
DO_LOGIT_SCALING = (logit_scaling != 0)
|
||||
DO_SOFTCAPPING : bool = bool(logit_softcapping != 0)
|
||||
DO_LOGIT_SCALING : bool = bool(logit_scaling != 0)
|
||||
|
||||
BLOCK_SIZE : int
|
||||
num_warps : int
|
||||
if n_chunks == 1:
|
||||
# For small vocabs <= 65336 like Llama, Mistral
|
||||
BLOCK_SIZE, num_warps = calculate_settings(vocab_size)
|
||||
|
|
@ -325,11 +340,13 @@ class Fast_CrossEntropyLoss(torch.autograd.Function):
|
|||
@staticmethod
|
||||
def backward(ctx, dlosses):
|
||||
logits, logsumexp, labels = ctx.saved_tensors
|
||||
n_rows : int
|
||||
vocab_size : int
|
||||
n_rows, vocab_size = logits.shape
|
||||
|
||||
BLOCK_SIZE = 4096
|
||||
BLOCK_SIZE : int = 4096
|
||||
div, mod = divmod(vocab_size, BLOCK_SIZE)
|
||||
n_blocks = div + (mod != 0)
|
||||
n_blocks : int = div + (mod != 0)
|
||||
|
||||
_cross_entropy_backward[(n_rows, n_blocks,)](
|
||||
logits, logits.stride(0),
|
||||
|
|
@ -342,14 +359,13 @@ class Fast_CrossEntropyLoss(torch.autograd.Function):
|
|||
SOFTCAP = ctx.logit_softcapping,
|
||||
DO_LOGIT_SCALING = ctx.DO_LOGIT_SCALING,
|
||||
LOGIT_SCALE = ctx.logit_scaling,
|
||||
num_warps = 8,
|
||||
num_warps = 8,
|
||||
)
|
||||
return logits, None, None, None,
|
||||
pass
|
||||
pass
|
||||
|
||||
|
||||
@torch._disable_dynamo
|
||||
def fast_cross_entropy_loss(
|
||||
logits,
|
||||
labels,
|
||||
|
|
@ -377,96 +393,12 @@ def fast_cross_entropy_loss(
|
|||
n_items = torch.count_nonzero(labels != -100)
|
||||
return loss.sum() / n_items
|
||||
pass
|
||||
|
||||
|
||||
from transformers.models.llama.modeling_llama import (
|
||||
LlamaForCausalLM,
|
||||
CausalLMOutputWithPast,
|
||||
Optional,
|
||||
Union,
|
||||
Cache,
|
||||
List,
|
||||
Tuple,
|
||||
)
|
||||
|
||||
# Transformers 4.47 need Unpack, KwargsForCausalLM
|
||||
try:
|
||||
from transformers.models.llama.modeling_llama import Unpack, KwargsForCausalLM
|
||||
except:
|
||||
pass
|
||||
if (Version(torch.__version__) < Version("2.4.0")) and \
|
||||
not hasattr(fast_cross_entropy_loss, "__wrapped__"):
|
||||
fast_cross_entropy_loss = torch._disable_dynamo(fast_cross_entropy_loss)
|
||||
pass
|
||||
|
||||
import inspect, re
|
||||
function = inspect.getsource(LlamaForCausalLM.forward)
|
||||
function = function.split("\n")
|
||||
i = re.match(r"[ ]{1,}", function[0]).span(0)[1]
|
||||
function = [x[i:] for x in function]
|
||||
function = "\n".join(function)
|
||||
function = function[function.find("def forward"):]
|
||||
replacement = """ loss = None
|
||||
logit_softcapping = getattr(self.config, "final_logit_softcapping", 0)
|
||||
logit_scaling = getattr(self.config, "logit_scale", 0)
|
||||
if labels is not None:
|
||||
shift_logits = logits
|
||||
if not hasattr(self, "extra_ignored_labels"):
|
||||
# Fixes https://github.com/unslothai/unsloth/issues/10
|
||||
self.extra_ignored_labels = torch.full((self.max_seq_length, 1), -100, device = "cuda:0")
|
||||
pass
|
||||
|
||||
shift_labels = torch.hstack((labels[..., 1:], self.extra_ignored_labels[:labels.shape[0]]))
|
||||
loss = fast_cross_entropy_loss(
|
||||
logits = shift_logits,
|
||||
labels = shift_labels,
|
||||
logit_softcapping = logit_softcapping,
|
||||
logit_scaling = logit_scaling,
|
||||
n_items = kwargs.get("num_items_in_batch", None) or kwargs.get("n_items", None),
|
||||
)
|
||||
else:
|
||||
if logit_scaling != 0:
|
||||
if logits.requires_grad:
|
||||
logits = logit_scaling * logits
|
||||
else:
|
||||
logits *= logit_scaling
|
||||
pass
|
||||
pass
|
||||
if logit_softcapping != 0:
|
||||
if logits.requires_grad:
|
||||
logits = (1.0 / logit_softcapping) * logits
|
||||
logits = torch.tanh(logits)
|
||||
logits = logit_softcapping * logits
|
||||
else:
|
||||
logits *= (1.0 / logit_softcapping)
|
||||
torch.tanh(logits, out = logits)
|
||||
logits *= logit_softcapping
|
||||
pass
|
||||
pass
|
||||
pass
|
||||
"""
|
||||
function = \
|
||||
function[:function.find(" loss = None")] + \
|
||||
replacement + \
|
||||
function[ function.find(" if not return_dict"):]
|
||||
function = function.replace("logits = logits.float()", "\n")
|
||||
# Missed spaces
|
||||
function = function.split("\n")
|
||||
# Not the first one though!
|
||||
function = [function[0]] + [" "*4 + x for x in function[1:]]
|
||||
function = "\n".join(function)
|
||||
function = f"class Unsloth_LlamaForCausalLM(LlamaForCausalLM):\n"\
|
||||
f" {function}\n"
|
||||
exec(function, globals())
|
||||
del function, replacement, inspect, re
|
||||
|
||||
|
||||
def patch_llama_for_causal_lm():
|
||||
import transformers.models.llama.modeling_llama
|
||||
transformers.models.llama.modeling_llama.LlamaForCausalLM = Unsloth_LlamaForCausalLM
|
||||
return
|
||||
pass
|
||||
|
||||
|
||||
def unpatch_llama_for_causal_lm():
|
||||
import transformers.models.llama.modeling_llama
|
||||
transformers.models.llama.modeling_llama.LlamaForCausalLM = LlamaForCausalLM
|
||||
return
|
||||
# Patch CE Losses in transformers
|
||||
def patch_loss_functions():
|
||||
_patch_loss_functions(fast_cross_entropy_loss)
|
||||
pass
|
||||
|
|
|
|||
|
|
@ -17,6 +17,9 @@ import triton
|
|||
import triton.language as tl
|
||||
import torch
|
||||
from .utils import calculate_settings
|
||||
from unsloth_zoo.patching_utils import (
|
||||
patch_layernorm,
|
||||
)
|
||||
|
||||
|
||||
@triton.jit
|
||||
|
|
@ -162,27 +165,6 @@ def fast_layernorm(layernorm, X):
|
|||
pass
|
||||
|
||||
|
||||
from torch.nn import LayerNorm
|
||||
class Unsloth_LayerNorm(LayerNorm):
|
||||
def forward(self, X):
|
||||
return fast_layernorm(self, X)
|
||||
pass
|
||||
pass
|
||||
|
||||
|
||||
def patch_layernorm():
|
||||
import torch.nn
|
||||
torch.nn.LayerNorm = Unsloth_LayerNorm
|
||||
return
|
||||
pass
|
||||
|
||||
|
||||
def unpatch_layernorm():
|
||||
import torch.nn
|
||||
torch.nn.LayerNorm = LayerNorm
|
||||
return
|
||||
pass
|
||||
|
||||
|
||||
def test_layernorm(
|
||||
dim = 1024, eps = 1e-5, dtype = torch.float16,
|
||||
|
|
|
|||
|
|
@ -53,14 +53,14 @@ def _rms_layernorm_forward(
|
|||
pass
|
||||
|
||||
|
||||
@triton.heuristics({"GEMMA": lambda args: args["GEMMA"],})
|
||||
@triton.heuristics({"GEMMA": lambda args: bool(args["GEMMA"]),})
|
||||
@triton.jit
|
||||
def _rms_layernorm_backward(
|
||||
dY, dY_row_stride,
|
||||
X, X_row_stride,
|
||||
W, W_row_stride,
|
||||
r, r_row_stride,
|
||||
dW, dW_row_stride,
|
||||
# dW, dW_row_stride,
|
||||
n_cols, eps,
|
||||
GEMMA : tl.constexpr,
|
||||
BLOCK_SIZE : tl.constexpr,
|
||||
|
|
@ -130,11 +130,15 @@ pass
|
|||
|
||||
class Fast_RMS_Layernorm(torch.autograd.Function):
|
||||
@staticmethod
|
||||
def forward(ctx, X, W, eps, gemma = False):
|
||||
def forward(ctx, X : torch.Tensor, W : torch.Tensor, eps : float, gemma : bool = False):
|
||||
shape = X.shape
|
||||
dim = shape[-1]
|
||||
dim : int = shape[-1]
|
||||
X = X.view(-1, dim)
|
||||
n_rows : int
|
||||
n_cols : int
|
||||
n_rows, n_cols = X.shape
|
||||
BLOCK_SIZE : int
|
||||
num_warps : int
|
||||
BLOCK_SIZE, num_warps = calculate_settings(n_cols)
|
||||
|
||||
Y = torch.empty((n_rows, n_cols), dtype = X.dtype, device = "cuda:0")
|
||||
|
|
@ -159,20 +163,22 @@ class Fast_RMS_Layernorm(torch.autograd.Function):
|
|||
pass
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, dY):
|
||||
def backward(ctx, dY : torch.Tensor):
|
||||
shape = dY.shape
|
||||
dim = shape[-1]
|
||||
dim : int = shape[-1]
|
||||
dY = dY.view(-1, dim)
|
||||
X, W, r = ctx.saved_tensors
|
||||
n_rows : int
|
||||
n_cols : int
|
||||
n_rows, n_cols = dY.shape
|
||||
dW = X
|
||||
# dW = X
|
||||
|
||||
_rms_layernorm_backward[(n_rows,)](
|
||||
dY, dY.stride(0),
|
||||
X, X .stride(0),
|
||||
W, W .stride(0),
|
||||
r, r .stride(0),
|
||||
dW, dW.stride(0),
|
||||
# dW, dW.stride(0),
|
||||
n_cols, ctx.eps,
|
||||
GEMMA = ctx.GEMMA,
|
||||
BLOCK_SIZE = ctx.BLOCK_SIZE,
|
||||
|
|
@ -184,9 +190,11 @@ class Fast_RMS_Layernorm(torch.autograd.Function):
|
|||
pass
|
||||
|
||||
|
||||
def fast_rms_layernorm(layernorm, X, gemma = False):
|
||||
W = layernorm.weight
|
||||
eps = layernorm.variance_epsilon if \
|
||||
# [TODO] Unsure why RMS Layernorm is not torch.compiling properly
|
||||
@torch.compiler.disable
|
||||
def fast_rms_layernorm(layernorm, X : torch.Tensor, gemma : bool = False):
|
||||
W : torch.Tensor = layernorm.weight
|
||||
eps : float = layernorm.variance_epsilon if \
|
||||
hasattr(layernorm, "variance_epsilon") \
|
||||
else layernorm.eps
|
||||
out = Fast_RMS_Layernorm.apply(X, W, eps, gemma)
|
||||
|
|
|
|||
|
|
@ -16,9 +16,9 @@ import triton
|
|||
import triton.language as tl
|
||||
import torch
|
||||
from .utils import calculate_settings
|
||||
ROPE_GROUP_SIZE = 4
|
||||
ROPE_GROUP_SIZE : int = 4
|
||||
|
||||
@triton.heuristics({"BACKWARD_PASS": lambda args: args["BACKWARD_PASS"],})
|
||||
@triton.heuristics({"BACKWARD_PASS": lambda args: bool(args["BACKWARD_PASS"]),})
|
||||
@triton.jit
|
||||
def _rope_embedding(
|
||||
Q, Q_row_stride,
|
||||
|
|
@ -75,8 +75,14 @@ class Fast_RoPE_Embedding(torch.autograd.Function):
|
|||
@staticmethod
|
||||
def forward(ctx, Q, cos, sin):
|
||||
cos, sin = cos.squeeze(), sin.squeeze()
|
||||
batch : int
|
||||
seq_len : int
|
||||
n_heads : int
|
||||
head_dim : int
|
||||
batch, seq_len, n_heads, head_dim = Q.shape
|
||||
Q = Q.view(batch*seq_len, n_heads*head_dim)
|
||||
n_rows : int
|
||||
n_cols : int
|
||||
n_rows, n_cols = Q.shape
|
||||
assert(seq_len <= cos.shape[0])
|
||||
|
||||
|
|
@ -85,8 +91,10 @@ class Fast_RoPE_Embedding(torch.autograd.Function):
|
|||
BLOCK_SIZE, num_warps = calculate_settings(head_dim//2) # (head_dim//2)
|
||||
|
||||
# group_size = 4 # 4 or 8, too large group_size can hurt performance.
|
||||
div : int
|
||||
mod : int
|
||||
div, mod = divmod(n_heads, ROPE_GROUP_SIZE)
|
||||
n_groups = div + (mod != 0)
|
||||
n_groups : int = div + (mod != 0)
|
||||
|
||||
_rope_embedding[(n_rows, n_groups, )](
|
||||
Q, Q.stride(0),
|
||||
|
|
@ -108,9 +116,15 @@ class Fast_RoPE_Embedding(torch.autograd.Function):
|
|||
|
||||
@staticmethod
|
||||
def backward(ctx, dY):
|
||||
batch : int
|
||||
seq_len : int
|
||||
n_heads : int
|
||||
head_dim : int
|
||||
batch, seq_len, n_heads, head_dim = dY.shape
|
||||
dY = dY.reshape(batch*seq_len, n_heads*head_dim)
|
||||
# Must be reshape not view
|
||||
n_rows : int
|
||||
n_cols : int
|
||||
n_rows, n_cols = dY.shape
|
||||
|
||||
cos = ctx.cos
|
||||
|
|
@ -130,7 +144,8 @@ class Fast_RoPE_Embedding(torch.autograd.Function):
|
|||
pass
|
||||
pass
|
||||
|
||||
|
||||
# [TODO] Unsure why RoPE Embedding is not torch.compiling properly
|
||||
@torch.compiler.disable
|
||||
def fast_rope_embedding(Q, K, cos, sin):
|
||||
Q = Fast_RoPE_Embedding.apply(Q.transpose(1, 2), cos, sin).transpose(1, 2)
|
||||
K = Fast_RoPE_Embedding.apply(K.transpose(1, 2), cos, sin).transpose(1, 2)
|
||||
|
|
|
|||
|
|
@ -13,7 +13,7 @@
|
|||
# limitations under the License.
|
||||
|
||||
import triton
|
||||
MAX_FUSED_SIZE = 65536
|
||||
MAX_FUSED_SIZE : int = 65536
|
||||
next_power_of_2 = triton.next_power_of_2
|
||||
|
||||
# torch.cuda.amp.custom_fwd is deprecated >= 2.4
|
||||
|
|
@ -40,12 +40,12 @@ else:
|
|||
pass
|
||||
|
||||
|
||||
def calculate_settings(n):
|
||||
BLOCK_SIZE = next_power_of_2(n)
|
||||
def calculate_settings(n : int) -> (int, int,):
|
||||
BLOCK_SIZE : int = next_power_of_2(n)
|
||||
if BLOCK_SIZE > MAX_FUSED_SIZE:
|
||||
raise RuntimeError(f"Cannot launch Triton kernel since n = {n} exceeds "\
|
||||
f"the maximum CUDA blocksize = {MAX_FUSED_SIZE}.")
|
||||
num_warps = 4
|
||||
num_warps : int = 4
|
||||
if BLOCK_SIZE >= 32768: num_warps = 32
|
||||
elif BLOCK_SIZE >= 8192: num_warps = 16
|
||||
elif BLOCK_SIZE >= 2048: num_warps = 8
|
||||
|
|
|
|||
|
|
@ -12,7 +12,7 @@
|
|||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
__version__ = "2024.10.7"
|
||||
__version__ = "2024.11.1"
|
||||
|
||||
__all__ = [
|
||||
"prepare_model_for_kbit_training",
|
||||
|
|
@ -41,9 +41,17 @@ __all__ = [
|
|||
"torch_amp_custom_bwd",
|
||||
"accelerate_old_send_to_device",
|
||||
"accelerate_new_send_to_device",
|
||||
"patch_gradient_accumulation_fix",
|
||||
"patch_compiling_bitsandbytes",
|
||||
"patch_regional_compilation",
|
||||
"patch_layernorm",
|
||||
"patch_torch_compile",
|
||||
"patch_model_and_tokenizer",
|
||||
|
||||
"patch_unsloth_gradient_checkpointing",
|
||||
"unpatch_unsloth_gradient_checkpointing",
|
||||
"patch_gradient_checkpointing",
|
||||
"unpatch_gradient_checkpointing",
|
||||
"patch_gradient_accumulation_fix",
|
||||
]
|
||||
|
||||
import torch
|
||||
|
|
@ -54,6 +62,28 @@ 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 as _patch_tokenizer,
|
||||
)
|
||||
from unsloth_zoo.patching_utils import (
|
||||
patch_compiling_bitsandbytes,
|
||||
patch_layernorm,
|
||||
patch_torch_compile,
|
||||
patch_regional_compilation,
|
||||
patch_model_and_tokenizer,
|
||||
)
|
||||
from unsloth_zoo.gradient_checkpointing import (
|
||||
Unsloth_Offloaded_Gradient_Checkpointer,
|
||||
unsloth_offloaded_gradient_checkpoint,
|
||||
patch_unsloth_gradient_checkpointing,
|
||||
unpatch_unsloth_gradient_checkpointing,
|
||||
|
||||
Unsloth_Gradient_Checkpointer,
|
||||
unsloth_gradient_checkpoint,
|
||||
patch_gradient_checkpointing,
|
||||
unpatch_gradient_checkpointing,
|
||||
)
|
||||
|
||||
# =============================================
|
||||
# Disable some warnings which can get annoying
|
||||
warnings.filterwarnings(action = "ignore", category = UserWarning, module = "torch")
|
||||
|
|
@ -70,6 +100,18 @@ warnings.filterwarnings(action = "ignore", category = RuntimeWarning, module = "
|
|||
# Stop "Special tokens have been added in the vocabulary, ..."
|
||||
import logging
|
||||
logging.getLogger("transformers.tokenization_utils_base").setLevel(logging.CRITICAL+1)
|
||||
|
||||
# Ignore logging messages
|
||||
class HideLoggingMessage(logging.Filter):
|
||||
def __init__(self, text): self.text = text
|
||||
def filter(self, x): return not x.getMessage().startswith(self.text)
|
||||
pass
|
||||
|
||||
# The speedups for torchdynamo mostly come wih GPU Ampere or higher and which is not detected here.
|
||||
from transformers.training_args import logger as transformers_training_args_logger
|
||||
transformers_training_args_logger.addFilter(HideLoggingMessage("The speedups"))
|
||||
del transformers_training_args_logger
|
||||
|
||||
# =============================================
|
||||
|
||||
# =============================================
|
||||
|
|
@ -129,7 +171,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
|
||||
|
|
@ -333,7 +374,8 @@ pass
|
|||
|
||||
# =============================================
|
||||
# Torch compile settings
|
||||
|
||||
UNSLOTH_COMPILE_DEBUG = "UNSLOTH_COMPILE_DEBUG" in os.environ
|
||||
UNSLOTH_COMPILE_MAXIMUM = "UNSLOTH_COMPILE_MAXIMUM" in os.environ
|
||||
# Just remove max_autotune_gemm warning
|
||||
import functools
|
||||
@functools.lru_cache(None)
|
||||
|
|
@ -345,47 +387,27 @@ def is_big_gpu(index):
|
|||
return True
|
||||
import torch._inductor.utils
|
||||
torch._inductor.utils.is_big_gpu = is_big_gpu
|
||||
patch_torch_compile(debug = UNSLOTH_COMPILE_DEBUG, O3 = UNSLOTH_COMPILE_MAXIMUM)
|
||||
|
||||
|
||||
# Torch compile arguments
|
||||
torch_compile_arguments = [
|
||||
"config.dce = True",
|
||||
"config.memory_planning = True",
|
||||
"config.memory_pool = 'combined'",
|
||||
"config.coordinate_descent_tuning = True",
|
||||
"config.max_autotune_gemm = False", # GEMM is unnecessary
|
||||
"config.autotune_multi_device = False",
|
||||
"config.max_autotune_gemm_backends = 'TRITON,ATEN,CPP'", # Not much faster
|
||||
"config.aggressive_fusion = False", # Careful changes results!
|
||||
"config.cuda.enable_cuda_lto = True",
|
||||
"config.cuda.use_fast_math = True",
|
||||
"config.cuda.compile_opt_level = '-O2'",
|
||||
]
|
||||
# Torch dynamo arguments
|
||||
torch_dynamo_arguments = [
|
||||
"config.accumulated_cache_size_limit = 1024", # Bump up a bit from 256
|
||||
"config.suppress_errors = True", # Supress errors for now
|
||||
"config.do_not_emit_runtime_asserts = True",
|
||||
"config.cache_size_limit = 1024", # Flex Attention
|
||||
"config.inline_inbuilt_nn_modules = True", # Torch 2.5 Regional recompilation
|
||||
]
|
||||
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,
|
||||
"shape_padding" : True,
|
||||
"trace.enabled" : False, # Output Triton kernel outputs!
|
||||
"trace.enabled" : UNSLOTH_COMPILE_DEBUG,
|
||||
"triton.cudagraphs" : False,
|
||||
}
|
||||
|
||||
import accelerate
|
||||
def torch_compile_kwargs(*args, **kwargs):
|
||||
print("Unsloth: Enabled auto compiling")
|
||||
return {"dynamic" : True, "fullgraph" : False, "options" : torch_compile_options,}
|
||||
pass
|
||||
|
||||
accelerate.utils.dataclasses.TorchDynamoPlugin.to_kwargs = torch_compile_kwargs
|
||||
accelerate.utils.TorchDynamoPlugin.to_kwargs = torch_compile_kwargs
|
||||
accelerate.accelerator.TorchDynamoPlugin.to_kwargs = torch_compile_kwargs
|
||||
del accelerate
|
||||
|
||||
# =============================================
|
||||
|
||||
def prepare_model_for_kbit_training(
|
||||
|
|
@ -455,137 +477,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__})
|
||||
|
||||
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
|
||||
|
||||
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 possible_pad_token is None 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:
|
||||
model.config.update({"pad_token_id" : tokenizer.pad_token_id})
|
||||
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:
|
||||
model.config.update({"pad_token_id" : tokenizer.pad_token_id})
|
||||
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.
|
||||
|
|
@ -618,6 +509,7 @@ if Version(peft_version) < Version("0.12.0"):
|
|||
)
|
||||
pass
|
||||
pass
|
||||
|
||||
# =============================================
|
||||
|
||||
import psutil
|
||||
|
|
@ -678,7 +570,9 @@ def get_statistics():
|
|||
# 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
|
||||
# You can disable this by setting UNSLOTH_DISABLE_STATISTICS
|
||||
import os
|
||||
if "UNSLOTH_DISABLE_STATISTICS" in os.environ: return
|
||||
from huggingface_hub.utils import disable_progress_bars, enable_progress_bars, are_progress_bars_disabled
|
||||
disabled = False
|
||||
if not are_progress_bars_disabled():
|
||||
|
|
@ -710,139 +604,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
|
||||
old_checkpoint = torch.utils.checkpoint
|
||||
def patch_gradient_checkpointing():
|
||||
torch.utils.checkpoint = unsloth_offloaded_gradient_checkpoint
|
||||
pass
|
||||
|
||||
def unpatch_gradient_checkpointing():
|
||||
torch.utils.checkpoint = old_checkpoint
|
||||
pass
|
||||
|
||||
|
||||
# =============================================
|
||||
# Fixes Bitsandbytes to remove missing warnings
|
||||
from transformers.utils.quantization_config import BitsAndBytesConfig, QuantizationMethod
|
||||
|
|
@ -1189,6 +950,7 @@ def patch_gradient_accumulation_fix(Trainer):
|
|||
# Fixes gradient accumulation
|
||||
import inspect
|
||||
if hasattr(Trainer, "get_batch_samples"):
|
||||
if Trainer.get_batch_samples.__name__ == "_unsloth_get_batch_samples": return
|
||||
if \
|
||||
not inspect.getsource(Trainer.get_batch_samples).strip()\
|
||||
.endswith("return batch_samples, num_items_in_batch"):
|
||||
|
|
@ -1215,6 +977,7 @@ def patch_gradient_accumulation_fix(Trainer):
|
|||
pass
|
||||
|
||||
# Also fix up loss scaling ie negate loss *= self.args.gradient_accumulation_steps
|
||||
if Trainer.training_step.__name__ == "_unsloth_training_step": return
|
||||
if "num_items_in_batch" not in inspect.signature(Trainer.training_step).parameters: return
|
||||
|
||||
function = inspect.getsource(Trainer.training_step)
|
||||
|
|
@ -1243,3 +1006,11 @@ def patch_gradient_accumulation_fix(Trainer):
|
|||
exec(function, globals())
|
||||
Trainer.training_step = _unsloth_training_step
|
||||
pass
|
||||
|
||||
|
||||
def patch_tokenizer(model, tokenizer):
|
||||
model, tokenizer = _patch_tokenizer(model, tokenizer)
|
||||
if model is not None:
|
||||
model.config.update({"unsloth_version" : __version__})
|
||||
return model, tokenizer
|
||||
pass
|
||||
|
|
|
|||
|
|
@ -339,60 +339,9 @@ class FastGemmaModel(FastLlamaModel):
|
|||
|
||||
|
||||
@staticmethod
|
||||
def post_patch(model):
|
||||
# Patch model for Gemma
|
||||
layers = model.model.layers
|
||||
|
||||
# Torch.compile fails on embedding matrix??
|
||||
# Workaround randomnly fixes it for torch versions < 2.2
|
||||
model.model.embed_tokens = torch.nn.Embedding.from_pretrained(model.model.embed_tokens.weight)
|
||||
model.config.update({"unsloth_version" : __version__})
|
||||
|
||||
# We also do this for the lm_head
|
||||
lm_head = torch.nn.Linear(1, 1, bias = None)
|
||||
del lm_head.weight
|
||||
lm_head.weight = model.lm_head.weight
|
||||
lm_head.in_features = lm_head.weight.shape[1]
|
||||
lm_head.out_features = lm_head.weight.shape[0]
|
||||
model.lm_head = lm_head
|
||||
|
||||
# Gemma has tied weights! This means lm_head == embed_tokens
|
||||
if model.model.embed_tokens.weight.data_ptr() != model.lm_head.weight.data_ptr():
|
||||
lm_head = torch.nn.Linear(1, 1, bias = None)
|
||||
del lm_head.weight
|
||||
lm_head.weight = model.model.embed_tokens.weight
|
||||
lm_head.in_features = lm_head.weight.shape[1]
|
||||
lm_head.out_features = lm_head.weight.shape[0]
|
||||
model.lm_head = lm_head
|
||||
pass
|
||||
|
||||
# Also patch all dtypes - BnB seems to not allocate the correct type?
|
||||
# BnB default dtype seems to be float16!
|
||||
correct_dtype = lm_head.weight.dtype
|
||||
|
||||
for name, module in model.named_modules():
|
||||
if isinstance(module, (Bnb_Linear4bit, Peft_Linear4bit)):
|
||||
weight = module.weight
|
||||
quant_state = weight.quant_state
|
||||
|
||||
if type(quant_state) is list:
|
||||
# BnB seems to have float16 as default!
|
||||
module.weight.quant_state[2] = correct_dtype # Cast to correct dtype
|
||||
else:
|
||||
# https://github.com/TimDettmers/bitsandbytes/pull/763/files
|
||||
quant_state.dtype = correct_dtype
|
||||
pass
|
||||
pass
|
||||
# Downcast RoPE embedding to correct data type
|
||||
# RoPE must be done in float32 for Gemma
|
||||
# if (name.endswith("rotary_emb") or hasattr(module, "cos_cached")) \
|
||||
# and (module.cos_cached.dtype != correct_dtype):
|
||||
|
||||
# module.cos_cached = module.cos_cached.to(correct_dtype)
|
||||
# module.sin_cached = module.sin_cached.to(correct_dtype)
|
||||
# pass
|
||||
# pass
|
||||
pass
|
||||
def post_patch(model, tokenizer):
|
||||
# Gemma does not downcast RoPE
|
||||
model, tokenizer = patch_model_and_tokenizer(model, tokenizer, downcast_rope = False)
|
||||
|
||||
# Add 1 to weight
|
||||
# return output * (1 + self.weight)
|
||||
|
|
@ -425,6 +374,6 @@ class FastGemmaModel(FastLlamaModel):
|
|||
for _ in range(3):
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
return model
|
||||
return model, tokenizer
|
||||
pass
|
||||
pass
|
||||
|
|
|
|||
|
|
@ -490,60 +490,9 @@ class FastGemma2Model(FastLlamaModel):
|
|||
|
||||
|
||||
@staticmethod
|
||||
def post_patch(model):
|
||||
# Patch model for Gemma
|
||||
layers = model.model.layers
|
||||
|
||||
# Torch.compile fails on embedding matrix??
|
||||
# Workaround randomnly fixes it for torch versions < 2.2
|
||||
model.model.embed_tokens = torch.nn.Embedding.from_pretrained(model.model.embed_tokens.weight)
|
||||
model.config.update({"unsloth_version" : __version__})
|
||||
|
||||
# We also do this for the lm_head
|
||||
lm_head = torch.nn.Linear(1, 1, bias = None)
|
||||
del lm_head.weight
|
||||
lm_head.weight = model.lm_head.weight
|
||||
lm_head.in_features = lm_head.weight.shape[1]
|
||||
lm_head.out_features = lm_head.weight.shape[0]
|
||||
model.lm_head = lm_head
|
||||
|
||||
# Gemma has tied weights! This means lm_head == embed_tokens
|
||||
if model.model.embed_tokens.weight.data_ptr() != model.lm_head.weight.data_ptr():
|
||||
lm_head = torch.nn.Linear(1, 1, bias = None)
|
||||
del lm_head.weight
|
||||
lm_head.weight = model.model.embed_tokens.weight
|
||||
lm_head.in_features = lm_head.weight.shape[1]
|
||||
lm_head.out_features = lm_head.weight.shape[0]
|
||||
model.lm_head = lm_head
|
||||
pass
|
||||
|
||||
# Also patch all dtypes - BnB seems to not allocate the correct type?
|
||||
# BnB default dtype seems to be float16!
|
||||
correct_dtype = lm_head.weight.dtype
|
||||
|
||||
for name, module in model.named_modules():
|
||||
if isinstance(module, (Bnb_Linear4bit, Peft_Linear4bit)):
|
||||
weight = module.weight
|
||||
quant_state = weight.quant_state
|
||||
|
||||
if type(quant_state) is list:
|
||||
# BnB seems to have float16 as default!
|
||||
module.weight.quant_state[2] = correct_dtype # Cast to correct dtype
|
||||
else:
|
||||
# https://github.com/TimDettmers/bitsandbytes/pull/763/files
|
||||
quant_state.dtype = correct_dtype
|
||||
pass
|
||||
pass
|
||||
# Downcast RoPE embedding to correct data type
|
||||
# RoPE must be done in float32 for Gemma
|
||||
# if (name.endswith("rotary_emb") or hasattr(module, "cos_cached")) \
|
||||
# and (module.cos_cached.dtype != correct_dtype):
|
||||
|
||||
# module.cos_cached = module.cos_cached.to(correct_dtype)
|
||||
# module.sin_cached = module.sin_cached.to(correct_dtype)
|
||||
# pass
|
||||
# pass
|
||||
pass
|
||||
def post_patch(model, tokenizer):
|
||||
# Gemma does not downcast RoPE
|
||||
model, tokenizer = patch_model_and_tokenizer(model, tokenizer, downcast_rope = False)
|
||||
|
||||
# Add 1 to weight
|
||||
# return output * (1 + self.weight)
|
||||
|
|
@ -576,6 +525,6 @@ class FastGemma2Model(FastLlamaModel):
|
|||
for _ in range(3):
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
return model
|
||||
return model, tokenizer
|
||||
pass
|
||||
pass
|
||||
|
|
|
|||
|
|
@ -57,8 +57,6 @@ from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING
|
|||
from transformers import set_seed as transformers_set_seed
|
||||
from peft import LoraConfig, TaskType, get_peft_model as _get_peft_model
|
||||
from peft import PeftModelForCausalLM
|
||||
from bitsandbytes.nn import Linear4bit as Bnb_Linear4bit
|
||||
from peft.tuners.lora import Linear4bit as Peft_Linear4bit
|
||||
from ..save import patch_saving_functions
|
||||
import re, os, inspect, math, sys
|
||||
try:
|
||||
|
|
@ -1518,6 +1516,7 @@ class FastLlamaModel:
|
|||
pass
|
||||
# Return old flag
|
||||
os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = old_hf_transfer
|
||||
os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
|
||||
|
||||
model_patcher.pre_patch()
|
||||
get_statistics() # For debugging - we use a download counter to see if environments are not breaking
|
||||
|
|
@ -1621,7 +1620,7 @@ class FastLlamaModel:
|
|||
)
|
||||
|
||||
model, tokenizer = patch_tokenizer(model, tokenizer)
|
||||
model = model_patcher.post_patch(model)
|
||||
model, tokenizer = model_patcher.post_patch(model, tokenizer)
|
||||
|
||||
# Patch up QKV / O and MLP
|
||||
for idx, layer in enumerate(model.model.layers):
|
||||
|
|
@ -1797,93 +1796,15 @@ class FastLlamaModel:
|
|||
internal_model = internal_model.model
|
||||
pass
|
||||
internal_model._saved_temp_tokenizer = tokenizer
|
||||
|
||||
# Also fix torch_dtype
|
||||
internal_model = model
|
||||
while hasattr(internal_model, "model"):
|
||||
if hasattr(internal_model, "config"):
|
||||
if internal_model.config.torch_dtype == "float32":
|
||||
internal_model.config.torch_dtype = torch.float32
|
||||
elif internal_model.config.torch_dtype == "bfloat16":
|
||||
internal_model.config.torch_dtype = torch.bfloat16
|
||||
elif internal_model.config.torch_dtype == "float16":
|
||||
internal_model.config.torch_dtype = torch.float16
|
||||
pass
|
||||
pass
|
||||
internal_model = internal_model.model
|
||||
pass
|
||||
if hasattr(internal_model, "config"):
|
||||
if internal_model.config.torch_dtype == "float32":
|
||||
internal_model.config.torch_dtype = torch.float32
|
||||
elif internal_model.config.torch_dtype == "bfloat16":
|
||||
internal_model.config.torch_dtype = torch.bfloat16
|
||||
elif internal_model.config.torch_dtype == "float16":
|
||||
internal_model.config.torch_dtype = torch.float16
|
||||
pass
|
||||
pass
|
||||
|
||||
return model, tokenizer
|
||||
pass
|
||||
|
||||
|
||||
@staticmethod
|
||||
def post_patch(model):
|
||||
# Patch model
|
||||
layers = model.model.layers
|
||||
|
||||
# Torch.compile fails on embedding matrix??
|
||||
# Workaround randomnly fixes it for torch versions < 2.
|
||||
model.set_input_embeddings(torch.nn.Embedding.from_pretrained(model.get_input_embeddings().weight))
|
||||
model.config.update({"unsloth_version" : __version__})
|
||||
|
||||
# We also do this for the lm_head
|
||||
lm_head = torch.nn.Linear(1, 1, bias = None)
|
||||
del lm_head.weight
|
||||
lm_head.weight = model.get_output_embeddings().weight
|
||||
lm_head.in_features = lm_head.weight.shape[1]
|
||||
lm_head.out_features = lm_head.weight.shape[0]
|
||||
model.lm_head = lm_head
|
||||
|
||||
# Also patch all dtypes - BnB seems to not allocate the correct type?
|
||||
# BnB default dtype seems to be float16!
|
||||
correct_dtype = lm_head.weight.dtype
|
||||
|
||||
for name, module in model.named_modules():
|
||||
if isinstance(module, (Bnb_Linear4bit, Peft_Linear4bit)):
|
||||
weight = module.weight
|
||||
quant_state = weight.quant_state
|
||||
|
||||
if type(quant_state) is list:
|
||||
# BnB seems to have float16 as default!
|
||||
module.weight.quant_state[2] = correct_dtype # Cast to correct dtype
|
||||
else:
|
||||
# https://github.com/TimDettmers/bitsandbytes/pull/763/files
|
||||
quant_state.dtype = correct_dtype
|
||||
pass
|
||||
pass
|
||||
# Downcast RoPE embedding to correct data type
|
||||
if (name.endswith("rotary_emb") or hasattr(module, "cos_cached")):
|
||||
|
||||
if hasattr(module, "cos_cached") and \
|
||||
(module.cos_cached.dtype != correct_dtype):
|
||||
|
||||
module.cos_cached = module.cos_cached.to(correct_dtype)
|
||||
module.sin_cached = module.sin_cached.to(correct_dtype)
|
||||
|
||||
elif hasattr(module, "short_cos_cached") and \
|
||||
(module.short_cos_cached.dtype != correct_dtype):
|
||||
|
||||
module.short_cos_cached = module.short_cos_cached.to(correct_dtype)
|
||||
module.short_sin_cached = module.short_sin_cached.to(correct_dtype)
|
||||
pass
|
||||
pass
|
||||
pass
|
||||
|
||||
# Clear deleted GPU items
|
||||
for _ in range(3):
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
return model
|
||||
def post_patch(model, tokenizer):
|
||||
model, tokenizer = patch_model_and_tokenizer(model, tokenizer, downcast_rope = True)
|
||||
return model, tokenizer
|
||||
pass
|
||||
|
||||
|
||||
|
|
@ -1910,6 +1831,11 @@ class FastLlamaModel:
|
|||
):
|
||||
transformers_set_seed(random_state)
|
||||
|
||||
if type(r) is not int:
|
||||
raise TypeError(f"Unsloth: Rank of {str(r)} must be an integer.")
|
||||
if r <= 0:
|
||||
raise TypeError(f"Unsloth: Rank of {str(r)} must be larger than 0.")
|
||||
|
||||
if isinstance(model, PeftModelForCausalLM):
|
||||
# Check if exactly the same and then pass through!
|
||||
assert(hasattr(model, "peft_config"))
|
||||
|
|
|
|||
|
|
@ -49,6 +49,7 @@ __all__ = [
|
|||
keynames = "\n" + "\n".join(os.environ.keys())
|
||||
IS_COLAB_ENVIRONMENT = "\nCOLAB_" in keynames
|
||||
IS_KAGGLE_ENVIRONMENT = "\nKAGGLE_" in keynames
|
||||
KAGGLE_TMP = "/tmp"
|
||||
del keynames
|
||||
|
||||
# Weights
|
||||
|
|
@ -447,13 +448,20 @@ def unsloth_save_model(
|
|||
if push_to_hub and "/" in save_directory:
|
||||
|
||||
# +1 solves absolute path issues
|
||||
username = save_directory[:save_directory.find("/")]
|
||||
new_save_directory = save_directory[save_directory.find("/")+1:]
|
||||
|
||||
logger.warning_once(
|
||||
f"Unsloth: You are pushing to hub, but you passed your HF username = {username}.\n"\
|
||||
f"We shall truncate {save_directory} to {new_save_directory}"
|
||||
)
|
||||
new_save_directory = save_directory
|
||||
username = new_save_directory[:new_save_directory.find("/")]
|
||||
new_save_directory = new_save_directory[new_save_directory.find("/")+1:]
|
||||
if IS_KAGGLE_ENVIRONMENT:
|
||||
new_save_directory = os.path.join(KAGGLE_TMP, new_save_directory[new_save_directory.find("/")+1:])
|
||||
logger.warning_once(
|
||||
"Unsloth: You are pushing to hub in Kaggle environment.\n"\
|
||||
f"To save memory, we shall move {save_directory} to {new_save_directory}"
|
||||
)
|
||||
else:
|
||||
logger.warning_once(
|
||||
f"Unsloth: You are pushing to hub, but you passed your HF username = {username}.\n"\
|
||||
f"We shall truncate {save_directory} to {new_save_directory}"
|
||||
)
|
||||
|
||||
save_pretrained_settings["save_directory"] = new_save_directory
|
||||
tokenizer_save_settings ["save_directory"] = new_save_directory
|
||||
|
|
@ -507,6 +515,10 @@ def unsloth_save_model(
|
|||
f"{round(max_ram/1024/1024/1024, 2)} out of "\
|
||||
f"{round(psutil.virtual_memory().total/1024/1024/1024, 2)} RAM for saving.")
|
||||
|
||||
# Move temporary_location to /tmp in Kaggle
|
||||
if IS_KAGGLE_ENVIRONMENT:
|
||||
temporary_location = os.path.join(KAGGLE_TMP, temporary_location)
|
||||
|
||||
# Max directory for disk saving
|
||||
if not os.path.exists(temporary_location):
|
||||
os.makedirs(temporary_location)
|
||||
|
|
@ -708,7 +720,7 @@ def unsloth_save_model(
|
|||
print("Done.")
|
||||
|
||||
if push_to_hub and hasattr(model, "config"):
|
||||
print(f"Saved merged model to https://huggingface.co/{username}/{save_directory.lstrip('/')}")
|
||||
print(f"Saved merged model to https://huggingface.co/{username}/{save_directory.lstrip('/').split('/')[-1]}")
|
||||
pass
|
||||
|
||||
save_pretrained_settings["state_dict"] = None
|
||||
|
|
@ -1108,14 +1120,17 @@ def save_to_gguf(
|
|||
# Check if quantization succeeded!
|
||||
if not os.path.isfile(final_location):
|
||||
if IS_KAGGLE_ENVIRONMENT:
|
||||
raise RuntimeError(
|
||||
f"Unsloth: Quantization failed for {final_location}\n"\
|
||||
"You are in a Kaggle environment, which might be the reason this is failing.\n"\
|
||||
"Kaggle only provides 20GB of disk space. Merging to 16bit for 7b models use 16GB of space.\n"\
|
||||
"This means using `model.{save_pretrained/push_to_hub}_merged` works, but\n"\
|
||||
"`model.{save_pretrained/push_to_hub}_gguf will use too much disk space.\n"\
|
||||
"I suggest you to save the 16bit model first, then use manual llama.cpp conversion."
|
||||
)
|
||||
if not Path(final_location).resolve().is_relative_to(Path('/tmp').resolve()):
|
||||
raise RuntimeError(
|
||||
f"Unsloth: Quantization failed for {final_location}\n"\
|
||||
"You are in a Kaggle environment, which might be the reason this is failing.\n"\
|
||||
"Kaggle only provides 20GB of disk space in the working directory.\n"\
|
||||
"Merging to 16bit for 7b models use 16GB of space.\n"\
|
||||
"This means using `model.{save_pretrained/push_to_hub}_merged` works, but\n"\
|
||||
"`model.{save_pretrained/push_to_hub}_gguf will use too much disk space.\n"\
|
||||
"You can try saving it to the `/tmp` directory for larger disk space.\n"\
|
||||
"I suggest you to save the 16bit model first, then use manual llama.cpp conversion."
|
||||
)
|
||||
else:
|
||||
raise RuntimeError(
|
||||
f"Unsloth: Quantization failed for {final_location}\n"\
|
||||
|
|
@ -1156,14 +1171,17 @@ def save_to_gguf(
|
|||
# Check if quantization succeeded!
|
||||
if not os.path.isfile(final_location):
|
||||
if IS_KAGGLE_ENVIRONMENT:
|
||||
raise RuntimeError(
|
||||
f"Unsloth: Quantization failed for {final_location}\n"\
|
||||
"You are in a Kaggle environment, which might be the reason this is failing.\n"\
|
||||
"Kaggle only provides 20GB of disk space. Merging to 16bit for 7b models use 16GB of space.\n"\
|
||||
"This means using `model.{save_pretrained/push_to_hub}_merged` works, but\n"\
|
||||
"`model.{save_pretrained/push_to_hub}_gguf will use too much disk space.\n"\
|
||||
"I suggest you to save the 16bit model first, then use manual llama.cpp conversion."
|
||||
)
|
||||
if not Path(final_location).resolve().is_relative_to(Path('/tmp').resolve()):
|
||||
raise RuntimeError(
|
||||
f"Unsloth: Quantization failed for {final_location}\n"\
|
||||
"You are in a Kaggle environment, which might be the reason this is failing.\n"\
|
||||
"Kaggle only provides 20GB of disk space in the working directory.\n"\
|
||||
"Merging to 16bit for 7b models use 16GB of space.\n"\
|
||||
"This means using `model.{save_pretrained/push_to_hub}_merged` works, but\n"\
|
||||
"`model.{save_pretrained/push_to_hub}_gguf will use too much disk space.\n"\
|
||||
"You can try saving it to the `/tmp` directory for larger disk space.\n"\
|
||||
"I suggest you to save the 16bit model first, then use manual llama.cpp conversion."
|
||||
)
|
||||
else:
|
||||
raise RuntimeError(
|
||||
"Unsloth: Quantization failed! You might have to compile llama.cpp yourself, then run this again.\n"\
|
||||
|
|
|
|||
|
|
@ -64,6 +64,7 @@ IGNORED_TOKENIZER_NAMES = frozenset(
|
|||
keynames = "\n" + "\n".join(os.environ.keys())
|
||||
IS_COLAB_ENVIRONMENT = "\nCOLAB_" in keynames
|
||||
IS_KAGGLE_ENVIRONMENT = "\nKAGGLE_" in keynames
|
||||
KAGGLE_TMP = "/tmp"
|
||||
del keynames
|
||||
|
||||
|
||||
|
|
@ -470,8 +471,12 @@ def _load_correct_tokenizer(
|
|||
cache_dir = "huggingface_tokenizers_cache",
|
||||
fix_tokenizer = True,
|
||||
):
|
||||
if IS_COLAB_ENVIRONMENT or IS_KAGGLE_ENVIRONMENT:
|
||||
if IS_COLAB_ENVIRONMENT:
|
||||
cache_dir = cache_dir
|
||||
elif IS_KAGGLE_ENVIRONMENT:
|
||||
# /tmp of Kaggle seems has a 80GB limit!
|
||||
# Let's utilize them
|
||||
cache_dir = os.path.join(KAGGLE_TMP, cache_dir)
|
||||
else:
|
||||
cache_dir = None
|
||||
pass
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue