From 8a14fe3f1e7c92085c35f7292fa946d74543eb35 Mon Sep 17 00:00:00 2001 From: Edd <68678137+Erland366@users.noreply.github.com> Date: Thu, 7 Nov 2024 00:16:02 +0400 Subject: [PATCH] Fix: cast logits to float32 in cross_entropy_forward to prevent errors (#1254) * Fix: cast logits to float32 in cross_entropy_forward to prevent errors * Update cross_entropy_loss.py --------- Co-authored-by: Daniel Han --- unsloth/kernels/cross_entropy_loss.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py index 0c07035097..cc3dbb1d87 100644 --- a/unsloth/kernels/cross_entropy_loss.py +++ b/unsloth/kernels/cross_entropy_loss.py @@ -84,12 +84,12 @@ def _cross_entropy_forward( logsumexp = c + tl.log(tl.sum(tl.exp(logits - c), 0)) if label_idx != -100: - x = tl.load(logits_ptr + label_idx) + x = tl.load(logits_ptr + label_idx).to(tl.float32) # Go logit scaling for Cohere: t * x if DO_LOGIT_SCALING: x = LOGIT_SCALE * x # Do logit softcapping for Gemma 2: t * tanh(1/t * x) if DO_SOFTCAPPING: x = SOFTCAP * triton_tanh(x / SOFTCAP) - loss = logsumexp - x.to(tl.float32) + loss = logsumexp - x else: loss = 0.0 tl.store(logsumexp_ptr, logsumexp) @@ -170,7 +170,7 @@ def _chunked_cross_entropy_forward( if DO_LOGIT_SCALING: x = LOGIT_SCALE * x # Do logit softcapping for Gemma 2: t * tanh(1/t * x) if DO_SOFTCAPPING: x = SOFTCAP * triton_tanh(x / SOFTCAP) - loss = -1.0 * x.to(tl.float32) + loss = -1.0 * x else: loss = 0.0 tl.store(loss_ptr, loss)