Inference on Apple Silicon

chaidiscovery / chai-lab

Chai-1, SOTA model for biomolecular structure prediction

https://www.chaidiscovery.com

Other

1.05k stars 134 forks source link

Inference on Apple Silicon #33

Open DomML opened 2 weeks ago

DomML commented 2 weeks ago

Would it be possible to run chai on apple silicon devices ? I replaced cuda:0 with mps:0, however I still have the AssertionError: Torch not compiled with CUDA enabled error. I believe it come from the feature_embedding.pt2 file, but I'm not sure and Idk how to solve the issue. Best, Dominique MIAS.

popfido commented 3 days ago

It's because every module of the model, as you can see from the source code published, is a jit compiled version by Torch Dynamo . Since you have mps device that probably not support the compiled module which was compiled by the default backend inductor of torch jit, you will receive this AssertionError.

I think currently there is no way for mps device to run current compiled module.

Maybe they should release another version of compiled checkpoint with Triton.