diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py index ca26a112bc..9c0a1bd1c1 100644 --- a/studio/install_llama_prebuilt.py +++ b/studio/install_llama_prebuilt.py @@ -3085,6 +3085,46 @@ def detect_host() -> HostInfo: ) +def _normalize_forwarded_gfx(value: str | None) -> str | None: + """Extract a single gfx token from a forwarded --rocm-gfx / env value. + setup.sh/setup.ps1 already picked the active GPU, so take the token as-is + without re-applying visible-device selection. Ignore anything malformed.""" + if not value: + return None + m = re.search(r"gfx[1-9][0-9a-z]{2,3}", value.lower()) + return m.group(0) if m else None + + +def _apply_host_overrides( + host: HostInfo, + *, + override_has_rocm: bool = False, + override_rocm_gfx: str | None = None, + force_cpu: bool = False, +) -> HostInfo: + """Fold setup.sh/setup.ps1's forwarded detection into the host profile. + A forwarded gfx (--rocm-gfx or UNSLOTH_ROCM_GFX_ARCH) is authoritative and + implies ROCm: the installer's own hipinfo/amd-smi probe can miss the arch on + amd-smi-only hosts or when setup inferred it from the GPU name, leaving + rocm_gfx_target None and no lemonade prebuilt selected. force_cpu is the + opposite explicit signal (arm64 Linux GPU host whose source build failed): + drop GPU attributes so the CPU prebuilt for this OS/arch is selected.""" + if force_cpu: + return dataclasses_replace( + host, + has_usable_nvidia = False, + has_physical_nvidia = False, + has_rocm = False, + rocm_gfx_target = None, + ) + gfx = _normalize_forwarded_gfx(override_rocm_gfx) + if gfx: + return dataclasses_replace(host, has_rocm = True, rocm_gfx_target = gfx) + if override_has_rocm and not host.has_rocm: + return dataclasses_replace(host, has_rocm = True) + return host + + def pick_windows_cuda_runtime(host: HostInfo) -> str | None: if not host.driver_cuda_version: return None @@ -6460,22 +6500,16 @@ def install_prebuilt( *, simple_policy: bool = False, override_has_rocm: bool = False, + override_rocm_gfx: str | None = None, force_cpu: bool = False, ) -> None: host = detect_host() - if override_has_rocm and not host.has_rocm: - host = dataclasses_replace(host, has_rocm = True) - if force_cpu: - # Explicit CPU fallback: drop GPU attributes so the CPU prebuilt for this - # OS/arch is selected. setup.sh uses this for arm64 Linux GPU hosts whose - # source build failed, where no arm64 CUDA prebuilt exists anywhere. - host = dataclasses_replace( - host, - has_usable_nvidia = False, - has_physical_nvidia = False, - has_rocm = False, - rocm_gfx_target = None, - ) + host = _apply_host_overrides( + host, + override_has_rocm = override_has_rocm, + override_rocm_gfx = override_rocm_gfx, + force_cpu = force_cpu, + ) choice: AssetChoice | None = None try: with install_lock(install_lock_path(install_dir)): @@ -6620,6 +6654,16 @@ def parse_args() -> argparse.Namespace: "so the HIP llama.cpp prebuilt is selected even when hipinfo is not on PATH." ), ) + parser.add_argument( + "--rocm-gfx", + default = os.environ.get("UNSLOTH_ROCM_GFX_ARCH"), + help = ( + "Forward the AMD gfx target (e.g. gfx1151) that setup.ps1/setup.sh " + "resolved, so the lemonade HIP prebuilt is selected even when the " + "installer's own hipinfo/amd-smi probe cannot report it. Implies " + "--has-rocm. Defaults to the UNSLOTH_ROCM_GFX_ARCH environment variable." + ), + ) parser.add_argument( "--cpu-fallback", action = "store_true", @@ -6751,6 +6795,7 @@ def main() -> int: published_release_tag = args.published_release_tag or "", simple_policy = args.simple_policy, override_has_rocm = args.has_rocm, + override_rocm_gfx = args.rocm_gfx, force_cpu = args.cpu_fallback, ) return EXIT_SUCCESS diff --git a/studio/setup.ps1 b/studio/setup.ps1 index 70e2fee334..28ad512690 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -2483,6 +2483,12 @@ if ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") { ) if ($HasROCm) { $prebuiltArgs += "--has-rocm" + # Forward the resolved gfx arch so the lemonade HIP prebuilt is picked + # even when the installer's own probe cannot report it (amd-smi-only + # hosts, name-inferred arch). + if ($script:ROCmGfxArch) { + $prebuiltArgs += @("--rocm-gfx", $script:ROCmGfxArch) + } } if ($env:UNSLOTH_LLAMA_RELEASE_TAG) { $prebuiltArgs += @("--published-release-tag", $env:UNSLOTH_LLAMA_RELEASE_TAG) diff --git a/studio/setup.sh b/studio/setup.sh index 77013f0e6d..b5c60030c3 100755 --- a/studio/setup.sh +++ b/studio/setup.sh @@ -857,6 +857,12 @@ else if [ -n "${UNSLOTH_LLAMA_RELEASE_TAG:-}" ]; then _PREBUILT_CMD+=(--published-release-tag "$UNSLOTH_LLAMA_RELEASE_TAG") fi + # Forward the gfx arch resolved above so the lemonade HIP prebuilt is picked + # even when the installer's own probe cannot report it (amd-smi-only hosts, + # name-inferred arch). Implies --has-rocm on the installer side. + if [ -n "${_setup_gfx:-}" ]; then + _PREBUILT_CMD+=(--rocm-gfx "$_setup_gfx") + fi _PREBUILT_LOG="$(mktemp)" set +e if _is_verbose; then diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py index ec99e9a8d9..5b6586f9c9 100644 --- a/tests/studio/install/test_rocm_support.py +++ b/tests/studio/install/test_rocm_support.py @@ -35,6 +35,8 @@ AssetChoice = prebuilt_mod.AssetChoice PrebuiltFallback = prebuilt_mod.PrebuiltFallback resolve_upstream_asset_choice = prebuilt_mod.resolve_upstream_asset_choice runtime_patterns_for_choice = prebuilt_mod.runtime_patterns_for_choice +_apply_host_overrides = prebuilt_mod._apply_host_overrides +_normalize_forwarded_gfx = prebuilt_mod._normalize_forwarded_gfx # install_python_stack.py _STACK_PATH = PACKAGE_ROOT / "studio" / "install_python_stack.py" @@ -2598,5 +2600,93 @@ class TestHipSdkInstalledButDeviceInaccessible: assert "GPU not ROCm-accessible" in source +# ============================================================================= +# TEST: --rocm-gfx forwarding -- setup.sh/setup.ps1 hand their resolved gfx arch +# to install_llama_prebuilt.py so the lemonade HIP prebuilt is selected even when +# the installer's own hipinfo/amd-smi probe cannot report it. +# ============================================================================= + +_SETUP_SH_PATH = PACKAGE_ROOT / "studio" / "setup.sh" + + +class TestNormalizeForwardedGfx: + """A forwarded gfx string is reduced to a single clean gfx token.""" + + def test_plain_token(self): + assert _normalize_forwarded_gfx("gfx1151") == "gfx1151" + + def test_uppercase_normalized(self): + assert _normalize_forwarded_gfx("GFX1151") == "gfx1151" + + def test_extracts_from_noise(self): + assert _normalize_forwarded_gfx("gcnArchName: gfx942") == "gfx942" + + def test_malformed_is_ignored(self): + assert _normalize_forwarded_gfx("not-a-gpu") is None + + def test_empty_and_none(self): + assert _normalize_forwarded_gfx("") is None + assert _normalize_forwarded_gfx(None) is None + + +class TestApplyHostOverrides: + """Forwarded ROCm detection is folded into the host profile correctly.""" + + def test_forwarded_gfx_fills_empty_probe(self): + # amd-smi-only / name-inferred host: installer probe found no gfx. + host = rocm_host(rocm_gfx_target = None) + out = _apply_host_overrides(host, override_rocm_gfx = "gfx1151") + assert out.has_rocm is True + assert out.rocm_gfx_target == "gfx1151" + + def test_forwarded_gfx_implies_rocm(self): + # A CPU-looking host with a forwarded gfx is an AMD host. + out = _apply_host_overrides(cpu_host(), override_rocm_gfx = "gfx1200") + assert out.has_rocm is True + assert out.rocm_gfx_target == "gfx1200" + + def test_forwarded_gfx_is_authoritative(self): + # setup already applied visible-device selection; its value wins. + host = rocm_host(rocm_gfx_target = "gfx1100") + out = _apply_host_overrides(host, override_rocm_gfx = "gfx1151") + assert out.rocm_gfx_target == "gfx1151" + + def test_has_rocm_only_keeps_probe_gfx(self): + out = _apply_host_overrides(cpu_host(), override_has_rocm = True) + assert out.has_rocm is True + assert out.rocm_gfx_target is None + + def test_malformed_forwarded_gfx_falls_back_to_has_rocm(self): + out = _apply_host_overrides( + cpu_host(), override_has_rocm = True, override_rocm_gfx = "junk" + ) + assert out.has_rocm is True + assert out.rocm_gfx_target is None + + def test_no_overrides_leaves_host_unchanged(self): + host = nvidia_host() + assert _apply_host_overrides(host) is host + + +class TestRocmGfxForwarding: + """setup.sh / setup.ps1 forward their resolved gfx; the installer accepts it.""" + + def test_installer_exposes_rocm_gfx_arg(self): + source = _PREBUILT_PATH.read_text(encoding = "utf-8") + assert '"--rocm-gfx"' in source + # Defaults to the env override so a standalone run still works. + assert 'os.environ.get("UNSLOTH_ROCM_GFX_ARCH")' in source + + def test_setup_sh_forwards_rocm_gfx(self): + source = _SETUP_SH_PATH.read_text(encoding = "utf-8") + assert "--rocm-gfx" in source + assert '"$_setup_gfx"' in source + + def test_setup_ps1_forwards_rocm_gfx(self): + source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") + assert "--rocm-gfx" in source + assert "$script:ROCmGfxArch" in source + + if __name__ == "__main__": pytest.main([__file__, "-v"])