Speed up checkpoint loading time at scale using Orbax on JAX

Imagine training a new AI / ML model like Gemma 3 or Llama 3.3 across hundreds of powerful accelerators like TPUs or GPUs to achieve a scientific breakthrough. You might have a team of powerful computers working in sync, constantly learning and refining. But every so often, they need to save their progress — a "checkpoint" — and then pick up from this known state in the case of an interruption.  

With the traditional approach, each device independently reads the same checkpoint from a central storage like Google Cloud Storage (GCS), resulting in duplicate data transfers. When GCS bandwidth of a project is fully utilized, it causes significant delays before training even begins. This bottleneck isn't just an inconvenience; it cuts productivity and increases cost (remember that you’re paying for all the accelerators that are waiting while the checkpoint is being saved or restored). 

Today, we'll explore how you ...