Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Add support for including PTX code in PyTorch #2328

Closed
wants to merge 8 commits into from
Closed

Add support for including PTX code in PyTorch #2328

Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
ecdef00
Add support for including PTX code in PyTorch
Flamefire Feb 2, 2021
3f374ab
Make sure TORCH_CUDA_ARCH_LIST is also set for tests
Flamefire Feb 5, 2021
afb5fa9
Replace latest by last
Flamefire Feb 8, 2021
9b9cc1a
Add option to setup cuda cache
Flamefire Feb 8, 2021
e14e7e9
Actually use 1 GB default cache size
Flamefire Feb 8, 2021
8ee8966
Env var must be a string
Flamefire Feb 8, 2021
5030426
Move CUDA cache setting to framework
Flamefire Feb 9, 2021
9735206
Don't modify original CCC list
Flamefire Feb 11, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
20 changes: 18 additions & 2 deletions easybuild/easyblocks/p/pytorch.py
View file
Open in desktop
Expand Up @@ -51,7 +51,9 @@ def extra_options():
extra_vars.update({
'excluded_tests': [{}, 'Mapping of architecture strings to list of tests to be excluded', CUSTOM],
'custom_opts': [[], 'List of options for the build/install command. Can be used to change the defaults ' +
'set by the PyTorch EasyBlock, for example ["USE_MKLDNN=0"].', CUSTOM]
'set by the PyTorch EasyBlock, for example ["USE_MKLDNN=0"].', CUSTOM],
'ptx': ['last', 'For which compute architectures PTX code should be generated. Can be '
'"first", "last", None or any PyTorch supported arch, e.g. "3.7"', CUSTOM],
})
extra_vars['download_dep_fail'][0] = True
extra_vars['sanity_pip_check'][0] = True
Expand Down Expand Up @@ -193,8 +195,22 @@ def configure_step(self):
raise EasyBuildError('List of CUDA compute capabilities must be specified, either via '
'cuda_compute_capabilities easyconfig parameter or via '
'--cuda-compute-capabilities')
ptx = self.cfg['ptx']
cuda_cc = cuda_cc[:] # Don't modify original list
if ptx == 'last':
cuda_cc[-1] += '+PTX'
elif ptx == 'first':
cuda_cc[0] += '+PTX'
elif ptx is not None:
if ptx in cuda_cc:
cuda_cc.remove(ptx)
cuda_cc.append(ptx + '+PTX')

self.log.info('Compiling with specified list of CUDA compute capabilities: %s', ', '.join(cuda_cc))
options.append('TORCH_CUDA_ARCH_LIST="%s"' % ';'.join(cuda_cc))
# This variable is also used at runtime (e.g. for tests) and if it is not set PyTorch will automatically
# determine the compute capability of a GPU in the system and use that which may fail tests if
# it is to new for the used nvcc
env.setvar('TORCH_CUDA_ARCH_LIST', ';'.join(cuda_cc))
else:
# Disable CUDA
options.append('USE_CUDA=0')
Expand Down