I found a way to resolve the issue - shard the dataset before doing the last transformation and save_to_disk. That way resulting shards are not copies of dataset.arrow and indices file. If you are curious to have a look, here it is create_pretraining_dataset.py 78