Reorder validate and checkpoint in train (#1542)

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

Author: wesleytruong

torchtitan/train.py

@@ -575,17 +575,17 @@ def train(self):
                     logger.warning("Ran out of data; last step was canceled.")
                     break
 
+                self.checkpointer.save(
+                    self.step, last_step=(self.step == job_config.training.steps)
+                )
+
                 # Run validation if validator is available
                 if (
                     self.job_config.validation.enabled
                     and self.validator.should_validate(self.step)
                 ):
                     self.validator.validate(self.model_parts, self.step)
 
-                self.checkpointer.save(
-                    self.step, last_step=(self.step == job_config.training.steps)
-                )
-
                 # signal the profiler that the next profiling step has started
                 if torch_profiler:
                     torch_profiler.step()