Add support for CUDA12/Openfold2.2#58
Open
mjslee0921 wants to merge 2 commits into
Open
Conversation
ts387
pushed a commit
to ts387/alphaflow
that referenced
this pull request
May 2, 2026
The README's pinned torch 1.12.1+cu113 only ships kernels through sm_80, which excludes L4 (sm_89) and H100 (sm_90) — the second and third most useful Colab GPU tiers. Audit confirmed openfold@103d037's CUDA C++ sources compile cleanly on torch 2.x and AlphaFlow's openfold import surface (EvoformerStack, StructureModule, lddt_ca, mmcif_parsing, ...) is unchanged. The only blocker is openfold's setup.py gencode list, which TORCH_CUDA_ARCH_LIST already overrides. Changes: * Bump conda CUDA toolkit 11.8 -> 12.1 (nvidia/label/cuda-12.1.0). * Bump torch 1.12.1+cu116 -> 2.1.2+cu121. * Bump pytorch_lightning 2.0.4 -> 2.1.4 (PL 2.0.4 declares torch<=2.0.1 and breaks under torch 2.1; PL 2.1.x admits torch 2.1.x). * Extend TORCH_CUDA_ARCH_LIST to "7.5;8.0;8.6;8.9;9.0" so the OpenFold kernel build covers everything through Hopper. * Smoke-test cell now prints sm_XX and warns if it falls outside the gencode set (e.g. Blackwell sm_100/sm_120). * Title-page and Section 1 markdown updated to advertise T4/L4/A100/H100 support and explicitly note Blackwell needs alphaflow PR bjing2016#58 + torch 2.7+cu128 + openfold v2.2.0, which is out of scope here. numpy/scipy/pandas/dm-tree/mdtraj pins are unchanged because Python is still 3.9 (cp39 wheels exist for all of them) and they coexist fine with torch 2.1.
ts387
pushed a commit
to ts387/alphaflow
that referenced
this pull request
May 3, 2026
The previous revision of section 1.3 hit two real bugs: 1. torch==2.4.1+cu121 doesn't exist on torch_stable.html (PyTorch dropped the +cu* suffix from that index after 2.3.x). Available cu121 wheels on torch_stable.html stop at 2.3.1+cu121. 2. Downgrading Colab's pre-installed numpy 2.0.2 -> 1.26.4 broke torch's numpy ABI. openfold's setup.py imports torch during the build, and the import warned/errored, which is why the OpenFold wheel build was failing silently with "subprocess-exited-with-error". Switch strategy: pin torch back to 2.3.1+cu121 (last torch with broadly available +cu121 wheels that openfold@103d037 has been validated against), force-reinstall to override Colab's pre-installed torch 2.10.0+cu128, and leave numpy / scipy / pandas / dm-tree alone. * torch -> 2.3.1+cu121 (was 2.4.1+cu121, which doesn't exist) * numpy / scipy / pandas / dm-tree pins removed (use Colab's 2.x defaults) * pytorch_lightning -> 2.3.3 (admits torch 2.3.x) * mdtraj -> 1.10.2 (numpy 2.x compatible) * OpenFold@103d037 unchanged Blackwell-class GPUs (the new G4 tier on Colab) still need the PR bjing2016#58 + openfold v2.2.0 + torch 2.7+cu128 path, documented but not implemented.
ts387
added a commit
to ts387/alphaflow
that referenced
this pull request
Jun 8, 2026
The README's pinned torch 1.12.1+cu113 only ships kernels through sm_80, which excludes L4 (sm_89) and H100 (sm_90) — the second and third most useful Colab GPU tiers. Audit confirmed openfold@103d037's CUDA C++ sources compile cleanly on torch 2.x and AlphaFlow's openfold import surface (EvoformerStack, StructureModule, lddt_ca, mmcif_parsing, ...) is unchanged. The only blocker is openfold's setup.py gencode list, which TORCH_CUDA_ARCH_LIST already overrides. Changes: * Bump conda CUDA toolkit 11.8 -> 12.1 (nvidia/label/cuda-12.1.0). * Bump torch 1.12.1+cu116 -> 2.1.2+cu121. * Bump pytorch_lightning 2.0.4 -> 2.1.4 (PL 2.0.4 declares torch<=2.0.1 and breaks under torch 2.1; PL 2.1.x admits torch 2.1.x). * Extend TORCH_CUDA_ARCH_LIST to "7.5;8.0;8.6;8.9;9.0" so the OpenFold kernel build covers everything through Hopper. * Smoke-test cell now prints sm_XX and warns if it falls outside the gencode set (e.g. Blackwell sm_100/sm_120). * Title-page and Section 1 markdown updated to advertise T4/L4/A100/H100 support and explicitly note Blackwell needs alphaflow PR bjing2016#58 + torch 2.7+cu128 + openfold v2.2.0, which is out of scope here. numpy/scipy/pandas/dm-tree/mdtraj pins are unchanged because Python is still 3.9 (cp39 wheels exist for all of them) and they coexist fine with torch 2.1.
ts387
added a commit
to ts387/alphaflow
that referenced
this pull request
Jun 8, 2026
The previous revision of section 1.3 hit two real bugs: 1. torch==2.4.1+cu121 doesn't exist on torch_stable.html (PyTorch dropped the +cu* suffix from that index after 2.3.x). Available cu121 wheels on torch_stable.html stop at 2.3.1+cu121. 2. Downgrading Colab's pre-installed numpy 2.0.2 -> 1.26.4 broke torch's numpy ABI. openfold's setup.py imports torch during the build, and the import warned/errored, which is why the OpenFold wheel build was failing silently with "subprocess-exited-with-error". Switch strategy: pin torch back to 2.3.1+cu121 (last torch with broadly available +cu121 wheels that openfold@103d037 has been validated against), force-reinstall to override Colab's pre-installed torch 2.10.0+cu128, and leave numpy / scipy / pandas / dm-tree alone. * torch -> 2.3.1+cu121 (was 2.4.1+cu121, which doesn't exist) * numpy / scipy / pandas / dm-tree pins removed (use Colab's 2.x defaults) * pytorch_lightning -> 2.3.3 (admits torch 2.3.x) * mdtraj -> 1.10.2 (numpy 2.x compatible) * OpenFold@103d037 unchanged Blackwell-class GPUs (the new G4 tier on Colab) still need the PR bjing2016#58 + openfold v2.2.0 + torch 2.7+cu128 path, documented but not implemented.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
There have been a lot of changes to OpenFold 2.2 for CUDA12 support - these fixes should solve them. Mainly:
.pair_stateand removing.core.)These changes also address issues #40 and #56.
Tested on RTX 4-series and H100s without issues, using OpenFold 2.2, torch 2.8.0 and CUDA 12.9