diff --git a/installer.py b/installer.py index 9e7e7e4b2..e61165a1a 100644 --- a/installer.py +++ b/installer.py @@ -699,7 +699,7 @@ def install_rocm_zluda(): if sys.platform == "win32" and (not args.use_zluda) and (device is not None) and (device.therock is not None) and not installed("rocm"): check_python(supported_minors=[11, 12, 13], reason='ROCm-Windows: python==3.11/3.12/3.13 required') - install(f"rocm[devel,libraries] --index-url https://rocm.nightlies.amd.com/{device.therock}") + install("rocm-sdk-devel --index-url https://rocm.nightlies.amd.com/whl-multi-arch") rocm.refresh() msg = f'ROCm: version={rocm.version}' @@ -733,9 +733,16 @@ def install_rocm_zluda(): else: # TODO rocm: switch to pytorch source when it becomes available if device is None: log.error('ROCm: no agent found - make sure that graphics driver is installed and up to date') - if isinstance(rocm.environment, rocm.PythonPackageEnvironment): + if device is not None and device.therock is not None: check_python(supported_minors=[11, 12, 13], reason='ROCm-Windows: python==3.11/3.12/3.13 required') - torch_command = os.environ.get('TORCH_COMMAND', f'torch torchvision --index-url https://rocm.nightlies.amd.com/{device.therock}') + # Extract device-specific package family from therock path (e.g., 'amd-torch-device-gfx1030' from 'whl-multi-arch/amd-torch-device-gfx1030') + torch_family = device.therock.rsplit('/', 1)[-1] + torchvision_family = torch_family.replace('amd-torch-device-', 'amd-torchvision-device-') + # Use device-specific index for torch/torchvision, with root index as fallback for torchaudio and other packages + torch_command = os.environ.get('TORCH_COMMAND', f'{torch_family} {torchvision_family} torchaudio --index-url https://rocm.nightlies.amd.com/{device.therock} --extra-index-url https://rocm.nightlies.amd.com/whl-multi-arch') + elif isinstance(rocm.environment, rocm.PythonPackageEnvironment): + check_python(supported_minors=[11, 12, 13], reason='ROCm-Windows: python==3.11/3.12/3.13 required') + torch_command = os.environ.get('TORCH_COMMAND', 'torch torchvision torchaudio --index-url https://rocm.nightlies.amd.com/whl-multi-arch') else: check_python(supported_minors=[12], reason='ROCm-Windows: preview python==3.12 required') # torch 2.8.0a0 is the last version with rocm 6.4 support diff --git a/modules/rocm.py b/modules/rocm.py index 15387d377..2ecdcd29c 100644 --- a/modules/rocm.py +++ b/modules/rocm.py @@ -133,28 +133,15 @@ class Agent: @property def therock(self) -> str | None: - if (self.gfx_version & 0xFFF0) == 0x1200: - return "v2/gfx120X-all" - if (self.gfx_version & 0xFFF0) == 0x1100: - return "v2/gfx110X-all" - if self.gfx_version == 0x1150: - return "v2-staging/gfx1150" - if self.gfx_version == 0x1151: - return "v2/gfx1151" - if self.gfx_version == 0x1152: - return "v2-staging/gfx1152" - if self.gfx_version == 0x1153: - return "v2-staging/gfx1153" - if self.gfx_version in (0x1030, 0x1031, 0x1032, 0x1034,): - return "v2-staging/gfx103X-dgpu" - #if (self.gfx_version & 0xFFF0) == 0x1010: - # return "gfx101X-dgpu" - #if (self.gfx_version & 0xFFF0) == 0x900: - # return "gfx90X-dcgpu" - #if (self.gfx_version & 0xFFF0) == 0x940: - # return "gfx94X-dcgpu" - #if self.gfx_version == 0x950: - # return "gfx950-dcgpu" + if self.gfx_version is None: + return None + gfx = self.name if self.name.startswith("gfx") else f"gfx{self.gfx_version:04x}" + if self.gfx_version & 0xFFF0 in (0x1200, 0x1100): + return f"whl-multi-arch/amd-torch-device-{gfx}" + if self.gfx_version in (0x1150, 0x1151, 0x1152, 0x1153): + return f"whl-multi-arch/amd-torch-device-{gfx}" + if self.gfx_version in (0x1030, 0x1031, 0x1032, 0x1033, 0x1034, 0x1035, 0x1036): + return f"whl-multi-arch/amd-torch-device-{gfx}" return None def get_gfx_version(self) -> str | None: diff --git a/scripts/rocm/rocm_mgr.py b/scripts/rocm/rocm_mgr.py index 3f679e0ff..feb7e925b 100644 --- a/scripts/rocm/rocm_mgr.py +++ b/scripts/rocm/rocm_mgr.py @@ -152,7 +152,10 @@ def _get_libs_pkg() -> str: agent = _rocm.Agent(i) therock = agent.therock if therock: - suffix = therock.split("/")[-1].replace("-", "_") + suffix = therock.rsplit("/", 1)[-1] + if suffix.startswith("amd-torch-device-"): + suffix = suffix[len("amd-torch-device-"):] + suffix = suffix.replace("-", "_") folder = "_rocm_sdk_libraries_" + suffix if folder in candidates: _libs_pkg_cache = folder