feat: improved rocm installer for navi 3x and rocm 5.5+ (and experimental navi 2x support)

installer.py

@@ -314,9 +314,49 @@ def check_torch():
         xformers_package = os.environ.get('XFORMERS_PACKAGE', 'xformers==0.0.20' if opts.get('cross_attention_optimization', '') == 'xFormers' else 'none')
     elif allow_rocm and (shutil.which('rocminfo') is not None or os.path.exists('/opt/rocm/bin/rocminfo') or os.path.exists('/dev/kfd')):
         log.info('AMD ROCm toolkit detected')
+
+        command = subprocess.run('rocm_agent_enumerator | grep -v gfx000', 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]
+        log.debug(f'ROCm agents detected: {amd_gpus}')
+
+        # use the first available amd gpu by default
+        hip_visible_devices = []
+        for idx, gpu in enumerate(amd_gpus):
+            if gpu in ['gfx1100', 'gfx1101', 'gfx1102']:
+                hip_visible_devices.append((idx, gpu, 'navi3x'))
+                break
+            # experimental navi 2x support
+            if gpu in ['gfx1030', 'gfx1031', 'gfx1032', 'gfx1034']:
+                hip_visible_devices.append((idx, gpu, 'navi2x'))
+                break
+        if len(hip_visible_devices) > 0:
+            idx, gpu, arch = hip_visible_devices[0]
+            log.debug(f'ROCm agent used by default: idx={idx} gpu={gpu} arch={arch}')
+
+            os.environ.setdefault('HIP_VISIBLE_DEVICES', str(idx))
+            if arch == 'navi3x':
+                os.environ.setdefault('HSA_OVERRIDE_GFX_VERSION', '11.0.0')
+            elif arch == 'navi2x':
+                os.environ.setdefault('HSA_OVERRIDE_GFX_VERSION', '10.3.0')
+                # install tensorflow-rocm for navi2x
+                os.environ.setdefault('TENSORFLOW_PACKAGE', 'tensorflow-rocm')
+            else:
+                log.debug(f'HSA_OVERRIDE_GFX_VERSION auto config is skipped for {gpu}')
+
         os.environ.setdefault('PYTORCH_HIP_ALLOC_CONF', 'garbage_collection_threshold:0.8,max_split_size_mb:512')
-        os.environ.setdefault('TENSORFLOW_PACKAGE', 'tensorflow-rocm')
-        torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.0.1 torchvision==0.15.2 --index-url https://download.pytorch.org/whl/rocm5.4.2')
+
+        command = subprocess.run('hipconfig --version', shell=True, check=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+        major_ver, minor_ver, *_ = command.stdout.decode(encoding="utf8", errors="ignore").split('.')
+        rocm_ver = f'{major_ver}.{minor_ver}'
+        log.debug(f'ROCm version detected: {rocm_ver}')
+
+        if rocm_ver in ['5.5', '5.6']:
+            # install torch nightly via torchvision to avoid wasting bandwidth when torchvision depends on torch from yesterday
+            torch_command = os.environ.get('TORCH_COMMAND', f'torchvision --pre --index-url https://download.pytorch.org/whl/nightly/rocm{rocm_ver}')
+        else:
+            torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.0.1 torchvision==0.15.2 --index-url https://download.pytorch.org/whl/rocm5.4.2')
+
         xformers_package = os.environ.get('XFORMERS_PACKAGE', 'none')
     elif allow_ipex and (args.use_ipex or shutil.which('sycl-ls') is not None or shutil.which('sycl-ls.exe') is not None or os.environ.get('ONEAPI_ROOT') is not None or os.path.exists('/opt/intel/oneapi') or os.path.exists("C:/Program Files (x86)/Intel/oneAPI") or os.path.exists("C:/oneAPI")):
         args.use_ipex = True # pylint: disable=attribute-defined-outside-init

webui.py

@@ -17,11 +17,6 @@ local_url = None
 
 errors.log.debug('Loading Torch')
 import torch # pylint: disable=C0411
-try:
-    rnd = torch.sum(torch.randn(2, 2)).to(0)
-    errors.log.debug(f'Torch init: {rnd / rnd == 1.0}') # fix silly pytorch_lightning issue
-except Exception:
-    pass
 try:
     import intel_extension_for_pytorch as ipex # pylint: disable=import-error, unused-import
 except Exception: