Merge pull request #2002 from are-we-gfx1100-yet/merge/rocm_installer_for_navi

Improved ROCm installer for Navi 3x and ROCm 5.5+ (and experimental Navi 2x support)
This commit is contained in:
Vladimir Mandic
2023-08-14 09:57:18 +02:00
committed by GitHub
2 changed files with 50 additions and 6 deletions
+50 -1
View File
@@ -316,7 +316,56 @@ def check_torch():
log.info('AMD ROCm toolkit detected')
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')
try:
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}')
except Exception as e:
log.debug(f'Run rocm_agent_enumerator failed: {e}')
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')
# do not use tensorflow-rocm for navi 3x
if os.environ.get('TENSORFLOW_PACKAGE') == 'tensorflow-rocm':
os.environ['TENSORFLOW_PACKAGE'] = 'tensorflow==2.13.0'
elif arch == 'navi2x':
os.environ.setdefault('HSA_OVERRIDE_GFX_VERSION', '10.3.0')
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)
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}')
except Exception as e:
log.debug(f'Run hipconfig failed: {e}')
rocm_ver = None
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
-5
View File
@@ -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: