returnn 1.20231229.22441__tar.gz → 1.20240102.20330__tar.gz

{returnn-1.20231229.22441/returnn.egg-info → returnn-1.20240102.20330}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20231229.22441
+Version: 1.20240102.20330
 Summary: The RWTH extensible training framework for universal recurrent neural networks
 Home-page: https://github.com/rwth-i6/returnn/
 Author: Albert Zeyer

returnn-1.20240102.20330/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20240102.020330'
+long_version = '1.20240102.020330+git.90ab8f0'

{returnn-1.20231229.22441 → returnn-1.20240102.20330}/returnn/torch/engine.py

@@ -221,14 +221,39 @@ class Engine(EngineBase):
         print(
             f"Starting training at epoch {self._start_epoch}, global train step {self.global_train_step}", file=log.v3
         )
-        self.epoch = self._start_epoch - 1
+        start_epoch = self._start_epoch
+        self.epoch = start_epoch - 1
         self._check_epoch_missing_eval()
         while self.epoch + 1 <= self._final_epoch:
             self.epoch += 1
             self._epoch_mp_shared.value = self.epoch
 
-            self.init_train_epoch()
-            self.train_epoch()
+            epoch_start_global_train_step = self.global_train_step
+            epoch_num_tries = 0
+            while True:
+                epoch_num_tries += 1
+                if epoch_num_tries > 1:
+                    print(f"Retry of train epoch {self.epoch}, global train step {self.global_train_step}", file=log.v3)
+                self.init_train_epoch()
+                try:
+                    self.train_epoch()
+                except torch.cuda.OutOfMemoryError as exc:
+                    print("torch.cuda.OutOfMemoryError:", exc, file=log.v1)
+                    # Some heuristics when to retry.
+                    if (
+                        epoch_num_tries <= 3  # heuristic
+                        and self.epoch > start_epoch  # heuristic
+                        and self.global_train_step > epoch_start_global_train_step  # heuristic
+                        and self._save_model_epoch_interval == 1  # cannot load most recent epoch
+                        and not self._torch_distributed_ctx  # need to fix correct sync, all workers need to restart ep
+                    ):
+                        print("Retry after OOM", file=log.v3)
+                        self._maybe_reset_dev_memory_caches(force=True)
+                        self._reset_train_epoch()
+                        assert self.global_train_step == epoch_start_global_train_step
+                        continue
+                    raise
+                break
 
         print(f"Finished training at epoch {self.epoch}, global train step {self.global_train_step}", file=log.v3)
 
@@ -249,8 +274,14 @@ class Engine(EngineBase):
             }
         )
 
-    def _maybe_reset_dev_memory_caches(self):
-        if not self._reset_dev_memory_caches:
+    def _reset_train_epoch(self):
+        epoch = self.epoch
+        self._load_model()
+        self._load_optimizer()
+        assert epoch == self.epoch
+
+    def _maybe_reset_dev_memory_caches(self, *, force: bool = False):
+        if not force and not self._reset_dev_memory_caches:
             return
         gc.collect()
         torch.cuda.empty_cache()
@@ -357,7 +388,8 @@ class Engine(EngineBase):
             if self._torch_distributed_ctx:
                 self._torch_distributed_ctx.step_after_param_update(module=self._pt_model, epoch_step_idx=step_idx)
 
-            elapsed_computation_time += time.time() - step_begin_time
+            step_duration = time.time() - step_begin_time
+            elapsed_computation_time += step_duration
 
             accumulated_losses_dict += losses_dict
             accumulated_inv_norm_factors_dict += inv_norm_factors_dict
@@ -365,6 +397,7 @@ class Engine(EngineBase):
                 f"ep {self.epoch} train",
                 step=step_idx,
                 eval_info=dict(losses_dict / inv_norm_factors_dict),
+                step_duration=step_duration,
                 log_memory_usage_device=self._device if self._log_memory_usage else None,
             )
 
@@ -969,6 +1002,7 @@ def _print_process(
     report_prefix: str,
     step: int,
     eval_info: Optional[Dict[str, Any]] = None,
+    step_duration: Optional[float] = None,
     log_memory_usage_device: Optional[str] = None,
 ):
     """
@@ -977,6 +1011,7 @@ def _print_process(
     :param report_prefix:
     :param step:
     :param eval_info:
+    :param step_duration:
     :param log_memory_usage_device: if given, will log memory usage (peak allocated memory)
     :return: nothing, will be printed to log
     """
@@ -990,6 +1025,8 @@ def _print_process(
                 info += [
                     f"mem_usage:{log_memory_usage_device} {util.human_bytes_size(torch.cuda.max_memory_allocated(dev))}"
                 ]
+        if step_duration is not None:
+            info += ["%.3f sec/step" % step_duration]
         print(", ".join(filter(None, info)), file=log.v5)
 
 

{returnn-1.20231229.22441 → returnn-1.20240102.20330/returnn.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20231229.22441
+Version: 1.20240102.20330
 Summary: The RWTH extensible training framework for universal recurrent neural networks
 Home-page: https://github.com/rwth-i6/returnn/
 Author: Albert Zeyer

returnn-1.20231229.22441/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20231229.022441'
-long_version = '1.20231229.022441+git.b5031e3'