Reorder validate and checkpoint in train

[wesleytruong]

If validation and checkpoint occur on the same training step, do checkpointing first

[tianyu-l]

could you say more about why?

[wesleytruong]

could you say more about why?

It may be a minor use case in real world where validation and checkpoint don't often occur on same step, but I think it could happen, especially since some people may set the behavior to occur this way.

I think it would be a good change since it makes training more robust against errors in validation. If validation fails on the same step then we would lose all our pre-checkpoint progress. Additionally, hidden problems could arise in validation that may interfere with checkpointing, such as I had a problem before where validation and checkpoint on the same step would cause a missing fqn in state dict error during save, but saving on a different step would work. This was likely because validation has to change pp schedule to only do forward and may not always correctly reset pp schedule or model to its original state. Moving checkpoint.save to before validate could help to isolate against errors such as this and make one less thing to worry about.

[tianyu-l]

Sounds good to me. That's very thoughtful!

If this is the case, it's less necessary to manually call reshard on FSDP modules, as it was originally there to make checkpoint saving happy. But we can keep it there for now.