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-11-03 21:18:26 -08:00
commit f6d50c78ac
1 changed files with 9 additions and 9 deletions
Download patch file Download diff file Expand all files Collapse all files

18
unsloth/kernels/cross_entropy_loss.py
View file

@ -37,10 +37,10 @@ def _cross_entropy_forward(
logsumexp_ptr ,
labels_ptr ,
VOCAB_SIZE ,
BLOCK_SIZE : tl.constexpr,
DO_SOFTCAPPING ,
BLOCK_SIZE : tl.constexpr(tl.int32),
DO_SOFTCAPPING : tl.constexpr(tl.int1),
SOFTCAP ,
DO_LOGIT_SCALING ,
DO_LOGIT_SCALING : tl.constexpr(tl.int1),
LOGIT_SCALE ,
):
"""
@ -111,10 +111,10 @@ def _chunked_cross_entropy_forward(
labels_ptr ,
VOCAB_SIZE ,
N_CHUNKS ,
BLOCK_SIZE : tl.constexpr,
DO_SOFTCAPPING ,
BLOCK_SIZE : tl.constexpr(tl.int32),
DO_SOFTCAPPING : tl.constexpr(tl.int1),
SOFTCAP ,
DO_LOGIT_SCALING ,
DO_LOGIT_SCALING : tl.constexpr(tl.int1),
LOGIT_SCALE ,
):
"""
@ -194,10 +194,10 @@ def _cross_entropy_backward(
logsumexp_ptr ,
labels_ptr ,
VOCAB_SIZE ,
BLOCK_SIZE : tl.constexpr,
DO_SOFTCAPPING ,
BLOCK_SIZE : tl.constexpr(tl.int32),
DO_SOFTCAPPING : tl.constexpr(tl.int1),
SOFTCAP ,
DO_LOGIT_SCALING ,
DO_LOGIT_SCALING : tl.constexpr(tl.int1),
LOGIT_SCALE ,
):
"""