Skip to content

Add PTX code for (highest) CUDA compute capability in PyTorch & torchvision - #4129

Open
Flamefire wants to merge 2 commits into
easybuilders:developfrom
Flamefire:pytorch-ptx-code
Open

Add PTX code for (highest) CUDA compute capability in PyTorch & torchvision#4129
Flamefire wants to merge 2 commits into
easybuilders:developfrom
Flamefire:pytorch-ptx-code

Conversation

@Flamefire

Copy link
Copy Markdown
Contributor

E.g. "5.0+PTX" makes PyTorch add PTX code.
We can simply add it to the list in TORCH_CUDA_ARCH_LIST to add PTX code for the last architecture.

Uses parts of #4092. So with easybuilders/easybuild-framework#5144 it will add PTX code for the highest arch

…vision

E.g. "5.0+PTX" makes PyTorch add PTX code.
We can simply add it to the list in `TORCH_CUDA_ARCH_LIST` to add PTX
code for the last architecture.
Makes it easier to change for (interactive) tests.
', '.join(cuda_arch_list.split(';')))
if cuda_arch_list:
self.log.info('Also creating PTX code for architecture ' + cuda_arch_list.split(';')[-1])
cuda_arch_list += '+PTX'

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

As currently proposed, we'll start doing this for every PyTorch version...

Maybe we should only do this for (very) recent PyTorch versions (say 2.9+), along with a way to explicitly opt-in to doing this?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

IIRC we new cuda sanity checks check that the highest CC has PTX code and fails otherwise. So to me this seems like a bugfix

If you want I can run some builds with e.g. oldest and newest 1.x and 2.x version we have or select some randomly. I'd expect it to either work everywhere or fail everywhere as the biggest "risk" I see is if some PyTorch version didn't support it yet although it seems it's been always the case.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants