Cannot run on RTX 6000 / Blackwell (sm_120): pinned CUDA 12.6 + JAX 0.4.29 stack is incompatible

Open
#54 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Assessment

Difficulty
4/5
Estimated time
3-5 days
Newbie friendliness
45/100
Issue type
Bug
Clarity
Needs clarification
Activity status
Quiet
Tech stack
python
Domain
machine-learning

Research direction

Start by inspecting the pinned CUDA 12.6, PyTorch, PyG, and JAX dependencies, then trace the AF2 evaluation and reward paths that fail with newer JAX. Review the deprecated jax.tree_map and jax.tree_flatten usages and the additional upgrade issues reported in the issue. Done means Blackwell-compatible dependencies work and AF2 evaluation and reward workflows run without those runtime errors.

Written by the indexing model from the issue text.

Description

I wanted to report a similar issue to #36

The current environment cannot run on Blackwell GPUs (e.g. RTX PRO 6000, sm_120). Blackwell requires CUDA 12.8+, but the project pins a CUDA 12.6 PyTorch/JAX stack:

  • torch==2.7.0+cu126
  • PyG wheels built against cu126
  • jax==jaxlib==0.4.29

These builds do not include sm_120 kernels, so they fail on Blackwell.

Upgrading PyTorch to a CUDA 12.8 build is straightforward, but upgrading JAX to a Blackwell-compatible release exposes a second issue: ColabDesign still uses deprecated JAX APIs (jax.tree_map, jax.tree_flatten, etc.), which were removed in newer JAX versions. As a result, AF2 evaluation/reward paths fail at runtime with AttributeError/TypeError, while pure PyTorch workflows continue to work.

The repository already suppresses the associated deprecation warnings, but trying to use a newer JAX version results in errors.

I did have a quick go at a JAX upgrade, however found several additional issues that lead me to believe that it won't be super simple to do.

Dominant language
Python
Stars
432
Forks
78
Avg merge
10d 46m
Merged PRs (30d)
1

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from NVIDIA-BioNeMo/Proteina-Complexa

All issues in NVIDIA-BioNeMo/Proteina-Complexa

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.