ref: decouple apex second attemp part 3/n (#4055)

pytorch_lightning/accelerators/base_accelerator.py

@@ -72,33 +72,14 @@ def process_dataloader(self, dataloader):
         return dataloader
 
     def backward(self, closure_loss, optimizer, opt_idx, *args, **kwargs):
-
-        # scale loss for 16 bit
         if self.trainer.precision == 16:
-            if self.trainer.amp_backend == AMPType.NATIVE:
-                closure_loss = self.trainer.scaler.scale(closure_loss)
-            else:
-                closure_loss = amp.scale_loss(closure_loss, optimizer)
-
-            # 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
-        closure_loss.backward(*args, **kwargs)
-
-        # 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()
+            closure_loss = self.trainer.precision_connector.backend.backward(closure_loss, optimizer, *args, **kwargs)
+        else:
+            # do backward pass
+            closure_loss.backward(*args, **kwargs)
+
+            # once backward has been applied, release graph
+            closure_loss = closure_loss.detach()
         return closure_loss
 
     def optimizer_step(self, optimizer, batch_idx, opt_idx, lambda_closure):

pytorch_lightning/plugins/apex.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 from typing import List, Tuple
 from torch.optim.optimizer import Optimizer
+from pytorch_lightning.utilities.distributed import rank_zero_warn
+from pytorch_lightning.utilities import AMPType
 
 try:
     from apex import amp
@@ -34,6 +36,28 @@ def training_step(self, fx, args):
         output = fx(args)
         return output
 
+    def backward(self, closure_loss, optimizer, *args, **kwargs):
+        closure_loss = amp.scale_loss(closure_loss, optimizer)
+
+        # enter apex context
+        self.trainer.dev_debugger.track_event('AMP', str(AMPType.APEX))
+        context = closure_loss
+        closure_loss = closure_loss.__enter__()
+
+        # do backward pass
+        closure_loss.backward(*args, **kwargs)
+
+        # exit amp context
+        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
+
     def configure_apex(
         self,
         amp: object,

pytorch_lightning/plugins/native_amp.py

@@ -23,6 +23,16 @@ def __init__(self, trainer):
     def connect(self, model, optimizers):
         return model, optimizers
 
+    def backward(self, closure_loss, optimizer, *args, **kwargs):
+        closure_loss = self.trainer.scaler.scale(closure_loss)
+
+        # do backward pass
+        closure_loss.backward(*args, **kwargs)
+
+        # once backward has been applied, release graph
+        closure_loss = closure_loss.detach()
+        return closure_loss
+
     def training_step(self, fx, args):
         with torch.cuda.amp.autocast():
             output = fx(*args)