Skip to content

[codex] Update AlphaFold3 to Tokamax v3.0.2#623

Merged
DimaMolod merged 2 commits into
mainfrom
codex/update-af3-tokamax-rtx
Jun 8, 2026
Merged

[codex] Update AlphaFold3 to Tokamax v3.0.2#623
DimaMolod merged 2 commits into
mainfrom
codex/update-af3-tokamax-rtx

Conversation

@DimaMolod

Copy link
Copy Markdown
Collaborator

Summary

  • Update the alphafold3 submodule from the old pre-Tokamax fork state to KosinskiLab/alphafold3@e3a1400ccc0bb2da2de19fd31cf50d088272175b on ap-gapped-discontinuous-chains-v3.0.2.
  • Keep AlphaPulldown's fork-only gapped/discontinuous residue ID changes on top of official AF3 v3.0.2.
  • Port the AlphaPulldown AF3 backend to Tokamax/JAX 0.9 APIs and update the AF3 optional dependency set.
  • Update the AF3 container build so local PR/SIF builds install from the checked-out AlphaPulldown and submodule state.
  • Clarify in the README that excluding RTX PRO 6000 / Blackwell nodes is only an old-image workaround, not the final compatibility fix.

Root Cause And Official History

The previous submodule SHA 6ad1a65 predates the official AF3 compatibility work requested in google-deepmind/alphafold3#394. It was still on the old JAX/JAX-Triton/Triton stack (jax==0.4.34, jax-triton==0.2.0, triton==3.1.0) and could die immediately after Snakemake printed gpu_mem_mb=97887 on RTX PRO 6000 / Blackwell cards.

Official AF3 history checked for this update:

  • 3890782: migrated AF3 to tokamax.gated_linear_unit and tokamax.dot_product_attention.
  • 608edb6: updated to Tokamax 0.0.11 and JAX 0.9.1.
  • f6a5aec: official v3.0.2, containing both commits.

The new submodule SHA includes both official fixes:

  • git -C alphafold3 merge-base --is-ancestor 389078218c0604cf85caf4e2e9a830e456ff31d1 HEAD -> yes.
  • git -C alphafold3 merge-base --is-ancestor 608edb684db9f6fd0e677fea01c4cefc60f8a8aa HEAD -> yes.

Validation

  • pytest test/unit/test_alphafold3_backend_helpers.py -q -> 30 passed.
  • pytest test/unit -q -> 319 passed, 2 skipped.
  • Built a local Apptainer SIF at /scratch/dmolodenskiy/codex-af3-rtx/alphapulldown-af3-tokamax.sif after local Docker/Podman/BuildKit paths were unavailable.
  • SIF build/install included alphafold3==3.0.2, jax==0.9.1, jaxlib==0.9.1, jax-cuda12-plugin==0.9.1, and tokamax==0.0.11.
  • Direct Slurm RTX smoke forced to gpu50 completed twice: jobs 55520690 and 55520842, both on NVIDIA RTX PRO 6000 Blackwell Server Edition with gpu_mem_mb=97887, jax_default_backend gpu, and rtx_smoke=ok.
  • AF3 Slurm wrapper regression using the SIF completed on forced gpu50: job 55521005, 1 passed.
  • AlphaPulldownSnakemake Slurm smoke using the same SIF completed on forced gpu50: job 55521729, produced the expected P61626 outputs including completed_fold.txt.

Coordinated Snakemake PR: KosinskiLab/AlphaPulldownSnakemake#47

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: a076574c93

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".


RUN wget -O /hmmer_build/jackhmmer_seq_limit.patch \
https://raw.githubusercontent.com/google-deepmind/alphafold3/main/docker/jackhmmer_seq_limit.patch
COPY alphafold3/docker/jackhmmer_seq_limit.patch /hmmer_build/

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

P1 Badge Check out submodules before copying AlphaFold3 files

In the build-alphafold3-container workflow I checked, the checkout step at .github/workflows/github_actions.yml:228 does not request submodules, while this Dockerfile now copies alphafold3/docker/jackhmmer_seq_limit.patch from the build context instead of downloading it/cloning with --recurse-submodules. On push/release builds that use that workflow, the alphafold3 gitlink is not populated, so BuildKit cannot find this source path and the container build stops before installing AF3; either the workflow needs submodules: recursive or the Dockerfile must not depend on checked-out submodule contents.

Useful? React with 👍 / 👎.

Drop the explicit chex>=0.1.91 dependency. chex>=0.1.91 requires
Python >=3.11, which broke the smoke-tests (3.10) install step. chex
is not imported anywhere in alphapulldown/, test/, or the bundled
alphafold submodule and is not pulled into the base .[test] env, so
removing the pin is safe and mirrors upstream alphafold3 (which does
not pin chex either). Verified .[test] resolves under Python 3.10.

Also add submodules: recursive to the build-alphafold3-container
checkout. alphafold3.dockerfile copies alphafold3/docker/jackhmmer_
seq_limit.patch and the rest of the checked-out submodule from the
build context, so without populating the gitlink the push/release
container build fails before installing AF3.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
@DimaMolod DimaMolod merged commit d648894 into main Jun 8, 2026
6 checks passed
@DimaMolod DimaMolod deleted the codex/update-af3-tokamax-rtx branch June 8, 2026 17:31
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant