You need to specifically use the save function PyTorch XLA gives you to avoid error, so to use accelerator.save. See the doc on saving 2 for mor context but basically:
# Wait for all the processes to arrive
accelerator.wait_for_everyone()
# Save the real model, not the DDP wrapper
unwrapped_model = accelerator.unwrap_model(model)
# Use accelerator.save to save
model.save_pretrained(local_folder, save_function=accelerate.save)

See also this script 2 for another example.