Skips DDP parameter sync (#4301)

* ddp no-sync

* Update pytorch_lightning/trainer/training_loop.py

Co-authored-by: ananthsub <[email protected]>

* Update training_loop.py

* factor __enter__ and __exit__ out to separate context manager

* delete _updated_model_last_step

Co-authored-by: justusschock <[email protected]>
Co-authored-by: Teddy Koker <[email protected]>
Co-authored-by: ananthsub <[email protected]>
Co-authored-by: chaton <[email protected]>
Co-authored-by: Rohit Gupta <[email protected]>
(cherry picked from commit bbd81df)

Author: 6 people

pytorch_lightning/trainer/training_loop.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import subprocess
+from contextlib import contextmanager
 from copy import copy, deepcopy
 
 import numpy as np
@@ -656,6 +656,7 @@ def run_training_batch(self, batch, batch_idx, dataloader_idx):
         # checks if backward or backward + optimizer step (via closure)
         accumulation_done = self._accumulated_batches_reached()
         is_final_batch = self._num_training_batches_reached()
+        should_accumulate = not (accumulation_done or is_final_batch)
 
         # lightning module hook
         splits = self.tbptt_split_batch(batch)
@@ -676,13 +677,17 @@ def run_training_batch(self, batch, batch_idx, dataloader_idx):
                     model = self.trainer.get_model()
                     model.toggle_optimizer(optimizer, opt_idx)
 
-                if not (accumulation_done or is_final_batch):
+                if should_accumulate:
                     # For gradient accumulation
 
                     # -------------------
                     # calculate loss (train step + train step end)
                     # -------------------
-                    self.training_step_and_backward(split_batch, batch_idx, opt_idx, optimizer, self.trainer.hiddens)
+
+                    # perform dpp sync only when performing optimizer_step
+                    with self.block_ddp_sync_behaviour():
+                        self.training_step_and_backward(split_batch, batch_idx, opt_idx, optimizer, self.trainer.hiddens)
+
                     batch_outputs = self._process_closure_result(
                         batch_callback_metrics=batch_callback_metrics,
                         batch_log_metrics=batch_log_metrics,
@@ -696,7 +701,6 @@ def run_training_batch(self, batch, batch_idx, dataloader_idx):
                 # gradient update with accumulated gradients
 
                 else:
-
                     if self.automatic_optimization:
 
                         def train_step_and_backward_closure():
@@ -761,6 +765,13 @@ def train_step_and_backward_closure():
         )
         return result
 
+    @contextmanager
+    def block_ddp_sync_behaviour(self):
+        if isinstance(self.trainer.model, torch.nn.parallel.DistributedDataParallel):
+            yield from self.trainer.model.no_sync()
+        else:
+            yield
+
     def _process_closure_result(
         self, batch_callback_metrics: list, batch_log_metrics: list, batch_outputs: list, opt_idx: int
     ) -> list: