.map() and distributed training
State: closed | Author: VictorSanh (CONTRIBUTOR)

Hi,
I have a question regarding distributed training and the `.map` call on a dataset.

I have a local dataset "my_custom_dataset" that I am loading with `datasets = load_from_disk(dataset_path=my_path)`.
`dataset` is then tokenized:
```python
datasets = load_from_disk(dataset_path=my_path)

[...]

def tokenize_function(examples):
    return tokenizer(examples[text_column_name])

logger.info("Mapping dataset to tokenized dataset.")
tokenized_datasets = datasets.map(
    tokenize_function,
    batched=True,
    num_proc=preprocessing_num_workers,
    remove_columns=column_names,
    load_from_cache_file=True,
)
```
I am using 31 workers (`preprocessing_num_workers=31`) and thus it creates 31 `cache*.arrow` files in `my_path/train` (there is only a train split).
When I relaunch the script, the map is tokenization is skipped in favor of loading the 31 previously cached files, and that's perfect.

Everything so far was done by launching a **single process script**.
I now launch the same training script in **distributed mode** (`pytorch -m torch.distributed.launch --nproc_per_node 2`). However, once it reaches the map call, it re-does the tokenization... instead of loading the 31 cached files. 

I tried adding the `cache_file_name` argument: `cache_file_name={"train": my_path/one_of_the_arrow_file}`, but I can't give the 31 cached files, so it probably isn't the right way to do it.

**My question: what is the best way to load cached files if they were pre-processed and dumped in multiple arrow files?** It seems automatically handled for single processes but fails on distributed training.

- I am following the same structure as the examples of transformers (more specifically [run_clm.py](https://github.com/huggingface/transformers/blob/master/examples/language-modeling/run_clm.py) in my case)
- I am using 1.5.0 version of datasets if that matters.

Comments: Hi, one workaround would be to save the mapped(tokenized in your case) file using `save_to_disk`, and having each process load this file using `load_from_disk`. This is what I am doing, and in this case, I turn off the ability to automatically load from the cache.

Also, multiprocessing the map function seems to be slower at the moment (#1992), hope this helps you.

Thanks @hwijeen for the workaround, feels a bit prototypical but it works! (it seems files are written twice then though)

(I haven't observed slowness using multiprocessed map function but I could be wrong)

To my understanding, files are written twice anyhow(one after load_dataset, another aftet map). It's just that you now have it at the location where you can see, whereas it was secretlely saved at caching folder(.cache/huggingface/datasets by default)! Correct me if I'm wrong!

Slowness in multiprocessing has been observed in certain environments but not others. We're investigating ;)

So to answer my initial question, I was just doing something stupid as I was not re-giving the `preprocessing_num_workers` arguments when launching the distributed training (and it was then set to `None`). I initially thought the hash was computed only with the `tokenize_function` but it's all arguments. Thanks @lhoestq for clarifying!

This cache process isn't really consistent. I just changed `per_device_train_batch_size` of training script and now it rebuilding the dataset cache!!!! Why?

Hi ! A `map` function is recomputed if the code changes or if any of the variables it uses changes. Can you check that your function doesn't use `per_device_train_batch_size` or any variable that contains `per_device_train_batch_size` ?

My code is actually a transformer's example for training t5, I modified a bit:

https://github.com/puraminy/transformers/blob/4b40877132eedb566043f83de8f1d29a84d71430/examples/flax/language-modeling/run_t5_mlm_flax.py#L614

No, it doesn't use `per_device_train_batch_size`. I remember it worked for several times and then for no reason or various reasons like the above it started to build the cache again, as if it had an expiration date (maybe), or maybe I had changed the code! 

So, to get rid of these problems I saved cache with a name (was forced to not use multiple_processes, because otherwise it generates multiple files) and then I load it from this cache file.