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:33:26 -08:00
commit 539a1406bd
1 changed files with 13 additions and 13 deletions
Download patch file Download diff file Expand all files Collapse all files

26
unsloth/kernels/cross_entropy_loss.py
View file

@ -32,16 +32,16 @@ from unsloth_zoo.loss_utils import (
@triton.jit
def _cross_entropy_forward(
logits_ptr ,
logits_row_stride ,
logits_row_stride : tl.constexpr(tl.int64),
loss_ptr ,
logsumexp_ptr ,
labels_ptr ,
VOCAB_SIZE ,
VOCAB_SIZE : tl.constexpr(tl.int32),
BLOCK_SIZE : tl.constexpr(tl.int32),
DO_SOFTCAPPING : tl.constexpr(tl.int1),
SOFTCAP : tl.constexpr(tl.float32),
SOFTCAP ,
DO_LOGIT_SCALING : tl.constexpr(tl.int1),
LOGIT_SCALE : tl.constexpr(tl.float32),
LOGIT_SCALE ,
):
"""
Cross Entropy Loss = 1/n sum [ -yi log(Pi) ]
@ -105,17 +105,17 @@ pass
@triton.jit
def _chunked_cross_entropy_forward(
logits_ptr ,
logits_row_stride ,
logits_row_stride : tl.constexpr(tl.int64),
loss_ptr ,
logsumexp_ptr ,
labels_ptr ,
VOCAB_SIZE ,
N_CHUNKS ,
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),
SOFTCAP ,
DO_LOGIT_SCALING : tl.constexpr(tl.int1),
LOGIT_SCALE : tl.constexpr(tl.float32),
LOGIT_SCALE ,
):
"""
256K vocab divided in 4 chunks
@ -188,17 +188,17 @@ pass
@triton.jit
def _cross_entropy_backward(
logits_ptr ,
logits_row_stride ,
logits_row_stride : tl.constexpr(tl.int64),
dloss_ptr ,
dloss_row_stride ,
logsumexp_ptr ,
labels_ptr ,
VOCAB_SIZE ,
VOCAB_SIZE : tl.constexpr(tl.int32),
BLOCK_SIZE : tl.constexpr(tl.int32),
DO_SOFTCAPPING : tl.constexpr(tl.int1),
SOFTCAP : tl.constexpr(tl.float32),
SOFTCAP ,
DO_LOGIT_SCALING : tl.constexpr(tl.int1),
LOGIT_SCALE : tl.constexpr(tl.float32),
LOGIT_SCALE ,
):
"""
CE_i = -y log(P) = y * (log[sum(exp(x))] - x)