Improving training time for Marian MT model with the Trainer 
## Environment info
- Platform: Linux-4.19.0-14-cloud-amd64-x86_64-with-debian-10.7
- Python version: 3.7.9
- PyTorch version (GPU): 1.7.1
- Using GPU in script?: Yes (2 Tesla V100 GPUs with 16160MiB memory)
- CUDA Version: 11.0 
- transformers: 4.3.2
- datasets: 1.3.0

### Who can help
@lhoestq @sgugger @sshleifer 

## Information
The model I am using is a [MarianMTModel](https://huggingface.co/transformers/model_doc/marian.html#marianmtmodel) to train a machine translation model for Italian to Dutch. In order to so I perform the following steps:

- Split the data
1. The dataset originally consists of 2 parts, 1 text file containing Italian sentences and 1 text file containing the corresponding Dutch sentences. (42,940,499 sentences). These text files are combined in a pandas dataframe with a source and target column and is written to a csv file. Based on this [issue](https://github.com/huggingface/datasets/issues/610#issuecomment-691672919), which solved out-of-memory issues, I split the dataframe in 1000 csv chunks.  
2. The csv files are loaded with the load_dataset method (which results in a 7.5G csv-train.arrow file) from the datasets library as follows:
```python
train_files = glob.glob(data_folder + 'shards/data_chunk_train_*') # list of the individual csv files
train_dataset = load_dataset('csv', data_files=train_files, split='train', cache_dir=cache_folder, keep_in_memory=True)
``` 
I use the keep_in_memory=True here to hopefully make things faster during training. 

- Tokenization
1. At first, I to batch tokenized all the sentences with the map function. However, this resulted in 16 * 32G cache-files and gave a training time of 2000 hours. So I changed this to use the set_transform method in the latest release from datasets as follows: 
``` python
def encode(example):
	return tokenizer.prepare_seq2seq_batch(src_texts=example['source'], tgt_texts=example['target'], padding='max_length', max_length=512)
train_dataset.set_transform(encode)
```

I use the max_length here to tokenize every sentence to the same size.
The tokenizer is a MarianTokenizer (where the spm files and vocab are trained with sentencepiece) and is defined as follows:
```python
tokenizer = MarianTokenizer(vocab='tokenizer/vocab.json', source_spm='tokenizer/source.model',
target_spm='tokenizer/target.model', source_lang='it', target_lang='nl', model_max_length=512)
```

- Model
1. A MarianMTModel is configured with the following MarianConfig (same config as the pretrained MarianMT models):
``` python
model = MarianConfig(decoder_layers=6, encoder_layers=6, d_model= 512, decoder_attention_heads=8,
decoder_ffn_dim=2048, decoder_layerdrop=0.0, encoder_attention_heads=8, encoder_ffn_dim=2048, 
encoder_layerdrop=0.0, max_position_embeddings=512)

model = MarianMTModel(configuration)
``` 

- Training
1. After I loaded the data, created the MarianTokenizer and configured the MarianMTmodel I started training with the Trainer. I used the following trainingArguments:
``` python
training_args = TrainingArguments(num_train_epochs=3,  per_device_train_batch_size=12,  
per_device_eval_batch_size=12, warmup_steps=100,  weight_decay=0.01, logging_dir='./logs', logging_steps=5000, 
save_steps=10000, disable_tqdm=False, logging_first_step=True, fp16=True, remove_unused_columns = False)
```

I was not able to increase the batch size, as this gave out of memory errors on the GPU. And finally started training:
```python
trainer = Trainer(
	model=model,                         # the instantiated 🤗 Transformers model to be trained
	args=training_args,                  # training arguments, defined above
	train_dataset=train_dataset,         # training dataset
	eval_dataset=dev_dataset # contains 50,000 samples
)
trainer.train()
```

## Expected behavior
I expect the model to train within a reasonable amount of time (i.e. a couple of days). However, the training process is going to take about 500 hours:
0 %|                          | 219/5367564 [01:18<531:01:44,  2.81it/s]

I was wondering if this is expected behaviour that it takes that long to train. Could you please give me any suggestions on how to modify this to make it faster? 
+1 I found the same problem. The bottleneck seems to be huggingface/datasets. Hence I switched back to use the old customized dataset, which was way more faster.<|||||>Hi ! There's currently an issue in huggingface/datasets that makes iterating through the dataset slow if your dataset is big.
We're working on a fix and we'll do a new release soon to address this :)

I'll ping you when the fix is ready if you want to try it out !
edit: https://github.com/huggingface/datasets/pull/2122 fixed it<|||||>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.<|||||>@gyin-ai could you share your custom solution? I'm running into the same problem