Memory usage during model load

[Turakar]

Thank you for publishing this model!

I noticed that loading esm1b_t33_650M_UR50S() has a high intermediate RAM usage, while only some of that is used after loading. This high intermediate RAM usage may lead to problems when loading the model on Google Colab, if you already have some other things sitting in RAM. I wonder, if there is some room for optimizing this loading process and avoiding the high intermediate RAM usage?

[tomsercu]

Hmm I see, the ESM-1b model checkpoint is 7.3GB which will have to fit in RAM.
On top of my head that would be due to 650M params * 4 bytes per float * 3. The 3 is because adam optimizer state holds on to first moment and second moment of the gradients. So if you'd strip the optimizer state from the checkpoint ahead of time you could cut checkpoint size and hence RAM usage with a factor 3.
The rest is all default pytorch checkpoint loading, no easy savings there.

[zhoubay]

Hi there, how could I strip the optimizer state, since the first moment and second moment is the nature feature of Adam optimizer? Change an optimizer? Thank you ahead of time!

[Turakar]

Workaround:
Step 1: Load the state dict on a machine with enough memory and delete irrelevant keys. In my case, the files were stored to ~/.cache/torch/hub/checkpoints/.

state_dict = torch.load("<path>.pt")
if "optimizer_history" in state_dict:
    del state_dict["optimizer_history"]
if "last_optimizer_state" in state_dict:
    del state_dict["last_optimizer_state"]
torch.save(state_dict, "<path>-inference.pt")

Step 2: Copy the regression weights file from <path>-contact-regression.pt to <path>-inference-contact-regression.pt.
Step 3: Move the files to the machine with low memory.
Step 4: Load it there using:

esm.pretrained.load_model_and_alphabet_local("<path>-inference.pt")

@tomsercu what do you think about adding this to the ESM library and providing the stripped-down versions via Facebook's Servers? I can contribute the code as an PR, or just copy it from here.