Call finish_checkpoint_resume to flush optimizer step (#2884)

Summary:
X-link: pytorch/torchrec#2244

Pull Request resolved: #2884

`iter` number is a critical counter in computing how much decay is needed together with `prev_iter`. However, `iter` is not checkpointed since it is not in `split_optimizer_state`. Instead, it relies on external calls of `set_optimizer_step` to flush the counter, which relies on explict calls of `finish_checkpoint_resume` to trigger `flush_state` in `FullSyncOptimizer`. This chain of actions is not properly set up in existing code, which causes the issue that after a preemption, `iter` gets reset to zero.

Reviewed By: yuxihu

Differential Revision: D60155781

fbshipit-source-id: 51b85fe10fb196df52bfdfc22424e6ab6f8c5bfb

Author: Wang Zhou

fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_training.py

@@ -2125,11 +2125,12 @@ def _set_learning_rate(self, lr: float) -> float:
         self.optimizer_args = self.optimizer_args._replace(learning_rate=lr)
         return 0.0
 
-    @torch.jit.export
+    @torch.jit.ignore
     def set_optimizer_step(self, step: int) -> None:
         """
         Sets the optimizer step.
         """
+        self.log(f"set_optimizer_step from {self.iter[0]} to {step}")
         if self.optimizer == OptimType.NONE:
             raise NotImplementedError(
                 f"Setting optimizer step is not supported for {self.optimizer}"