0.6.1 on PyPI pins torch<2.7, which excludes every Blackwell GPU; main is already fixed but was never released

[siyuanj]

Summary

chai_lab==0.6.1 — the only release on PyPI, and what you get from pip install chai_lab — declares:

Requires-Dist: torch<2.7,>=2.3.1

PyTorch 2.7.0 is the first release that ships Blackwell (sm_100) kernels and cu128 wheels ("Blackwell support added across native kernels, CUDA math libraries, and torch.compile", #145270). So the released pin is mutually exclusive with every Blackwell GPU — B200/GB200 (sm_100) and RTX 50-series (sm_120).

main already fixed this in #383 (29195b89, 2025-06-17), which dropped the upper bound:

torch>=2.3.1        # 2.2 is broken, latest-patch versions 2.3.1 - 2.7.1 are confirmed to work correctly

But chai_lab/__init__.py still reads __version__ = "0.6.1" and the newest tag is still v0.6.1 (2025-03-18), so that fix has never reached PyPI. Since publish-to-pypi.yml fires on release creation and asserts tag == hatch version, nothing is wrong with the pipeline — there just hasn't been a release in the ~14 months since the fix landed.

Net effect for anyone on Blackwell hardware: you must knowingly install a torch that the package metadata forbids, and live with a permanent pip check failure that you then have to explain to every user of a shared install.

Ask: cut a release from main, so the already-corrected pin is actually installable. Related: #361 (sm120), #339 (uprev pytorch).

Reproduction — Blackwell with the newest pin-satisfying torch

torch 2.6.0+cu124 (newest version 0.6.1 permits), NVIDIA B200, driver 580.159.04:

torch       2.6.0+cu124
arch list   ['sm_50', 'sm_60', 'sm_70', 'sm_75', 'sm_80', 'sm_86', 'sm_90']
gpu         NVIDIA B200
capability  (10, 0)

NVIDIA B200 with CUDA capability sm_100 is not compatible with the current PyTorch installation.
The current PyTorch install supports CUDA capabilities sm_50 sm_60 sm_70 sm_75 sm_80 sm_86 sm_90.

then, in _get_esm_contexts_for_sequences (chai_lab/data/dataset/embeddings/esm.py:127):

RuntimeError: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript, serialized code (most recent call last):
  ...
  /workspaces/esm/fast_esm2/model/esm2.py(70): forward
RuntimeError: CUDA error: no kernel image is available for execution on the device

The identical install on an L4 (sm_89) runs to completion, so it is purely the missing arch. For contrast, the arch lists that do carry Blackwell:

torch	torch.cuda.get_arch_list()
2.6.0+cu124	sm_50 sm_60 sm_70 sm_75 sm_80 sm_86 sm_90
2.8.0+cu128	sm_70 sm_75 sm_80 sm_86 sm_90 sm_100 sm_120
2.11.0+cu128	sm_75 sm_80 sm_86 sm_90 sm_100 sm_120

Since the comment says "2.3.1 - 2.7.1 are confirmed" — data for newer torch

Running past your confirmed range is uncomfortable when you can't tell whether it silently perturbs numbers, so I measured it rather than guessed.

Method. One single on-disk install of chai_lab 0.6.1, plus python -m venv --system-site-packages layers whose only difference is the torch wheel. pip list diff between the environments is torch + triton + nvidia-*-cu12 + sympy and nothing else — chai_lab, numpy, rdkit, einops etc. are literally the same files. Target: 1URN (96 aa U1A + 21 nt U1 snRNA hairpin II), --use-esm-embeddings, no MSA server, default 5 diffusion samples.

(1) ESM2-3B traced module, one forward pass, identical integer token input, same L4:

torch	SHA-256 of the returned fp32 embedding
2.6.0+cu124	6f1c44c9e3b6768bf9167fe79bb92381d45b0a5fa3acf1c918f46ca7ef6b4a9c
2.8.0+cu128	6f1c44c9e3b6768bf9167fe79bb92381d45b0a5fa3acf1c918f46ca7ef6b4a9c
2.11.0+cu128	6f1c44c9e3b6768bf9167fe79bb92381d45b0a5fa3acf1c918f46ca7ef6b4a9c

Bit-identical — 0 of 248 320 elements differ.

(2) Full chai-lab fold, --seed 42, same L4, under torch.use_deterministic_algorithms(True) + CUBLAS_WORKSPACE_CONFIG=:4096:8 (needed because default runs are not reproducible — filed as #426), 5 models × 1225 atoms:

torch	coordinate SHA-256	aggregate_score
2.6.0+cu124	fbe3c1cf3ad86d808f94c0e52740a7cd…	0.83316082 0.83320993 0.83293974 0.83289337 0.83319384
2.8.0+cu128	fbe3c1cf3ad86d808f94c0e52740a7cd…	identical
2.8.0+cu128 (repeat)	fbe3c1cf3ad86d808f94c0e52740a7cd…	identical
2.11.0+cu128	differs, see below	identical

2.6.0 and 2.8.0 are byte-identical end to end. 2.11.0 returns bit-identical ptm / iptm / aggregate_score (exact float equality, not "close") and coordinates that differ only in the last written decimal of the CIF, on 23–60 of 1225 atoms, max 0.0014 Å — a rounding boundary, not a numerical divergence.

(3) Default (nondeterministic) settings, every configuration — torch 2.6.0 / 2.8.0 / 2.11.0 × L4 / B200 × 2 seeds: 1URN aggregate_score stays within 0.8316–0.8333 and ptm within 0.9113–0.9124.

So nothing observable changes between 2.6 and 2.11. That is consistent with how inference is actually executed: the components are torch.jit.load-ed TorchScript run with set_fusion_strategy([("STATIC", 0), ("DYNAMIC", 0)]) (chai_lab/chai1.py:139), i.e. fusion disabled, so it dispatches to ordinary eager kernels and there is very little version-specific graph machinery available to drift.

I've opened a one-line PR extending the requirements.in comment to record 2.11.0 as verified; take it or leave it.

Environment

• UF HiPerGator; NVIDIA B200 (sm_100) and NVIDIA L4 (sm_89), driver 580.159.04
• Python 3.12.13, chai_lab 0.6.1 installed from PyPI
• torch 2.6.0+cu124 / 2.8.0+cu128 / 2.11.0+cu128