From 8a8dcd48ddcf8a4c77557e6b0d5ade7bae60f4c7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=87=91=E9=BB=84=E8=89=B2=E8=91=A1=E8=90=84=E7=90=83?=
 =?UTF-8?q?=E5=90=9B=E5=90=9B?= <yueyuan@amd.com>
Date: Sun, 1 Mar 2026 15:59:17 +0800
Subject: [PATCH] fix(ROCm): Comprehensive RDNA GPU support - fix Gemma3 NaN &
 add is_rdna() (#4109)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* fix(ROCm): comprehensive RDNA GPU support - fix Gemma3 NaN & add is_rdna()

- Add is_rdna() detection for RDNA3/3.5/RDNA4 consumer GPUs (gfx11xx, gfx1151, gfx12xx)
- Disable torch.compile for Gemma3 on HIP to fix NaN loss (fixes #3385, #4029)
- Export is_cdna/is_rdna from kernels for downstream use
- Import is_rdna into cross_entropy_loss for future RDNA-specific tuning

Tested on AMD Radeon PRO W7900 (gfx1100) with ROCm 7.1:
  ✓ Gemma3-1B: loss 3.37→3.25 (no NaN)
  ✓ Llama-3.2-1B: loss 2.44→2.37 (no NaN)
  ✓ Qwen2.5-1.5B: loss 1.89→1.85 (no NaN)
  ✓ RMS LayerNorm Triton kernel: bf16/fp16 PASSED
  ✓ Cross Entropy Loss Triton kernel: 32K/256K vocab PASSED

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Address review: scope compile disable to RDNA only, use partial mode, remove unused import

Changes based on Daniel's review:
1. (HIGH) Replace DEVICE_TYPE=='hip' with is_rdna() to avoid disabling
   torch.compile on CDNA GPUs (MI250X/MI300X/MI350) where it works fine
2. (MEDIUM) Use 'partial' instead of '1' for UNSLOTH_COMPILE_DISABLE to
   only disable model forward compilation while keeping loss compilation,
   matching the existing Sesame pattern
3. (LOW) Remove unused is_rdna import from cross_entropy_loss.py (F401)

* Remove redundant is_cdna/is_rdna exports from kernels/__init__.py

These functions are imported directly from .utils where needed
(e.g. cross_entropy_loss.py, loader.py). No external code imports
them from the unsloth.kernels namespace.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 unsloth/kernels/utils.py | 9 +++++++++
 unsloth/models/loader.py | 8 ++++++++
 2 files changed, 17 insertions(+)

diff --git a/unsloth/kernels/utils.py b/unsloth/kernels/utils.py
index 0b13c04532..b1211ae345 100644
--- a/unsloth/kernels/utils.py
+++ b/unsloth/kernels/utils.py
@@ -86,6 +86,15 @@ def is_cdna():
     )
 
 
+@functools.lru_cache(1)
+def is_rdna():
+    """Detect RDNA consumer/workstation GPUs (RDNA3, RDNA3.5, RDNA4)."""
+    if not is_hip():
+        return False
+    arch = triton.runtime.driver.active.get_current_target().arch
+    return arch.startswith("gfx1") and not is_cdna()
+
+
 def calculate_settings(
     n: int,
 ) -> (
diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py
index 711476b759..7bb2a24e00 100644
--- a/unsloth/models/loader.py
+++ b/unsloth/models/loader.py
@@ -1131,6 +1131,14 @@ class FastModel(FastBaseModel):
             # Set norms to float32 since anyways they get upcasted to float32
             # common in both gemma-3 and gemma-3n
             os.environ["UNSLOTH_HIGH_PRECISION_LAYERNORM"] = "1"
+            # ROCm/HIP: Gemma3 compiled forward produces NaN on RDNA GPUs
+            # (gfx1100, gfx1101, gfx1102, gfx1150, gfx1151, etc.).
+            # Disable torch.compile for model forward; loss compilation is fine.
+            # See https://github.com/unslothai/unsloth/issues/3385
+            from unsloth.kernels.utils import is_rdna
+
+            if is_rdna():
+                os.environ["UNSLOTH_COMPILE_DISABLE"] = "partial"
         # Cohere
         elif "cohere2" in model_types_all and transformers_version < Version(
             "4.50.0.dev0"