From e23ef99c033474de786d4947dce1152f9db5e9d1 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 3 Nov 2024 21:38:19 -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 19b384a3bd..8b2204436b 100644 --- a/unsloth/kernels/cross_entropy_loss.py +++ b/unsloth/kernels/cross_entropy_loss.py @@ -38,9 +38,9 @@ def _cross_entropy_forward( labels_ptr , VOCAB_SIZE : tl.constexpr(tl.int32), BLOCK_SIZE : tl.constexpr(tl.int32), - DO_SOFTCAPPING : tl.constexpr(tl.int1), + DO_SOFTCAPPING , SOFTCAP , - DO_LOGIT_SCALING : tl.constexpr(tl.int1), + DO_LOGIT_SCALING , LOGIT_SCALE , ): """ @@ -112,9 +112,9 @@ def _chunked_cross_entropy_forward( VOCAB_SIZE : tl.constexpr(tl.int32), N_CHUNKS : tl.constexpr(tl.int32), BLOCK_SIZE : tl.constexpr(tl.int32), - DO_SOFTCAPPING : tl.constexpr(tl.int1), + DO_SOFTCAPPING , SOFTCAP , - DO_LOGIT_SCALING : tl.constexpr(tl.int1), + DO_LOGIT_SCALING , LOGIT_SCALE , ): """ @@ -195,9 +195,9 @@ def _cross_entropy_backward( labels_ptr , VOCAB_SIZE : tl.constexpr(tl.int32), BLOCK_SIZE : tl.constexpr(tl.int32), - DO_SOFTCAPPING : tl.constexpr(tl.int1), + DO_SOFTCAPPING , SOFTCAP , - DO_LOGIT_SCALING : tl.constexpr(tl.int1), + DO_LOGIT_SCALING , LOGIT_SCALE , ): """