Honestly, I haven鈥檛 tried this code, just found it when researching to answer your question. It鈥檚 possible that there is a bug in datasets?
We are going to add this feature to all HF Trainer-based examples [examples] [distributed] process datasets.map only on main process in 路 Issue #12345 路 huggingface/transformers 路 GitHub so watch that Issue, I will try to work on it shortly and see what needs to be done to make it work.