- add a post-step callback hook for checkpoint saves - preserve updated model, optimizer, and scheduler state - cover checkpoint ordering with a regression test