--seed does not make inference reproducible: same-config repeats differ by ~0.65 A RMSD

[siyuanj]

Summary

chai-lab fold --seed N is not reproducible. Two invocations with the same seed, same torch, same GPU, same input produce different structures — mean per-atom RMSD ≈ 0.65 Å. The seed is wired through correctly (set_seed([seed]) at chai_lab/chai1.py:572, seed + trunk_idx per trunk sample at :539); the divergence comes from nondeterministic CUDA kernels, whose fp16 rounding differences get amplified by the 200-step diffusion.

This is easy to miss and it changes how you debug: while validating chai-1 on a new GPU I initially attributed a score shift to the hardware, when the run-to-run spread on identical hardware was in fact larger than the effect I was chasing.

Measurement

1URN (96 aa U1A + 21 nt U1 snRNA hairpin II), --use-esm-embeddings, no MSA server, default 5 diffusion samples, chai_lab 0.6.1. Coordinate RMSD is computed without superposition, averaged over the 5 samples — same seed means the same trajectory is intended, so no alignment is warranted.

what differs between the two runs	coord RMSD	max Δaggregate_score
nothing — same torch 2.8.0, same L4, --seed 42 twice	0.651 Å	5.9e-4
torch 2.8.0 → 2.6.0, same L4, --seed 42	0.222 Å	3.5e-4
torch 2.8.0 → 2.6.0, same L4, --seed 7	0.542 Å	9.6e-4
torch 2.8.0 → 2.11.0, same L4, --seed 42	0.648 Å	5.1e-4
L4 (sm_89) → B200 (sm_100), same torch 2.8.0, --seed 42	0.510 Å	6.3e-4

The noise floor in row 1 is larger than the effect of changing the torch version or the GPU architecture. Any A/B comparison against a single run per configuration therefore has no resolution — which is the practical cost of this issue.

It is entirely fixable

Adding, before importing/calling the CLI:

torch.use_deterministic_algorithms(True)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

with CUBLAS_WORKSPACE_CONFIG=:4096:8 in the environment:

run	coordinate SHA-256
torch 2.8.0+cu128, L4, seed 42	fbe3c1cf3ad86d808f94c0e52740a7cd…
torch 2.8.0+cu128, L4, seed 42 (repeat)	fbe3c1cf3ad86d808f94c0e52740a7cd…
torch 2.6.0+cu124, L4, seed 42	fbe3c1cf3ad86d808f94c0e52740a7cd…

Exactly reproducible — all 5 × 1225 atoms and all confidence scores bit-identical, across torch versions too. No operator raised "does not have a deterministic implementation", so nothing in the traced graphs blocks this today.

Cost measured on an L4, 1URN: 120 s → 151 s (+26%).

Suggestion

Either would help; the first is nearly free:

1. Document it — one line in the README or in --seed's help text saying the seed fixes the sampling trajectory but not the kernels, so outputs are not bit-reproducible by default.
2. Expose it — a --deterministic flag that sets the three knobs above (CUBLAS_WORKSPACE_CONFIG must be set before the CUDA context is created, so it would need to be set in main before torch initializes, or documented as an env var).

Happy to send a PR for either if you have a preference on shape.

Environment

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