add extra methods to wsd scheduler

Author: Chris-hughes10

pytorch_accelerated/callbacks.py

@@ -834,12 +834,7 @@ def on_training_run_start(self, trainer, **kwargs):
     def on_train_step_end(self, trainer, step: int, **kwargs):
         """Handle checkpoint saving and progress logging"""
 
-        # Calculate global step accounting for distributed training and gradient accumulation
-        total_steps = (
-            (trainer.run_history.current_epoch - 1)
-            * trainer.run_config.num_update_steps_per_epoch
-            + step // trainer.run_config.gradient_accumulation_steps
-        )
+        total_steps = trainer.scheduler.get_current_step()
 
         # Skip if we've already saved at this step
         if total_steps == self.last_checkpoint_step:

pytorch_accelerated/schedulers/wsd_scheduler.py

@@ -217,6 +217,23 @@ def get_decay_info(self) -> List[dict]:
         """
         return self.checkpoint_decay_info
 
+    def get_current_step(self) -> int:
+        """Get the current step count of the scheduler.
+        
+        Returns:
+            int: The current number of optimizer updates completed
+        """
+        return self._num_updates
+
+    def get_current_phase_info(self) -> dict:
+        """Get phase information for the current step.
+        
+        Returns:
+            dict: Phase information containing period_start, period_end, 
+                decay_steps, and pre_decay_step for current position
+        """
+        return self.get_phase_info(self._num_updates)
+
     @lru_cache(maxsize=1)
     def _get_checkpoint_info(self, num_updates):
         """Get information about the current checkpoint period."""