updates

Author: Chris-hughes10

pytorch_accelerated/schedulers/wsd_scheduler.py

@@ -359,7 +359,7 @@ def load_state_dict(self, state_dict: dict):
     def create_scheduler_fn(
         cls,
         total_num_epochs: int = TrainerPlaceholderValues.NUM_EPOCHS,
-        num_update_steps_per_epoch: int = TrainerPlaceholderValues.NUM_LOCAL_UPDATE_STEPS_PER_EPOCH,
+        num_update_steps_per_epoch: int = TrainerPlaceholderValues.PER_PROCESS_NUM_UPDATE_STEPS_PER_EPOCH,
         num_warmup_epochs: int = None,
         decay_phase_ratio: float = 0.1,
         lr_min: float = 1e-6,

pytorch_accelerated/trainer.py

@@ -808,6 +808,7 @@ def _run_train_epoch(self, train_dl):
             self,
         )
 
+        # updates across all processes
         max_total_update_steps = self.run_config.max_num_train_steps
 
         updates_completed = (
@@ -849,9 +850,11 @@ def _run_train_epoch(self, train_dl):
                     + (step + 1) // self.run_config.gradient_accumulation_steps
                 )
 
+                global_process_updates = process_updates * self._accelerator.num_processes
+
                 if (
                     self.run_config.max_num_train_steps is not None
-                    and process_updates >= max_total_update_steps
+                    and global_process_updates >= max_total_update_steps
                 ):
                     reached_max_steps = True
                     # Synchronize reached_max_steps across processes
@@ -860,7 +863,7 @@ def _run_train_epoch(self, train_dl):
                             [reached_max_steps], device=self.device
                         )
                         reached_max_steps = (
-                            self.gather(reached_max_steps_tensor).all().item()
+                            self.gather(reached_max_steps_tensor).any().item()
                         )
                     break
 

test/test_trainer.py

@@ -1,5 +1,3 @@
-from pathlib import Path
-from tempfile import TemporaryFile, TemporaryDirectory
 from unittest.mock import MagicMock, Mock, call
 import pytest
 
@@ -402,6 +400,8 @@ def test_max_steps_stops_training_correctly():
     )
 
 
+
+
 def test_max_steps_overrides_num_epochs():
     """Test that max_num_train_steps takes precedence over num_epochs"""