From 7d5111a40c13de36aa675d19fcb9c6c6a9deb5de Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 4 Nov 2024 00:05:51 -0800 Subject: [PATCH] constexpr --- unsloth/kernels/cross_entropy_loss.py | 12 ++++++------ unsloth/models/_utils.py | 2 +- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py index 9cf7ddc36d..d1b8ae8275 100644 --- a/unsloth/kernels/cross_entropy_loss.py +++ b/unsloth/kernels/cross_entropy_loss.py @@ -39,9 +39,9 @@ def _cross_entropy_forward( VOCAB_SIZE : tl.constexpr, BLOCK_SIZE : tl.constexpr, DO_SOFTCAPPING , - SOFTCAP : tl.constexpr, + SOFTCAP , DO_LOGIT_SCALING , - LOGIT_SCALE : tl.constexpr, + LOGIT_SCALE , ): """ Cross Entropy Loss = 1/n sum [ -yi log(Pi) ] @@ -113,9 +113,9 @@ def _chunked_cross_entropy_forward( N_CHUNKS : tl.constexpr, BLOCK_SIZE : tl.constexpr, DO_SOFTCAPPING , - SOFTCAP : tl.constexpr, + SOFTCAP , DO_LOGIT_SCALING , - LOGIT_SCALE : tl.constexpr, + LOGIT_SCALE , ): """ 256K vocab divided in 4 chunks @@ -196,9 +196,9 @@ def _cross_entropy_backward( VOCAB_SIZE : tl.constexpr, BLOCK_SIZE : tl.constexpr, DO_SOFTCAPPING , - SOFTCAP : tl.constexpr, + SOFTCAP , DO_LOGIT_SCALING , - LOGIT_SCALE : tl.constexpr, + LOGIT_SCALE , ): """ CE_i = -y log(P) = y * (log[sum(exp(x))] - x) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index dd37d26ae4..5c099548d6 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -366,7 +366,7 @@ pass # ============================================= # Torch compile settings -UNSLOTH_COMPILE_DEBUG = "UNSLOTH_COMPILE_DEBUG" in os.environ +UNSLOTH_COMPILE_DEBUG = True #"UNSLOTH_COMPILE_DEBUG" in os.environ UNSLOTH_COMPILE_MAXIMUM = "UNSLOTH_COMPILE_MAXIMUM" in os.environ # Just remove max_autotune_gemm warning import functools