I could also add a with no_sync(not should_step) to fit_batch, the way it's implemented it will only not sync if DDP is used and should_step is False, which is the desired behaviour, what do you think @pplantinga ?

Hmm, that's a great question! On the one hand, I imagine in nearly every case this would have the desired behavior. On the other hand, if it wasn't the desired behavior, it might be tricky to debug. Why don't you go ahead and add it for now, and we'll get feedback from one more person before we merge it. @mravanelli @TParcollet @Gastron do you have opinions about this question?

Hum, I have no idea what No_sync will do outside of DDP though.

It won't do anything, it will just be an empty context manager.

Then let's have it only for DDP.

Then let's have it only for DDP.

I'm not sure exactly what you mean, there isn't an easy way to apply this only when DDP is active. Code without:

outputs = self.compute_forward(batch, Stage.TRAIN)
loss = self.compute_objectives(outputs, batch, Stage.TRAIN)
(loss / self.grad_accumulation_factor).backward()

Code with:

outputs = self.compute_forward(batch, Stage.TRAIN)
loss = self.compute_objectives(outputs, batch, Stage.TRAIN)
with self.no_sync(not should_step):
  (loss / self.grad_accumulation_factor).backward()

I think this is a good addition because it speeds up computation by 10-20% in cases where DDP and grad accumulation are used.

What's our conclusion here?

I'd say lets go ahead with the self.no_sync(not should_step): feature

Feel free to merge @pplantinga

Okay great, will add the final touches today/tomorrow