Support grain data checkpoint for elastic training

[aireenmei]

Description

• migrate RemoteIterator to colocated python class
• Add checkpointing logic to RemoteIterator, so data iterator in the colocated sidecar writes checkpoint to the checkpoint path, prevent sending data iterator state to the controller
• Add grain.ElasticIterator support controlled by flag grain_use_elastic_iterator, can be used with or without elastic training, with or without colocated_python. This class allows recovering checkpoint with a dynamic scale (up or down), with the following limitations, future work will loose the limitations:
  (1). Only arrayrecord files are supported, parquet or tfrecord are not supported
  (2). Does not support many-to-one transformations, including packing, filtering
  (3). Does not support mixing datasets

Tests

Tested on Pathways saving and restoring data iterator checkpoints with different # of v5e-32 slices
jobset.yaml

With colocated_python

colocated_python_data_input=true colocated_python_checkpointing=true grain_use_elastic_iterator=true, checkpoints in gs://aireenmei-multipod/pathways-v5e/20260424_grain_elastic_2/aireen-pathways-v5e/checkpoints:

• Start with 1 slice, set steps=25, checkpoint at step 0, 10, 20, 24: log (log confirms Num devices: 32),
• Resume with 2 slices, set steps=45, checkpoint at step 30, 40, 44: log(log confirms Num devices: 64),
• Resume with 1 slice, set steps=65, checkpoint at step 50, 60, 64: log(log confirms Num devices: 32)

Without colocated_python

colocated_python_data_input=false colocated_python_checkpointing=false grain_use_elastic_iterator=true, checkpoints in gs://aireenmei-multipod/pathways-v5e/20260424_grain_elastic_3/aireen-pathways-v5e/checkpoints:

• Start with 1 slice, set steps=25, checkpoint at step 0, 10, 20, 24: log (log confirms Num devices: 32),
• Resume with 2 slices, set steps=45, checkpoint at step 30, 40, 44: log(log confirms Num devices: 64),
• Resume with 1 slice, set steps=65, checkpoint at step 50, 60, 64: log(log confirms Num devices: 32)

Index verification

Inspecting the checkpoints in gs://aireenmei-multipod/pathways-v5e/20260424_grain_elastic_2/aireen-pathways-v5e/checkpoints/{step}/iter/process_0.json

• Step 20: "global_next_index": 672
• Step 24: "global_next_index": 800
• Step 30: "global_next_index": 1184
• Step 20-24 is on 1 slice, 32 devices, batch_size=32, matches (800-672) / 32 = 4 steps
• Step 24-30 is on 2 slices, 64 devices, batch_size=64, matches (1184-800) / 64 = 6 steps

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 22.28916% with 129 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...rc/maxtext/input_pipeline/multihost_dataloading.py	20.28%	55 Missing ⚠️
src/maxtext/common/checkpointing.py	17.07%	25 Missing and 9 partials ⚠️
...rc/maxtext/input_pipeline/grain_data_processing.py	31.11%	26 Missing and 5 partials ⚠️
...rc/maxtext/input_pipeline/data_processing_utils.py	22.22%	4 Missing and 3 partials ⚠️
src/maxtext/input_pipeline/tfds_data_processing.py	0.00%	2 Missing ⚠️

📢 Thoughts on this report? Let us know!

[github-actions]

🤖 Hi @aireenmei, 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]

🤖 I'm sorry @aireenmei, but I was unable to process your request. Please see the logs for more details.

[github-actions]

🤖 Hi @aireenmei, 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]

🤖 I'm sorry @aireenmei, but I was unable to process your request. Please see the logs for more details.

[lukebaumann]

How are you confirming that the dataset is resumed from the correct index?

[lukebaumann]

We want ElasitcIterator to also be a RemoteIterator. Is that happening?

[aireenmei]

We want to support ElasitcIterator in both regular (Pathways, mcJAX) and colocated python environments. Because ElasitcIterator allows flexible chip counts, it has use case for regular customers who may change chip counts while keeping training progress. This specific path is for non-colocated cases (including Pathways and mcJAX), the Pathways + colocated case is handled by RemoteIteratorWrapper in the lines above. Let me improve the comments

[SujeethJinesh]

Can you pull out what this actually is just for easier readability?

[aireenmei]

I refactored and removed "[0]". Hopefully better readability now

[SujeethJinesh]

Do need any specialization?

[aireenmei]

Could you explain specialization?

[aireenmei]

How are you confirming that the dataset is resumed from the correct index?

Good question. I just added a section "Index verification" in PR description.