From 6eb409ace3e369f082d4fcf89717bfee6b50d13e Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 3 Nov 2024 23:53:03 -0800 Subject: [PATCH] int64 --- unsloth/kernels/cross_entropy_loss.py | 20 ++++++++++---------- unsloth/models/_utils.py | 13 ++++++++++++- 2 files changed, 22 insertions(+), 11 deletions(-) diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py index b4780db278..11e582711d 100644 --- a/unsloth/kernels/cross_entropy_loss.py +++ b/unsloth/kernels/cross_entropy_loss.py @@ -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 diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 837b4849f1..a260466915 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -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