Persistent OOM Errors

[cooper-svajda]

Hi there,

Been trying to get MF setup as a pre-screen for later AF runs on single sequence, de novo protein designs. I initially tried to get it running on an AWS instance where I had AF and other tools running, a g4dn.xlarge that had 16GB of RAM and VRAM on a T4 GPU.

This was giving me OOM errors on the CPU side so I decided to size up to a g6.2xlarge with 32 GB of RAM and 24 GB VRAM on a L4 TPU. But even after this size up I'm seeing OOM issues with an apparent RAM consumption of 30+ GB. This was with the 12L model.

Is there a known issue on memory overflow that I can alleviate somehow?

Thanks!

[jwohlwend]

Specifically CPU memory?

[ikalvet]

Can confirm that this model requires a lot of memory. Regardless of whether it's on a GPU or CPU-only node, it requires >32GB of system (CPU) memory to run, otherwise it throws OOM errors. Is this by design, or could something be wrong with our setup?
Testing with 24g memory on a CPU, the code threw OOM at this point in create_model():

    # Initialize alphabet
    _, alphabet = load_model_and_alphabet(hparams["esm_model_name"])

Otherwise, a wonderful model, and if it gets enough memory it really is super fast! And seems to work well for pre-screening sequences for a more expensive AF3 run.