Guard torch.compile on ROCm when triton_key is missing (#3923)

* Guard torch.compile on ROCm when triton_key missing

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

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

* Update unsloth/import_fixes.py

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

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

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

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

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

* Tighten ROCm Triton import handling

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

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

---------

Co-authored-by: Rachel Li <rachelliqx07@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

unsloth/__init__.py

@@ -127,6 +127,7 @@ from .import_fixes import (
     fix_vllm_aimv2_issue,
     fix_vllm_guided_decoding_params,
     fix_vllm_pdl_blackwell,
+    fix_rocm_triton_key_error,
     ignore_logger_messages,
     patch_ipykernel_hf_xet,
     patch_trackio,
@@ -141,6 +142,7 @@ fix_xformers_performance_issue()
 fix_vllm_aimv2_issue()
 fix_vllm_guided_decoding_params()
 fix_vllm_pdl_blackwell()
+fix_rocm_triton_key_error()
 ignore_logger_messages()
 patch_ipykernel_hf_xet()
 patch_trackio()
@@ -154,6 +156,7 @@ del fix_xformers_performance_issue
 del fix_vllm_aimv2_issue
 del fix_vllm_guided_decoding_params
 del fix_vllm_pdl_blackwell
+del fix_rocm_triton_key_error
 del ignore_logger_messages
 del patch_ipykernel_hf_xet
 del patch_trackio

unsloth/import_fixes.py

@@ -666,6 +666,39 @@ def fix_huggingface_hub():
         )
 
 
+def fix_rocm_triton_key_error():
+    """
+    ROCm + torch.compile can fail if Triton lacks `triton_key`.
+    Disable Inductor/compile only on ROCm when that symbol is missing.
+    """
+    try:
+        import torch
+    except (ImportError, ModuleNotFoundError):
+        return
+
+    if not getattr(torch.version, "hip", None):
+        return
+
+    try:
+        import triton
+    except (ImportError, ModuleNotFoundError):
+        return
+
+    try:
+        from triton.runtime import triton_key  # noqa: F401
+
+        return
+    except ImportError:
+        pass
+
+    os.environ.setdefault("TORCHINDUCTOR_DISABLE", "1")
+    os.environ.setdefault("TORCH_COMPILE_DISABLE", "1")
+    logger.info(
+        "Unsloth: ROCm detected and Triton lacks triton_key; "
+        "disabling torch.compile/Inductor to avoid backend crash."
+    )
+
+
 def fix_vllm_pdl_blackwell():
     """
     Fix vLLM PDL (Programmatic Dependent Launch) bug on Blackwell GPUs (SM100).