fix out of resources issue for llama3.2 sft on amd gpu (#3455)

Co-authored-by: Xun Wang <xunwang2@amd.com>

unsloth/kernels/cross_entropy_loss.py

@@ -21,6 +21,7 @@ from .utils import (
     triton_tanh,
     triton_cast,
     torch_gpu_device,
+    is_cdna,
 )
 from transformers.models.llama.modeling_llama import logger
 from packaging.version import Version
@@ -332,7 +333,7 @@ class Fast_CrossEntropyLoss(torch.autograd.Function):
                     SOFTCAP          = logit_softcapping,
                     DO_LOGIT_SCALING = DO_LOGIT_SCALING,
                     LOGIT_SCALE      = logit_scaling,
-                    num_warps        = 32,
+                    num_warps        = 32 if not is_cdna() else 16,
                 )
             # logsumexp(chunked_logsumexp) - x
             # Do the -x separately

unsloth/kernels/utils.py

@@ -71,6 +71,14 @@ else:
 pass
 
 
+def is_hip():
+    return triton.runtime.driver.active.get_current_target().backend == "hip"
+
+
+def is_cdna():
+    return is_hip() and triton.runtime.driver.active.get_current_target().arch in ('gfx940', 'gfx941', 'gfx942')
+
+
 def calculate_settings(n : int) -> (int, int,):
     BLOCK_SIZE : int = next_power_of_2(n)
     if BLOCK_SIZE > MAX_FUSED_SIZE: