diff --git a/unsloth/__init__.py b/unsloth/__init__.py index d3093cf4c0..0b819a546e 100644 --- a/unsloth/__init__.py +++ b/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 diff --git a/unsloth/import_fixes.py b/unsloth/import_fixes.py index 4f88808c2a..89fd152857 100644 --- a/unsloth/import_fixes.py +++ b/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).