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 01:23:41 -07:00
commit 8942b2fc76
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

@ -92,10 +92,10 @@ def _cross_entropy_forward(
pass
# @triton.heuristics({
# "DO_SOFTCAPPING": lambda args: args["DO_SOFTCAPPING" ],
# "DO_LOGIT_SCALING": lambda args: args["DO_LOGIT_SCALING"],
# })
@triton.heuristics({
"DO_SOFTCAPPING": lambda args: args["DO_SOFTCAPPING" ],
"DO_LOGIT_SCALING": lambda args: args["DO_LOGIT_SCALING"],
})
@triton.jit
def _chunked_cross_entropy_forward(
logits_ptr ,
@ -106,9 +106,9 @@ def _chunked_cross_entropy_forward(
VOCAB_SIZE ,
N_CHUNKS ,
BLOCK_SIZE : tl.constexpr,
DO_SOFTCAPPING ,
DO_SOFTCAPPING : tl.constexpr(tl.int1),
SOFTCAP ,
DO_LOGIT_SCALING ,
DO_LOGIT_SCALING : tl.constexpr(tl.int1),
LOGIT_SCALE ,
):
"""
@ -189,9 +189,9 @@ def _cross_entropy_backward(
labels_ptr ,
VOCAB_SIZE ,
BLOCK_SIZE : tl.constexpr,
DO_SOFTCAPPING ,
DO_SOFTCAPPING : tl.constexpr(tl.int1),
SOFTCAP ,
DO_LOGIT_SCALING ,
DO_LOGIT_SCALING : tl.constexpr(tl.int1),
LOGIT_SCALE ,
):
"""
@ -347,7 +347,7 @@ class Fast_CrossEntropyLoss(torch.autograd.Function):
SOFTCAP = ctx.logit_softcapping,
DO_LOGIT_SCALING = ctx.DO_LOGIT_SCALING,
LOGIT_SCALE = ctx.logit_scaling,
num_warps = 8,
num_warps = 8,
)
return logits, None, None, None,
pass