ref: inner train loop (intermediate step) 1/n (Lightning-AI#3359)

Author: williamFalcon

pytorch_lightning/accelerators/base_backend.py

@@ -1,6 +1,7 @@
 import torch
 from typing import Any
 from pytorch_lightning.utilities.apply_func import move_data_to_device
+from pytorch_lightning.utilities import AMPType, rank_zero_warn
 
 
 class Accelerator(object):
@@ -31,3 +32,36 @@ def validation_step_end(self, output):
 
     def process_dataloader(self, dataloader):
         return dataloader
+
+    def backward(self, closure_loss, optimizer, opt_idx):
+        model_ref = self.trainer.get_model()
+
+        # scale loss for 16 bit
+        if self.trainer.precision == 16:
+            closure_loss = model_ref.amp_scale_loss(
+                closure_loss,
+                optimizer,
+                opt_idx,
+                amp_backend=self.trainer.amp_backend
+            )
+
+            # enter amp context
+            if self.trainer.amp_backend == AMPType.APEX:
+                self.trainer.dev_debugger.track_event('AMP', str(AMPType.APEX))
+                context = closure_loss
+                closure_loss = closure_loss.__enter__()
+
+        # do backward pass
+        model_ref.backward(self, closure_loss, optimizer, opt_idx)
+
+        # exit amp context
+        if self.trainer.precision == 16 and self.trainer.amp_backend == AMPType.APEX:
+            a, b, c = None, None, None
+            error = context.__exit__(a, b, c)
+            if error:
+                rank_zero_warn(a, b, c)
+                raise Exception('apex unscale error')
+
+        # once backward has been applied, release graph
+        closure_loss = closure_loss.detach()
+        return closure_loss

pytorch_lightning/accelerators/cpu_backend.py

@@ -14,7 +14,7 @@
 import torch
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from pytorch_lightning.accelerators.base_backend import Accelerator
-from pytorch_lightning.utilities import AMPType
+from pytorch_lightning.utilities import AMPType, rank_zero_warn
 
 
 class CPUBackend(Accelerator):

pytorch_lightning/accelerators/tpu_backend.py

@@ -220,3 +220,14 @@ def __setup_tpu_training(self, model: LightningModule, trainer):
         log.info(f'INIT TPU local core: {trainer.tpu_local_core_rank},'
                  f' global rank: {trainer.tpu_global_core_rank}'
                  f' with XLA_USE_BF16={os.environ.get("XLA_USE_BF16")}')
+
+    def backward(self, closure_loss, optimizer, opt_idx):
+        model_ref = self.trainer.get_model()
+
+        # do backward pass
+        model_ref.backward(self, closure_loss, optimizer, opt_idx)
+
+        # detach after backward
+        closure_loss = closure_loss.detach()
+
+        return closure_loss

pytorch_lightning/trainer/training_loop.py

@@ -988,36 +988,16 @@ def optimizer_closure(self, split_batch, batch_idx, opt_idx, optimizer, hiddens)
         untouched_loss = closure_loss.detach().clone()
 
         # backward pass
-        model_ref = self.get_model()
         with self.profiler.profile('model_backward'):
-            # scale loss for 16 bit
-            if self.precision == 16 and not self.on_tpu:
-                closure_loss = model_ref.amp_scale_loss(closure_loss, optimizer, opt_idx, amp_backend=self.amp_backend)
-
-                # enter amp context
-                if self.amp_backend == AMPType.APEX:
-                    self.dev_debugger.track_event('AMP', str(AMPType.APEX))
-                    context = closure_loss
-                    closure_loss = closure_loss.__enter__()
-
-            # do backward pass
-            model_ref.backward(self, closure_loss, optimizer, opt_idx)
-
-            # exit amp context
-            if self.precision == 16 and self.amp_backend == AMPType.APEX and not self.on_tpu:
-                a, b, c = None, None, None
-                error = context.__exit__(a, b, c)
-                if error:
-                    rank_zero_warn(a, b, c)
-                    raise Exception('apex unscale error')
-
-            # once backward has been applied, release graph
-            closure_loss = closure_loss.detach()
-
-            if is_result_obj:
-                training_step_output.detach()
-            else:
-                training_step_output.batch_loss = training_step_output.batch_loss.detach()
+            closure_loss = self.accelerator_backend.backward(closure_loss, optimizer, opt_idx)
+
+        # --------------------
+        # ON AFTER BACKWARD TODO
+        # --------------------
+        if is_result_obj:
+            training_step_output.detach()
+        else:
+            training_step_output.batch_loss = training_step_output.batch_loss.detach()
 
         if self.use_horovod:
             # Synchronize Horovod to ensure gradient manipulations (e.g., loss scaling) are valid