google-deepmind / google-deepmind/alphafold
bug: NAN in latest version of jax when using multimer+bfloat16
Open
- Dominant language
- Python
- Stars
- 14.9k
- Forks
- 2.9k
- PR merge metrics
- No merged PRs in 30d
Description
I think I tracked down the issue.
The issue comes from when the multimer model is run without templates. Zero features are appended to `msa_mask`. This creates a very large negative bias that bfloat16 can't seem to handle.
if you replace:
`bias = (1e9 * (msa_mask - 1.))[:, None, None, :]`
with:
`bias = (1e4 * (msa_mask - 1.))[:, None, None, :]`
No more NANS!
See:
https://github.com/deepmind/alphafold/blob/main/alphafold/model/modules.py#L764
Contributor guide
Assessment
This issue has not been assessed yet.