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

int64

Browse source
This commit is contained in:
Daniel Han 2024-11-03 23:53:03 -08:00
commit 6eb409ace3
2 changed files with 22 additions and 11 deletions
Download patch file Download diff file Expand all files Collapse all files

20
unsloth/kernels/cross_entropy_loss.py
View file

@ -36,8 +36,8 @@ def _cross_entropy_forward(
loss_ptr ,
logsumexp_ptr ,
labels_ptr ,
VOCAB_SIZE ,
BLOCK_SIZE : tl.constexpr(tl.int32),
VOCAB_SIZE : tl.constexpr,
BLOCK_SIZE : tl.constexpr,
DO_SOFTCAPPING ,
SOFTCAP ,
DO_LOGIT_SCALING ,
@ -64,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
logits_ptr += row_idx * logits_row_stride.to(tl.int64)
loss_ptr += row_idx
logsumexp_ptr += row_idx
labels_ptr += row_idx
@ -109,9 +109,9 @@ def _chunked_cross_entropy_forward(
loss_ptr ,
logsumexp_ptr ,
labels_ptr ,
VOCAB_SIZE ,
N_CHUNKS ,
BLOCK_SIZE : tl.constexpr(tl.int32),
VOCAB_SIZE : tl.constexpr,
N_CHUNKS : tl.constexpr,
BLOCK_SIZE : tl.constexpr,
DO_SOFTCAPPING ,
SOFTCAP ,
DO_LOGIT_SCALING ,
@ -143,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
logits_ptr += row_idx * logits_row_stride.to(tl.int64)
loss_ptr += row_idx
logsumexp_ptr += row_idx * N_CHUNKS + chunk_idx
labels_ptr += row_idx
@ -193,8 +193,8 @@ def _cross_entropy_backward(
dloss_row_stride ,
logsumexp_ptr ,
labels_ptr ,
VOCAB_SIZE ,
BLOCK_SIZE : tl.constexpr(tl.int32),
VOCAB_SIZE : tl.constexpr,
BLOCK_SIZE : tl.constexpr,
DO_SOFTCAPPING ,
SOFTCAP ,
DO_LOGIT_SCALING ,
@ -218,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
logits_ptr += row_idx * logits_row_stride.to(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

13
unsloth/models/_utils.py
View file

@ -92,6 +92,17 @@ 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.
import transformers.training_args.logger
transformers.training_args.logger.addFilter(HideLoggingMessage("The speedups"))
# =============================================
# =============================================
@ -380,7 +391,7 @@ torch_compile_options = {
import accelerate
def torch_compile_kwargs(*args, **kwargs):
print("Unsloth: Enabled auto compiling")
return {"dynamic" : True, "fullgraph" : False, "options" : torch_compile_options}
return {"dynamic" : True, "fullgraph" : False, "options" : torch_compile_options,}
pass
accelerate.utils.dataclasses.TorchDynamoPlugin.to_kwargs = torch_compile_kwargs