From 6fd28a4dce5f544d67bbbcbefc4060a439165a24 Mon Sep 17 00:00:00 2001 From: wangxunx <104552583+wangxunx@users.noreply.github.com> Date: Mon, 27 Oct 2025 12:20:47 +0800 Subject: [PATCH] fix cross entropy loss issue for small vocab size on amd gpu (#3503) --- unsloth/kernels/cross_entropy_loss.py | 1 + unsloth/models/mistral.py | 7 +------ 2 files changed, 2 insertions(+), 6 deletions(-) diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py index e0af458a96..d3b618582e 100644 --- a/unsloth/kernels/cross_entropy_loss.py +++ b/unsloth/kernels/cross_entropy_loss.py @@ -300,6 +300,7 @@ class Fast_CrossEntropyLoss(torch.autograd.Function): if n_chunks == 1: # For small vocabs <= 65336 like Llama, Mistral BLOCK_SIZE, num_warps = calculate_settings(vocab_size) + if is_cdna(): num_warps = num_warps // 2 logsumexp = torch.empty(n_rows, dtype = torch.float32, device = device) with torch_gpu_device(device): diff --git a/unsloth/models/mistral.py b/unsloth/models/mistral.py index b547739df2..faab2d30b1 100644 --- a/unsloth/models/mistral.py +++ b/unsloth/models/mistral.py @@ -298,12 +298,7 @@ def MistralForCausalLM_fast_forward( else: RETURN_LOGITS = os.environ.get("UNSLOTH_RETURN_LOGITS", "0") == "1" # < 1024 Normal Unsloth uses less VRAM! - if DEVICE_TYPE == "hip": - # [TODO] AMD GPUs fail on chunked_cross_entropy loss! - # RuntimeError: Triton Error [HIP]: Code: 1, Messsage: invalid argument - RETURN_LOGITS = False - elif bsz*q_len <= 1024: - RETURN_LOGITS = True + if bsz * q_len <= 1024: RETURN_LOGITS = True if not RETURN_LOGITS and labels is not None: n_items = kwargs.get("num_items_in_batch", None) or kwargs.get("n_items", None)