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 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,
                    )

                    self._updated_model_last_step = False

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

                else:
                    self._updated_model_last_step = True

                    if self.automatic_optimization:

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

Review comment (suggestion):
is `backward` in the DDP accelerator aware of `self.block_ddp_sync_behaviour()` ?  Otherwise will `require_backward_grad_sync=False` on the wrapped DDP model be ignored?

https://github.com/PyTorchLightning/pytorch-lightning/blob/41de4538aa0c187793709a93875e67666c2ddde8/pytorch_lightning/accelerators/accelerator.py#L89-L101

https://github.com/PyTorchLightning/pytorch-lightning/blob/41de4538aa0c187793709a93875e67666c2ddde8/pytorch_lightning/trainer/connectors/model_connector.py#L54-L57

Code after:
            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,
                                self.trainer.hiddens
                            )
                            return None if result is None else result.loss