From 9a68c7fec9cf624b6a87fabee1da62501c60dd28 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 3 Nov 2024 19:59:59 -0800 Subject: [PATCH] Update cross_entropy_loss.py --- unsloth/kernels/cross_entropy_loss.py | 15 ++++++--------- 1 file changed, 6 insertions(+), 9 deletions(-) diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py index d33cb78409..70b0f116d5 100644 --- a/unsloth/kernels/cross_entropy_loss.py +++ b/unsloth/kernels/cross_entropy_loss.py @@ -279,9 +279,6 @@ class Fast_CrossEntropyLoss(torch.autograd.Function): n_chunks = div + (mod != 0) losses = torch.empty(n_rows, dtype = torch.float32, device = "cuda:0") - DO_SOFTCAPPING : bool = bool(logit_softcapping != 0) - DO_LOGIT_SCALING : bool = bool(logit_scaling != 0) - if n_chunks == 1: # For small vocabs <= 65336 like Llama, Mistral BLOCK_SIZE, num_warps = calculate_settings(vocab_size) @@ -294,9 +291,9 @@ class Fast_CrossEntropyLoss(torch.autograd.Function): labels, VOCAB_SIZE = vocab_size, BLOCK_SIZE = BLOCK_SIZE, - DO_SOFTCAPPING = DO_SOFTCAPPING, + DO_SOFTCAPPING = bool(logit_softcapping != 0), SOFTCAP = logit_softcapping, - DO_LOGIT_SCALING = DO_LOGIT_SCALING, + DO_LOGIT_SCALING = bool(logit_scaling != 0), LOGIT_SCALE = logit_scaling, num_warps = num_warps, ) @@ -312,9 +309,9 @@ class Fast_CrossEntropyLoss(torch.autograd.Function): VOCAB_SIZE = vocab_size, N_CHUNKS = n_chunks, BLOCK_SIZE = MAX_FUSED_SIZE, - DO_SOFTCAPPING = DO_SOFTCAPPING, + DO_SOFTCAPPING = bool(logit_softcapping != 0), SOFTCAP = logit_softcapping, - DO_LOGIT_SCALING = DO_LOGIT_SCALING, + DO_LOGIT_SCALING = bool(logit_scaling != 0), LOGIT_SCALE = logit_scaling, num_warps = 32, ) @@ -326,9 +323,9 @@ class Fast_CrossEntropyLoss(torch.autograd.Function): pass ctx.save_for_backward(logits, logsumexp, labels) - ctx.DO_SOFTCAPPING = DO_SOFTCAPPING + ctx.DO_SOFTCAPPING = bool(logit_softcapping != 0) ctx.logit_softcapping = logit_softcapping - ctx.DO_LOGIT_SCALING = DO_LOGIT_SCALING + ctx.DO_LOGIT_SCALING = bool(logit_scaling != 0) ctx.logit_scaling = logit_scaling return losses pass