Inaccurate reduction in OPT_LEVEL 'O1'

I want to disable all-reduce during gradient accumulation. If my gradient accumulation is 2, I want to enable all reduce every other step. This will speed up my training.

I am using the following code with latest apex but the reduced gradients are coming inaccurate. Can I know what's wrong here?

from apex.parallel import DistributedDataParallel

    self.model = DistributedDataParallel(self.model, delay_allreduce=True)
    def _enable_or_disable_reduction(self, step):
            # https://github.com/NVIDIA/apex/blob/master/apex/parallel/distributed.py
            if (step + 1) % self.args.gradient_accumulation_steps == 0:
                self.logger.debug('Enabling all reduce')
                self.model.enable_allreduce()
            else:
                self.logger.debug('Disabling all reduce')
                self.model.disable_allreduce()
            pass
    def _train_step(self, batch, step):
        self.logger.debug(f'***** Training step {step} *****')
        """ Performs forward and backwards pass"""
        batch = [t.to(self.args.device) for t in batch] # move data to device
        batch, batch_len = Tasks.get_model_inputs(self.args, batch)
        self.sequences += batch_len
        self.logger.debug(self.get_param_info())
        output = self._forward(batch)
        loss = output[0]
        loss = loss / self.args.gradient_accumulation_steps

        self._enable_or_disable_reduction(step)
        self.logger.debug(f'step = {step}, _disable_allreduce = {self.model._disable_allreduce}')

        self._backward(loss)
        self.logger.debug(self.get_param_info())
        return loss

With this code reduction is not happening between 2 processes

Start with apex/parallel/distributed.py, the linked implementation, and inspect enable_allreduce(), disable_allreduce(), and their call order around backward(). Reproduce the provided two-process gradient-accumulation example and verify when reduction occurs and whether the resulting gradients are accurate; the issue does not name a test or a specific expected code change.