Repository: Lightning-AI/pytorch-lightning

Issue:
training_epoch_end log output gets combined with next epoch training
## 🐛 Bug

So, I put 'training_epoch_end' function in my LightningModule. I have it return this dictionary
`{'log':{'train_loss': tensor(0.3616, device='cuda:0'), 'epoch': 0}`

I check the `run_training_epoch_end` function in the Pytorch library, it looks like it is working normally as `log_epoch_metrics` is showing the 'log' part in the dictionary produced by 'training_epoch_end' function
`{'train_loss': tensor(0.3616, device='cuda:0'), 'epoch': 0}`

So, they send it off to the logger. But there is problem, it is trying to combine the dictionary above with the results from the training step fo the next epoch. When I check the variable `self._metrics_to_agg`, I get the following result. Of course, it is impossible to combine these dictionaries as they have different keys. I guess the main problem is that the code is combining the log results of `run_training_epoch_end` function with the results of the next training batch.
`[{'train_loss': tensor(0.3616, device='cuda:0'), 'epoch': 0}, {'loss': 0.48756, 'input': ..., 'ouput': ...}]`

Any ideas to solve this problem? I will appreciate your help! Here is the whole error stack:
```
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-5-a4569c0ab5cd> in <module>
      3 logging.getLogger('allennlp').setLevel(logging.INFO)
      4 
----> 5 train("/shared/labs/workflow/configs/ner_span_path_2.jsonnet", "/Dropbox/EvidLabs/atee/lab_data_enrichment_project/train/DataEnrichmentSpanPath/DataEnrichmentSpanPath__INV_M_span_path_42_43", AllenNlpLightningModule, tags=[])
      6 
      7 # train("/shared/labs/workflow/configs/candidate_span_grouper.jsonnet", "./DataEnrichmentBunchedCrfTaggerTest", AllenNlpLightningModule, tags=["Frozen"])

/code/evid-research/evid2/pytorch_lightning/trainer/config_file_trainer.py in train(config_file, project_dir, lightning_module, run_id, run_name_override, tags, dry_run)
    171     # 3 START TRAINING
    172     # ------------------------
--> 173     trainer.fit(model)
    174 
    175 

/code/pytorch-lightning/pytorch_lightning/trainer/trainer.py in fit(self, model, train_dataloader, val_dataloaders)
    971         # easier to avoid NCCL issues
    972         elif self.use_dp:
--> 973             self.dp_train(model)
    974 
    975         elif self.use_horovod:

/code/pytorch-lightning/pytorch_lightning/trainer/distrib_parts.py in dp_train(self, model)
    265         model = LightningDataParallel(model, device_ids=device_ids)
    266 
--> 267         self.run_pretrain_routine(model)
    268 
    269         model.forward = model_autocast_original_forward

/code/pytorch-lightning/pytorch_lightning/trainer/trainer.py in run_pretrain_routine(self, model)
   1154 
   1155         # CORE TRAINING LOOP
-> 1156         self.train()
   1157 
   1158     def test(

/code/pytorch-lightning/pytorch_lightning/trainer/training_loop.py in train(self)
    368                 # RUN TNG EPOCH
    369                 # -----------------
--> 370                 self.run_training_epoch()
    371 
    372                 if self.max_steps and self.max_steps <= self.global_step:

/code/pytorch-lightning/pytorch_lightning/trainer/training_loop.py in run_training_epoch(self)
    480             # SAVE METRICS TO LOGGERS
    481             # -----------------------------------------
--> 482             self.save_train_loop_metrics_to_loggers(batch_idx, batch_output)
    483 
    484             # progress global step according to grads progress

/code/pytorch-lightning/pytorch_lightning/trainer/training_loop.py in save_train_loop_metrics_to_loggers(self, batch_idx, batch_output)
    564             # logs user requested information to logger
    565             print("batch_output.batch_log_metrics", batch_output.batch_log_metrics.keys())
--> 566             self.log_metrics(batch_output.batch_log_metrics, batch_output.grad_norm_dic)
    567 
    568     def save_loggers_in_training_loop(self, batch_idx):

/code/pytorch-lightning/pytorch_lightning/trainer/logging.py in log_metrics(self, metrics, grad_norm_dic, step)
     72         if self.is_global_zero and self.logger is not None:
     73             print("scalar_metrics", scalar_metrics.keys())
---> 74             self.logger.agg_and_log_metrics(scalar_metrics, step=step)
     75             self.logger.save()
     76 

/code/pytorch-lightning/pytorch_lightning/loggers/base.py in agg_and_log_metrics(self, metrics, step)
    131         """
    132         print("metrics", metrics.keys())
--> 133         agg_step, metrics_to_log = self._aggregate_metrics(metrics=metrics, step=step)
    134 
    135         if metrics_to_log:

/code/pytorch-lightning/pytorch_lightning/loggers/base.py in _aggregate_metrics(self, metrics, step)
     88 
     89         # compute the metrics
---> 90         agg_step, agg_mets = self._reduce_agg_metrics()
     91 
     92         # as new step received reset accumulator

/code/pytorch-lightning/pytorch_lightning/loggers/base.py in _reduce_agg_metrics(self)
    109             print("self._agg_key_funcs", self._agg_key_funcs)
    110             print("self._agg_default_func", self._agg_default_func)
--> 111             agg_mets = merge_dicts(self._metrics_to_agg, self._agg_key_funcs, self._agg_default_func)
    112         return self._prev_step, agg_mets
    113 

/code/pytorch-lightning/pytorch_lightning/loggers/base.py in merge_dicts(dicts, agg_key_funcs, default_func)
    386             print("fn or default_func", fn or default_func)
    387             print("values_to_agg", values_to_agg)
--> 388             d_out[k] = (fn or default_func)(values_to_agg)
    389 
    390     return d_out

<__array_function__ internals> in mean(*args, **kwargs)

/venvs/dev3.7/lib/python3.7/site-packages/numpy/core/fromnumeric.py in mean(a, axis, dtype, out, keepdims)
   3333 
   3334     return _methods._mean(a, axis=axis, dtype=dtype,
-> 3335                           out=out, **kwargs)
   3336 
   3337 

/venvs/dev3.7/lib/python3.7/site-packages/numpy/core/_methods.py in _mean(a, axis, dtype, out, keepdims)
    149             is_float16_result = True
    150 
--> 151     ret = umr_sum(arr, axis, dtype, out, keepdims)
    152     if isinstance(ret, mu.ndarray):
    153         ret = um.true_divide(

TypeError: unsupported operand type(s) for +: 'dict' and 'dict'
```


### To Reproduce

Steps to reproduce the behavior:

1. Add a training_epoch_end function to your Lightning Module and run it. You can use mine in the code sample section if you want. The key is that the "log"  section of training_epoch_end dictionary output must be of a different format than the dictionary that contains the results of your training sample while training (like in the example I provided `{'train_loss': tensor(0.3616, device='cuda:0'), 'epoch': 0}` and `{'loss': 0.48756, 'input': ..., 'ouput': ...}` are different formats as they don't share the same keys).


#### Code sample
```
 def training_epoch_end(self, outputs):
        print("training_epoch_end")
        print(len(outputs))

        prefix = "train_"
        metric_modules = self.training_metric_modules

        # Handle loss:
        avg_loss = torch.stack([x['loss'] for x in outputs]).mean()
        result = {'log': {}, 'progress_bar': {}}
        result[prefix + 'loss'] = avg_loss
        result['log'][prefix + 'loss'] = avg_loss
        result['progress_bar'][prefix + 'loss'] = avg_loss

        # Add tracking variables
        result['log']['epoch'] = self.current_epoch

        return result
```


### Expected behavior

I thought I would be able to run training_epoch_end function with no combination with the training samples in the next epoch. I expected no error, like running validation_epoch_end.

### Environment
* CUDA:
        - GPU:
                - Tesla V100-SXM2-16GB
        - available:         True
        - version:           10.2
* Packages:
        - numpy:             1.18.5
        - pyTorch_debug:     False
        - pyTorch_version:   1.5.0
        - pytorch-lightning: 0.8.4
        - tensorboard:       2.2.2
        - tqdm:              4.46.1
* System:
        - OS:                Linux
        - architecture:
                - 64bit
                -
        - processor:         x86_64
        - python:            3.7.7
        - version:           #113-Ubuntu SMP Wed Jan 29 14:54:54 UTC 2020


Discussion:
Hi! thanks for your contribution!, great first issue!
An update to the situation: I think I found the cause of the error. It seems that the `self.global_step` wasn't incrementing after processing `training_epoch_end`. Basically, the counter would increment for validation_epoch_end.
validation_epoch_end -> step 547

And then for `training_epoch_end`
training_epoch_end -> step 548

But it wouldn't increment as we progress to the next epoch
First batch of next epoch -> step 548

This triggered this piece of code from `_aggregate_metrics`, putting the log output of `training_epoch_end` and `first batch of next epoch` together as explained in the initial problem statement
```
        if step == self._prev_step:
            self._metrics_to_agg.append(metrics)
            return step, None
```

My suggestion to the solution is to add `self.global_step += 1` to the function `run_training_epoch_end` like so:

```
    def run_training_epoch_end(self, epoch_output):
        model = self.get_model()
        if self.is_overridden('training_epoch_end', model=model):
            .......
            self.global_step += 1
```

Let me know if you foresee any issues or problems with this solution. If you would like me to submit a pull request, I would be happy to do so. 

@ameliatqy please submit a PR. good catch!

Resolving patch:
diff --git a/pytorch_lightning/callbacks/early_stopping.py b/pytorch_lightning/callbacks/early_stopping.py
--- a/pytorch_lightning/callbacks/early_stopping.py
+++ b/pytorch_lightning/callbacks/early_stopping.py
@@ -9,6 +9,7 @@
 
 import numpy as np
 import torch
+import torch.distributed as dist
 
 from pytorch_lightning import _logger as log
 from pytorch_lightning.callbacks.base import Callback
@@ -16,6 +17,14 @@
 
 torch_inf = torch.tensor(np.Inf)
 
+try:
+    import torch_xla
+    import torch_xla.core.xla_model as xm
+except ImportError:
+    XLA_AVAILABLE = False
+else:
+    XLA_AVAILABLE = True
+
 
 class EarlyStopping(Callback):
     r"""
@@ -138,17 +147,38 @@ def _run_early_stopping_check(self, trainer, pl_module):
 
         current = logs.get(self.monitor)
         if not isinstance(current, torch.Tensor):
-            current = torch.tensor(current)
+            current = torch.tensor(current, device=pl_module.device)
 
-        if self.monitor_op(current - self.min_delta, self.best_score):
+        if self.monitor_op(current - self.min_delta, self.best_score.to(pl_module.device)):
             self.best_score = current
             self.wait_count = 0
         else:
             self.wait_count += 1
-            if self.wait_count >= self.patience:
+            should_stop = self.wait_count >= self.patience
+
+            if bool(should_stop):
                 self.stopped_epoch = trainer.current_epoch
                 trainer.should_stop = True
 
+        # stop every ddp process if any world process decides to stop
+        self._stop_distributed_training(trainer, pl_module)
+
+    def _stop_distributed_training(self, trainer, pl_module):
+
+        # in ddp make sure all processes stop when one is flagged
+        if trainer.use_ddp or trainer.use_ddp2:
+            stop = torch.tensor(int(trainer.should_stop), device=pl_module.device)
+            dist.all_reduce(stop, op=dist.reduce_op.SUM)
+            dist.barrier()
+            trainer.should_stop = stop == trainer.world_size
+
+        # if trainer.use_tpu:
+        #     stop = torch.tensor(int(trainer.should_stop), device=pl_module.device)
+        #     xm.all_reduce('sum', [stop])
+        #     print(type(stop))
+        #     torch_xla.core.xla_model.rendezvous("pl.EarlyStoppingCallback.stop_distributed_training_check")
+        #     trainer.should_stop = stop.item() == trainer.world_size
+
     def on_train_end(self, trainer, pl_module):
         if self.stopped_epoch > 0 and self.verbose > 0:
             rank_zero_warn('Displayed epoch numbers by `EarlyStopping` start from "1" until v0.6.x,'


Test patch:
diff --git a/tests/models/test_gpu.py b/tests/models/test_gpu.py
--- a/tests/models/test_gpu.py
+++ b/tests/models/test_gpu.py
@@ -58,6 +58,30 @@ def test_multi_gpu_model(tmpdir, backend):
     memory.get_memory_profile('min_max')
 
 
+@pytest.mark.spawn
+@pytest.mark.parametrize("backend", ['dp', 'ddp', 'ddp2'])
+@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
+def test_multi_gpu_early_stop(tmpdir, backend):
+    """Make sure DDP works. with early stopping"""
+    tutils.set_random_master_port()
+
+    trainer_options = dict(
+        default_root_dir=tmpdir,
+        early_stop_callback=True,
+        max_epochs=50,
+        limit_train_batches=10,
+        limit_val_batches=10,
+        gpus=[0, 1],
+        distributed_backend=backend,
+    )
+
+    model = EvalModelTemplate()
+    # tutils.run_model_test(trainer_options, model)
+    trainer = Trainer(**trainer_options)
+    result = trainer.fit(model)
+    assert result
+
+
 @pytest.mark.spawn
 @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="test requires multi-GPU machine")
 def test_ddp_all_dataloaders_passed_to_fit(tmpdir):
diff --git a/tests/models/test_tpu.py b/tests/models/test_tpu.py
--- a/tests/models/test_tpu.py
+++ b/tests/models/test_tpu.py
@@ -19,6 +19,27 @@
     TPU_AVAILABLE = True
 
 
+@pytest.mark.skipif(not TPU_AVAILABLE, reason="test requires TPU machine")
+@pytest.mark.parametrize(['tpu_cores', 'expected_device'], [
+    pytest.param([1], 'xla:1'),
+    pytest.param([8], 'xla:8'),
+])
+def test_early_stop_checkpoints_on_tpu(tmpdir, tpu_cores, expected_device):
+    """Test if single TPU core training works"""
+    model = EvalModelTemplate()
+    trainer = Trainer(
+        early_stop_callback=True,
+        default_root_dir=tmpdir,
+        progress_bar_refresh_rate=0,
+        max_epochs=50,
+        limit_train_batches=10,
+        limit_val_batches=10,
+        tpu_cores=tpu_cores,
+    )
+    trainer.fit(model)
+    assert torch_xla._XLAC._xla_get_default_device() == expected_device
+
+
 @pytest.mark.skipif(not TPU_AVAILABLE, reason="test requires TPU machine")
 @pytest.mark.parametrize(['tpu_cores', 'expected_device'], [
     pytest.param([1], 'xla:1'),