b4d57c6d49682094efe22fbe2c03fa2c4973869f,fairseq/trainer.py,Trainer,train_step,#Trainer#Any#Any#,496

Before Change

        overflow = False
        try:
            if self.tpu and self.data_parallel_world_size > 1:
                import torch_xla.core.xla_model as xm

                gradients = xm._fetch_gradients(self.optimizer.optimizer)
                xm.all_reduce(
                    "sum", gradients, scale=1.0 / self.data_parallel_world_size
                )

            with torch.autograd.profiler.record_function("multiply-grads"):
                // multiply gradients by (// GPUs / sample_size) since DDP
                // already normalizes by the number of GPUs. Thus we get
                // (sum_of_gradients / sample_size).

After Change

        try:
            with torch.autograd.profiler.record_function("reduce-grads"):
                self.optimizer.all_reduce_grads(self.model)
                if utils.has_parameters(self.criterion):
                    self.optimizer.all_reduce_grads(self.criterion)

            with torch.autograd.profiler.record_function("multiply-grads"):
                // multiply gradients by (data_parallel_size / sample_size) since
                // DDP already normalizes by the number of data parallel workers.
                // Thus we get (sum_of_gradients / sample_size) at the end.

Instances