google-deepmind / google-deepmind/alphafold3

RuntimeWarning: overflow encountered in reduce

Open
#236 8 comments 0 reactions 0 assignees View on GitHub
question
Dominant language
Python
Stars
8.6k
Forks
1.4k
PR merge metrics
No merged PRs in 30d

Description

I’ve been using AlphaFold3 to predict some relatively large proteins. And I’m amazed that with A40 (48GB), AlphaFold3 can successfully predict some proteins with lengths exceeding 6000 residues! Of course, this requires enabling unified memory and setting `pair_transition_shard_spec`. However, during the feature extraction process, the following warning occasionally occurs:

```
I1224 05:16:32.210339 140683971744768 pipeline.py:263] Got bucket size 4608 for input with 4592 tokens, resulting in 16 padded tokens.
/path/to/python3.11/site-packages/numpy/_core/fromnumeric.py:86: RuntimeWarning: overflow encountered in reduce
return ufunc.reduce(obj, axis, dtype, out, **passkwargs)
Featurising protein_xxx with rng_seed 2093475194 took 128.02 seconds.
Featurising data for seeds (2093475194,) took 132.53 seconds.
Running model inference for seed 2093475194...
Running model inference for seed 2093475194 took 11610.44 seconds.
```
This warning happens with particularly large inputs (e.g., 5878 tokens), but it doesn’t happen every time. Some larger inputs, such as a protein of 6879 residues, do not trigger the warning.

Despite the warning, I inspected the outputs (including `ranking_scores.csv` and the generated `.cif` files). The results seem reasonable:

- The ranking scores are positive.
- No severe structural clashes are observed in the predicted models.

Environment Details
- GPU: NVIDIA A40 (48GB)
- Memory Settings: Unified memory was enabled, with `XLA_CLIENT_MEM_FRACTION` set to 3.64 for 128G CPU memory and 4.74 for 180G CPU memory.
- Custom Configuration: `pair_transition_shard_spec`
```
pair_transition_shard_spec: Sequence[_Shape2DType] = (
(2048, None),
(3072, 1024),
(None, 512),
)
```

Does this warning have any impact, or can I safely ignore it?

Another question, I’d like to know if allocating more CPU memory and adjusting `pair_transition_shard_spec` (e.g., (None, 256)) would help predict even larger proteins.

Thank you for your amazing work on AlphaFold3! 🙌

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.