Allow to choose the `writer_batch_size` when using `save_to_disk`

### Feature request

Add an argument in `save_to_disk` regarding batch size, which would be passed to `shard` and other methods.

### Motivation

The `Dataset.save_to_disk` method currently calls `shard` without passing a `writer_batch_size` argument, thus implicitly using the default value (1000). This can result in RAM saturation when using a lot of processes on long text sequences or other modalities, or for specific IO configs.

### Your contribution

I would be glad to submit a PR, as long as it does not imply extensive tests refactoring.