Cannot find checkpoint during Trainer._load_best_model when using deepspeed
Author: nathan-az
Pull request: no

### System Info

```
- `transformers` version: 4.37.0
- Platform: Linux-6.2.0-1017-aws-x86_64-with-glibc2.35
- Python version: 3.10.12
- Huggingface_hub version: 0.20.3
- Safetensors version: 0.4.2
- Accelerate version: 0.23.0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.1.2+cu121 (False)
- 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?: yes
```
Note the above was run in a container on a different instance from the job compute, but with the same docker image.

### Who can help?

@pacman100 @muellerzr 

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

Came across this using the SFT script in the [alignment handbook](https://github.com/huggingface/alignment-handbook).

I can add more information but I think the relevant info is as follows:
In terms of trainer args:
```yaml
load_best_model_at_end: true
num_train_epochs: 40
output_dir: /local_disk0/hf/outputs
overwrite_output_dir: true
resume_from_checkpoint: false
save_on_each_node: true
save_only_model: true
save_steps: 1
save_strategy: "epoch"
save_total_limit: 5
```
Note that I am attempting to only save 5 models, but to keep track of the best model and load it at the end for saving. I also set `save_only_model` to `true` as I don't currently care to be able to actually load a checkpoint for continued training, and suspect this is the problem.

Note that the output directory `/local_disk0/hf/outputs` is a directory path that _exists_ on each node, but is _not_ a shared filesystem/NFS (so each node contains its information in that path).

My setup is distributed multi-gpu and multi-node via pdsh.
```yaml
compute_environment: LOCAL_MACHINE
deepspeed_config:
  deepspeed_multinode_launcher: pdsh
  deepspeed_hostfile: {TRAIN_DIR}/hostfile
  deepspeed_config_file: {CONFIG_FILE}
  zero3_init_flag: true
distributed_type: DEEPSPEED
```
I've tried to clean up the stacktrace, since I'm getting multiple (it appears to be one per rank)
```
  File "/databricks/python3/lib/python3.10/site-packages/trl/trainer/sft_trainer.py", line 315, in train
    main()
  File "/local_disk0/.ephemeral_nfs/training/alignment-handbook/scripts/run_sft.py", line 164, in main
            output = super().train(*args, **kwargs)
train_result = trainer.train(resume_from_checkpoint=checkpoint)output = super().train(*args, **kwargs)
  File "/databricks/python3/lib/python3.10/site-packages/transformers/trainer.py", line 1539, in train
    output = super().train(*args, **kwargs)
  File "/databricks/python3/lib/python3.10/site-packages/transformers/trainer.py", line 1539, in train
    return inner_training_loop(
  File "/databricks/python3/lib/python3.10/site-packages/transformers/trainer.py", line 1972, in _inner_training_loop
    self._load_best_model()
  File "/databricks/python3/lib/python3.10/site-packages/transformers/trainer.py", line 2167, in _load_best_model
    deepspeed_load_checkpoint(self.model_wrapped, self.state.best_model_checkpoint)
  File "/databricks/python3/lib/python3.10/site-packages/transformers/integrations/deepspeed.py", line 408, in deepspeed_load_checkpoint
    raise ValueError(f"Can't find a valid checkpoint at {checkpoint_path}")
ValueError: Can't find a valid checkpoint at /local_disk0/hf/outputs/checkpoint-29
```

### Expected behavior

I expect the model to be loaded at the end, simply so that it can be saved (i.e. the motivation is to save the model with the best eval metric for use during inference).

The error message indicates the `checkpoint` could not be found. I suspect that maybe it is expecting a full checkpoint including parameters, optimiser states, etc. so that training can continue, but am unsure. If this is the cause, this might be more of a feature request? Since it would be good to have a way to keep track of and save just the parameters of the iteration with the best eval metric.

Comments:
Closing this.

The docs are clear that the issue is [save_only_model](https://huggingface.co/docs/transformers/main_classes/trainer#transformers.TrainingArguments.save_only_model): `Note that when this is true, you won’t be able to resume training from checkpoint`.

It would be nice to have a version of the model checkpointing that saves the model files using a distributed equivalent of `save_pretrained` which retains the best model, without the training information, but either way this does not appear to currently be a bug.

> Closing this.
> 
> The docs are clear that the issue is [save_only_model](https://huggingface.co/docs/transformers/main_classes/trainer#transformers.TrainingArguments.save_only_model): `Note that when this is true, you won’t be able to resume training from checkpoint`.
> 
> It would be nice to have a version of the model checkpointing that saves the model files using a distributed equivalent of `save_pretrained` which retains the best model, without the training information, but either way this does not appear to currently be a bug.

I'm not using，[save_only_model](https://huggingface.co/docs/transformers/main_classes/trainer#transformers.TrainingArguments.save_only_model):， but still got same issue,  its safetensor format. it doesn't contain .pt