CPU runs produce distorted structures with excessively long bonds and colliding atoms

[ozvoz]

Running Boltz on the CPU for large elongated structures results in models with with excessively long bonds and colliding atoms (see example below PDB ID 6TE1) . Running the same structures on the CPU in single sequence mode (no MSA) produces the same defects. Structures predicted on the GPU are fine.

System:
Ubuntu 22.04.3 LTS
Intel Core i9-12900K × 24
32 GiB RAM
AMD Radeon RX 6900 XT

Version: Boltz v2.2.1
Command line:
boltz predict boltzinput --out_dir boltzoutput --cache aimodels/boltz --accelerator cpu --no_kernels --output_format pdb --diffusion_samples 1 --seed 0
Input YAML file:

sequences:
    - protein:
        id: A
        sequence: PSGVEGAAFQSRLPHDRMTSQEAACFPDIISGPQQTQKVFLFIRNRTLQLWLDNPKIQLTFEATLQQLEAPYNSDTVLVHRVHSYLERHGLINFGIYKRIKPLPTKKTGKVIIIGSGVSGLAAARQLQSFGMDVTLLEARDRVGGRVATFRKGNYVADLGAMVVTGLGGNPMAVVSKQVNMELAKIKQKCPLYEANGQAVPKEKDEMVEQEFNRLLEATSYLSHQLDFNVLNNKPVSLGQALEVVIQLQEKHVKDEQIEHWKKIVKTQEELKELLNKMVNLKEKIKELHQQYKEASEVKPPRDITAEFLVKSKHRDLTALCKEYDELAETQGKLEEKLQELEANPPSDVYLSSRDRQILDWHFANLEFANATPLSTLSLKHWDQDDDFEFTGSHLTVRNGYSCVPVALAEGLDIKLNTAVRQVRYTASGCEVIAVNTRSTSQTFIYKCDAVLCTLPLGVLKQQPPAVQFVPPLPEWKTSAVQRMGFGNLNKVVLCFDRVFWDPSVNLFGHVGSTTASRGELFLFWNLYKAPILLALVAGEAAGIMENISDDVIVGRCLAILKGIFGSSAVPQPKETVVSRWRADPWARGSYSYVAAGSSGNDYDLMAQPITPGPSIPGAPQPIPRLFFAGEHTIRNYPATVHGALLSGLREAGRIADQFLGAMYTL
        msa: /tmp/boltzdata/6TE1_1.csv
    - protein:
        id: B
        sequence: RKPPKGMFLSQEDVEAVSANATAATTVLRQLDMELVSVKRQIQNIKQTNSALKEKLDGGIEPYRLPEVIQKCNARWTTEEQLLAVQAIRKYGRDFQAISDVIGNKSVVQVKNFFVNYRRRFNIDEVLQEWEAE
        msa: /tmp/boltzdata/6TE1_2.csv

MSA files: 6TE1_1.csv, 6TE1_2.csv
Output PDB file: https://github.com/ozvoz/filedump/blob/main/6TE1_model_0.pdb

[volgin]

We've identified and fixed this in boltz-community v2.9.1.

Root cause: Boltz-2 falls through to bf16-mixed precision when --accelerator cpu is used. On modern CPUs with AVX-512 BF16 support (Intel 12th gen+, recent AMD Zen 4), PyTorch actually executes bfloat16 operations — unlike older CPUs that silently fall back to float32. This produces drastically reduced numerical precision, leading to the distorted bond lengths and atom clashes you're seeing.

Fix: One line — extend the float32 precision guard to cover cpu in addition to mps:

# before
if accelerator == "mps":
    precision = 32

# after  
if accelerator in ("mps", "cpu"):
    precision = 32

In the meantime, as a workaround on the upstream version you can force it via:

PYTORCH_ENABLE_MPS_FALLBACK=0 boltz predict ... --accelerator cpu

...though installing boltz-community is the cleaner fix until this is patched upstream.

[ozvoz]

Setting the precision from bf16-mixed to 32, as proposed by @volgin, fixes the problem of incorrectly predicted atom positions and reduces the computing time for the given example above from 12 hours to 45 minutes (!) while maintaining comparable peak RAM usage.

Applied fix in main.py:

-        precision=32 if model == "boltz1" else "bf16-mixed",
+        precision=32 if model == "boltz1" or accelerator=="cpu" else "bf16-mixed",