load_model_and_alphabet_local not working with DataDistributedParallel?

[jasperhyp]

Hi, I was trying to use load_model_and_alphabet_local in a DDP script, but it keeps failing. It works with one GPU, and load_model_and_alphabet_hub also worked. The error message is:

Traceback (most recent call last):
  File "/home/workspace/PLM/TxPLM/txplm/model/esm.py", line 139, in __init__
    self.model, _ = esm.pretrained.load_model_and_alphabet_local(pretrained_weights_dir+'/esm2_t36_3B_UR50D.pt')
  File "/home/.conda/envs/plm/lib/python3.8/site-packages/esm/pretrained.py", line 70, in load_model_and_alphabet_local
    model_data = torch.load(str(model_location), map_location="cpu")
  File "/home/.conda/envs/plm/lib/python3.8/site-packages/torch/serialization.py", line 789, in load
    return _load(opened_zipfile, map_location, pickle_module, **pickle_load_args)
  File "/home/.conda/envs/plm/lib/python3.8/site-packages/torch/serialization.py", line 1131, in _load
    result = unpickler.load()
  File "/home/.conda/envs/plm/lib/python3.8/site-packages/torch/serialization.py", line 1101, in persistent_load
    load_tensor(dtype, nbytes, key, _maybe_decode_ascii(location))
  File "/home/.conda/envs/plm/lib/python3.8/site-packages/torch/serialization.py", line 1079, in load_tensor
    storage = zip_file.get_storage_from_record(name, numel, torch.UntypedStorage).storage().untyped()
OSError: [Errno 14] Bad address

Never met this error before, so I'm not sure if this is a PyTorch issue -- perhaps map_location shouldn't be set to CPU under DDP?

[tomsercu]

Can you provide a minimum reproducible example of the script calling esm.pretrained.load_model_and_alphabet_local()?
cc @nikitos9000