From 1aed9df52960fde0876faa50fd4016b76d3c19e4 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 3 Nov 2024 21:09:34 -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 17168b230a..15a148b00e 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, + VOCAB_SIZE , 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) ] @@ -109,13 +109,13 @@ def _chunked_cross_entropy_forward( loss_ptr , logsumexp_ptr , labels_ptr , - VOCAB_SIZE : tl.constexpr, - N_CHUNKS : tl.constexpr, + VOCAB_SIZE , + N_CHUNKS , 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 @@ -193,12 +193,12 @@ def _cross_entropy_backward( dloss_row_stride , logsumexp_ptr , labels_ptr , - VOCAB_SIZE : tl.constexpr, + VOCAB_SIZE , 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)