Race condition when loading models from local folders with custom code
Author: dakinggg
Pull request: no
Labels: name: WIP

name: Good Difficult Issue

### System Info

- `transformers` version: 4.34.0
- Platform: Linux-5.4.0-148-generic-x86_64-with-glibc2.31
- Python version: 3.10.13
- Huggingface_hub version: 0.17.3
- Safetensors version: 0.4.0
- Accelerate version: 0.20.3
- Accelerate config: 	not found
- PyTorch version (GPU?): 2.1.0+cu121 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed

### Who can help?

not sure who to tag, so starting with @arthurzucker

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

This is pretty difficult to reproduce, so I don't have an easy script or anything like that, but hopefully the symptom is enough to fix it. It seems that when I load a model with `transformers.AutoModel.from_pretrained(<my local folder>)`, and that model contains custom code, the custom code is copied into the transformers cache. That is fine, except there seems to be a possible race condition when doing this from multiple processes at the same time. So if I call `transformers.AutoModel.from_pretrained(<my local folder>)` from multiple processes at the same time, very occasionally, I end up with a module not found error for a module that very clearly exists. My guess is that the separate processes are overwriting each other and there is a brief moment where some file doesn't exist, while one of the processes is attempting to load it.

I worked around this by just creating the model first on rank 0 to prepopulate the transformers module cache, but should this operation be implemented in a safer way in transformers?

Thanks!

### Expected behavior

`from_pretrained` does not need to be run first on rank 0 in order to be safe.

Comments:
Hi @dakinggg, thanks for raising this issue! 

Could you give a bit more information about how you're running the code from multiple processes? 

Note: there was a recent update in the recent release `4.35` where we use safetensors by default when saving/loading models which might also affect this (although I doubt it). 

It is simply multiple parallel processes (one per GPU) each running the `from_pretrained` call at the same time to load the model in for that process. The processes are set up with the normal torch distributed environment variables.What other info would be helpful? I also would not expect safetensors to affect this since it's related to the code modules not the weights.

Hi @dakinggg, apologies for the delay on my side. I've been away the past two weeks. 

OK, so race condition is exactly what seems to be happening -> each process is trying to download the weights to the same location. Ultimately, you only want to try to download once. The easiest fix is making sure that the checkpoint is downloaded first before running the script. This means all the models will just load the cache in their process. 

Could you share a minimal code snippet which can reproduce the error (even if it's sporadic?) as well as the env variable values, how multiprocessing is being set up and any CLI arguments being used to launch the script. 

Hi Amy, sorry for the delay, haven't had time to write and test a minimal script, but will try to soon.

If you want to investigate in the meantime, the pseudocode for that script is
```
# outside of the main script
# this can be any model with custom code, we just need a local folder to load from
model = transformers.AutoModelForCausalLM.from_pretrained('mosaicml/mpt-7b', trust_remote_code=True)
model.save_pretrained('my-local-folder')
```

main script:
```
def main():
    config = transformers.AutoConfig.from_pretrained('my-local-folder', trust_remote_code=True)
    if local_rank() == 0:
        model = AutoModelForCausalLM.from_pretrained('my-local-folder', trust_remote_code=True, config=config)
    else:
        with init_empty_weights(include_buffers=False):
            model = AutoModelForCausalLM.from_config(config, trust_remote_code=True)
```

then run the script with whatever distributed launcher you like (accelerate, torchrun, etc). Likely requires a bunch of times to hit the issue, and I'm not sure whether the issue is hardware specific or not.

not stale

I can confirm this problem still exists when I load the model from a local folder (instead of downloading from internet). It's not very easy to reproduce because I get `AttributeError: module '......' has no attribute '....'` error occasionally.

@randxie OK - so the model is completely downloaded and the issue is happening when many models are trying to read from the same weights file at the same time? 


This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

not stale

In case it helps someone, it seems that this issue became more prevalent with the latest transformers release (4.42.x). I assume this has to do with changes to caching structure. Previously we only encountered this when loading an hf model from a local folder. Now we hit it also when loading from the hub with trust_remote_code=True (i.e. code from the repo, not from transformers).

Hey @Rocketknight1, this seems like an issue that has been going on for a while; do you mind taking a look at what we could do to fix it?

Hmmn, I'm not sure if this would be an easy fix at the `transformers` end - we could add some kind of lock for model downloading, but I think it could cause headaches in the (vastly more common) single-process case, especially if the lock becomes stale etc., so I'm hesitant to do that.

@dakinggg have you tried either adding a barrier so that one process fully loads the model before the others start, or just cloning the repo you want to a local dir before loading it?

Yes, the workaround is here (https://github.com/mosaicml/llm-foundry/blob/502eb12ff40c69b8a7d693ace8120057afd34338/llmfoundry/models/hf/hf_causal_lm.py#L294-L316). IMO `from_pretrained`/`from_config` should be made safe to run on multiple processes, but at a minimum, it would be good to put this somewhere prominently in the docs that they are not.

FYI, this has now impacted pretty much anywhere you load anything from huggingface. Datasets, models, configs, tokenizers. I assume as a result of all the hf hub caching changes.

Hi there, sorry for the long delay. We've finally been able to track down and fix the race-condition issue! You can find the fix in the last patch release `huggingface_hub==0.24.7`. More details in https://github.com/huggingface/huggingface_hub/releases/tag/v0.24.7

Thanks everyone for the patience, kudos to @regisss and @Wauplin !