From 66b807b6ddaaac86241f3200530e8403ea95b44c Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 3 Nov 2024 21:11:34 -0800 Subject: [PATCH] Update cross_entropy_loss.py --- unsloth/kernels/cross_entropy_loss.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py index 15a148b00e..bcdef31997 100644 --- a/unsloth/kernels/cross_entropy_loss.py +++ b/unsloth/kernels/cross_entropy_loss.py @@ -32,11 +32,11 @@ from unsloth_zoo.loss_utils import ( @triton.jit def _cross_entropy_forward( logits_ptr , - logits_row_stride : tl.constexpr, + logits_row_stride , loss_ptr , logsumexp_ptr , labels_ptr , - VOCAB_SIZE , + VOCAB_SIZE : tl.constexpr, BLOCK_SIZE : tl.constexpr, DO_SOFTCAPPING , SOFTCAP , @@ -105,12 +105,12 @@ pass @triton.jit def _chunked_cross_entropy_forward( logits_ptr , - logits_row_stride : tl.constexpr, + logits_row_stride , loss_ptr , logsumexp_ptr , labels_ptr , - VOCAB_SIZE , - N_CHUNKS , + VOCAB_SIZE : tl.constexpr, + N_CHUNKS : tl.constexpr, BLOCK_SIZE : tl.constexpr, DO_SOFTCAPPING , SOFTCAP , @@ -188,12 +188,12 @@ pass @triton.jit def _cross_entropy_backward( logits_ptr , - logits_row_stride : tl.constexpr, + logits_row_stride , dloss_ptr , dloss_row_stride , logsumexp_ptr , labels_ptr , - VOCAB_SIZE , + VOCAB_SIZE : tl.constexpr, BLOCK_SIZE : tl.constexpr, DO_SOFTCAPPING , SOFTCAP ,