From 88bf684c615f2d1b0e5d06ead94b6e7b5d14117b Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Sat, 24 Feb 2024 02:06:11 +1100 Subject: [PATCH] Update cross_entropy_loss.py --- unsloth/kernels/cross_entropy_loss.py | 21 ++++++++++++--------- 1 file changed, 12 insertions(+), 9 deletions(-) diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py index 4b16f0599b..4b4bda0b4e 100644 --- a/unsloth/kernels/cross_entropy_loss.py +++ b/unsloth/kernels/cross_entropy_loss.py @@ -99,16 +99,19 @@ def _large_cross_entropy_forward( # Maximum stops overflow lse = tl.log(tl.sum(tl.exp(logits - max_logits), 0)) + max_logits tl.store(lse_ptr, lse) - - if (label_idx != -100) and \ - (label_idx >= (col_idx+0)*BLOCK_SIZE) and \ - (label_idx < min((col_idx+1)*BLOCK_SIZE, n_cols)): - loss = tl.load(logits_ptr + label_idx).to(tl.float32) - lse = 0.0 - loss = lse - logits_label # We add the final logsumexp after a reduction - else: - loss = 0.0 + loss = 0.0 + # chained boolean operators (A or B or C) are not supported; use parentheses to split the chain. + if (label_idx != -100): + if (label_idx >= (col_idx+0)*BLOCK_SIZE) and \ + (label_idx < min((col_idx+1)*BLOCK_SIZE, n_cols)): + + loss = tl.load(logits_ptr + label_idx).to(tl.float32) + lse = 0.0 + loss = lse - logits_label # We add the final logsumexp after a reduction + pass + pass + tl.store(loss_ptr, loss) pass