trainer became extremely slow after reload dataset by `load_from_disk`
## Describe the bug
Due to memory problem, I need to save my tokenized datasets locally by CPU and reload it by multi GPU for running training script. However, after I reload it by `load_from_disk` and start training, the speed is extremely slow. It says I need about 1500 hours with 8 A100 cards. Before this, I can run the whole script in one day with a single A100 card.
Since I am try to pre-train a BERT, **my dataset is very large(29058165 rows)**

## Steps to reproduce the bug
```python
tokenized_datasets.save_to_disk(
        "/pathto/dataset"
    )
tokenized_datasets = load_from_disk(
        "/pathto/dataset"
    )
trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=tokenized_datasets["train"] if training_args.do_train else None,
        eval_dataset=tokenized_datasets["validation"]
        if training_args.do_eval
        else None,
        tokenizer=tokenizer,
        data_collator=data_collator,
    )
 train_result = trainer.train(resume_from_checkpoint=checkpoint)
```

## Expected results
Without the save and reload process, I only need about one day to run the whole script with one A100 card.

## Actual results
```
[INFO|trainer.py:1290] 2022-05-23 22:49:46,266 >> ***** Running training *****
[INFO|trainer.py:1291] 2022-05-23 22:49:46,266 >>   Num examples = 29058165
[INFO|trainer.py:1292] 2022-05-23 22:49:46,266 >>   Num Epochs = 5
[INFO|trainer.py:1293] 2022-05-23 22:49:46,266 >>   Instantaneous batch size per device = 16
[INFO|trainer.py:1294] 2022-05-23 22:49:46,266 >>   Total train batch size (w. parallel, distributed & accumulation) = 256
[INFO|trainer.py:1295] 2022-05-23 22:49:46,266 >>   Gradient Accumulation steps = 2
[INFO|trainer.py:1296] 2022-05-23 22:49:46,266 >>   Total optimization steps = 567540
  0%|          | 1/567540 [00:09<1544:49:04,  9.80s/it]
  0%|          | 2/567540 [00:17<1320:00:17,  8.37s/it]
  0%|          | 3/567540 [00:26<1393:10:17,  8.84s/it]
  0%|          | 4/567540 [00:34<1344:56:33,  8.53s/it]
  0%|          | 5/567540 [00:43<1359:36:12,  8.62s/it]
```

## Environment info
```
torch              1.11.0+cu113
torchaudio         0.11.0+cu113
torchvision        0.12.0+cu113
transformers       4.18.0
datasets           2.2.2
```