Race condition while saving the checkpoint of the model
Author: upperwal
Pull request: no

### System Info

- `transformers` version: 4.36.0
- Platform: Linux-5.4.0-80-generic-x86_64-with-glibc2.31
- Python version: 3.11.5
- Huggingface_hub version: 0.19.4
- Safetensors version: 0.4.1
- Accelerate version: 0.25.0
- Accelerate config: 	not found
- PyTorch version (GPU?): 2.0.1 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: Distributed


### Who can help?

@muellerz  @pacman100 



### Error
There is a race condition at `Trainer > _maybe_log_save_evaluate` in a `multi_gpu` setup with `DDP`. Race condition happens because `self.control.should_save` allows all nodes to enter `self._save_checkpoint()`.

All nodes (even non-root nodes) enter `self._save_checkpoint()` and creates `staging_output_dir` if not available to save the random number generator state using `_save_rng_state`. Then eventually one of these nodes reaches [`if staging_output_dir != output_dir`](https://github.com/huggingface/transformers/blob/ec43d6870aa1afb42a6d2b1b0a03743d3f9b3ce6/src/transformers/trainer.py#L2385) and renames the `staging_output_dir` to `output_dir`. Any node including the root node trying to save the model state in `staging_output_dir` will fail as that directory has been moved.



https://github.com/huggingface/transformers/blob/ec43d6870aa1afb42a6d2b1b0a03743d3f9b3ce6/src/transformers/trainer.py#L2276C1-L2278

```
Traceback (most recent call last):
  File "...py", line 123, in <module>
    main()
  File "....py", line 119, in main
    trainer.train(resume_from_checkpoint=ckpt_dir)
  File ".../lib/python3.11/site-packages/transformers/trainer.py", line 1537, in train
    return inner_training_loop(
           ^^^^^^^^^^^^^^^^^^^^
  File ".../lib/python3.11/site-packages/transformers/trainer.py", line 1914, in _inner_training_loop
    self._maybe_log_save_evaluate(tr_loss, model, trial, epoch, ignore_keys_for_eval)
  File ".../lib/python3.11/site-packages/transformers/trainer.py", line 2274, in _maybe_log_save_evaluate
    self._save_checkpoint(model, trial, metrics=metrics)
  File ".../lib/python3.11/site-packages/transformers/trainer.py", line 2350, in _save_checkpoint
    self.save_model(staging_output_dir, _internal_call=True)
  File ".../lib/python3.11/site-packages/transformers/trainer.py", line 2837, in save_model
    self._save(output_dir)
  File ".../lib/python3.11/site-packages/transformers/trainer.py", line 2893, in _save
    safetensors.torch.save_file(state_dict, os.path.join(output_dir, SAFE_WEIGHTS_NAME))
  File ".../lib/python3.11/site-packages/safetensors/torch.py", line 281, in save_file
    serialize_file(_flatten(tensors), filename, metadata=metadata)
RuntimeError: Parent directory ./tmp-checkpoint-500 does not exist.
    super().__init__(torch._C.PyTorchFileWriter(str(name)))
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
```

### Fix
Only global root process should rename the staging directory.
```py
if self.args.process_index == 0 and staging_output_dir != output_dir:
    os.rename(staging_output_dir, output_dir)
```

This will still have race condition incase `save_on_each_node == True` as root process might rename the directory when some non-root node is saving the state in `staging_output_dir`. Should ideally sync all the processes and then rename but that will halt some processes. Better write to the `output_dir` directly.


### 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

1. Setup the `Trainer` in DDP with multi-node
2. Enable saving checkpoints at some intervals
3. Train

### Expected behavior

Should see the following file in the checkpoint directory without `RuntimeError: Parent directory ./tmp-checkpoint-500 does not exist` error

1. model.safetensors  
2. rng_state_0.pth  
3. rng_state_1.pth  
4. rng_state_XXX.pth  
5. scheduler.pt             
6. tokenizer.json         
7. trainer_state.json
8. optimizer.pt       
9. special_tokens_map.json  
10. tokenizer_config.json  
11. training_args.bin

Comments:
@upperwal see https://github.com/huggingface/transformers/pull/28009

@thundergolfer Oh, awesome! Will close this issue.