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-05 14:43:49 -08:00
commit cf2a714d74
1 changed files with 6 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

10
unsloth/kernels/cross_entropy_loss.py
View file

@ -181,10 +181,10 @@ def _chunked_cross_entropy_forward(
pass
# @triton.heuristics({
# "DO_SOFTCAPPING": lambda args: bool(args["DO_SOFTCAPPING" ]),
# "DO_LOGIT_SCALING": lambda args: bool(args["DO_LOGIT_SCALING"]),
# })
@triton.heuristics({
"DO_SOFTCAPPING": lambda args: bool(args["DO_SOFTCAPPING" ]),
"DO_LOGIT_SCALING": lambda args: bool(args["DO_LOGIT_SCALING"]),
})
@triton.jit
def _cross_entropy_backward(
logits_ptr ,
@ -337,6 +337,8 @@ class Fast_CrossEntropyLoss(torch.autograd.Function):
return losses
pass
@torch.compiler.disable
@staticmethod
def backward(ctx, dlosses):
logits, logsumexp, labels = ctx.saved_tensors