Optimizations to cache `all_steps` to avoid repeated file accesses. A few extra minor optimizations for the save path to avoid duplicate work, but these should have a marginal impact compared to that of all_steps.
Optimizations to cache `all_steps` to avoid repeated file accesses. A few extra minor optimizations for the save path to avoid duplicate work, but these should have a marginal impact compared to that of all_steps.

PiperOrigin-RevId: 471606497

google/paxml,google/paxml