Thank you adilism and lewtun,
I tried both suggestions, but they did not work out for me .
Neither cleanup_cache_files() nor datasets.set_caching_enabled(False) do not seem to affect saving at all, which is curious. I still have 16Gb for each train, validation, and test in my wikipedia_rank_nocache folders.
The thing that kind of worked is to only create new ArrowDataset objects for validaiton and test. These two are small and fit into memory easily.
dataset_dict["test"] = datasets.arrow_dataset.Dataset.from_dict(dataset_dict_val["test"].data.to_pydict())
dataset_dict["validation"] = datasets.arrow_dataset.Dataset.from_dict(dataset_dict_val["validation"].data.to_pydict())
dataset_dict.save_to_disk("../data/wikipedia_rank_nocache")