Loading from checkpoint seems to hang indefinitely for Roberta
Author: elie-h
Pull request: no

## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->

- `transformers` version: 4.5.0
- Platform: Linux-5.8.0-48-generic-x86_64-with-glibc2.29
- Python version: 3.8.7
- PyTorch version (GPU?): 1.8.1+cu111 (True)
- Tensorflow version (GPU?): 2.4.1 (False)
- Using GPU in script?: Yes - RTX 3090
- Using distributed or parallel set-up in script?: No

Models:
- albert, bert, xlm: @LysandreJik
- Trainer: @sgugger 

## Information

Model I am using (Bert, XLNet ...): Roberta

The problem arises when using:
* [x] my own modified scripts: (give details below)

The tasks I am working on is:
* [ ] my own task or dataset


## To reproduce

I'm trying to resume training my Roberta model from a checkpoint. When the training initialises it seems to pick up the last checkpoint: 

```
  Continuing training from checkpoint, will skip to saved global_step
  Continuing training from epoch 0
  Continuing training from global step 660000
  Will skip the first 0 epochs then the first 660000 batches in the first epoch.
```

After that it just hangs, training does not start, no further logging and GPU utilisation is 0. I've left it for over 6 hours and still no progress.

I've tried both loading directly from a checkpoint and initialising the trainer with checkpoint=True: trainer.train("ml/models/araberto/checkpoint-660000") and trainer.train(checkpoint=True)

Code below:
```
from datasets import load_dataset
from datasets import ClassLabel, Value, Sequence

tokenizer = RobertaTokenizerFast.from_pretrained(output_path)
dataset = load_dataset('text', 
                            data_files={
                                "train":[str(x) for x in Path(f"{dataset_path}/train").glob("*.txt")], 
                                "test": str(Path(f"{dataset_path}/test.txt"))
                            })

def encode(batch):
    tokenized = tokenizer(batch.get("text", ""), padding="max_length", truncation=True, max_length=max_length)
    return tokenized

dataset.set_transform(encode)

import torch
if not torch.cuda.is_available():
    raise Exception("GPU not available")


from transformers import RobertaForMaskedLM
from transformers import RobertaConfig
from transformers import DataCollatorForLanguageModeling

config = RobertaConfig(
    vocab_size=vocab_size,
    max_position_embeddings=514,
    num_attention_heads=12,
    num_hidden_layers=6,
    type_vocab_size=1,
)

model = RobertaForMaskedLM(config=config)
model.num_parameters()

data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer, mlm=True, mlm_probability=0.15
)

from transformers import Trainer, TrainingArguments

training_args = TrainingArguments(
    output_dir=output_path,
    overwrite_output_dir=True,
    num_train_epochs=1,
    per_device_train_batch_size=48,
    save_steps=10_000,
    save_total_limit=2,
    remove_unused_columns=False,
    fp16=True,
    fp16_backend="amp"
)

trainer = Trainer(
    model=model,
    args=training_args,
    data_collator=data_collator,
    train_dataset=dataset["train"],
    eval_dataset=dataset["test"] 
)

%%time
torch.cuda.is_available()
torch.cuda.empty_cache()
trainer.train("ml/models/roberto/checkpoint-660000")

```

Debug logs:
```
Loading model from ml/models/roberto/checkpoint-660000).
loading configuration file ml/models/roberto/checkpoint-660000/config.json
Model config RobertaConfig {
  "architectures": [
    "RobertaForMaskedLM"
  ],
  "attention_probs_dropout_prob": 0.1,
  "bos_token_id": 0,
  "eos_token_id": 2,
  "gradient_checkpointing": false,
  "hidden_act": "gelu",
  "hidden_dropout_prob": 0.1,
  "hidden_size": 768,
  "initializer_range": 0.02,
  "intermediate_size": 3072,
  "layer_norm_eps": 1e-12,
  "max_position_embeddings": 514,
  "model_type": "roberta",
  "num_attention_heads": 12,
  "num_hidden_layers": 6,
  "pad_token_id": 1,
  "position_embedding_type": "absolute",
  "transformers_version": "4.5.0",
  "type_vocab_size": 1,
  "use_cache": true,
  "vocab_size": 51000
}

loading weights file ml/models/roberto/checkpoint-660000/pytorch_model.bin
All model checkpoint weights were used when initializing RobertaForMaskedLM.

All the weights of RobertaForMaskedLM were initialized from the model checkpoint at ml/models/roberto/checkpoint-660000.
If your task is similar to the task the model of the checkpoint was trained on, you can already use RobertaForMaskedLM for predictions without further training.
***** Running training *****
  Num examples = 808405026
  Num Epochs = 1
  Instantaneous batch size per device = 48
  Total train batch size (w. parallel, distributed & accumulation) = 48
  Gradient Accumulation steps = 1
  Total optimization steps = 16841772
  Continuing training from checkpoint, will skip to saved global_step
  Continuing training from epoch 0
  Continuing training from global step 660000
  Will skip the first 0 epochs then the first 660000 batches in the first epoch.
```

Iteration over the dataset seems fine:
```
x = dataset["train"][808405025]
print(x)

{'input_ids': [0, 14527, 606, 606, 503, 616, 13117, 1319, 7537, 93, 2506, 7712, 4897, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], 'attention_mask': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]}
```
## Expected behavior

Training to resume from the checkpoint 

Comments:
The script is not hanging, it is skipping the first 660,000 batches since you are resuming training from there which takes a lot of time. If you don't mind continuing training with the same data, you can use the option `ignore_data_skip=True` in your training arguments.

@eh-93 Do you remember how much time it took to train that checkpoint? 
@sgugger How about we add a progress bar to make the trainer more user-friendly? 

@sgugger got it, thanks
@cronoik - around 5 days to get to that checkpoint

After 28 hours the training resumed

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/master/CONTRIBUTING.md) are likely to be ignored.

That's not a good solution. A better solution that should be implemented is to pass some boolean to the CustomDataset, telling it that we are now in a 'skip' mode, so that the CustomDataset could prevent from doing expensive and unneeded steps in the skipping phase, like for example tokenize words. Does HF work on such solution?