Learning Rate doesn't anneal properly after resume_from_checkpoint
Author: jmzeng
Pull request: no

### System Info

I'm resuming a run with HF trainer using resume_from_checkpoint and loading from the checkpoint. 

The LR scheduler is set as cosine annealing with warmup. However, when I resumed, the learning rate didn't resemble cosine annealing at all. I'm finetuning a llama2 model with custom code. 


See image below:
<img width="786" alt="Screenshot 2023-11-28 at 9 35 17 AM" src="https://github.com/huggingface/transformers/assets/5641698/8328ea10-c284-4ad0-bfc9-9a22e7da8c4a">

The previous run's LR is below, so I was expecting it will continue from there:
<img width="786" alt="Screenshot 2023-11-28 at 9 41 44 AM" src="https://github.com/huggingface/transformers/assets/5641698/e42b17a1-0869-4eca-b3a2-92843cecf3b1">

Transformers info:

- `transformers` version: 4.34.0
- Platform: Linux-5.15.0-84-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.22.0
- 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
- Using GPU in script?: yes
- Using distributed or parallel set-up in script?: yes, multinode training



### Who can help?

@muellerzr and @pacman100

### 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

The steps to reach the bug was below:

1. Train a multinode model using HF Trainer with LLama2 with cosine annealing and warmup:
    learning_rate: 2.0e-5
    adam_beta1: 0.9
    adam_beta2: 0.95
    weight_decay: 0.0
    lr_scheduler_type: cosine
    warmup_ratio: 0.03
2. Resume the multinode trainer with default usage of resume_from_checkpoint

### Expected behavior

It should resume from where it left off in cosine annealing with warmup

Comments:
Moreover, I'm seeing that the epoch variable gets reset to 1.0 at about 0.45 epochs. Would this be a possible bug?

<img width="786" alt="Screenshot 2023-11-28 at 10 52 48 AM" src="https://github.com/huggingface/transformers/assets/5641698/aab54221-f955-48ee-9642-3e5abadaaec6">

In the logs:
```{'loss': 0.1561, 'learning_rate': 1.996214146675922e-05, 'epoch': 0.47}
734 {'loss': 0.1319, 'learning_rate': 1.996214146675922e-05, 'epoch': 0.47}
735 {'loss': 0.2349, 'learning_rate': 1.996214146675922e-05, 'epoch': 0.47}
736 {'loss': 0.3174, 'learning_rate': 1.996214146675922e-05, 'epoch': 0.47}
737 {'loss': 0.1351, 'learning_rate': 1.996214146675922e-05, 'epoch': 0.47}
738 {'loss': 0.1666, 'learning_rate': 1.996214146675922e-05, 'epoch': 0.47}
739  10%|▉         | 1072/11275 [11:21:26<188:55:54, 66.66s/it]
740 {'loss': 0.2105, 'learning_rate': 1.996214146675922e-05, 'epoch': 1.0}
741 {'loss': 0.1975, 'learning_rate': 1.996214146675922e-05, 'epoch': 1.0}
742 {'loss': 0.1663, 'learning_rate': 1.996214146675922e-05, 'epoch': 1.0}
743 {'loss': 0.2243, 'learning_rate': 1.996214146675922e-05, 'epoch': 1.0}
744 {'loss': 0.117, 'learning_rate': 1.996214146675922e-05, 'epoch': 1.0}
745 {'loss': 0.25, 'learning_rate': 1.996214146675922e-05, 'epoch': 1.0}
746 {'loss': 0.1695, 'learning_rate': 1.996214146675922e-05, 'epoch': 1.0}
```

cc @muellerzr


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.

Gently pinging @muellerzr 🤗 

Gentle ping @muellerzr 

@jmzeng how are you resuming from your checkpoint? Are you passing it in through the `TrainingArguments` or in `trainer.train()`. It needs to be done at the `trainer.train()` level for it to actually apply like it needs to.

I tried using your scheduler configuration on a small-scale test and I couldn't reproduce it/they were identical. Let me know how you set things up and we can debug further!

My repro:

```python
import evaluate
import numpy as np
from datasets import load_dataset
from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer,
    DataCollatorWithPadding,
    Trainer,
    TrainingArguments,
)
model_name = "distilbert-base-uncased"
dataset_name = "imdb"
metric = "accuracy"
id2label = {0: "NEGATIVE", 1: "POSITIVE"}
label2id = {v: k for k, v in id2label.items()}
dataset = load_dataset(dataset_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)
def tokenize_function(examples):
    return tokenizer(examples["text"], truncation=True)
tokenized_dataset = dataset.map(tokenize_function, batched=True)
data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
accuracy = evaluate.load(metric)
def compute_metrics(evaluation_preds):
    predictions, labels = evaluation_preds
    predictions = np.argmax(predictions, axis=1)
    return accuracy.compute(predictions=predictions, references=labels)

model = AutoModelForSequenceClassification.from_pretrained(
    model_name, num_labels=2, id2label=id2label, label2id=label2id
)
training_args = TrainingArguments(
    output_dir="results/sequence_classification",  # Where weights are stored
    learning_rate=2e-5,  # The learning rate during training
    adam_beta1=.9,
    adam_beta2=.95,
    weight_decay=0.0,
    lr_scheduler_type="cosine",
    warmup_ratio=0.03,
    per_device_train_batch_size=16,  # Number of samples per batch during training
    max_steps=200,  # How many iterations through the dataloaders should be done
    save_strategy="steps",  # When to try and save the best model (such as a step number or every iteration)
    save_steps=20,
    logging_steps=10,
    report_to="wandb",  # Where to report metrics (such as to TensorBoard)
    run_name="sequence-classification-annealing-resume-from-120",
    fp16=True,
)
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset["train"],
    eval_dataset=tokenized_dataset["test"],
    data_collator=data_collator,
    compute_metrics=compute_metrics,
)
trainer.train(resume_from_checkpoint="results/sequence_classification/checkpoint-120",)
# trainer.train()
```

![image](https://github.com/huggingface/transformers/assets/7831895/ac38f284-f5d2-46c0-a6a3-cfa0a5b0abfa)


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.