go-inoue (NONE) commented on 2020-08-28 13:29:49:
Here's a minimal working example to reproduce this issue.

Assumption:
- You have access to TPU.
- You have installed `transformers` and `nlp`.
- You have tokenizer files (`config.json`, `merges.txt`, `vocab.json`) under the directory named `model_name`.
- You have `xla_spawn.py` (Download from https://github.com/huggingface/transformers/blob/master/examples/xla_spawn.py).
- You have saved the following script as `prepare_cached_dataset.py`.

```python
#!/usr/bin/env python3
import argparse
from transformers import AutoTokenizer
from nlp import load_dataset

def main():
    parser = argparse.ArgumentParser(description='description')
    parser.add_argument('--tokenizer_name', type=str, help='Pretrained tokenizer name or path if not the same as model_name')
    parser.add_argument('--data_file', type=str, help='The input data file (a text file).')
    parser.add_argument('--cache_dir', type=str, default=None, help='Where do you want to store the pretrained models downloaded from s3')
    parser.add_argument('--block_size', type=int, default=-1, help='The training dataset will be truncated in block of this size for training')
    parser.add_argument('--tpu_num_cores', type=int, default=1, help='Number of TPU cores to use (1 or 8). For xla_apwan.py')
    model_args = parser.parse_args()
    
    tokenizer = AutoTokenizer.from_pretrained(model_args.tokenizer_name, cache_dir=model_args.cache_dir, use_fast=True)
    
    dataset = load_dataset("text", data_files=model_args.data_file, split="train")
    dataset = dataset.map(lambda ex: tokenizer(ex["text"], add_special_tokens=True,
                                truncation=True, max_length=model_args.block_size), batched=True)
    dataset.set_format(type='torch', columns=['input_ids'])

def _mp_fn(index):
    # For xla_spawn (TPUs)
    main()

if __name__ == "__main__":
    main()
```

- Run the following command. Replace `your_training_data` with some text file.

```bash
export TRAIN_DATA=your_training_data

python prepare_cached_dataset.py \
--tokenizer_name=model_name \
--block_size=512 \
--data_file=$TRAIN_DATA
```
- Check the cached directory.
```bash
ls -lha /home/*****/.cache/huggingface/datasets/text/default-e84dd29acc4ad9ef/0.0.0/447f2bcfa2a721a37bc8fdf23800eade1523cf07f7eada6fe661fe4d070d380d
total 132M
drwxr-xr-x 2 ***** ***** 4.0K Aug 28 13:08 .
drwxr-xr-x 3 ***** ***** 4.0K Aug 28 13:08 ..
-rw------- 1 ***** *****  99M Aug 28 13:08 cache-bfc7cb0702426d19242db5e8c079f04b.arrow
-rw-r--r-- 1 ***** *****  670 Aug 28 13:08 dataset_info.json
-rw-r--r-- 1 ***** *****    0 Aug 28 13:08 LICENSE
-rw-r--r-- 1 ***** *****  33M Aug 28 13:08 text-train.arrow
```

- Run the same script again. (The output should be just `Using custom data configuration default`.)
```
python prepare_cached_dataset.py \
--tokenizer_name=model_name \
--block_size=512 \
--data_file=$TRAIN_DATA
```
- Check the cached directory.
```bash
ls -lha /home/*****/.cache/huggingface/datasets/text/default-e84dd29acc4ad9ef/0.0.0/447f2bcfa2a721a37bc8fdf23800eade1523cf07f7eada6fe661fe4d070d380d
total 132M
drwxr-xr-x 2 ***** ***** 4.0K Aug 28 13:08 .
drwxr-xr-x 3 ***** ***** 4.0K Aug 28 13:08 ..
-rw------- 1 ***** *****  99M Aug 28 13:08 cache-bfc7cb0702426d19242db5e8c079f04b.arrow
-rw-r--r-- 1 ***** *****  670 Aug 28 13:20 dataset_info.json
-rw-r--r-- 1 ***** *****    0 Aug 28 13:20 LICENSE
-rw-r--r-- 1 ***** *****  33M Aug 28 13:08 text-train.arrow
```
- The cached file (`cache-bfc7cb0702426d19242db5e8c079f04b.arrow`) is reused.
- Now, run this script with `xla_spawn.py`. Ideally, it should reuse the cached file, however, you will see each process is creating a cache file again.

```bash
python xla_spawn.py --num_cores 8 \
prepare_cached_dataset.py \
--tokenizer_name=model_name \
--block_size=512 \
--data_file=$TRAIN_DATA
```

- Check the cached directory. There are two arrrow files.
```bash
ls -lha /home/*****/.cache/huggingface/datasets/text/default-e84dd29acc4ad9ef/0.0.0/447f2bcfa2a721a37bc8fdf23800eade1523cf07f7eada6fe661fe4d070d380d
total 230M
drwxr-xr-x 2 ***** ***** 4.0K Aug 28 13:25 .
drwxr-xr-x 3 ***** ***** 4.0K Aug 28 13:08 ..
-rw------- 1 ***** *****  99M Aug 28 13:08 cache-bfc7cb0702426d19242db5e8c079f04b.arrow
-rw------- 1 ***** *****  99M Aug 28 13:25 cache-e0e2313e49c8a110aafcc8133154c19a.arrow
-rw-r--r-- 1 ***** *****  670 Aug 28 13:24 dataset_info.json
-rw-r--r-- 1 ***** *****    0 Aug 28 13:24 LICENSE
-rw-r--r-- 1 ***** *****  33M Aug 28 13:08 text-train.arrow
```