diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 0883f56cea..23022f46e1 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -604,18 +604,49 @@ def _tilelang_importable() -> bool: return False +def _torch_has_hip() -> bool: + """True iff the installed torch is a HIP / ROCm build. + + We check `torch.version.hip` (non-None on ROCm wheels). This is the + reliable signal even on x86_64 Linux Strix Halo / MI300, where + `sys.platform` and `platform.machine()` look identical to a CUDA box. + + Importing torch here is acceptable in the worker subprocess context: + the next step after kernel installers is the model load, which + imports torch anyway. We swallow import errors so a missing torch + (extremely unusual at this point) is treated as "not HIP" and the + rest of the gate stack handles it. + """ + try: + import torch as _torch + return getattr(_torch.version, "hip", None) is not None + except Exception: + return False + + def _tilelang_platform_supported() -> bool: - """True iff the current platform has a tilelang 0.1.8 wheel. + """True iff the current platform has a usable tilelang 0.1.8 backend. tilelang publishes manylinux x86_64/aarch64 and macOS arm64 wheels plus a 93MB sdist; we never want the sdist on a Studio worker, so we restrict to Linux x86_64/aarch64 explicitly. + + Excludes HIP / ROCm torch builds: tilelang 0.1.8 has no HIP GEMM + instruction, so `_select_gemm_instruction` raises `Unsupported + target for gemm: hip` mid-compile during Qwen3.5 GDN backward. + Reported by h34v3nzc0dex on Strix Halo (gfx1151, ROCm 7.13). The + pip wheel installs fine and imports cleanly, but FLA's TileLang + dispatcher then crashes at first training step. See PR 5434. """ import platform as _platform if not sys.platform.startswith("linux"): return False - return _platform.machine().lower() in _TILELANG_SUPPORTED_LINUX_MACHINES + if _platform.machine().lower() not in _TILELANG_SUPPORTED_LINUX_MACHINES: + return False + if _torch_has_hip(): + return False + return True def _pip_install_cmd(*args: str) -> list[str]: @@ -854,6 +885,23 @@ def _install_fast_path_hooks(event_queue: Any, model_name: str) -> None: logger.info("Fast-path hooks disabled via env; using substring fallback") return + # Defensive: on HIP/ROCm torch builds, FLA's TileLang backend (when + # tilelang is installed for any reason — e.g. a stale CUDA env that + # was reused for ROCm) crashes mid-backward with + # "Unsupported target for gemm: hip" inside + # `tilelang.tileop.gemm._select_gemm_instruction`. The install gate + # in `_ensure_tilelang_backend_unconditional` prevents NEW installs + # on HIP; this env-var setdefault disables FLA's TileLang dispatch + # for already-installed tilelang too. Users can override by setting + # FLA_TILELANG=1 explicitly. Reported by h34v3nzc0dex on Strix Halo. + if _torch_has_hip() and os.environ.get("FLA_TILELANG") is None: + os.environ["FLA_TILELANG"] = "0" + logger.info( + "HIP/ROCm torch detected; setting FLA_TILELANG=0 to keep " + "FLA on the safe Triton path (tilelang 0.1.8 has no HIP " + "GEMM backend)" + ) + try: from transformers.utils import import_utils as _iu except Exception as exc: diff --git a/studio/backend/tests/test_training_worker_flash_attn.py b/studio/backend/tests/test_training_worker_flash_attn.py index 683e995be7..e4d08444bc 100644 --- a/studio/backend/tests/test_training_worker_flash_attn.py +++ b/studio/backend/tests/test_training_worker_flash_attn.py @@ -1221,3 +1221,91 @@ def test_run_training_process_eagerly_installs_causal_conv1d_in_normal_mode(): "branch, so SSM models that bypass is_causal_conv1d_available() still " "get the eager install" ) + + +# ─────────────────────────────────────────────────────────────────── +# HIP / ROCm regression coverage (h34v3nzc0dex Strix Halo report). +# tilelang 0.1.8 has no HIP GEMM backend; FLA's TileLang dispatch +# crashes mid-backward on AMD with "Unsupported target for gemm: hip". +# The fix: skip the install on HIP-built torch AND setdefault +# FLA_TILELANG=0 so already-installed tilelang doesn't get used either. +# ─────────────────────────────────────────────────────────────────── + + +def test_tilelang_platform_unsupported_on_hip_torch(monkeypatch): + """Strix Halo / MI300 with ROCm torch: linux + x86_64 looks + identical to a CUDA box at the OS level, so the platform check + must consult torch.version.hip explicitly. + """ + monkeypatch.setattr(worker, "_torch_has_hip", lambda: True) + assert worker._tilelang_platform_supported() is False + + +def test_tilelang_install_skipped_on_hip_torch(monkeypatch): + """End-to-end: the unconditional installer must not call pip on HIP torch.""" + monkeypatch.delenv(worker._TILELANG_SKIP_ENV, raising=False) + monkeypatch.setattr(worker, "_torch_has_hip", lambda: True) + run_mock = mock.Mock(return_value=mock.Mock(returncode=0, stdout="")) + monkeypatch.setattr(worker._sp, "run", run_mock) + monkeypatch.setattr(worker, "_send_status", lambda *a, **k: None) + + result = worker._ensure_tilelang_backend_unconditional(event_queue=[]) + + assert result is False + run_mock.assert_not_called() + + +def test_install_fast_path_hooks_sets_fla_tilelang_zero_on_hip(monkeypatch): + """When HIP torch is detected, hook installer must set + FLA_TILELANG=0 (via setdefault — respects user override) so any + PRE-EXISTING tilelang install isn't used by FLA's dispatcher. + """ + import os as _os + monkeypatch.delenv("FLA_TILELANG", raising=False) + monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising=False) + monkeypatch.setattr(worker, "_torch_has_hip", lambda: True) + monkeypatch.setattr(worker, "_ensure_flash_linear_attention_unconditional", lambda eq: True) + monkeypatch.setattr(worker, "_ensure_tilelang_backend_unconditional", lambda eq: True) + monkeypatch.setattr(worker, "_install_package_wheel_first", lambda **kw: True) + + worker._install_fast_path_hooks( + event_queue=_FakeQueue(), model_name="unsloth/Qwen3.5-2B" + ) + + assert _os.environ.get("FLA_TILELANG") == "0" + + +def test_install_fast_path_hooks_respects_user_fla_tilelang_override(monkeypatch): + """If the user explicitly set FLA_TILELANG (even on HIP), don't + overwrite — they may know they have a HIP-aware tilelang fork. + """ + import os as _os + monkeypatch.setenv("FLA_TILELANG", "1") + monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising=False) + monkeypatch.setattr(worker, "_torch_has_hip", lambda: True) + monkeypatch.setattr(worker, "_ensure_flash_linear_attention_unconditional", lambda eq: True) + monkeypatch.setattr(worker, "_ensure_tilelang_backend_unconditional", lambda eq: True) + monkeypatch.setattr(worker, "_install_package_wheel_first", lambda **kw: True) + + worker._install_fast_path_hooks( + event_queue=_FakeQueue(), model_name="unsloth/Qwen3.5-2B" + ) + + assert _os.environ["FLA_TILELANG"] == "1" + + +def test_install_fast_path_hooks_does_not_set_fla_tilelang_on_cuda(monkeypatch): + """CUDA path must NOT set FLA_TILELANG (tilelang is wanted there).""" + import os as _os + monkeypatch.delenv("FLA_TILELANG", raising=False) + monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising=False) + monkeypatch.setattr(worker, "_torch_has_hip", lambda: False) + monkeypatch.setattr(worker, "_ensure_flash_linear_attention_unconditional", lambda eq: True) + monkeypatch.setattr(worker, "_ensure_tilelang_backend_unconditional", lambda eq: True) + monkeypatch.setattr(worker, "_install_package_wheel_first", lambda **kw: True) + + worker._install_fast_path_hooks( + event_queue=_FakeQueue(), model_name="unsloth/Qwen3.5-2B" + ) + + assert _os.environ.get("FLA_TILELANG") is None