Repository: Lightning-AI/pytorch-lightning
File: pytorch_lightning/trainer/training_loop.py (Python)

Code before:

        for split_idx, split_batch in enumerate(splits):
            self.trainer.split_idx = split_idx

            # in manual optimization we loop over all optimizers at once
            optimizers = self.get_optimizers_iterable()
            if not self.automatic_optimization:
                optimizers = [optimizers[0]]

            # loop over optimizers
            for opt_idx, optimizer in optimizers:
                # make sure only the gradients of the current optimizer's parameters are calculated
                # in the training step to prevent dangling gradients in multiple-optimizer setup.
                if self.automatic_optimization and len(self.trainer.optimizers) > 1:
                    model = self.trainer.get_model()
                    model.toggle_optimizer(optimizer, opt_idx)

                if not (accumulation_done or is_final_batch):
                    # For gradient accumulation

                    # -------------------
                    # calculate loss (train step + train step end)
                    # -------------------

                    # no ddp sync at the beginning of forward or backward due to parameter changes in this or last step required
                    no_sync = self._updated_model_last_step and isinstance(self.trainer.model, torch.nn.parallel.DistributedDataParallel)
                    if no_sync:
                        self.trainer.model.no_sync.__enter__()

                    self.training_step_and_backward(split_batch, batch_idx, opt_idx, optimizer, self.trainer.hiddens)

                    if no_sync:
                        self.trainer.model.no_sync.__exit__()

                    batch_outputs = self._process_closure_result(
                        batch_callback_metrics=batch_callback_metrics,
                        batch_log_metrics=batch_log_metrics,
                        batch_outputs=batch_outputs,
                        opt_idx=opt_idx,
                    )

                    self._updated_model_last_step = False

                # ------------------------------
                # BACKWARD PASS
                # ------------------------------
                # gradient update with accumulated gradients

                else:
                    self._updated_model_last_step = True


Review comment (performance):
Great catch there !

Would you mind making this a bit more readable ? 
Like moving no_sync outside of main loop and hiding those __enter__() and __exit__() ?
Also, you don't need `self._updated_model_last_step` variable. It can derived from `self.trainer.batch_idx and self.trainer.accumulate_grad_batches`

Code after:
        # lightning module hook
        splits = self.tbptt_split_batch(batch)

        for split_idx, split_batch in enumerate(splits):
            self.trainer.split_idx = split_idx

            # in manual optimization we loop over all optimizers at once
            optimizers = self.get_optimizers_iterable()
            if not self.automatic_optimization:
                optimizers = [optimizers[0]]

            # loop over optimizers
            for opt_idx, optimizer in optimizers:
                # make sure only the gradients of the current optimizer's parameters are calculated
                # in the training step to prevent dangling gradients in multiple-optimizer setup.
                if self.automatic_optimization and len(self.trainer.optimizers) > 1:
                    model = self.trainer.get_model()
                    model.toggle_optimizer(optimizer, opt_idx)

                if should_accumulate:
                    # For gradient accumulation

                    # -------------------
                    # calculate loss (train step + train step end)
                    # -------------------

                    # 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,
                        batch_outputs=batch_outputs,
                        opt_idx=opt_idx,
                    )

                # ------------------------------
                # BACKWARD PASS
                # ------------------------------
                # gradient update with accumulated gradients

                else:
                    if self.automatic_optimization:

                        def train_step_and_backward_closure():
                            result = self.training_step_and_backward(
                                split_batch,
                                batch_idx,
                                opt_idx,
                                optimizer,