Repository: Lightning-AI/pytorch-lightning

Issue:
GPU memory show max/min only
Instead of logging GPU mem for all GPUs, probably better to just show the min/max usage? Maybe the flag can be modified?

```python
# now
Trainer(log_gpu_memory=False)  

# proposed
Trainer(log_gpu_memory='min_max')  

# and for all GPUs
Trainer(log_gpu_memory='all')  

```

@adefazio


Discussion:
Should it be by default `min_max` or empty string or `False`? 
default = None

Resolving patch:
diff --git a/pytorch_lightning/root_module/memory.py b/pytorch_lightning/root_module/memory.py
--- a/pytorch_lightning/root_module/memory.py
+++ b/pytorch_lightning/root_module/memory.py
@@ -178,6 +178,33 @@ def count_mem_items():  # pragma: no cover
     return nb_params, nb_tensors
 
 
+def get_memory_profile(mode):
+    """
+    'all' means return memory for all gpus
+    'min_max' means return memory for max and min
+    :param mode:
+    :return:
+    """
+    memory_map = get_gpu_memory_map()
+
+    if mode == 'min_max':
+        min_mem = 1000000
+        min_k = None
+        max_mem = 0
+        max_k = None
+        for k, v in memory_map:
+            if v > max_mem:
+                max_mem = v
+                max_k = k
+            if v < min_mem:
+                min_mem = v
+                min_k = k
+
+        memory_map = {min_k: min_mem, max_k: max_mem}
+
+    return memory_map
+
+
 def get_gpu_memory_map():
     """Get the current gpu usage.
 
@@ -196,6 +223,6 @@ def get_gpu_memory_map():
     gpu_memory = [int(x) for x in result.strip().split('\n')]
     gpu_memory_map = {}
     for k, v in zip(range(len(gpu_memory)), gpu_memory):
-        k = 'gpu_%i' % k
+        k = f'gpu_{k}'
         gpu_memory_map[k] = v
     return gpu_memory_map
diff --git a/pytorch_lightning/root_module/root_module.py b/pytorch_lightning/root_module/root_module.py
--- a/pytorch_lightning/root_module/root_module.py
+++ b/pytorch_lightning/root_module/root_module.py
@@ -92,16 +92,20 @@ def configure_optimizers(self):
         """
         raise NotImplementedError
 
-    def optimizer_step(self, epoch_nb, batch_nb, optimizer, optimizer_i):
+    def optimizer_step(self, epoch_nb, batch_nb, optimizer, optimizer_i, second_order_closure=None):
         """
         Do something instead of the standard optimizer behavior
         :param epoch_nb:
         :param batch_nb:
         :param optimizer:
         :param optimizer_i:
+        :param second_order_closure: closure for second order methods
         :return:
         """
-        optimizer.step()
+        if isinstance(optimizer, torch.optim.LBFGS):
+            optimizer.step(second_order_closure)
+        else:
+            optimizer.step()
 
         # clear gradients
         optimizer.zero_grad()
diff --git a/pytorch_lightning/trainer/trainer.py b/pytorch_lightning/trainer/trainer.py
--- a/pytorch_lightning/trainer/trainer.py
+++ b/pytorch_lightning/trainer/trainer.py
@@ -15,7 +15,7 @@
 from torch.optim.optimizer import Optimizer
 
 from pytorch_lightning.root_module.root_module import LightningModule
-from pytorch_lightning.root_module.memory import get_gpu_memory_map
+from pytorch_lightning.root_module import memory
 from pytorch_lightning.logging import TestTubeLogger
 from pytorch_lightning.trainer.trainer_io import TrainerIO
 from pytorch_lightning.pt_overrides.override_data_parallel import (
@@ -66,7 +66,7 @@ def __init__(self,
                  process_position=0,
                  nb_gpu_nodes=1,
                  gpus=None,
-                 log_gpu_memory=False,
+                 log_gpu_memory=None,
                  show_progress_bar=True,
                  overfit_pct=0.0,
                  track_grad_norm=-1,
@@ -98,7 +98,7 @@ def __init__(self,
         :param process_position: shown in the tqdm bar
         :param nb_gpu_nodes: number of GPU nodes
         :param gpus: int. (ie: 2 gpus) OR list to specify which GPUs [0, 1] or '0,1'
-        :param log_gpu_memory: Bool. If true, adds memory logs
+        :param log_gpu_memory: str. None, 'min_max', 'all'
         :param show_progress_bar: Bool. If true shows tqdm bar
         :param overfit_pct: float. uses this much of all datasets
         :param track_grad_norm: int. -1 no tracking. Otherwise tracks that norm
@@ -275,6 +275,8 @@ def __init_amp(self, use_amp):
             raise ModuleNotFoundError(msg)
 
     def __configure_accumulated_gradients(self, accumulate_grad_batches):
+        self.accumulate_grad_batches = None
+
         if isinstance(accumulate_grad_batches, dict):
             self.accumulation_scheduler = GradientAccumulationScheduler(accumulate_grad_batches)
         elif isinstance(accumulate_grad_batches, int):
@@ -704,14 +706,6 @@ def fit(self, model):
             else:
                 nb_gpus = self.nb_requested_gpus
                 nb_tasks = self.nb_slurm_tasks
-                msg = f"""
-                You requested {nb_gpus}s GPUs but launched {nb_tasks}s slurm tasks.
-                We will launch {nb_gpus}s processes for you.
-                We recommend you let slurm manage the processes by setting:
-                --ntasks-per-node={nb_gpus}s
-                If you're not using SLURM, ignore this message!
-                """
-                warnings.warn(msg)
                 mp.spawn(self.ddp_train, nprocs=self.num_gpus, args=(model, ))
 
         # 1 gpu or dp option triggers training using DP module
@@ -1086,8 +1080,8 @@ def run_training_epoch(self):
                 metrics = self.__training_tqdm_dict
 
                 # add gpu memory
-                if self.on_gpu and self.log_gpu_memory:
-                    mem_map = get_gpu_memory_map()
+                if self.on_gpu and self.log_gpu_memory is not None:
+                    mem_map = memory.get_memory_profile(mode=self.log_gpu_memory)
                     metrics.update(mem_map)
 
                 # add norms
@@ -1267,27 +1261,35 @@ def __run_training_batch(self, batch, batch_nb):
         # call training_step once per optimizer
         for opt_idx, optimizer in enumerate(self.optimizers):
 
-            # forward pass
-            loss, model_specific_tqdm_metrics = self.__training_forward(batch, batch_nb, opt_idx)
+            # wrap the forward step in a closure so second order methods work
+            def optimizer_closure():
+                # forward pass
+                output = self.__training_forward(batch, batch_nb, opt_idx)
+                closure_loss, model_specific_tqdm_metrics = output
 
-            # track metrics
-            self.__add_tqdm_metrics(model_specific_tqdm_metrics)
+                # track metrics
+                self.__add_tqdm_metrics(model_specific_tqdm_metrics)
 
-            # accumulate loss
-            # (if accumulate_grad_batches = 1 no effect)
-            loss = loss / self.accumulate_grad_batches
+                # accumulate loss
+                # (if accumulate_grad_batches = 1 no effect)
+                closure_loss = closure_loss / self.accumulate_grad_batches
 
-            # backward pass
-            if self.use_amp:
-                with amp.scale_loss(loss, optimizer) as scaled_loss:
-                    scaled_loss.backward()
-            else:
-                loss.backward()
+                # backward pass
+                if self.use_amp:
+                    with amp.scale_loss(closure_loss, optimizer) as scaled_loss:
+                        scaled_loss.backward()
+                else:
+                    closure_loss.backward()
+
+                # insert after step hook
+                if self.__is_function_implemented('on_after_backward'):
+                    model_ref = self.__get_model()
+                    model_ref.on_after_backward()
+
+                return closure_loss
 
-            # insert after step hook
-            if self.__is_function_implemented('on_after_backward'):
-                model_ref = self.__get_model()
-                model_ref.on_after_backward()
+            # calculate loss
+            loss = optimizer_closure()
 
             # nan grads
             if self.print_nan_grads:
@@ -1311,7 +1313,8 @@ def __run_training_batch(self, batch, batch_nb):
                 # calls .step(), .zero_grad()
                 # override function to modify this behavior
                 model = self.__get_model()
-                model.optimizer_step(self.current_epoch, batch_nb, optimizer, opt_idx)
+                model.optimizer_step(self.current_epoch, batch_nb,
+                                     optimizer, opt_idx, optimizer_closure)
 
                 # calculate running loss for display
                 self.running_loss.append(self.batch_loss_value)


Test patch:
diff --git a/pytorch_lightning/testing/lm_test_module_base.py b/pytorch_lightning/testing/lm_test_module_base.py
--- a/pytorch_lightning/testing/lm_test_module_base.py
+++ b/pytorch_lightning/testing/lm_test_module_base.py
@@ -115,11 +115,14 @@ def training_step(self, batch, batch_idx):
     # ---------------------
     def configure_optimizers(self):
         """
-        return whatever optimizers we want here
+        return whatever optimizers we want here.
         :return: list of optimizers
         """
         # try no scheduler for this model (testing purposes)
-        optimizer = optim.Adam(self.parameters(), lr=self.hparams.learning_rate)
+        if self.hparams.optimizer_name == 'lbfgs':
+            optimizer = optim.LBFGS(self.parameters(), lr=self.hparams.learning_rate)
+        else:
+            optimizer = optim.Adam(self.parameters(), lr=self.hparams.learning_rate)
 
         # test returning only 1 list instead of 2
         return optimizer
diff --git a/tests/test_logging.py b/tests/test_logging.py
--- a/tests/test_logging.py
+++ b/tests/test_logging.py
@@ -3,15 +3,20 @@
 import shutil
 
 import numpy as np
+import torch
 from pytorch_lightning import Trainer
 
 from pytorch_lightning.testing import LightningTestModel
-
 from .test_models import get_hparams, get_test_tube_logger, init_save_dir, clear_save_dir
 
+RANDOM_SEEDS = list(np.random.randint(0, 10000, 1000))
+
 
 def test_testtube_logger():
-    """verify that basic functionality of test tube logger works"""
+    """
+    verify that basic functionality of test tube logger works
+    """
+    reset_seed()
 
     hparams = get_hparams()
     model = LightningTestModel(hparams)
@@ -35,7 +40,11 @@ def test_testtube_logger():
 
 
 def test_testtube_pickle():
-    """Verify that pickling a trainer containing a test tube logger works"""
+    """
+    Verify that pickling a trainer containing a test tube logger works
+    """
+    reset_seed()
+
     hparams = get_hparams()
     model = LightningTestModel(hparams)
 
@@ -58,7 +67,11 @@ def test_testtube_pickle():
 
 
 def test_mlflow_logger():
-    """verify that basic functionality of mlflow logger works"""
+    """
+    verify that basic functionality of mlflow logger works
+    """
+    reset_seed()
+
     try:
         from pytorch_lightning.logging import MLFlowLogger
     except ModuleNotFoundError:
@@ -90,7 +103,11 @@ def test_mlflow_logger():
 
 
 def test_mlflow_pickle():
-    """verify that pickling trainer with mlflow logger works"""
+    """
+    verify that pickling trainer with mlflow logger works
+    """
+    reset_seed()
+
     try:
         from pytorch_lightning.logging import MLFlowLogger
     except ModuleNotFoundError:
@@ -115,3 +132,9 @@ def test_mlflow_pickle():
     pkl_bytes = pickle.dumps(trainer)
     trainer2 = pickle.loads(pkl_bytes)
     trainer2.logger.log_metrics({"acc": 1.0})
+
+
+def reset_seed():
+    SEED = RANDOM_SEEDS.pop()
+    torch.manual_seed(SEED)
+    np.random.seed(SEED)
diff --git a/tests/test_models.py b/tests/test_models.py
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -31,9 +31,11 @@
 from pytorch_lightning.logging import TestTubeLogger
 from examples import LightningTemplateModel
 
-SEED = 2334
-torch.manual_seed(SEED)
-np.random.seed(SEED)
+# generate a list of random seeds for each test
+ROOT_SEED = 1234
+torch.manual_seed(ROOT_SEED)
+np.random.seed(ROOT_SEED)
+RANDOM_SEEDS = list(np.random.randint(0, 10000, 1000))
 
 
 # ------------------------------------------------------------------------
@@ -44,6 +46,7 @@ def test_default_logger_callbacks_cpu_model():
     Test each of the trainer options
     :return:
     """
+    reset_seed()
 
     trainer_options = dict(
         max_nb_epochs=1,
@@ -63,6 +66,31 @@ def test_default_logger_callbacks_cpu_model():
     model.unfreeze()
 
 
+def test_lbfgs_cpu_model():
+    """
+    Test each of the trainer options
+    :return:
+    """
+    reset_seed()
+
+    trainer_options = dict(
+        max_nb_epochs=1,
+        gradient_clip_val=1.0,
+        overfit_pct=0.20,
+        print_nan_grads=True,
+        show_progress_bar=False,
+        train_percent_check=0.2,
+        val_percent_check=0.2
+    )
+
+    model, hparams = get_model(use_test_model=True, lbfgs=True)
+    run_model_test_no_loggers(trainer_options, model, hparams, on_gpu=False)
+
+    # test freeze on cpu
+    model.freeze()
+    model.unfreeze()
+
+
 def test_multi_gpu_model_ddp2():
     """
     Make sure DDP2 works
@@ -72,6 +100,9 @@ def test_multi_gpu_model_ddp2():
         return
 
     os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0])
+
+    reset_seed()
+
     model, hparams = get_model()
     trainer_options = dict(
         show_progress_bar=True,
@@ -94,6 +125,8 @@ def test_dp_resume():
     if not can_run_gpu_test():
         return
 
+    reset_seed()
+
     hparams = get_hparams()
     model = LightningTestModel(hparams)
 
@@ -174,6 +207,8 @@ def test_running_test_pretrained_model_ddp():
     if not can_run_gpu_test():
         return
 
+    reset_seed()
+
     hparams = get_hparams()
     model = LightningTestModel(hparams)
 
@@ -219,6 +254,8 @@ def test_running_test_pretrained_model_ddp():
 
 def test_running_test_after_fitting():
     """Verify test() on fitted model"""
+    reset_seed()
+
     hparams = get_hparams()
     model = LightningTestModel(hparams)
 
@@ -257,6 +294,8 @@ def test_running_test_after_fitting():
 
 
 def test_running_test_without_val():
+    reset_seed()
+
     """Verify test() works on a model with no val_loader"""
     class CurrentTestModel(LightningTestMixin, LightningTestModelBase):
         pass
@@ -298,6 +337,8 @@ class CurrentTestModel(LightningTestMixin, LightningTestModelBase):
 
 
 def test_running_test_pretrained_model():
+    reset_seed()
+
     """Verify test() on pretrained model"""
     hparams = get_hparams()
     model = LightningTestModel(hparams)
@@ -340,6 +381,8 @@ def test_running_test_pretrained_model():
 
 
 def test_running_test_pretrained_model_dp():
+    reset_seed()
+
     """Verify test() on pretrained model"""
     if not can_run_gpu_test():
         return
@@ -386,6 +429,8 @@ def test_running_test_pretrained_model_dp():
 
 
 def test_gradient_accumulation_scheduling():
+    reset_seed()
+
     """
     Test grad accumulation by the freq of optimizer updates
     """
@@ -466,6 +511,8 @@ def test_multi_gpu_model_ddp():
         return
 
     os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0])
+    reset_seed()
+
     model, hparams = get_model()
     trainer_options = dict(
         show_progress_bar=False,
@@ -480,6 +527,7 @@ def test_multi_gpu_model_ddp():
 
 
 def test_optimizer_return_options():
+    reset_seed()
 
     trainer = Trainer()
     model, hparams = get_model()
@@ -510,6 +558,8 @@ def test_optimizer_return_options():
 
 
 def test_single_gpu_batch_parse():
+    reset_seed()
+
     if not can_run_gpu_test():
         return
 
@@ -556,6 +606,7 @@ def test_early_stopping_cpu_model():
     Test each of the trainer options
     :return:
     """
+    reset_seed()
 
     stopping = EarlyStopping(monitor='val_loss')
     trainer_options = dict(
@@ -583,6 +634,8 @@ def test_no_val_module():
     Tests use case where trainer saves the model, and user loads it from tags independently
     :return:
     """
+    reset_seed()
+
     hparams = get_hparams()
 
     class CurrentTestModel(LightningTestModelBase):
@@ -629,6 +682,7 @@ def test_no_val_end_module():
     Tests use case where trainer saves the model, and user loads it from tags independently
     :return:
     """
+    reset_seed()
 
     class CurrentTestModel(LightningValidationStepMixin, LightningTestModelBase):
         pass
@@ -675,6 +729,8 @@ def test_simple_cpu():
     Verify continue training session on CPU
     :return:
     """
+    reset_seed()
+
     hparams = get_hparams()
     model = LightningTestModel(hparams)
 
@@ -702,6 +758,8 @@ def test_amp_single_gpu():
     Make sure DDP + AMP work
     :return:
     """
+    reset_seed()
+
     if not torch.cuda.is_available():
         warnings.warn('test_amp_gpu_ddp cannot run.'
                       'Rerun on a GPU node to run this test')
@@ -730,6 +788,8 @@ def test_no_amp_single_gpu():
     Make sure DDP + AMP work
     :return:
     """
+    reset_seed()
+
     if not torch.cuda.is_available():
         warnings.warn('test_amp_gpu_ddp cannot run.'
                       'Rerun on a GPU node to run this test')
@@ -759,6 +819,8 @@ def test_cpu_restore_training():
     Verify continue training session on CPU
     :return:
     """
+    reset_seed()
+
     hparams = get_hparams()
     model = LightningTestModel(hparams)
 
@@ -831,6 +893,8 @@ def test_amp_gpu_ddp():
 
     os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0])
 
+    reset_seed()
+
     hparams = get_hparams()
     model = LightningTestModel(hparams)
 
@@ -850,6 +914,8 @@ def test_cpu_slurm_save_load():
     Verify model save/load/checkpoint on CPU
     :return:
     """
+    reset_seed()
+
     hparams = get_hparams()
     model = LightningTestModel(hparams)
 
@@ -925,6 +991,8 @@ def assert_pred_same():
 
 
 def test_loading_meta_tags():
+    reset_seed()
+
     from argparse import Namespace
     hparams = get_hparams()
 
@@ -946,6 +1014,7 @@ def test_loading_meta_tags():
 
 
 def test_dp_output_reduce():
+    reset_seed()
 
     # test identity when we have a single gpu
     out = torch.rand(3, 1)
@@ -971,6 +1040,8 @@ def test_model_saving_loading():
     Tests use case where trainer saves the model, and user loads it from tags independently
     :return:
     """
+    reset_seed()
+
     hparams = get_hparams()
     model = LightningTestModel(hparams)
 
@@ -1026,6 +1097,8 @@ def test_model_saving_loading():
 
 
 def test_model_freeze_unfreeze():
+    reset_seed()
+
     hparams = get_hparams()
     model = LightningTestModel(hparams)
 
@@ -1045,6 +1118,8 @@ def test_amp_gpu_ddp_slurm_managed():
     os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0])
     os.environ['SLURM_LOCALID'] = str(0)
 
+    reset_seed()
+
     hparams = get_hparams()
     model = LightningTestModel(hparams)
 
@@ -1111,6 +1186,7 @@ def test_cpu_model_with_amp():
     Make sure model trains on CPU
     :return:
     """
+    reset_seed()
 
     trainer_options = dict(
         show_progress_bar=False,
@@ -1132,6 +1208,7 @@ def test_cpu_model():
     Make sure model trains on CPU
     :return:
     """
+    reset_seed()
 
     trainer_options = dict(
         show_progress_bar=False,
@@ -1151,6 +1228,7 @@ def test_all_features_cpu_model():
     Test each of the trainer options
     :return:
     """
+    reset_seed()
 
     trainer_options = dict(
         gradient_clip_val=1.0,
@@ -1174,6 +1252,8 @@ def test_single_gpu_model():
     Make sure single GPU works (DP mode)
     :return:
     """
+    reset_seed()
+
     if not torch.cuda.is_available():
         warnings.warn('test_single_gpu_model cannot run.'
                       ' Rerun on a GPU node to run this test')
@@ -1197,6 +1277,8 @@ def test_multi_gpu_none_backend():
     distributed_backend = None
     :return:
     """
+    reset_seed()
+
     if not can_run_gpu_test():
         return
 
@@ -1218,6 +1300,8 @@ def test_multi_gpu_model_dp():
     Make sure DP works
     :return:
     """
+    reset_seed()
+
     if not can_run_gpu_test():
         return
 
@@ -1242,6 +1326,8 @@ def test_amp_gpu_dp():
     Make sure DP + AMP work
     :return:
     """
+    reset_seed()
+
     if not can_run_gpu_test():
         return
 
@@ -1266,6 +1352,8 @@ def test_ddp_sampler_error():
 
     os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0])
 
+    reset_seed()
+
     hparams = get_hparams()
     model = LightningTestModel(hparams, force_remove_distributed_sampler=True)
 
@@ -1292,6 +1380,8 @@ def test_multiple_val_dataloader():
     Verify multiple val_dataloader
     :return:
     """
+    reset_seed()
+
     class CurrentTestModel(
         LightningValidationMultipleDataloadersMixin,
         LightningTestModelBase
@@ -1327,6 +1417,8 @@ def test_multiple_test_dataloader():
     Verify multiple test_dataloader
     :return:
     """
+    reset_seed()
+
     class CurrentTestModel(
         LightningTestMultipleDataloadersMixin,
         LightningTestModelBase
@@ -1447,9 +1539,11 @@ def get_hparams(continue_training=False, hpc_exp_number=0):
     return hparams
 
 
-def get_model(use_test_model=False):
+def get_model(use_test_model=False, lbfgs=False):
     # set up model with these hyperparams
     hparams = get_hparams()
+    if lbfgs:
+        setattr(hparams, 'optimizer_name', 'lbfgs')
 
     if use_test_model:
         model = LightningTestModel(hparams)
@@ -1554,5 +1648,11 @@ def can_run_gpu_test():
     return True
 
 
+def reset_seed():
+    SEED = RANDOM_SEEDS.pop()
+    torch.manual_seed(SEED)
+    np.random.seed(SEED)
+
+
 if __name__ == '__main__':
     pytest.main([__file__])