From 88cfe5eeb22691ce6d4d5052850f5dd37a3d7457 Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Mon, 26 Feb 2024 03:48:16 +1100 Subject: [PATCH] Update cross_entropy_loss.py --- unsloth/kernels/cross_entropy_loss.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py index 176f30e3f7..9ffd7c7ad8 100644 --- a/unsloth/kernels/cross_entropy_loss.py +++ b/unsloth/kernels/cross_entropy_loss.py @@ -157,13 +157,16 @@ def _cross_entropy_backward( col_offsets = block_idx*BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) mask = col_offsets < VOCAB_SIZE label_idx = tl.load(labels_ptr + row_idx).to(tl.int32) - + + if label_idx != -100: + dloss = tl.load(dloss_ptr) + else: + dloss = 0.0 logits = tl.load(logits_ptr + col_offsets, mask = mask, other = -float("inf")).to(tl.float32) lse = tl.load(lse_ptr + row_idx) probs = tl.exp(logits - lse) - probs = tl.where(col_offsets == label_idx, probs - 1.0, probs) - dloss = tl.load(dloss_ptr) if label_idx != -100 else 0.0 + probs = tl.where(col_offsets == label_idx, probs - 1.0, probs) tl.store(logits_ptr + col_offsets, dloss * probs, mask = mask) pass