Allow sharding data files by row when data file count is too small

[aireenmei]

Description

Currently, in the scenario when data file count < data loading host count, the following paths have issues handling:

• For tfds, it will have all the hosts reading all the files, impacting perf. On a large scale, when a file has hundreds reader, it will hang or crash the workload. (b/505191017)
• For grain, with parquet or tfrecord formats, each data loading host requires at least 1 data file. If this requirement is not met, it fails with error Training stopped: load_next_batch() failed with <class 'IndexError'> exception: (list index out of range). (example log)

After this PR, the new behavior is:

• Allow data file split between hosts (shard data file by row). In the case of data_file_count < data_loading_host_count, each host gets data_file_count / data_loading_host_count files.

Test

Unit test

Added ComputeFileShardingNormalCaseTest

End to end test

Grain

Tests (dataset contains 1 files, on 2x v5e-32, 16 data loading hosts). Test 2 runs to make sure checkpointing works.

• TfRecord: screenshot
• Parquet: screenshot

TFDS

(checkpoint not supported, 1 run)
Dataset contains 8 shards, test on 4x v5e-32, 32 data loading hosts.

• log, also checked the warning log screenshot

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.

[codecov]

Codecov Report

❌ Patch coverage is 50.81967% with 30 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/input_pipeline/input_pipeline_utils.py	48.48%	17 Missing ⚠️
...rc/maxtext/input_pipeline/grain_data_processing.py	50.00%	3 Missing and 4 partials ⚠️
src/maxtext/input_pipeline/tfds_data_processing.py	57.14%	4 Missing and 2 partials ⚠️

📢 Thoughts on this report? Let us know!

[igorts-git]

Please add a unit test for compute_file_sharding. Otherwise LGTM.

[aireenmei]

Please add a unit test for compute_file_sharding. Otherwise LGTM.

Good suggestion, unit tests added