From 10b50c84df97c52d644ee3ff4c6a5bbfdcb35f7a Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Sun, 17 May 2026 08:43:24 +0000 Subject: [PATCH] studio: skip tilelang on HIP / ROCm torch (Strix Halo crash report) h34v3nzc0dex tested PR 5434 on Strix Halo (gfx1151, ROCm 7.13, torch 2.11.0+rocm7.13.0) and hit a hard regression: File ".../fla/ops/common/backends/tilelang/__init__.py", line 92, in chunk_bwd_dqkwg File ".../tilelang/jit/kernel.py", line 137, in __init__ File ".../tilelang/tileop/gemm/__init__.py", line 143, in _select_gemm_instruction tvm.error.InternalError: Check failed: (0) is false: Unsupported target for gemm: hip -keys=hip,gpu -mcpu=gfx1151 ... `tilelang==0.1.8` ships no HIP GEMM instruction; `_select_gemm_instruction` raises at lower-time, not import-time. So: - pip install succeeds - `import tilelang` succeeds - `TileLangBackend.is_available()` returns True - FLA's dispatcher picks TileLang for `chunk_bwd_dqkwg` - training subprocess dies at first GDN backward, no graceful fallback The PR's existing platform gate (`_tilelang_platform_supported`) checked only `sys.platform == "linux"` and `platform.machine()`, both of which look identical on a ROCm box. Fix has two layers: 1. INSTALL GATE: new `_torch_has_hip()` helper checks `torch.version.hip is not None`. `_tilelang_platform_supported` now returns False on HIP torch, so the install never fires. 2. RUNTIME GATE: even with the install skipped, a user could have tilelang already present (e.g. venv carried over from a CUDA box). `_install_fast_path_hooks` now calls `os.environ.setdefault("FLA_TILELANG", "0")` when HIP is detected, which is the env-var FLA's `TileLangBackend` already honors. Users who know they have a HIP-aware tilelang fork can override by setting `FLA_TILELANG=1` explicitly. This costs nothing on CUDA (the gate is a no-op when `torch.version.hip is None`), and removes the crash for AMD users. The benchmark numbers in the PR description (1.43x on B200 sm_100) are not affected. The other halves of the PR are confirmed working on gfx1151 by the same report: - `flash-linear-attention 0.5.0` runs at production scale (B=1 T=8192 H=16 K=128 V=128 and others) with no patches. - `causal-conv1d` runs at the shapes the fast-path gate cares about. (A separate Ubuntu 24.04 `--gcc-install-dir` build workaround is needed for the source-build path; that mirrors bbf004c's llama.cpp fix and is out of scope here.) Tests added: - test_tilelang_platform_unsupported_on_hip_torch - test_tilelang_install_skipped_on_hip_torch - test_install_fast_path_hooks_sets_fla_tilelang_zero_on_hip - test_install_fast_path_hooks_respects_user_fla_tilelang_override - test_install_fast_path_hooks_does_not_set_fla_tilelang_on_cuda Total 50 passing (was 45). --- studio/backend/core/training/worker.py | 52 ++++++++++- .../tests/test_training_worker_flash_attn.py | 88 +++++++++++++++++++ 2 files changed, 138 insertions(+), 2 deletions(-) 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