Cache the checkpoint saving plan

[Leahlijuan]

Currently we need to generate plan to save the checkpoint every time, while since the model structure is keeping the same throughout the training, the plan stays the same. As the generate_plan step is blocking time which would slow down the training, if we can cache the plan, we can save a significant amount of time.