diff --git a/installer.py b/installer.py index 32e8cdb88..447fadd80 100644 --- a/installer.py +++ b/installer.py @@ -452,28 +452,23 @@ def check_onnx(): def install_rocm_zluda(torch_command): + from modules import rocm + + if not rocm.is_installed: + log.warning('Could not find ROCm toolkit installed.') + log.info('Using CPU-only torch') + return os.environ.get('TORCH_COMMAND', 'torch torchvision') + check_python(supported_minors=[10, 11], reason='ROCm or ZLUDA backends require Python 3.10 or 3.11') - is_windows = platform.system() == 'Windows' log.info('AMD ROCm toolkit detected') os.environ.setdefault('PYTORCH_HIP_ALLOC_CONF', 'garbage_collection_threshold:0.8,max_split_size_mb:512') # if not is_windows: # os.environ.setdefault('TENSORFLOW_PACKAGE', 'tensorflow-rocm') try: - if is_windows: - command = subprocess.run('hipinfo', shell=True, check=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE) - amd_gpus = command.stdout.decode(encoding="utf8", errors="ignore").split('\n') - amd_gpus = [x.split(' ')[-1].strip() for x in amd_gpus if x.startswith('gcnArchName:')] - elif os.environ.get('WSL_DISTRO_NAME', None) is not None: # WSL does not have 'rocm_agent_enumerator' - command = subprocess.run('rocminfo', shell=True, check=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE) - amd_gpus = command.stdout.decode(encoding="utf8", errors="ignore").split('\n') - amd_gpus = [x.strip().split(" ")[-1] for x in amd_gpus if x.startswith(' Name:') and "CPU" not in x] - else: - command = subprocess.run('rocm_agent_enumerator', shell=True, check=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE) - amd_gpus = command.stdout.decode(encoding="utf8", errors="ignore").split('\n') - amd_gpus = [x for x in amd_gpus if x and x != 'gfx000'] - log.debug(f'ROCm agents detected: {amd_gpus}') + amd_gpus = rocm.get_agents() + log.info(f'ROCm agents detected: {amd_gpus}') except Exception as e: - log.debug(f'ROCm agent enumerator failed: {e}') + log.warning(f'ROCm agent enumerator failed: {e}') amd_gpus = [] hip_visible_devices = [] # use the first available amd gpu by default @@ -484,6 +479,7 @@ def install_rocm_zluda(torch_command): if gpu in ['gfx1030', 'gfx1031', 'gfx1032', 'gfx1034']: # experimental navi 2x support hip_visible_devices.append((idx, gpu, 'navi2x')) break + hip_found_device = len(hip_visible_devices) > 0 if hip_found_device: idx, gpu, arch = hip_visible_devices[0] @@ -498,65 +494,59 @@ def install_rocm_zluda(torch_command): else: log.debug(f'HSA_OVERRIDE_GFX_VERSION auto config is skipped for {gpu}') - try: - command = subprocess.run('hipconfig --version', shell=True, check=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE) - arr = command.stdout.decode(encoding="utf8", errors="ignore").split('.') - rocm_ver = f'{arr[0]}.{arr[1]}' if len(arr) >= 2 else None - log.debug(f'ROCm version detected: {rocm_ver}') - except Exception as e: - log.debug(f'ROCm hipconfig failed: {e}') - rocm_ver = None + log.info(f'ROCm version detected: {rocm.version}') - if args.use_zluda: - log.warning("ZLUDA support: experimental") - error = None - from modules import zluda_installer - try: - if args.reinstall_zluda: - zluda_installer.uninstall() - zluda_path = zluda_installer.get_path() - zluda_installer.install(zluda_path) - zluda_installer.make_copy(zluda_path) - except Exception as e: - error = e - log.warning(f'Failed to install ZLUDA: {e}') - if error is None: + if sys.platform == "win32": + if args.use_zluda: + log.warning("ZLUDA support: experimental") + error = None + from modules import zluda_installer try: - zluda_installer.load(zluda_path) - torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.3.0 torchvision --index-url https://download.pytorch.org/whl/cu118') - log.info(f'Using ZLUDA in {zluda_path}') + if args.reinstall_zluda: + zluda_installer.uninstall() + zluda_path = zluda_installer.get_path() + zluda_installer.install(zluda_path) + zluda_installer.make_copy(zluda_path) except Exception as e: error = e - log.warning(f'Failed to load ZLUDA: {e}') - if error is not None: + log.warning(f'Failed to install ZLUDA: {e}') + if error is None: + try: + zluda_installer.load(zluda_path) + torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.3.0 torchvision --index-url https://download.pytorch.org/whl/cu118') + log.info(f'Using ZLUDA in {zluda_path}') + except Exception as e: + error = e + log.warning(f'Failed to load ZLUDA: {e}') + if error is not None: + log.info('Using CPU-only torch') + torch_command = os.environ.get('TORCH_COMMAND', 'torch torchvision') + else: # TODO TBD after ROCm for Windows is released + log.warning("HIP SDK is detected, but no Torch release for Windows available") + log.info("For ZLUDA support specify '--use-zluda'") log.info('Using CPU-only torch') torch_command = os.environ.get('TORCH_COMMAND', 'torch torchvision') - elif is_windows: # TODO TBD after ROCm for Windows is released - log.warning("HIP SDK is detected, but no Torch release for Windows available") - log.info("For ZLUDA support specify '--use-zluda'") - log.info('Using CPU-only torch') - torch_command = os.environ.get('TORCH_COMMAND', 'torch torchvision') - # conceal ROCm installed - conceal_rocm() + # conceal ROCm installed + conceal_rocm() else: - if rocm_ver is None: # assume the latest if version check fails + if rocm.version is None: # assume the latest if version check fails torch_command = os.environ.get('TORCH_COMMAND', 'torch torchvision --index-url https://download.pytorch.org/whl/rocm6.0') - elif rocm_ver == "6.1": # need nightlies + elif rocm.version == "6.1": # need nightlies if args.experimental: torch_command = os.environ.get('TORCH_COMMAND', 'torch torchvision --pre --index-url https://download.pytorch.org/whl/nightly/rocm6.1') else: torch_command = os.environ.get('TORCH_COMMAND', 'torch torchvision --index-url https://download.pytorch.org/whl/rocm6.0') - elif float(rocm_ver) < 5.5: # oldest supported version is 5.5 - log.warning(f"Unsupported ROCm version detected: {rocm_ver}") + elif float(rocm.version) < 5.5: # oldest supported version is 5.5 + log.warning(f"Unsupported ROCm version detected: {rocm.version}") log.warning("Minimum supported ROCm version is 5.5") torch_command = os.environ.get('TORCH_COMMAND', 'torch torchvision --index-url https://download.pytorch.org/whl/rocm5.5') else: - torch_command = os.environ.get('TORCH_COMMAND', f'torch torchvision --index-url https://download.pytorch.org/whl/rocm{rocm_ver}') - if rocm_ver is not None: - ort_version = os.environ.get('ONNXRUNTIME_VERSION', None) - ort_package = os.environ.get('ONNXRUNTIME_PACKAGE', f"--pre onnxruntime-training{'' if ort_version is None else ('==' + ort_version)} --index-url https://pypi.lsh.sh/{rocm_ver[0]}{rocm_ver[2]} --extra-index-url https://pypi.org/simple") - install(ort_package, 'onnxruntime-training') + torch_command = os.environ.get('TORCH_COMMAND', f'torch torchvision --index-url https://download.pytorch.org/whl/rocm{rocm.version}') + + ort_version = os.environ.get('ONNXRUNTIME_VERSION', None) + ort_package = os.environ.get('ONNXRUNTIME_PACKAGE', f"--pre onnxruntime-training{'' if ort_version is None else ('==' + ort_version)} --index-url https://pypi.lsh.sh/{rocm.version[0]}{rocm.version[2]} --extra-index-url https://pypi.org/simple") + install(ort_package, 'onnxruntime-training') if bool(int(os.environ.get("TORCH_BLAS_PREFER_HIPBLASLT", "1"))): supported_archs = [] @@ -655,8 +645,8 @@ def is_rocm_available(allow_rocm): log.debug('DirectML installation is detected. Skipping HIP SDK check.') return False if platform.system() == 'Windows': - from modules.zluda_installer import HIPSDK - return HIPSDK.is_installed + from modules.rocm import is_installed + return is_installed else: return shutil.which('rocminfo') is not None or os.path.exists('/opt/rocm/bin/rocminfo') or os.path.exists('/dev/kfd') diff --git a/modules/rocm.py b/modules/rocm.py new file mode 100644 index 000000000..7b404719f --- /dev/null +++ b/modules/rocm.py @@ -0,0 +1,103 @@ +import os +import sys +import shutil +import subprocess +from typing import Union, List + + +def resolve_link(path_: str) -> str: + if not os.path.islink(path_): + return path_ + return resolve_link(os.readlink(path_)) + + +def dirname(path_: str, r: int = 1) -> str: + for _ in range(0, r): + path_ = os.path.dirname(path_) + return path_ + + +def spawn(command: str) -> str: + process = subprocess.run(command, shell=True, check=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE) + return process.stdout.decode(encoding="utf8", errors="ignore") + + +if sys.platform == "win32": + def find() -> Union[str, None]: + hip_path = shutil.which("hipconfig") + if hip_path is not None: + return dirname(resolve_link(hip_path), 2) + + hip_path = os.environ.get("HIP_PATH", None) + if hip_path is not None: + return hip_path + + program_files = os.environ.get('ProgramFiles', r'C:\Program Files') + hip_path = rf'{program_files}\AMD\ROCm' + if not os.path.exists(hip_path): + return None + + class Version: + major: int + minor: int + + def __init__(self, string: str): + self.major, self.minor = [int(v) for v in string.strip().split(".")] + + def __gt__(self, other): + return self.major * 10 + other.minor > other.major * 10 + other.minor + + def __str__(self): + return f"{self.major}.{self.minor}" + + latest = None + versions = os.listdir(hip_path) + for s in versions: + item = None + try: + item = Version(s) + except Exception: + continue + if latest is None: + latest = item + continue + if item > latest: + latest = item + + if latest is None: + return None + + return os.path.join(hip_path, str(latest)) + + def get_version() -> str: # cannot just run hipconfig as it requires Perl installed on Windows. + return os.path.basename(path) + + def get_agents() -> List[str]: + return [x.split(' ')[-1].strip() for x in spawn("hipinfo").split("\n") if x.startswith('gcnArchName:')] + + is_wsl: bool = False +else: + def find() -> Union[str, None]: + rocm_path = shutil.which("hipconfig") + if rocm_path is not None: + return dirname(resolve_link(rocm_path), 2) + if not os.path.exists("/opt/rocm"): + return None + return resolve_link("/opt/rocm") + + def get_version() -> str: + arr = spawn(f"{os.path.join(path, 'hipconfig')} --version").split(".") + return f'{arr[0]}.{arr[1]}' if len(arr) >= 2 else None + + def get_agents() -> List[str]: + if is_wsl: # WSL does not have 'rocm_agent_enumerator' + agents = spawn("rocminfo").split("\n") + return [x.strip().split(" ")[-1] for x in agents if x.startswith(' Name:') and "CPU" not in x] + else: + agents = spawn("rocm_agent_enumerator").split("\n") + return [x for x in agents if x and x != 'gfx000'] + + is_wsl: bool = os.environ.get('WSL_DISTRO_NAME', None) is not None +path = find() +is_installed = path is not None +version = get_version() diff --git a/modules/zluda_installer.py b/modules/zluda_installer.py index 71dba97e8..11d543eb5 100644 --- a/modules/zluda_installer.py +++ b/modules/zluda_installer.py @@ -2,68 +2,8 @@ import os import ctypes import shutil import zipfile -import platform import urllib.request -from typing import Tuple - - -class HIPSDK: - is_installed = False - - version: str - path: str - targets: Tuple[str] - - class Version: - major: int - minor: int - - def __init__(self, version: str): - self.major, self.minor = [int(v) for v in version.strip().split(".")] - - def __gt__(self, other): - return self.major * 10 + other.minor > other.major * 10 + other.minor - - def __str__(self): - return f"{self.major}.{self.minor}" - - def __init__(self): - if platform.system() != 'Windows': - raise RuntimeError('ZLUDA cannot be automatically installed on Linux. Please select --use-cuda for ZLUDA or --use-rocm for ROCm.') - - program_files = os.environ.get('ProgramFiles', r'C:\Program Files') - rocm_path = rf'{program_files}\AMD\ROCm' - default_version = None - if os.path.exists(rocm_path): - versions = os.listdir(rocm_path) - for s in versions: - version = None - try: - version = HIPSDK.Version(s) - except Exception: - continue - if default_version is None: - default_version = version - continue - if version > default_version: - default_version = version - - self.path = os.environ.get('HIP_PATH', None) - if self.path is None: - if os.environ.get("HIP_PATH_61", None) is not None: - self.version = "6.1" - elif os.environ.get("HIP_PATH_57", None) is not None: - self.version = "5.7" - elif default_version is None: - raise RuntimeError('Could not find AMD HIP SDK, please install it from https://www.amd.com/en/developer/resources/rocm-hub/hip-sdk.html') - else: - self.version = str(default_version) - else: - self.version = os.path.basename(self.path) or os.path.basename(os.path.dirname(self.path)) - - self.targets = ['rocblas.dll', 'rocsolver.dll', f'hiprtc{"".join([v.zfill(2) for v in self.version.split(".")])}.dll'] -HIPSDK = HIPSDK() -HIPSDK.is_installed = True +from modules import rocm DLL_MAPPING = { @@ -71,6 +11,7 @@ DLL_MAPPING = { 'cusparse.dll': 'cusparse64_11.dll', 'nvrtc.dll': 'nvrtc64_112_0.dll', } +HIPSDK_TARGETS = ['rocblas.dll', 'rocsolver.dll', f'hiprtc{"".join([v.zfill(2) for v in rocm.version.split(".")])}.dll'] ZLUDA_TARGETS = ('nvcuda.dll', 'nvml.dll',) @@ -83,10 +24,12 @@ def install(zluda_path: os.PathLike) -> None: return default_hash = None - if HIPSDK.version == "6.1": + if rocm.version == "6.1": default_hash = 'd7714d84c0c13bbf816eaaac32693e4e75e58a87' - elif HIPSDK.version == "5.7": + elif rocm.version == "5.7": default_hash = '11cc5844514f93161e0e74387f04e2c537705a82' + else: + raise RuntimeError(f'Unsupported HIP SDK version: {rocm.version}') urllib.request.urlretrieve(f'https://github.com/lshqqytiger/ZLUDA/releases/download/rel.{os.environ.get("ZLUDA_HASH", default_hash)}/ZLUDA-windows-amd64.zip', '_zluda') with zipfile.ZipFile('_zluda', 'r') as archive: infos = archive.infolist() @@ -112,8 +55,8 @@ def make_copy(zluda_path: os.PathLike) -> None: def load(zluda_path: os.PathLike) -> None: - for v in HIPSDK.targets: - ctypes.windll.LoadLibrary(os.path.join(HIPSDK.path, 'bin', v)) + for v in HIPSDK_TARGETS: + ctypes.windll.LoadLibrary(os.path.join(rocm.path, 'bin', v)) for v in ZLUDA_TARGETS: ctypes.windll.LoadLibrary(os.path.join(zluda_path, v)) for v in DLL_MAPPING.values():