diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py index f0193c74d8..13d90baafe 100644 --- a/unsloth/kernels/cross_entropy_loss.py +++ b/unsloth/kernels/cross_entropy_loss.py @@ -181,10 +181,10 @@ def _chunked_cross_entropy_forward( pass -@triton.heuristics({ - "DO_SOFTCAPPING": lambda args: bool(args["DO_SOFTCAPPING" ]), - "DO_LOGIT_SCALING": lambda args: bool(args["DO_LOGIT_SCALING"]), -}) +# @triton.heuristics({ +# "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 , @@ -345,6 +345,8 @@ class Fast_CrossEntropyLoss(torch.autograd.Function): n_rows, vocab_size = logits.shape BLOCK_SIZE : int = 4096 + div : int + mod : int div, mod = divmod(vocab_size, BLOCK_SIZE) n_blocks : int = div + (mod != 0) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index b5c17434c7..bb004d588d 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -393,7 +393,7 @@ UNSLOTH_COMPILE_DEBUG = os.environ.get("UNSLOTH_COMPILE_DEBUG", "0") == "1" UNSLOTH_COMPILE_MAXIMUM = os.environ.get("UNSLOTH_COMPILE_MAXIMUM", "0") == "1" patch_torch_compile( debug = UNSLOTH_COMPILE_DEBUG, - O3 = UNSLOTH_COMPILE_MAXIMUM, + O3 = UNSLOTH_COMPILE_MAXIMUM, ) torch_compile_options = {