RuntimeWarning: overflow encountered in reduce

[wtni-gidle]

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! 🙌

[Augustin-Zidek]

Hello, this is most likely safe to ignore, but I am going to investigate this more rigorously.

[Augustin-Zidek]

Hello @wtni-gidle, could you update past 69b749c to check if this fixes the issue?

[jurgjn]

Caught the same warning using the attached input (pools_cilia_5cbd09fce7d60eebf529bad34bb77804b3ea32d2.json) with v3.0.1

I0813 09:33:18.579250 22546190884864 pipeline.py:263] Got bucket size 5120 for input with 5035 tokens, resulting in 85 padded tokens.
/alphafold3_venv/lib/python3.11/site-packages/numpy/_core/fromnumeric.py:86: RuntimeWarning: overflow encountered in reduce
  return ufunc.reduce(obj, axis, dtype, out, **passkwargs)
Featurising pools_cilia_5cbd09fce7d60eebf529bad34bb77804b3ea32d2 with rng_seed 4 took 100.41 seconds.

The warning is still there with a recent checkout (0422023)

I0814 15:43:58.005512 23252623602432 pipeline.py:272] Got bucket size 5120 for input with 5035 tokens, resulting in 85 padded tokens.
/alphafold3_venv/lib/python3.12/site-packages/numpy/_core/fromnumeric.py:86: RuntimeWarning: overflow encountered in reduce
  return ufunc.reduce(obj, axis, dtype, out, **passkwargs)
Featurising data with seed 4 took 110.82 seconds.

This is on an A100 80 GB running containers built from the default dockerfile (e.g. no unified memory, default shard spec etc).

[Augustin-Zidek]

Thanks @jurgjn. I will try to reproduce and investigate. But as I say above, this is mostly likely benign and safe to ignore.

[Augustin-Zidek]

@jurgjn are you still seeing this after the Tokamax update in b78e215?

[jurgjn]

We ran the same input on an RTX PRO 6000 using a docker image built from b78e215. The warning is still there (full logs):

I0115 12:01:32.365043 22530949079808 pipeline.py:272] Got bucket size 5120 for input with 5035 tokens, resulting in 85 padded tokens.
/alphafold3_venv/lib/python3.12/site-packages/numpy/_core/fromnumeric.py:83: RuntimeWarning: overflow encountered in reduce
  return ufunc.reduce(obj, axis, dtype, out, **passkwargs)
Featurising data with seed 4 took 91.11 seconds.

There are additional warnings about the tokamax autotuning cache not being found:

I0115 12:03:03.277778 22530949079808 cache.py:95] Autotuning cache file not found: /alphafold3_venv/lib/python3.12/site-packages/tokamax/data/autotuning/nvidia_rtx_pro_6000_blackwell_server_edition/pallas_triton_gated_linear_unit.json
I0115 12:03:03.447485 22530949079808 cache.py:95] Autotuning cache file not found: /alphafold3_venv/lib/python3.12/site-packages/tokamax/data/autotuning/nvidia_rtx_pro_6000_blackwell_server_edition/pallas_triton_flash_attention.json
/alphafold3_venv/lib/python3.12/site-packages/alphafold3/model/network/featurization.py:104: UserWarning: Explicitly requested dtype int64 requested in broadcasted_iota is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/jax-ml/jax#current-gotchas for more.
  iota = jax.lax.broadcasted_iota(jnp.int64, logits.shape, axis)

The same tokamak warnings also show up with unit tests (full logs for b78e215 and ebfe70a).

[Augustin-Zidek]

Thanks for the quick reply, that is useful to know.

There are additional warnings about the tokamax autotuning cache not being found:

Yeah, this is a known issue and should be fixed soon on Tokamax side. Once it is, I will update the Tokamax dependency.