unified memory

[rmarabini]

Hi,

Is openfold able to use "unified memory"?

When using alphafold adding the following variables

export TF_FORCE_UNIFIED_MEMORY='1' export XLA_PYTHON_CLIENT_MEM_FRACTION='4.0' export TF_FORCE_GPU_ALLOW_GROWTH=true

increased the size of the sequence that can be processed (of course, the use of unified memory decreases speed) I have notice that this same approach does not seem to work for openfold.

In summary the question is: Can I force openfold to use unified memory? How?

[gahdritz]

OpenFold is based on PyTorch, not TensorFlow, and doesn't recognize those flags. PyTorch does not have CUDA Unified Memory support. However, we do have a bunch of optimizations that make inference on sequences of length 4700 possible on a single A100 (check out the "long_sequence_inference" config preset). Beyond that point, the accuracy of OF/AF isn't that good anyway.

[rmarabini]

Thanks for the info.