In such scenarios you ideally want to do all the pre-processing separately before you do the training, so at training time you only need to read data.

This is what I tried initially. I tried to save the preprocessed data to the disk using save_to_disk(). But my dataset is extremely large, and even saving the dataset within a reasonable time period was not possible given that my hard disk space is limited as well as my RAM.
So I tried using the torch.distributed.barrier() as you pointed me to.
test_dataset = Dataset.from_dict(test_dataset)

if training_args.local_rank > 0:
    logger.info('Waiting for main process to perform the mapping')
    torch.distributed.barrier()

test_dataset = test_dataset.map(
    tokenize_function,
    batched=True,
    num_proc=data_args.preprocessing_num_workers,
    cache_file_name='/data/Jason/.cache/test_dataset.cache',
    load_from_cache_file=True if training_args.local_rank > 0 else False,
)

if training_args.local_rank == 0:
    logger.info('Loading results from the main process')
    torch.distributed.barrier()

I can see that the main process does process and save the cache as seen below.

saved_cache763×280 50.7 KB

However, once the main process reaches the barrier and other non-main processes resume, these non-main processes do not load the cache saved by the main process. Is there something that I’m doing wrong in my code?
Thanks.