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:31:18 -08:00
commit a54954fe15
1 changed files with 6 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

12
unsloth/kernels/cross_entropy_loss.py
View file

@ -39,9 +39,9 @@ def _cross_entropy_forward(
VOCAB_SIZE ,
BLOCK_SIZE : tl.constexpr(tl.int32),
DO_SOFTCAPPING : tl.constexpr(tl.int1),
SOFTCAP ,
SOFTCAP : tl.constexpr(tl.float32),
DO_LOGIT_SCALING : tl.constexpr(tl.int1),
LOGIT_SCALE ,
LOGIT_SCALE : tl.constexpr(tl.float32),
):
"""
Cross Entropy Loss = 1/n sum [ -yi log(Pi) ]
@ -113,9 +113,9 @@ def _chunked_cross_entropy_forward(
N_CHUNKS ,
BLOCK_SIZE : tl.constexpr(tl.int32),
DO_SOFTCAPPING : tl.constexpr(tl.int1),
SOFTCAP ,
SOFTCAP : tl.constexpr(tl.float32),
DO_LOGIT_SCALING : tl.constexpr(tl.int1),
LOGIT_SCALE ,
LOGIT_SCALE : tl.constexpr(tl.float32),
):
"""
256K vocab divided in 4 chunks
@ -196,9 +196,9 @@ def _cross_entropy_backward(
VOCAB_SIZE ,
BLOCK_SIZE : tl.constexpr(tl.int32),
DO_SOFTCAPPING : tl.constexpr(tl.int1),
SOFTCAP ,
SOFTCAP : tl.constexpr(tl.float32),
DO_LOGIT_SCALING : tl.constexpr(tl.int1),
LOGIT_SCALE ,
LOGIT_SCALE : tl.constexpr(tl.float32),
):
"""
CE_i = -y log(P) = y * (log[sum(exp(x))] - x)