1.3GB dataset creates over 107GB of cache file!
Author: DarshanDeshpande
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.4.0 dev0
- Platform: Google Colab
- Python version: 3.6
- PyTorch version (GPU?): 1.7
- Tensorflow version (GPU?): None
- Using GPU in script?: None. Colab TPU is used
- Using distributed or parallel set-up in script?:  Using default ```run_mlm.py``` script

### Who can help
@sgugger

<!-- Your issue will be replied to more quickly if you can figure out the right person to tag with @
 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- albert, bert, xlm: @LysandreJik
- blenderbot, bart, marian, pegasus, encoderdecoder,  t5: @patrickvonplaten, @patil-suraj
- longformer, reformer, transfoxl, xlnet: @patrickvonplaten
- fsmt: @stas00
- funnel: @sgugger
- gpt2: @patrickvonplaten, @LysandreJik
- rag: @patrickvonplaten, @lhoestq
- tensorflow: @jplu

Library:

- benchmarks: @patrickvonplaten
- deepspeed: @stas00
- ray/raytune: @richardliaw, @amogkam
- text generation: @patrickvonplaten
- tokenizers: @n1t0, @LysandreJik
- trainer: @sgugger
- pipelines: @LysandreJik

Documentation: @sgugger

HF projects:

- nlp datasets: [different repo](https://github.com/huggingface/nlp)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Examples:

- maintained examples (not research project or legacy): @sgugger, @patil-suraj
- research_projects/bert-loses-patience: @JetRunner
- research_projects/distillation: @VictorSanh

 -->

## Information

Model I am using (Bert, XLNet ...): DistilBert

The problem arises when using:
* [X] the official example scripts: (give details below)
* [ ] my own modified scripts: (give details below)

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [X] my own task or dataset: (give details below)

## To reproduce

Steps to reproduce the behavior:
```
!python /content/transformers/examples/xla_spawn.py --num_cores 8 /content/transformers/examples/language-modeling/run_mlm.py 
--model_type distilbert --config_name /content/TokenizerFiles \
--tokenizer_name /content/drive/TokenizerFiles \
--train_file Corpus.txt \
--mlm_probability 0.15 \ 
--output_dir "/content/TrainingCheckpoints" \
--do_train \
--per_device_train_batch_size 32 \
--save_steps 500 --disable_tqdm False \
--line_by_line True \
--max_seq_length 128 \
--pad_to_max_length True \
--cache_dir /content/cache_dir --save_total_limit 2
```
The script ends up creating more than 107GB of cache files only with 54% processing done which crashes the Colab environment
This means that 200+ GB of space is required to cache and preprocess a mere 1GB file. Am I doing something wrong here? I ran the same script a few days ago and it didn't give me any such "Out of disk space" error. Because I wanted to use the TPU, I changed pad_to_max_length=True [(10192)](https://github.com/huggingface/transformers/issues/10192) . That's all I changed and it does this. Let me know if anyone requires any more data to help me out with this

<!-- If you have code snippets, error messages, stack traces please provide them here as well.
     Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting
     Do not use screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.-->

## Expected behavior

<!-- A clear and concise description of what you would expect to happen. -->
The dataset should cache in a minimum amount of disk space. It currently occupies over 150-200x the space of the actual dataset

Comments:
cc @lhoestq 

Related to https://github.com/huggingface/datasets/issues/861
Maybe on-the-fly tokenization can help.
Or if we stick to having the tokenization in the preprocessing, at least reduce the precision of the integers stored on disk and maybe do the padding on the fly.

@lhoestq Are there any minor changes that could fix this temporarily?  Will changing the map function to set transform as mentioned [here](https://github.com/huggingface/datasets/issues/1825) help?

Currently the Trainer doesn't handle `set_transform` but this will be supported soon.

Another think you could try is specify `features=` in the parameters of the map function to specify the precision of the integers that are written on disk. For example
```python
from datasets import Features, Sequence, Value

features = Features({
    "input_ids": Sequence(Value("int32")),
    "token_type_ids": Sequence(Value("bool")),
    "attention_mask": Sequence(Value("bool")),
    "special_tokens_mask": Sequence(Value("bool")),
})
tokenized_datasets = datasets.map(
    tokenize_function,
    batched=True,
    num_proc=data_args.preprocessing_num_workers,
    remove_columns=[text_column_name],
    load_from_cache_file=not data_args.overwrite_cache,
    features=features,
)
```
The tokenization will still be done during the preprocessing and store the tokenized texts on disk, but this time it will take much less space since you'll store int32 and booleans instead of int64 by default.

@lhoestq Nope. I get a casting error as attached below
```
Exception in device=TPU:4: Could not convert 1 with type int: tried to convert to boolean
Traceback (most recent call last):
  File "/opt/conda/lib/python3.7/site-packages/torch_xla/distributed/xla_multiprocessing.py", line 330, in _mp_start_fn
    _start_fn(index, pf_cfg, fn, args)
  File "/opt/conda/lib/python3.7/site-packages/torch_xla/distributed/xla_multiprocessing.py", line 324, in _start_fn
    fn(gindex, *args)
  File "/kaggle/working/run_mlm_custom.py", line 461, in _mp_fn
    main()
  File "/kaggle/working/run_mlm_custom.py", line 355, in main
    features=features
  File "/opt/conda/lib/python3.7/site-packages/datasets/dataset_dict.py", line 386, in map
    for k, dataset in self.items()
  File "/opt/conda/lib/python3.7/site-packages/datasets/dataset_dict.py", line 386, in <dictcomp>
    for k, dataset in self.items()
  File "/opt/conda/lib/python3.7/site-packages/datasets/arrow_dataset.py", line 1140, in map
    update_data=update_data,
  File "/opt/conda/lib/python3.7/site-packages/datasets/arrow_dataset.py", line 167, in wrapper
    out: Union["Dataset", "DatasetDict"] = func(self, *args, **kwargs)
  File "/opt/conda/lib/python3.7/site-packages/datasets/fingerprint.py", line 312, in wrapper
    out = func(self, *args, **kwargs)
  File "/opt/conda/lib/python3.7/site-packages/datasets/arrow_dataset.py", line 1411, in _map_single
    writer.write_batch(batch)
  File "/opt/conda/lib/python3.7/site-packages/datasets/arrow_writer.py", line 279, in write_batch
    pa_table = pa.Table.from_pydict(typed_sequence_examples)
  File "pyarrow/table.pxi", line 1474, in pyarrow.lib.Table.from_pydict
  File "pyarrow/array.pxi", line 322, in pyarrow.lib.asarray
  File "pyarrow/array.pxi", line 222, in pyarrow.lib.array
  File "pyarrow/array.pxi", line 110, in pyarrow.lib._handle_arrow_array_protocol
  File "/opt/conda/lib/python3.7/site-packages/datasets/arrow_writer.py", line 100, in __arrow_array__
    out = pa.array(self.data, type=type)
pyarrow.lib.ArrowInvalid: Could not convert 1 with type int: tried to convert to boolean
Traceback (most recent call last):
  File "transformers/examples/xla_spawn.py", line 85, in <module>
    main()
  File "transformers/examples/xla_spawn.py", line 81, in main
    xmp.spawn(mod._mp_fn, args=(), nprocs=args.num_cores)
  File "/opt/conda/lib/python3.7/site-packages/torch_xla/distributed/xla_multiprocessing.py", line 395, in spawn
    start_method=start_method)
  File "/opt/conda/lib/python3.7/site-packages/torch/multiprocessing/spawn.py", line 157, in start_processes
    while not context.join():
  File "/opt/conda/lib/python3.7/site-packages/torch/multiprocessing/spawn.py", line 112, in join
    (error_index, exitcode)
Exception: process 5 terminated with exit code 17
```

Fixed this using 
features = Features({
            "input_ids": Sequence(Value("int32")),
            "attention_mask": Sequence(Value("int32")),
            "special_tokens_mask": Sequence(Value("int32")),
        })
but it still ends up taking a lot of space. Runs out of space again on kaggle notebook with 19GB free

Looks like PyArrow doesn't know how to convert 1 to True ^^'
I created an issue on Apache Arrow's JIRA [here](https://issues.apache.org/jira/browse/ARROW-11646) to track the issue.
Also I think we can add support for uint16 in the pytorch integration of `datasets` to reduce the size even more. Feel free to open an issue on the `datasets` repo about this if you want.

In the end I think the best solution is probably to not do padding during preprocessing and do it on the fly in the Trainer. This can use `dataset.set_transform` or a data collator. As long as padding is done before sending the data to the TPU it should be good.

@lhoestq 
You said in the comment before that ```set-transform``` support hasn't been added to the trainer yet. How exactly do I do it on the fly then?

Yes you're right, we first need to make the Trainer support `set_transform`, sorry if it was confusing.

`Trainer` in master completely supports `set_transform`. If there are some columns removed that should not be, you just have to set the training arguments `remove_unused_columns` to `False` for the time being.

Nice ! so one could try to replace the .map with .set_transform(tokenize_function)

> `Trainer` in master completely supports `set_transform`. If there are some columns removed that should not be, you just have to set the training arguments `remove_unused_columns` to `False` for the time being.

I tried changing remove_unused_columns to False but it gives me an error during the Trainer call. The set_transform function returns a NoneType object and so the Trainer complains of getting a None instead of training data. The run_mlm_custom.py file below is the same run_mlm file, just with set_transform instead of map (you can have a look at it [here](https://drive.google.com/file/d/1--ijV3UK-Rq9TnzkcWFAjkk2A-0j_XWK/view?usp=sharing))

```
Traceback (most recent call last):
  File "run_mlm_custom.py", line 452, in <module>
    main()
  File "run_mlm_custom.py", line 397, in main
    train_dataset=tokenized_datasets["train"] if training_args.do_train else None,  # print(tokenized_datasets) gives None
TypeError: 'NoneType' object is not subscriptable
```

Do you have any template code for passing the data to the trainer?

Indeed `set_transform` is in-place. For example you can do
```python
dataset.set_transform(tokenize)  # return None, but sets the transform of the current dataset
```

If you want to use a non in-place function like what was doing map, you can do
```python
dataset = dataset.with_transform(tokenize)  # return a new dataset object with the specified transform
```

Also I'm not a big fan of having two functions that does the same thing (except one is in-place) so we might deprecate one or the other in the future. I guess the second one is more convenient and is more aligned with the other Dataset functions. Let me know what you think

Returning the dataset is more intuitive I feel. Anyway, this is some really good news. I will try to modify the script and make it work. If it does then maybe, if you want, I can clean the code and create a pull request for the same.

@DarshanDeshpande did it work for you after you made the changes? 
I have the same issue, trying to train a roberta mlm on 1.3G of text data on cloud TPU and got the no space on device error (the code works with 300M of data though). I changed the run_mlm.py code based on your PR to do tokenization on the fly, but now I get this error:

```
[INFO|trainer.py:946] 2021-03-16 17:51:34,455 >> ***** Running training *****
[INFO|trainer.py:947] 2021-03-16 17:51:34,456 >>   Num examples = 14602056
[INFO|trainer.py:948] 2021-03-16 17:51:34,456 >>   Num Epochs = 2
[INFO|trainer.py:949] 2021-03-16 17:51:34,456 >>   Instantaneous batch size per device = 8
[INFO|trainer.py:950] 2021-03-16 17:51:34,456 >>   Total train batch size (w. parallel, distributed & accumulation) = 256
[INFO|trainer.py:951] 2021-03-16 17:51:34,456 >>   Gradient Accumulation steps = 4
[INFO|trainer.py:952] 2021-03-16 17:51:34,456 >>   Total optimization steps = 114078
  0%|                                                                                                                                                                                       | 0/114078 [00:00<?, ?it/s]Exception in thread Thread-2:
Traceback (most recent call last):
  File "/anaconda3/envs/torch-xla-1.7/lib/python3.6/threading.py", line 916, in _bootstrap_inner
    self.run()
  File "/anaconda3/envs/torch-xla-1.7/lib/python3.6/threading.py", line 864, in run
    self._target(*self._args, **self._kwargs)
  File "/anaconda3/envs/torch-xla-1.7/lib/python3.6/site-packages/torch_xla/distributed/parallel_loader.py", line 141, in _loader_worker
    _, data = next(data_iter)
  File "/anaconda3/envs/torch-xla-1.7/lib/python3.6/site-packages/torch/utils/data/dataloader.py", line 435, in __next__
    data = self._next_data()
  File "/anaconda3/envs/torch-xla-1.7/lib/python3.6/site-packages/torch/utils/data/dataloader.py", line 475, in _next_data
    data = self._dataset_fetcher.fetch(index)  # may raise StopIteration
  File "/anaconda3/envs/torch-xla-1.7/lib/python3.6/site-packages/torch/utils/data/_utils/fetch.py", line 44, in fetch
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "/anaconda3/envs/torch-xla-1.7/lib/python3.6/site-packages/torch/utils/data/_utils/fetch.py", line 44, in <listcomp>
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "/anaconda3/envs/torch-xla-1.7/lib/python3.6/site-packages/datasets/arrow_dataset.py", line 1130, in __getitem__
    format_kwargs=self._format_kwargs,
  File "/anaconda3/envs/torch-xla-1.7/lib/python3.6/site-packages/datasets/arrow_dataset.py", line 1117, in _getitem
    pa_subtable, key, formatter=formatter, format_columns=format_columns, output_all_columns=output_all_columns
  File "/anaconda3/envs/torch-xla-1.7/lib/python3.6/site-packages/datasets/formatting/formatting.py", line 375, in format_table
    return formatter(pa_table, query_type=query_type)
  File "/anaconda3/envs/torch-xla-1.7/lib/python3.6/site-packages/datasets/formatting/formatting.py", line 173, in __call__
    return self.format_row(pa_table)
  File "/anaconda3/envs/torch-xla-1.7/lib/python3.6/site-packages/datasets/formatting/formatting.py", line 239, in format_row
    formatted_batch = self.format_batch(pa_table)
  File "/anaconda3/envs/torch-xla-1.7/lib/python3.6/site-packages/datasets/formatting/formatting.py", line 268, in format_batch
    return self.transform(batch)
  File "/home/aida_delfan/pretrain/run_mlm.py", line 363, in tokenize_function
    return_special_tokens_mask=True,
  File "/anaconda3/envs/torch-xla-1.7/lib/python3.6/site-packages/transformers/tokenization_utils_base.py", line 2266, in __call__
    **kwargs,
  File "/anaconda3/envs/torch-xla-1.7/lib/python3.6/site-packages/transformers/tokenization_utils_base.py", line 2451, in batch_encode_plus
    **kwargs,
  File "/anaconda3/envs/torch-xla-1.7/lib/python3.6/site-packages/transformers/models/gpt2/tokenization_gpt2_fast.py", line 163, in _batch_encode_plus
    return super()._batch_encode_plus(*args, **kwargs)
  File "/anaconda3/envs/torch-xla-1.7/lib/python3.6/site-packages/transformers/tokenization_utils_fast.py", line 411, in _batch_encode_plus
    for key in tokens_and_encodings[0][0].keys():
IndexError: list index out of range

```
wondering if you saw a similar error and if there is a fix for it.

I have transformers 4.4.0 and datasets 1.4.0.

Thanks!

@aidad This looks like a tokenizer issue. The script in the PR works for me. Try checking your tokenizer files or altering the script to use the Roberta tokenizer specifically instead of the included AutoTokenizer