Wanted to let you know that binder generation (complexa design) crashes early in the forward pass for me:
RuntimeError: CUDA driver error: invalid argument
community_models/openfold/data/data_transforms.py:980 atom37_to_torsion_angles
via nn/feature_factory/seq_feats.py:401 _get_sidechain_angles
Dug into it a bit: the target coords/mask come in as float64, and atom37_to_torsion_angles runs torch.prod on the float64 mask. Turns out FP64 torch.prod is broken on my box, so it falls over — even though the model runs in float32 anyway.
Quick repro:
import torch
torch.prod(torch.rand(4, 2, device="cuda", dtype=torch.float64), dim=-1) # boom
torch.prod(torch.rand(4, 2, device="cuda", dtype=torch.float32), dim=-1) # fine
Casting the target features to float32 before the transform fixes it and generation runs fine.
Is there any reason these are computed in float64? Seems like float32 would be a bit faster and avoid leaning on FP64 reduction kernels, which are crippled on consumer cards and clearly a bit flaky across torch/driver combos. Happy to be told this is just my setup.
Env: RTX 3090, driver 610.43.02 / CUDA 13.3, torch 2.7.0+cu126, repo UV env, single-pass PDL1 (02_PDL1), batch 2. Probably a torch-cu126 vs CUDA-13 driver thing on my end, but figured the float32 change might be worth it regardless.
Wanted to let you know that binder generation (complexa design) crashes early in the forward pass for me:
Dug into it a bit: the target coords/mask come in as float64, and atom37_to_torsion_angles runs torch.prod on the float64 mask. Turns out FP64 torch.prod is broken on my box, so it falls over — even though the model runs in float32 anyway.
Quick repro:
Casting the target features to float32 before the transform fixes it and generation runs fine.
Is there any reason these are computed in float64? Seems like float32 would be a bit faster and avoid leaning on FP64 reduction kernels, which are crippled on consumer cards and clearly a bit flaky across torch/driver combos. Happy to be told this is just my setup.
Env: RTX 3090, driver 610.43.02 / CUDA 13.3, torch 2.7.0+cu126, repo UV env, single-pass PDL1 (02_PDL1), batch 2. Probably a torch-cu126 vs CUDA-13 driver thing on my end, but figured the float32 change might be worth it regardless.