Checkpoint tfds data iterator

[mattdonati]

Enable deterministic training with preemption when using tfds pipeline by checkpointing data iterator.

Creates a checkpoint handler for data iterator that implements orbax.checkpoint.CheckpointHandler, similar to https://github.com/google/grain/blob/main/grain/_src/python/checkpoint_handlers.py. Handler utilizes tf.train.Checkpoint to save and restore iterator.

Makes checkpointing the data iterator optional, since this method will save large checkpoints. Adds a bool flag to base.yml

Async checkpointing is handled at the level of the orbax checkpoint manager.

Updates input pipeline description to reflect option to checkpoint tfds iterator.

[aireenmei]

Thanks for using and contributing to our repo!
I would like to understand the use case of this feature. You must have checked out our doc and know about the Grain pipeline which is optimized for ckpt support of data iterators. It's efficient and will save very small data iterator ckpt containing only indices, therefore it's the recommended way for use cases that need check-pointing data iterators. We would like to hear if there's any difficulty in adopting Grain for your use case?

And I tested the branch with this command on v4-8 python3 MaxText/train.py MaxText/configs/base.yml steps=20 per_device_batch_size=8.0 learning_rate=3e-4 enable_checkpointing=true base_output_directory=gs://aireenmei-multipod/tfds_ckpt dataset_path=gs://maxtext_dataset tfds_iter_checkpointing=True run_name=$(date +%m%d-$H%M) checkpoint_period=10 and got this error (no error if tfds_iter_checkpointing=False): https://gist.github.com/aireenmei/42a4c4e0dd8caed0b7ce8182f5ca8292

[aireenmei]

Since it saves large checkpoint, it would be better to set the default to False

[aireenmei]

Could you add more details on what contributes to the size of the checkpoint? Would be good to provide a way to estimate the size to help the users make plans