Repository: Lightning-AI/pytorch-lightning

Issue:
Batchsize and learning rate scheduler
I was wondering if I need to adjust the batch size when using TPUs. I had a memory error when trying to run a (image 256x256x3) batch of size 128, which works perfectly fine on GPUs. 

Furthermore, do I need to adjust my custom learning rate scheduler which on GPUs run every batch (not just epoch):
```
def configure_optimizers(self):
        optimizer = get_optimizer(self)
        one_cycle_scheduler = OneCycleScheduler(optimizer, EPOCHS * steps_per_epoch // GRAD_ACCUMULATE)
        scheduler = {'scheduler': one_cycle_scheduler, "interval": "step"}

        return [optimizer], [scheduler]
```


pl version: 0.8.5
system: kaggle


Discussion:
Hi! thanks for your contribution!, great first issue!
@sachinruk Can you list the parameters being passed to the trainer in both the cases?
```
early_stopping = EarlyStopping('val_loss')
trainer = Trainer(tpu_cores=8, max_epochs=1, val_check_interval=0.5, early_stop_callback=early_stopping, accumulate_grad_batches=GRAD_ACCUMULATE)
trainer.fit(model)
```
where `GRAD_ACCUMULATE = 1`.  Only difference is for gpus, I have `gpus=1` inside the `Trainer` instead. Let me know if there is anything else needed. Oh and its a TPUv3-8 on kaggle if that makes a difference.
@sachinruk I couldn't find any known bugs where the batch size doesn't translate directly between a GPU and TPU because ideally, a TPUv3-8 should have more memory. Is the memory error an OOM? Can you post the error message and code to replicate this?
I think setting batch_size=128 will actually be 128*8 here since it's running on all tpu-cores. I don't think TPUv3-8 can handle this.
https://pytorch-lightning.readthedocs.io/en/latest/trainer.html#tpu-cores
ok I've made the kernel public: https://www.kaggle.com/sachin/torch-lightning-3-tpu (which shows the error). Hopefully this tells you what i'm doing wrong. I reduced the batchsize to 16 but that doesn't seem to have helped. Mind you it might be another problem. It did work on GPU without issues so I'm a bit stumped on this.
Hey people, would love to hear your thoughts on this. Is the fact that I'm using a 'weird' lr_scheduler possibly the issue?
Can't clone the notebook now since the competition is going to end soon. Although I don't think there is a problem with lr_scheduler there. If you still think it is then I suggest remove it and check.
@sachinruk any updates on this one?
I think there is some issue with earlystopping when using TPUs, because the error is coming there only.

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
@@ -229,7 +229,7 @@ def _stop_distributed_training(self, trainer, pl_module):
 
         if trainer.use_tpu:
             stop = torch.tensor(int(trainer.should_stop), device=pl_module.device, dtype=torch.int32)
-            stop = xm.mesh_reduce("stop_signal", stop, torch.cat)
+            stop = xm.mesh_reduce("stop_signal", stop, sum)
             torch_xla.core.xla_model.rendezvous("pl.EarlyStoppingCallback.stop_distributed_training_check")
             trainer.should_stop = int(stop.item()) == trainer.world_size
 
diff --git a/pytorch_lightning/core/step_result.py b/pytorch_lightning/core/step_result.py
--- a/pytorch_lightning/core/step_result.py
+++ b/pytorch_lightning/core/step_result.py
@@ -299,6 +299,8 @@ def __str__(self):
     def __copy__(self):
         newone = type(self)()
         for k, v in self.items():
+            if isinstance(v, torch.Tensor):
+                v = v.detach()
             newone[k] = copy(v)
         return newone
 
@@ -843,6 +845,6 @@ def write_dict(self, predictions_dict, filename='predictions.pt'):
 
 def weighted_mean(result, weights):
     weights = weights.to(result.device)
-    numerator = torch.dot(result.float(), weights.t().float())
+    numerator = torch.dot(result.float(), weights.transpose(-1, 0).float())
     result = numerator / weights.sum().float()
     return result


Test patch:
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
@@ -4,7 +4,7 @@
 from torch.utils.data import DataLoader
 
 import tests.base.develop_pipelines as tpipes
-from pytorch_lightning import Trainer
+from pytorch_lightning import Trainer, seed_everything
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from tests.base import EvalModelTemplate
 from tests.base.datasets import TrialMNIST
@@ -241,3 +241,36 @@ def test_exception_when_no_tpu_found(tmpdir):
 def test_distributed_backend_set_when_using_tpu(tmpdir, tpu_cores):
     """Test if distributed_backend is set to `tpu` when tpu_cores is not None"""
     assert Trainer(tpu_cores=tpu_cores).distributed_backend == 'tpu'
+
+
+@pytest.mark.skipif(not TPU_AVAILABLE, reason="test requires TPU machine")
+@pl_multi_process_test
+def test_result_obj_on_tpu(tmpdir):
+    seed_everything(1234)
+    os.environ['PL_DEV_DEBUG'] = '1'
+
+    batches = 5
+    epochs = 2
+
+    model = EvalModelTemplate()
+    model.training_step = model.training_step_result_obj
+    model.training_step_end = None
+    model.training_epoch_end = None
+    model.validation_step = model.validation_step_result_obj
+    model.validation_step_end = None
+    model.validation_epoch_end = None
+    model.test_step = model.test_step_result_obj
+    model.test_step_end = None
+    model.test_epoch_end = None
+
+    trainer_options = dict(
+        default_root_dir=tmpdir,
+        max_epochs=epochs,
+        early_stop_callback=True,
+        row_log_interval=2,
+        limit_train_batches=batches,
+        weights_summary=None,
+        tpu_cores=8
+    )
+
+    tpipes.run_model_test(trainer_options, model, on_gpu=False, with_hpc=False)