Dataset.map() execute twice when in PyTorch DDP mode
Opened by wanghaoyucn (closed)

### Describe the bug

I use `torchrun --standalone --nproc_per_node=2 train.py` to start training. And write the code following the [docs](https://huggingface.co/docs/datasets/process#distributed-usage). The trick about using `torch.distributed.barrier()` to only execute map at the main process doesn't always work. When I am training model, it will map twice. When I am running a test for dataset and dataloader (just print the batches), it can work. Their code about loading dataset are same.

And on another server with 30 CPU cores, I use 2 GPUs and it can't work neither. 

I have tried to use `rank` and `local_rank` to check, they all didn't make sense.

### Steps to reproduce the bug

use `torchrun --standalone --nproc_per_node=2 train.py` or `torchrun --standalone train.py` to run

This is my code:

```python
            if args.distributed and world_size > 1:
                if args.local_rank > 0:
                    print(f"Rank {args.rank}: Gpu {args.gpu} waiting for main process to perform the mapping", force=True)
                    torch.distributed.barrier()
                print("Mapping dataset")
                dataset = dataset.map(lambda x: cut_reorder_keys(x, num_stations_list=args.num_stations_list, is_pad=True, is_train=True), num_proc=8, desc="cut_reorder_keys")
                dataset = dataset.map(lambda x: random_shift(x, shift_range=(-160, 0), feature_scale=16), num_proc=8, desc="random_shift")
                dataset_test = dataset_test.map(lambda x: cut_reorder_keys(x, num_stations_list=args.num_stations_list, is_pad=True, is_train=False), num_proc=8, desc="cut_reorder_keys")
                if args.local_rank == 0:
                    print("Mapping finished, loading results from main process")
                    torch.distributed.barrier()
```


### Expected behavior

Only the main process will execute `map`, while the sub process will load cache from disk.

### Environment info

server with 64 CPU cores (AMD Ryzen Threadripper PRO 5995WX 64-Cores) and 2 RTX 4090

- `python==3.9.16`
- `datasets==2.13.1`
- `torch==2.0.1+cu117`
- `22.04.1-Ubuntu`

server with 30 CPU cores (Intel(R) Xeon(R) Platinum 8375C CPU @ 2.90GHz) and 2 RTX 4090

- `python==3.9.0`
- `datasets==2.13.1`
- `torch==2.0.1+cu117`
- `Ubuntu 20.04`

Comments:
Sorry for asking a duplicate question about `num_proc`, I searched the forum and find the solution.

But I still can't make the trick with `torch.distributed.barrier()` to only map at the main process work. The [post on forum]( https://discuss.huggingface.co/t/slow-processing-with-map-when-using-deepspeed-or-fairscale/7229/7) didn't help.

If it does the `map` twice then it means the hash of your map function is not some same between your two processes.

Can you make sure your map functions have the same hash in different processes ?

```python
from datasets.fingerprint import Hasher

print(Hasher.hash(lambda x: cut_reorder_keys(x, num_stations_list=args.num_stations_list, is_pad=True, is_train=True)))
print(Hasher.hash(lambda x: random_shift(x, shift_range=(-160, 0), feature_scale=16)))
```

You can also set the fingerprint used to reload the resulting dataset by passing `new_finegrprint=` in `map`, see https://huggingface.co/docs/datasets/v2.13.1/en/about_cache#the-cache. This will force the different processes to use the same fingerprint used to locate the resulting dataset in the cache.

Thanks for help! I find the fingerprint between processes don't have same hash:
```
Rank 0: Gpu 0 cut_reorder_keys fingerprint c7f47f40e9a67657
Rank 0: Gpu 0 random_shift fingerprint 240a0ce79831e7d4

Rank 1: Gpu 1 cut_reorder_keys fingerprint 20edd3d9cf284001
Rank 1: Gpu 1 random_shift fingerprint 819f7c1c18e7733f
```
But my functions only process the example one by one and don't need rank or other arguments. After all it can work in the test for dataset and dataloader.
I'll try to set `new_fingerprint` to see if it works and figure out the reason of different hash.

I finally figure it out. The fingerprint of the function will change if other key-value pairs change in `args` even the `args.num_stations_list` is not changed.

```python
lambda x: cut_reorder_keys(x, num_stations_list=args.num_stations_list, is_pad=True, is_train=True)
```

My `args` contains the key `rank` which refers the rank of its GPU, so the fingerprints change among the GPUs.
I use `partial` in `functools` to generate a partial function that fixs the argument `num_stations_list=args.num_stations_list`, and the fingerprint of this partial function keeps among the GPUs. Finally I can reuse the mapped cache.