From a54954fe153ca164e50edd30b4badabf146867c6 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 3 Nov 2024 21:31:18 -0800 Subject: [PATCH] Update cross_entropy_loss.py --- unsloth/kernels/cross_entropy_loss.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py index 4ebe69d565..7939758fda 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 , BLOCK_SIZE : tl.constexpr(tl.int32), DO_SOFTCAPPING : tl.constexpr(tl.int1), - SOFTCAP , + SOFTCAP : tl.constexpr(tl.float32), DO_LOGIT_SCALING : tl.constexpr(tl.int1), - LOGIT_SCALE , + LOGIT_SCALE : tl.constexpr(tl.float32), ): """ Cross Entropy Loss = 1/n sum [ -yi log(Pi) ] @@ -113,9 +113,9 @@ def _chunked_cross_entropy_forward( N_CHUNKS , BLOCK_SIZE : tl.constexpr(tl.int32), DO_SOFTCAPPING : tl.constexpr(tl.int1), - SOFTCAP , + SOFTCAP : tl.constexpr(tl.float32), DO_LOGIT_SCALING : tl.constexpr(tl.int1), - LOGIT_SCALE , + LOGIT_SCALE : tl.constexpr(tl.float32), ): """ 256K vocab divided in 4 chunks @@ -196,9 +196,9 @@ def _cross_entropy_backward( VOCAB_SIZE , BLOCK_SIZE : tl.constexpr(tl.int32), DO_SOFTCAPPING : tl.constexpr(tl.int1), - SOFTCAP , + SOFTCAP : tl.constexpr(tl.float32), DO_LOGIT_SCALING : tl.constexpr(tl.int1), - LOGIT_SCALE , + LOGIT_SCALE : tl.constexpr(tl.float32), ): """ CE_i = -y log(P) = y * (log[sum(exp(x))] - x)