load_from_checkpoints does not work on fine-tuning Llama 13B
Author: raj-swype
Pull request: no

### System Info

- `transformers` version: 4.28.0.dev0
- Platform: Linux-5.4.0-137-generic-x86_64-with-glibc2.27
- Python version: 3.10.8
- Huggingface_hub version: 0.13.3
- Safetensors version: not installed
- PyTorch version (GPU?): 2.0.0+cu117 (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?: yes

### Who can help?

@sgugger 

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

I'm trying to finetune Llama 13B on a custom dataset.

Training code:

```python
 model = transformers.LlamaForCausalLM.from_pretrained(
        model_args.model_name_or_path,
        cache_dir=training_args.cache_dir,
        use_auth_token=HF_TOKEN
    )

    tokenizer = transformers.LlamaTokenizer.from_pretrained(
        model_args.model_name_or_path,
        cache_dir=training_args.cache_dir,
        model_max_length=training_args.model_max_length,
        padding_side="right",
        use_auth_token=HF_TOKEN,
        use_fast=False,
    )

training_args = TrainingArguments("test", evaluation_strategy="steps", eval_steps=100, do_eval=True)
trainer = Trainer(tokenizer=tokenizer, args=training_args, **data_module, compute_metrics=compute_metrics, model_init=model_init)
```

Script command:

```sh
python3 -m torch.distributed.run \
    --nproc_per_node=8 \
    --master_port=5102 /root/deepshard/training/train.py \
    --model_name_or_path swype/deepshard-13B-raw \
    --data_path /root/deepshard/datasets/data/train.jsonl \
    --eval_path /root/deepshard/datasets/data/test.jsonl \
    --search False \
    --fp16 True \
    --output_dir /root/deepshard/training/finetuned \
    --num_train_epochs 3 \
    --per_device_train_batch_size 2 \
    --per_device_eval_batch_size 2 \
    --gradient_accumulation_steps 4 \
    --evaluation_strategy "steps" \
    --eval_steps 1000 \
    --save_strategy "steps" \
    --save_steps 1000 \
    --learning_rate 2e-6 \
    --weight_decay 0. \
    --warmup_ratio 0.03 \
    --resume_from_checkpoint True \
    --lr_scheduler_type "cosine" \
    --logging_steps 10 \
    --fsdp "full_shard auto_wrap" \
    --fsdp_transformer_layer_cls_to_wrap 'LlamaDecoderLayer'
```

I also try setting `resume_from_checkpoint` to the checkpoint folder, but the script still starts from epoch 0 and starts training from scratch rather than from the checkpoint.


### Expected behavior

The `Trainer` should continue from the last saved checkpoint.

Comments:
I'm not too sure how to reconcile the command you typed with the code you wrote, as the code you are showing doesn't accept all the arguments you are passing to the launching command.

@raj-swype As a debugging step, do you experience the same issue if you remove the `compute_metrics` argument from the `Trainer(...)` initialization? It is causing my Trainer to hang.

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.

I'm confused what metrics are you trying to compute for a text generation model ?, i thought you only need to run `prediction_loss_only=True`, cause usual metrics doesn't really work with this type of models