Title: trainer became extremely slow after reload dataset by `load_from_disk`
Author: conan1024hao (NONE)
State: open
Created: 2022-05-23T14:04:37Z
Labels: bug

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

Comments:
I tried to make the dataset much more smaller (100000 rows) , then the speed became `33.88it/s` from`8.62s/it`.  It's nearly 200 times... Do you have any idea? Thank you!

Similar issue: https://github.com/huggingface/transformers/issues/8818

I changed `RandomSampler` to  `SequentialSampler` in the `trainer.py`, but the speed didn't become faster.

I changed
```
tokenized_datasets = load_from_disk(
        "/pathto/dataset"
    )
```
to
```
tokenized_datasets = load_from_disk(
        "/pathto/dataset", keep_in_memory=True
    )
```
and obtained normal speed. It's seems that the problem is on the os's speed limit.

Hi ! Currently `save_to_disk` saves one big Arrow file, which causes some slow downs. This has been discussed in #3735 and we'll implement sharding pretty soon to solve this

For now you can try splitting and saving your dataset in several Arrow files. Then you can load them one by one and use `concatenate_datasets` to have your big dataset again and hopefully with a better speed

Any update on fixing this? The issue still seems to be present.