ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Update cross_entropy_loss.py

Browse source
This commit is contained in:
Daniel Han 2024-10-31 02:00:04 -07:00
commit f3bd5d6d33
1 changed files with 7 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

14
unsloth/kernels/cross_entropy_loss.py
View file

@ -30,8 +30,8 @@ def _cross_entropy_forward(
loss_ptr ,
logsumexp_ptr ,
labels_ptr ,
VOCAB_SIZE : tl.constexpr,
BLOCK_SIZE : tl.constexpr,
VOCAB_SIZE : tl.constexpr(tl.int32),
BLOCK_SIZE : tl.constexpr(tl.int32),
DO_SOFTCAPPING : tl.constexpr(tl.int1),
SOFTCAP : tl.constexpr(tl.float32),
DO_LOGIT_SCALING : tl.constexpr(tl.int1),
@ -103,9 +103,9 @@ def _chunked_cross_entropy_forward(
loss_ptr ,
logsumexp_ptr ,
labels_ptr ,
VOCAB_SIZE ,
N_CHUNKS ,
BLOCK_SIZE : tl.constexpr,
VOCAB_SIZE : tl.constexpr(tl.int32),
N_CHUNKS : tl.constexpr(tl.int32),
BLOCK_SIZE : tl.constexpr(tl.int32),
DO_SOFTCAPPING : tl.constexpr(tl.int1),
SOFTCAP : tl.constexpr(tl.float32),
DO_LOGIT_SCALING : tl.constexpr(tl.int1),
@ -187,8 +187,8 @@ def _cross_entropy_backward(
dloss_row_stride ,
logsumexp_ptr ,
labels_ptr ,
VOCAB_SIZE ,
BLOCK_SIZE : tl.constexpr,
VOCAB_SIZE : tl.constexpr(tl.int32),
BLOCK_SIZE : tl.constexpr(tl.int32),
DO_SOFTCAPPING : tl.constexpr(tl.int1),
SOFTCAP : tl.constexpr(tl.float32),
DO_LOGIT_SCALING : tl.constexpr(tl.int1),