Incorporate Orbax emergency replicator checkpoint manager

[xuefgu]

Description

Incorporate Orbax emergency replicator checkpoint manager.

It is a feature under test and does not affect any other usage of MaxText.

Tests

Manually tested and observed that the behavior is WAI: local checkpoints are saved on each slice and no GCS checkpoints are saved.

To-be-tested: The restore path with the relevant GKE component. We agreed that we would first check in the code for the save path and leave E2E for later.

Checklist

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

• I have performed a self-review of my code.
• 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.