From c190577837b66a7eeca012f47b4ee8240eeb7fdf Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 3 Nov 2024 21:29:19 -0800 Subject: [PATCH] Update cross_entropy_loss.py --- unsloth/kernels/cross_entropy_loss.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py index f92ddbd05d..4ebe69d565 100644 --- a/unsloth/kernels/cross_entropy_loss.py +++ b/unsloth/kernels/cross_entropy_loss.py @@ -36,12 +36,12 @@ def _cross_entropy_forward( loss_ptr , logsumexp_ptr , labels_ptr , - VOCAB_SIZE : tl.constexpr(tl.int32), + VOCAB_SIZE , BLOCK_SIZE : tl.constexpr(tl.int32), DO_SOFTCAPPING : tl.constexpr(tl.int1), - SOFTCAP : tl.constexpr(tl.float32), + SOFTCAP , DO_LOGIT_SCALING : tl.constexpr(tl.int1), - LOGIT_SCALE : tl.constexpr(tl.float32), + LOGIT_SCALE , ): """ Cross Entropy Loss = 1/n sum [ -yi log(Pi) ] @@ -109,13 +109,13 @@ def _chunked_cross_entropy_forward( loss_ptr , logsumexp_ptr , labels_ptr , - VOCAB_SIZE : tl.constexpr(tl.int32), - N_CHUNKS : tl.constexpr(tl.int32), + VOCAB_SIZE , + N_CHUNKS , BLOCK_SIZE : tl.constexpr(tl.int32), DO_SOFTCAPPING : tl.constexpr(tl.int1), - SOFTCAP : tl.constexpr(tl.float32), + SOFTCAP , DO_LOGIT_SCALING : tl.constexpr(tl.int1), - LOGIT_SCALE : tl.constexpr(tl.float32), + LOGIT_SCALE , ): """ 256K vocab divided in 4 chunks @@ -193,12 +193,12 @@ def _cross_entropy_backward( dloss_row_stride , logsumexp_ptr , labels_ptr , - VOCAB_SIZE : tl.constexpr(tl.int32), + VOCAB_SIZE , BLOCK_SIZE : tl.constexpr(tl.int32), DO_SOFTCAPPING : tl.constexpr(tl.int1), - SOFTCAP : tl.constexpr(tl.float32), + SOFTCAP , DO_LOGIT_SCALING : tl.constexpr(tl.int1), - LOGIT_SCALE : tl.constexpr(tl.float32), + LOGIT_SCALE , ): """ CE_i = -y log(P) = y * (log[sum(exp(x))] - x)