Regression: HF JSONL training halts early with

### Bug report

@aireenmei

I keep getting an error related to running out of shards. Training is halted early, and gives the error:
“Run out of shards on host 0, shard 256 is not available”
followed by StopIteration → “You may have run out of training data.”

Setup:
- MaxText @ main (Aug 2025)
- TPU v5-32 (8 hosts)
- dataset_type=hf, streaming JSONL from GCS
- 256 shards matched by: hf_train_files='gs://mybucket/train*.jsonl'
(files named train_aa.jsonl … train_jo.jsonl)
- Command: `python -m MaxText.train ... hf_path=json hf_data_dir= hf_train_files='gs://…/train*.jsonl' ...`

I think this is a regression, because earlier this did run until none of the workers had any data.

### Logs/Output

_No response_

### Environment Information

_No response_

### Additional Context

_No response_