learning rate resets on resumption from checkpoint
Author: agneet42
Pull request: no

### System Info

- `transformers` version: 4.26.1
- Platform: Linux-3.10.0-1160.71.1.el7.x86_64-x86_64-with-glibc2.17
- Python version: 3.9.12
- Huggingface_hub version: 0.13.0
- PyTorch version (GPU?): 1.13.1+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  @stas00 @ArthurZucker 

### 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

``Steps to reproduce the behaviour : 

1. Training Script -- 

```
if __name__ == "__main__":

    model_path = '/checkpoint-38000'
    
    model = AutoModelForCausalLM.from_pretrained(model_path)
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    tokenizer.pad_token_id = tokenizer.eos_token_id

    train_path = '/train/*'
    train_data = glob(train_path)
    
    val_path = 'val/*'
    val_data = glob(val_path)

    dataset = load_dataset("json", data_files = {"train": train_data, "validation" : val_data})
    dataset = dataset.map(transform, batched=True, remove_columns = ["id" ,"tokens"])
    
    train_dataset = dataset["train"]
    val_dataset = dataset["validation"]
    
    parser = HfArgumentParser(TrainingArguments)
    parser.add_argument("--model_name_or_dir")
    
    training_args, args = parser.parse_args_into_dataclasses()
    transformers.logging.set_verbosity_debug()
    
    trainer = Trainer(
        model,
        training_args,
        train_dataset=train_dataset,
        eval_dataset=val_dataset,
        tokenizer=tokenizer,
        data_collator=DataCollatorForTokenClassification(tokenizer, padding='longest'),
        compute_metrics=None, 
        callbacks = [TensorBoardCallback()]      
    )
    if trainer.is_world_process_zero():
        print(dataset)
    
    trainer.pop_callback(MLflowCallback)
    
    if training_args.do_train:
        if trainer.is_world_process_zero():
            print("Training...")

        start = time.time()
        trainer.train(model_path=model_path)
        mlflow.log_metric(
            "time/epoch", (time.time() - start) / 60 / training_args.num_train_epochs
        )
```

2. Params --

```export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5

accelerate launch train_cerebras_checkpoint.py \
--output_dir output_dir \
--num_train_epochs 30 \
--do_train --per_device_train_batch_size 10 \
--fsdp "full_shard auto_wrap" \
--fsdp_transformer_layer_cls_to_wrap "GPT2Block" \
--logging_steps 1 \
--save_strategy "steps" \
--save_steps 2000 \
--fp16 \
--gradient_checkpointing true
```

3. trainer_state.json (only last few epochs shown)--

```
{
      "epoch": 13.47,
      "learning_rate": 5.270405303307255e-06,
      "loss": 2.4699,
      "step": 37990
    },
    {
      "epoch": 13.47,
      "learning_rate": 5.2691867124856815e-06,
      "loss": 2.4614,
      "step": 37991
    },
    {
      "epoch": 13.47,
      "learning_rate": 5.267968121664108e-06,
      "loss": 2.3527,
      "step": 37992
    },
    {
      "epoch": 13.47,
      "learning_rate": 5.266749530842534e-06,
      "loss": 2.42,
      "step": 37993
    },
    {
      "epoch": 13.47,
      "learning_rate": 5.2655309400209605e-06,
      "loss": 2.6322,
      "step": 37994
    },
    {
      "epoch": 13.47,
      "learning_rate": 5.264312349199386e-06,
      "loss": 2.566,
      "step": 37995
    },
    {
      "epoch": 13.47,
      "learning_rate": 5.263093758377812e-06,
      "loss": 2.5026,
      "step": 37996
    },
    {
      "epoch": 13.47,
      "learning_rate": 5.261875167556239e-06,
      "loss": 2.6096,
      "step": 37997
    },
    {
      "epoch": 13.47,
      "learning_rate": 5.260656576734664e-06,
      "loss": 2.7513,
      "step": 37998
    },
    {
      "epoch": 13.47,
      "learning_rate": 5.2594379859130905e-06,
      "loss": 2.5066,
      "step": 37999
    },
    {
      "epoch": 13.48,
      "learning_rate": 5.258219395091516e-06,
      "loss": 2.7268,
      "step": 38000
    }
```

4. Learning rate after resumption --
```
'loss': 2.4654, 'learning_rate': 2.754905437352246e-05, 'epoch': 13.48}                                                      
{'loss': 2.8525, 'learning_rate': 2.754905437352246e-05, 'epoch': 13.48}
{'loss': 2.9781, 'learning_rate': 2.7548463356973997e-05, 'epoch': 13.48}
{'loss': 2.8067, 'learning_rate': 2.7547872340425534e-05, 'epoch': 13.48}

```

### Expected behavior

1. Learning rate should resume from the last stored LR.
2.  The loss seems higher post resumption as compared to before (Maybe due to LR?)
3. I also notice that the time per iteration has almost double (19s to 38s), even though I am using the same number of GPUs

Comments:
When doing `model_path = '=/checkpoint-38000'` you are not resuming training from the checkpoint, you are starting a new fresh training with the model of this checkpoint.

To resume training from a checkpoint, you need to use the [`resume_from_checkpoint` argument](https://huggingface.co/docs/transformers/main_classes/trainer#transformers.Trainer.train.resume_from_checkpoint) of `Trainer.train`.

@sgugger not sure what I am missing : 

Code : 

```
model_path = 'output_dir/checkpoint-38000'
    
    model = AutoModelForCausalLM.from_pretrained(model_path)
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    tokenizer.pad_token_id = tokenizer.eos_token_id

    train_path = '/train/*'
    train_data = glob(train_path)
    
    val_path = '/val/*'
    val_data = glob(val_path)

    dataset = load_dataset("json", data_files = {"train": train_data, "validation" : val_data})
    dataset = dataset.map(transform, batched=True, remove_columns = ["id" ,"tokens"])
    
    train_dataset = dataset["train"]
    val_dataset = dataset["validation"]
    
    print('Training data length', len(train_dataset))
    print('Validation data length', len(val_dataset))
    
    parser = HfArgumentParser(TrainingArguments)
    parser.add_argument("--model_name_or_dir")
    
    training_args, args = parser.parse_args_into_dataclasses()
    transformers.logging.set_verbosity_debug()
    
    trainer = Trainer(
        args=training_args,
        model=model,
        tokenizer = tokenizer,
        train_dataset=train_dataset,
        eval_dataset=val_dataset,
        data_collator=DataCollatorForTokenClassification(tokenizer, padding='longest'),
        compute_metrics=None, 
        callbacks = [TensorBoardCallback()]      
    )
    if trainer.is_world_process_zero():
        print(dataset)
    
    trainer.pop_callback(MLflowCallback)
    
    if training_args.do_train:
        if trainer.is_world_process_zero():
            print("Training...")

        start = time.time()
        trainer.train(resume_from_checkpoint=True)
        mlflow.log_metric(
            "time/epoch", (time.time() - start) / 60 / training_args.num_train_epochs
        )
```

Script : 
`
accelerate launch train_cerebras_checkpoint.py \
--resume_from_checkpoint True \
--output_dir /output_dir \
--num_train_epochs 30 \
--do_train --per_device_train_batch_size 10 \
--fsdp "full_shard auto_wrap" \
--fsdp_transformer_layer_cls_to_wrap "GPT2Block" \
--logging_steps 1 \
--save_strategy "steps" \
--save_steps 2000 \
--fp16 \
--gradient_checkpointing true
`
In the logs, I see the following --

`Continuing training from checkpoint, will skip to saved global_step`

However the learning rate still resets. I expect it to be around `e-06` but it is at `e-05`
I can confirm that `output_dir` contains `checkpoint-38000`

I also debugged and noticed that the execution goes through here - https://github.com/huggingface/transformers/blob/v4.26.1/src/transformers/trainer.py#L2333

Furthermore, I checked the last_lr from my optimizer and it seems to be as I see in the training_state.json : 
```
path = '/output_dir/checkpoint-38000/scheduler.pt'
x = torch.load(path)
x['_last_lr']
[5.258219395091516e-06, 5.258219395091516e-06]
```

Not able to understand why therefore the LR starts from `e-05` when I resume from checkpoint.
@sgugger

@agneet42 
Is there a typo here:
```
model_path = 'output_dir/checkpoint-38000'
```
Should be below or are you running under root path?
```
model_path = '/output_dir/checkpoint-38000'
```

@sgugger Hi, I'm using `Trainer.train(last_checkpoint_path)` but still got lr reset. Here is some info might help.

- checkpoint folder files
```text
- config.json
- generation_config.json
- optimizer.pt
- pytorch_model.bin
- rng_state.pth
- scaler.pt
- trainer_state.json
- training_args.bin
```

- Code: Load from checkpoint
```python

args.from_checkpoint = "./checkpoint-30000"


# Define the training arguments
training_args = TrainingArguments(
    output_dir=args.save_dir,
    overwrite_output_dir=True,
    num_train_epochs=2,
    per_device_train_batch_size=12,
    per_device_eval_batch_size=12,
    gradient_accumulation_steps=3,
    evaluation_strategy='steps',
    eval_steps=40000,
    save_steps=10000,
    logging_steps=100,
    learning_rate=5e-5,
    warmup_steps=1000,
    fp16=True,
    logging_dir='./logs'
)

# Create a trainer instance
trainer = Trainer(
    model=model,
    args=training_args, 
    train_dataset=train_dataset, 
    eval_dataset=eval_dataset, 
    data_collator=lambda data: {'input_ids': torch.stack([f[0] for f in data]),
                                'attention_mask': torch.stack([f[1] for f in data]),
                                'labels': torch.stack([f[0] for f in data])})

# Fine-tune the model
trainer.train(args.from_checkpoint)
```

Log: Previous training logs
```text
{'loss': 0.1625, 'learning_rate': 3.67605680426363e-06, 'epoch': 0.93}                                                   
{'loss': 0.1627, 'learning_rate': 3.3450296269323714e-06, 'epoch': 0.94}                                                 
{'loss': 0.1618, 'learning_rate': 3.0140024496011123e-06, 'epoch': 0.94}                                                 
{'loss': 0.1617, 'learning_rate': 2.6829752722698537e-06, 'epoch': 0.95}                                                 
{'loss': 0.1623, 'learning_rate': 2.3519480949385943e-06, 'epoch': 0.95}                                                 
{'loss': 0.1608, 'learning_rate': 2.0209209176073357e-06, 'epoch': 0.96}
```

Log: After load from checkpoint logs
```text
{'loss': 0.1715, 'learning_rate': 2.631964570647042e-05, 'epoch': 0.96}                                         
{'loss': 0.1757, 'learning_rate': 2.6238236347650525e-05, 'epoch': 0.97}                                        
{'loss': 0.1758, 'learning_rate': 2.6156826988830634e-05, 'epoch': 0.97}                                        
{'loss': 0.1766, 'learning_rate': 2.6075417630010747e-05, 'epoch': 0.97}                                        
{'loss': 0.1754, 'learning_rate': 2.5994008271190857e-05, 'epoch': 0.98}                                        
{'loss': 0.1789, 'learning_rate': 2.5912598912370966e-05, 'epoch': 0.98}                                        
{'loss': 0.1797, 'learning_rate': 2.583118955355108e-05, 'epoch': 0.98}                                         
{'loss': 0.18, 'learning_rate': 2.5750594288319384e-05, 'epoch': 0.99}
```



@wmhcqw I still don't have any code I can reproduce on this issue. To be able to reproduce the code needs to include the dataset/model creation as whole as the creation of the checkpoint from which you are then resuming training.

Resuming training is tested in our CI and there is no issue of learning rate resetting there, so the examples of this situation we have on our side work. To debug what is particular to the bug you are encountering, I need to be able to reproduce it.

@sgugger I think I've found the problem. Here's the sample code to reproduce this issue.

File: train.py
```python
import os
import argparse

from tqdm import tqdm

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset
from transformers import GPTNeoConfig, GPTNeoForCausalLM, GPT2Tokenizer, TrainingArguments, Trainer

import random
from random import randint


class DummyDataset(Dataset):
    
    def __init__(self, words, tokenizer):
        self.tokenizer = tokenizer
        self.data = {
            "input_ids": [],
            "attention_mask": []
        }
        for word in tqdm(words):
            res = tokenizer(word)
            self.data["input_ids"].append(torch.LongTensor(res["input_ids"]))
            self.data["attention_mask"].append(torch.LongTensor(res["attention_mask"]))
        
    def __len__(self):
        return len(self.data['input_ids'])
    
    def __getitem__(self, idx):
        input_ids = self.data['input_ids'][idx]
        attention_mask = self.data['attention_mask'][idx]
        return input_ids, attention_mask
    

if __name__ == "__main__":
    tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
    # print(tokenizer("Hello World"))
    
    names=["We","I","They","He","She","Jack","Jim","Rose","You"]
    verbs=["was", "is", "are", "were"]
    nouns=["playing a game", "watching television", "talking", "dancing", "speaking", "playing basketball", "eating dinner"]

    random.seed(42)
    train_sens = []
    for i in range(1000):
        train_sens.append(names[randint(0,len(names)-1)]+" "+verbs[randint(0,len(verbs)-1)]+" "+nouns[randint(0,len(nouns)-1)])
    
    eval_sens = []
    for i in range(100):
        eval_sens.append(names[randint(0,len(names)-1)]+" "+verbs[randint(0,len(verbs)-1)]+" "+nouns[randint(0,len(nouns)-1)])
    
    train_dataset = DummyDataset(train_sens, tokenizer)
    eval_dataset = DummyDataset(eval_sens, tokenizer)
    
    config = GPTNeoConfig(
        vocab_size=len(tokenizer.get_vocab()),
        n_positions=1024,
        n_ctx=2048,
        n_embd=768,
        n_layer=1,
        n_head=1,
        intermediate_size=3072
    )
    model = GPTNeoForCausalLM(config).cuda()
    
    training_args = TrainingArguments(
        output_dir="./dummpy_model",
        overwrite_output_dir=True,
        num_train_epochs=1,  # 2
        per_device_train_batch_size=1,
        per_device_eval_batch_size=1,
        gradient_accumulation_steps=1,
        evaluation_strategy='steps',
        eval_steps=1000,
        save_steps=100,
        logging_steps=10,
        learning_rate=5e-5,
        warmup_steps=10,
        fp16=True,
        logging_dir='./logs'
    )
    
    trainer = Trainer(
    model=model,
    args=training_args, 
    train_dataset=train_dataset, 
    eval_dataset=eval_dataset, 
    data_collator=lambda data: {'input_ids': torch.stack([f[0] for f in data]),
                                'attention_mask': torch.stack([f[1] for f in data]),
                                'labels': torch.stack([f[0] for f in data])})

    trainer.train()
    # trainer.train("./dummpy_model/checkpoint-1000")
    
```

Steps to reproduce:

1. run `python train.py`, and you will get a checkpoint folder named 'checkpoint-1000'
2. change the `train.py` code
    - change the training_args num_train_epcohs from 1->2 (**This is the reason, the total steps changed.**)
    - comment trainer.train()
    - uncomment trainer.train("./dummpy_model/checkpoint-1000")
3. run `python train.py` again, training resume from step 1000 but with reset lr.

Logs:

Step1. `python train.py`
```text
{'loss': 1.6059, 'learning_rate': 2.7878787878787883e-05, 'epoch': 0.45}                                        
{'loss': 1.6617, 'learning_rate': 2.7373737373737374e-05, 'epoch': 0.46}                                        
{'loss': 1.4485, 'learning_rate': 2.686868686868687e-05, 'epoch': 0.47}                                         
{'loss': 1.5028, 'learning_rate': 2.636363636363636e-05, 'epoch': 0.48}                                         
{'loss': 1.5889, 'learning_rate': 2.585858585858586e-05, 'epoch': 0.49}                                         
{'loss': 1.3763, 'learning_rate': 2.5353535353535356e-05, 'epoch': 0.5}                                         
{'loss': 1.5049, 'learning_rate': 2.4848484848484847e-05, 'epoch': 0.51} 
...                                
{'loss': 1.2957, 'learning_rate': 6.060606060606061e-07, 'epoch': 0.99}                                         
{'loss': 1.406, 'learning_rate': 1.0101010101010101e-07, 'epoch': 1.0}                                          
{'eval_loss': 1.3103933334350586, 'eval_runtime': 5.7138, 'eval_samples_per_second': 17.502, 'eval_steps_per_second': 17.502, 'epoch': 1.0}                                                                                     
{'train_runtime': 497.58, 'train_samples_per_second': 2.01, 'train_steps_per_second': 2.01, 'train_loss': 1.8285735349655152, 'epoch': 1.0} 
```

Step3. `python train.py`
```text
{'loss': 1.5102, 'learning_rate': 2.492462311557789e-05, 'epoch': 1.01}                                         
{'loss': 1.3294, 'learning_rate': 2.4673366834170854e-05, 'epoch': 1.02}                                        
{'loss': 1.511, 'learning_rate': 2.442211055276382e-05, 'epoch': 1.03}                                          
{'loss': 1.4702, 'learning_rate': 2.4170854271356786e-05, 'epoch': 1.04}                                        
{'loss': 1.5096, 'learning_rate': 2.391959798994975e-05, 'epoch': 1.05}                                         
{'loss': 1.4866, 'learning_rate': 2.3668341708542715e-05, 'epoch': 1.06}                                        
{'loss': 1.4644, 'learning_rate': 2.3417085427135678e-05, 'epoch': 1.07}                                        
{'loss': 1.2187, 'learning_rate': 2.3165829145728644e-05, 'epoch': 1.08}                                        
{'loss': 1.5627, 'learning_rate': 2.291457286432161e-05, 'epoch': 1.09}                                         
{'loss': 1.5059, 'learning_rate': 2.2663316582914573e-05, 'epoch': 1.1} 
```

As you can see, the learning rate is reset to where epoch 0.5, so I think the learning rate is not saved but the steps, and learning rate is calculated according to the total steps.

This is not a bug. I thought the learning rate is saved but I was wrong. Changing training arguments before loading from checkpoint is not an expected behaviour, I think.


> 2. change the train.py code change the training_args num_train_epcohs from 1->2 (This is the reason, the total steps changed.)

You cannot change a single a training argument when resuming training and expect the training to resume properly. This is the source of the bug, not the `Trainer` itself.

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.

@sgugger I am facing related problem, though not pertaining to learning rate but relates to resuming the training with change in GPUs available. I have written the detailed issue [here (huggingface forums)](https://discuss.huggingface.co/t/skipped-batches-do-not-consider-distributed-training/43832). Would really appreciate if you can help! 

Thanks and regards!

There used to be a bug in huggingface that hf loses control of resuming lr_scheduler when using deepspeed. The newest version have fixed it. Ref: https://github.com/huggingface/transformers/issues/24656#issuecomment-1733069714


> Resuming training is tested in our CI and there is no issue of learning rate resetting there, so the examples of this situation we have on our side work. To debug what is particular to the bug you are encountering, I need to be able to reproduce it.

@sgugger Have you considered deepspeed when designing the CI?