Querying examples from big datasets is slower than small datasets
My workaround is to shard the dataset into splits in my ssd disk and feed the data in different training sessions. But it is a bit of a pain when we need to reload the last training session with the rest of the split with the Trainer in transformers.

I mean, when I split the training and then reloads the model and optimizer, it not gets the correct global_status of the optimizer, so I need to hardcode some things. I'm planning to open an issue in transformers and think about it.
```
from datasets import load_dataset

book_corpus = load_dataset("bookcorpus", split="train[:25%]")
wikicorpus = load_dataset("wikicorpus", split="train[:25%]")
openwebtext = load_dataset("openwebtext", split="train[:25%]")

big_dataset = datasets.concatenate_datasets([wikicorpus, openwebtext, book_corpus])
big_dataset.shuffle(seed=42)
big_dataset = big_dataset.map(encode, batched=True, num_proc=20, load_from_cache_file=True, writer_batch_size=5000)
big_dataset.set_format(type='torch', columns=["text", "input_ids", "attention_mask", "token_type_ids"])


training_args = TrainingArguments(
    output_dir="./linear_bert",
    overwrite_output_dir=True,
    per_device_train_batch_size=71,
    save_steps=500,
    save_total_limit=10,
    logging_first_step=True,
    logging_steps=100,
    gradient_accumulation_steps=9,
    fp16=True,
    dataloader_num_workers=20,
    warmup_steps=24000,
    learning_rate=0.000545205002870214,
    adam_epsilon=1e-6,
    adam_beta2=0.98,
    weight_decay=0.01,
    max_steps=138974,  # the total number of steps after concatenating 100% datasets
    max_grad_norm=1.0,
)

trainer = Trainer(
    model=model,
    args=training_args,
    data_collator=data_collator,
    train_dataset=big_dataset,
    tokenizer=tokenizer))
```

I do one training pass with the total steps of this shard and I use len(bbig)/batchsize to stop the training (hardcoded in the trainer.py) when I pass over all the examples in this split.

Now Im working, I will edit the comment with a more elaborated answer when I left the work.
After some experiments with bookcorpus I noticed that querying examples from big datasets is slower than small datasets.
For example
```python
from datasets import load_dataset

b1 = load_dataset("bookcorpus", split="train[:1%]")
b50 = load_dataset("bookcorpus", split="train[:50%]")
b100 = load_dataset("bookcorpus", split="train[:100%]")

%timeit _ = b1[-1]                                                                     
# 12.2 µs ± 70.4 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

%timeit _ = b50[-1]                                                                    
# 92.5 µs ± 1.24 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

%timeit _ = b100[-1]                                                                      
# 177 µs ± 3.13 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

```

It looks like the time to fetch the example increases with the size of the dataset.

This is maybe due to the use of the Arrow streaming format to store the data on disk. I guess pyarrow needs to iterate through the file as a stream to find the queried sample.

Maybe switching to the Arrow IPC file format could help fixing this issue.

Indeed according to the [documentation](https://arrow.apache.org/docs/format/Columnar.html?highlight=arrow1#ipc-file-format), it's identical to the streaming format except that it contains the memory offsets of each sample, which could fix the issue:
> We define a “file format” supporting random access that is build with the stream format. The file starts and ends with a magic string ARROW1 (plus padding). What follows in the file is identical to the stream format. At the end of the file, we write a footer containing a redundant copy of the schema (which is a part of the streaming format) plus memory offsets and sizes for each of the data blocks in the file. This enables random access any record batch in the file. See File.fbs for the precise details of the file footer.

cc @gaceladri since it can help speed up your training when this one is fixed.
Querying examples from big datasets is slower than small datasets 
 After some experiments with bookcorpus I noticed that querying examples from big datasets is slower than small datasets.
For example
```python
from datasets import load_dataset

b1 = load_dataset("bookcorpus", split="train[:1%]")
b50 = load_dataset("bookcorpus", split="train[:50%]")
b100 = load_dataset("bookcorpus", split="train[:100%]")

%timeit _ = b1[-1]                                                                     
# 12.2 µs ± 70.4 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

%timeit _ = b50[-1]                                                                    
# 92.5 µs ± 1.24 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

%timeit _ = b100[-1]                                                                      
# 177 µs ± 3.13 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

```

It looks like the time to fetch the example increases with the size of the dataset.

This is maybe due to the use of the Arrow streaming format to store the data on disk. I guess pyarrow needs to iterate through the file as a stream to find the queried sample.

Maybe switching to the Arrow IPC file format could help fixing this issue.

Indeed according to the [documentation](https://arrow.apache.org/docs/format/Columnar.html?highlight=arrow1#ipc-file-format), it's identical to the streaming format except that it contains the memory offsets of each sample, which could fix the issue:
> We define a “file format” supporting random access that is build with the stream format. The file starts and ends with a magic string ARROW1 (plus padding). What follows in the file is identical to the stream format. At the end of the file, we write a footer containing a redundant copy of the schema (which is a part of the streaming format) plus memory offsets and sizes for each of the data blocks in the file. This enables random access any record batch in the file. See File.fbs for the precise details of the file footer.

cc @gaceladri since it can help speed up your training when this one is fixed. 
 My workaround is to shard the dataset into splits in my ssd disk and feed the data in different training sessions. But it is a bit of a pain when we need to reload the last training session with the rest of the split with the Trainer in transformers.

I mean, when I split the training and then reloads the model and optimizer, it not gets the correct global_status of the optimizer, so I need to hardcode some things. I'm planning to open an issue in transformers and think about it.
```
from datasets import load_dataset

book_corpus = load_dataset("bookcorpus", split="train[:25%]")
wikicorpus = load_dataset("wikicorpus", split="train[:25%]")
openwebtext = load_dataset("openwebtext", split="train[:25%]")

big_dataset = datasets.concatenate_datasets([wikicorpus, openwebtext, book_corpus])
big_dataset.shuffle(seed=42)
big_dataset = big_dataset.map(encode, batched=True, num_proc=20, load_from_cache_file=True, writer_batch_size=5000)
big_dataset.set_format(type='torch', columns=["text", "input_ids", "attention_mask", "token_type_ids"])


training_args = TrainingArguments(
    output_dir="./linear_bert",
    overwrite_output_dir=True,
    per_device_train_batch_size=71,
    save_steps=500,
    save_total_limit=10,
    logging_first_step=True,
    logging_steps=100,
    gradient_accumulation_steps=9,
    fp16=True,
    dataloader_num_workers=20,
    warmup_steps=24000,
    learning_rate=0.000545205002870214,
    adam_epsilon=1e-6,
    adam_beta2=0.98,
    weight_decay=0.01,
    max_steps=138974,  # the total number of steps after concatenating 100% datasets
    max_grad_norm=1.0,
)

trainer = Trainer(
    model=model,
    args=training_args,
    data_collator=data_collator,
    train_dataset=big_dataset,
    tokenizer=tokenizer))
```

I do one training pass with the total steps of this shard and I use len(bbig)/batchsize to stop the training (hardcoded in the trainer.py) when I pass over all the examples in this split.

Now Im working, I will edit the comment with a more elaborated answer when I left the work.