Boost Training Goodput: How Continuous Checkpointing Optimizes Reliability in Orbax and MaxText

Summary

Orbax and MaxText introduce continuous checkpointing to balance training reliability and performance when hardware failures or preemption events interrupt long-running jobs. Fixed checkpoint intervals create a trade-off: infrequent saves risk substantial lost work, while frequent saves can block or bottleneck training, particularly over unstable networks. MaxText enables asynchronous checkpointing that attempts a new save only after the previous background request is fulfilled; its continuous mode overrides checkpoint_period, and max_num_checkpoints_to_keep can limit storage use. A llama-3.1-70B benchmark on two v5p-128 slices found smaller P50 checkpoint intervals, but average training step time increased because of more frequent device-to-host transfers; the analysis estimated resource conservation against mean-time-between-failure assumptions. Orbax also supports minimum intervals and customizable preservation policies, while co-locating storage with training is recommended because cross-metro bandwidth can reduce checkpointing speed and reliability.

Takeaways

1. 1

   Continuous checkpointing produced smaller P50 checkpoint intervals than checkpointing every 100 steps in the described llama-3.1-70B benchmark, but increased average training step time because of more frequent device-to-host transfers.

2. 2

   In MaxText, continuous checkpointing overrides checkpoint_period and attempts a save after the previous background saving request is fulfilled; max_num_checkpoints_to_keep limits retained storage.

3. 3

   Orbax supports minimum intervals and custom preservation logic, while co-locating storage with the training cluster helps avoid cross-metro bandwidth degrading checkpoint speed and reliability.