Repository: rwth-i6/returnn

Issue:
CUDA out of memory issue when resuming Torch model training
When running with multi-gpu, once the training job is interrupted, when restarting, when it loads the existing checkpoint, I always get` OutOfMemoryError: CUDA out of memory` at the first training step (the same GPUs are used). But such error does not happen when I run training with single GPU.

The possible reason could be that with multi-GPU training, we save the checkpoint when `rank==0`. Thus when loading the checkpoint, instead of loading the checkpoints to each device respectively, they are all loaded to device 0. 



Discussion:

Resolving patch:
diff --git a/returnn/torch/engine.py b/returnn/torch/engine.py
--- a/returnn/torch/engine.py
+++ b/returnn/torch/engine.py
@@ -6,6 +6,7 @@
 from typing import Optional, Union, Callable, Dict
 from contextlib import nullcontext
 
+import gc
 import os
 import numpy
 import torch
@@ -187,7 +188,9 @@ def init_train_from_config(
             self._ddp_pt_model = self._torch_distributed_class(
                 self._pt_model, device_ids=get_device_ids(), **self._torch_distributed_options
             )
-        self._updater = Updater(self.config, self._pt_model, self.learning_rate)
+        self._updater = Updater(
+            config=self.config, network=self._pt_model, device=self._device, initial_learning_rate=self.learning_rate
+        )
         self._updater.create_optimizer()
         if self._start_epoch > 1:
             self._load_optimizer()
@@ -481,7 +484,7 @@ def _load_model(self):
         if model_epoch_filename:
             filename = model_epoch_filename + util.get_model_filename_postfix()
             print("Load model %s" % (filename,), file=log.v4)
-            checkpoint_state = torch.load(filename)
+            checkpoint_state = torch.load(filename, map_location=self._device)
             if epoch is None:
                 epoch = checkpoint_state["epoch"]
             step = checkpoint_state["step"]
@@ -577,6 +580,9 @@ def _load_model(self):
                     assert not missing_prefix_keys, f"Missing keys and ignore_missing=False: {missing_prefix_keys}"
                 print(f"Missing keys: {missing_keys}", file=log.v4)
 
+        # https://github.com/rwth-i6/returnn/issues/1345
+        del checkpoint_state
+        gc.collect()
         self._pt_model.to(self._device)
 
         if model_epoch_filename and is_training:
diff --git a/returnn/torch/updater.py b/returnn/torch/updater.py
--- a/returnn/torch/updater.py
+++ b/returnn/torch/updater.py
@@ -5,6 +5,7 @@
 
 from __future__ import annotations
 
+import gc
 import torch
 import typing
 from typing import Any, Set, Dict
@@ -80,7 +81,7 @@ class Updater(object):
     Wraps a torch.optim.Optimizer, and extends it by some further functionality.
     """
 
-    def __init__(self, config, network, initial_learning_rate=1.0):
+    def __init__(self, *, config, network, device, initial_learning_rate=1.0):
         """
         :param returnn.config.Config config: config defining the training conditions.
         :param torch.nn.Module network: PyTorch Module defining the network.
@@ -89,6 +90,7 @@ def __init__(self, config, network, initial_learning_rate=1.0):
         self.config = config
         self.learning_rate = initial_learning_rate
         self.network = network
+        self._device = device
         self.optimizer = None  # type: typing.Optional[torch.optim.Optimizer]
 
     def set_learning_rate(self, value):
@@ -122,8 +124,11 @@ def load_optimizer(self, filename):
         :param str filename: File from which to load the optimizer state.
         """
         print("Load optimizer %s" % filename, file=log.v4)
-        optimizer_state = torch.load(filename)
+        optimizer_state = torch.load(filename, map_location=self._device)
         self.optimizer.load_state_dict(optimizer_state)
+        # https://github.com/rwth-i6/returnn/issues/1345
+        del optimizer_state
+        gc.collect()
 
     def save_optimizer(self, filename):
         """


Test patch: