Add utility to reshard TFDS dataset for large scale training

[RissyRan]

Description

Add a utility to reshard TFDS dataset for large scale training, b/508024540.

• Eliminates overhead of dynamic re-sharding during training by pre-sharding data to match the host count.
• Tests show a reduction in training loss (from 4.075 to 3.582). This is because the TFDS fallback logic (ds.shard) performs a "logical shard" that effectively reconstructs the un-shuffled sequential order of the original dataset across the global batch. By physically re-sharding into dedicated files, each host can stream contiguous records into its local shuffle buffer. This ensures that when host data is combined, the global batch contains a truly randomized distribution of data categories, leading to higher-quality gradients and better model performance.

Tests

Local conversation - example log

Before this change:

• Warning from loading both training and eval dataset, and there is huge performance hit for each training steps
• end-to-end training logs by loading real dataset, and performance is Tokens/s/device: 80-90 per step link

/deps/src/maxtext/input_pipeline/tfds_data_processing.py:64: UserWarning: WARNING: Inefficient dataloading. Your c4/en:3.0.5 contains 128shards, smaller than dataloading_host_count=2048. This is known to lead to inefficient dataloading.

/deps/src/maxtext/input_pipeline/tfds_data_processing.py:64: UserWarning: WARNING: Inefficient dataloading. Your c4/en:3.0.5 contains 64shards, smaller than dataloading_host_count=2048. This is known to lead to inefficient dataloading

After this change:

• We reshard the dataset based on host count
• No warning is found in logs: link
• end-to-end training logs by loading real dataset, and performance is Tokens/s/device: 600-700 per step link

For training loss:

• before the change, initial loss: 5.819
• after the change, initial loss: 5.816
• observe the training efficiency after re-sharding, the llm_loss at step 77 reduce from 4.075 (before resharding) to 3.582 (after resharding). Both train/validation loss are aligned.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[github-actions]

🤖 Hi @shuningjin, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[github-actions]

## 📋 Review Summary

This Pull Request introduces a valuable utility for resharding TFDS datasets, which is critical for maintaining efficient dataloading performance in large-scale training environments (e.g., matching shard count to host count). The implementation is well-structured, leverages multiprocessing effectively, and correctly updates the dataset metadata.

🔍 General Feedback

• The use of raw_bytes reading and the round-robin distribution strategy into target shards is efficient and helps ensure a balanced distribution of data.
• The automatic adjustment of num_workers to be a factor of num_shards is a clever way to maintain deterministic shard ownership per worker and avoid write collisions.
• Overall, the code is high quality and follows good practices for data processing scripts.

[github-actions]

## 📋 Review Summary (Addendum)

Providing one additional piece of feedback that was missed in the initial submission regarding dataset naming conventions.

[shuningjin]

Approve to unblock.