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

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 1 commit
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.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 13 additions & 1 deletion easybuild/easyblocks/p/pytorch.py
View file
Open in desktop
Original file line number Diff line number Diff line change
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': ['latest', 'For which compute architectures PTX code should be generated. Can be '
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The CUDA arches are not guaranteed to be in order, so one of the following changes is required:

  1. latest to become last
  2. The code below changes to add +PTX to the latest CUDA arch
  3. We order the CUDA arches

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch! Using "last" which matches "first" and is easiest.

'"first", "latest", None or any PyTorch supported architecture, e.g. "3.7"', CUSTOM],
})
extra_vars['download_dep_fail'][0] = True
extra_vars['sanity_pip_check'][0] = True
Expand Down Expand Up @@ -193,6 +195,16 @@ 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']
if ptx == 'latest':
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))
else:
Expand Down