The common use case of this is to set it to False every couple of iterations, so I think passing a default value for the constructor is not that useful. Instead, we can assume the default is False, and that users will have to call a function model.accumulate_grads(enable: bool) every iteration to control what to do.

Adding some background for this PR: sometimes, a batch of inputs is too large to fit into GPU, which has to be split into several micro-batches. However, to improve efficiency, it would be helpful to only apply params/gradients sync at original batch boundaries instead of micro-batch boundaries. Hence, @chenyangyu1988 would like to control when to skip/apply sync.

@apaszke Could you elaborate on how to creating a wrapper model which takes a list of sub-batches to forward, runs the backward on each of them separately, and only then returns?

For the latest DistributedDataParallel, we need to call prepare_for_backward in its forward function, I am not sure how to do this with context_manager because when exit the context, all forward and backward already finished, there is no easy way to kick off gradients sync with current design.

The context wrapper approach won't be able to overlap any reduction with the final backward pass, because we only know that we want to reduce again after the call to backward has already returned.

Either way, we'll need some way to toggle this behavior on the DistributedDataParallel class, because it controls how it executes its forward function. This doesn't need to be a construction time parameter because you would toggle it anyway. A setter function would be nice though. @mrshenli and I just chatted about this and I think there are two approaches possible here. One that doesn't look nice and is easy to screw, and one that is better looking and harder to screw up. @apaszke I think your first suggestion maps to this second approach, can you confirm? This approach does require you to add the loss computation to your forward function, so the call to backward can happen in line.

# Alternative 1: simple but ugly and easy to screw up.
# Doesn't deal with edge condition here.
for index, batch in enumerate(batches):
  model.accumulate_grads(index % 5 != 4) # True for [0, 3], False for [4]
  model(batch).backward()

# Alternative 2: wrapper (TODO: find better name)
class AccumulateGradientsAcrossBatches(nn.Module):
  def __init__(self, model, batch_size):
    super().__init__()
    self.model = model
    self.batch_size = batch_size
    
  def forward(self, megabatch):
    batches = torch.chunk(megabatch, self.batch_size)
    self.model.accumulate_grads(True)
    for batch in batches[0:-1]:
      self.model(batch).backward()
    self.model.accumulate_grads(False)
    return self.model(batches[-1])

model = AccumulateGradientsAcrossBatches(model, 32)
for megabatch in batches:
  model(megabatch).backward()

@soumith @apaszke Does @pietern's suggestion look good to you?

@chenyangyu1988 If we all have consensus on that, I think you only need to change accumulate_grads from a construction time argument to become a function, which should be sufficient to unblock you. We will add AccumulateGradientsAcrossBatches later (maybe under a different name).

@mrshenli Sounds great, I already make it a function rather than a construction time argument :)

Ok, I think my suggestion with the wrapper model might not work. Still, those things absolutely shouldn't be setters, because that's just too error prone. If something, they should be context managers like the one I proposed.

An acceptable alternative which doesn't involve a context manager would be to add a parameter to the constructor which allows one to disable automatic gradient syncs, and then trigger them manually every time they are needed. So the code would look like that:

model = DistributedDataParallel(..., auto_sync=False)

def train(...):
  for batch in ...:
    for subbatch in batch:
      model(subbatch).backward()
    model.sync_grads()

@apaszke The DDP overlaps the gradients sync with the backward call to get more performance gain, so we probably need some ways to make it happen in the last backward call.

@apaszke Can you elaborate why not a setter here, considering we have existing ones for e.g. train(bool)? I agree they can be error prone, cause problems w.r.t. boundary conditions, etc, etc.

I like your proposed alternative, but we need to find an easy and robust way to incorporate the overlapping of backward with gradient reduction. Before the last call to backward we need to let DDP know it needs to reduce the grads as they are computed. This could look something like this:

model = DistributedDataParallel(..., auto_sync=False)

def train(...):
  for batch in ...:
    for index, subbatch in enumerate(batch):
      output = model(subbatch)
      if index == len(batch) - 1:
        model.sync_grads(output)
      output.backward()

We could opt to store a copy of the model output in DDP so that you don't have to pass it back to sync_grads, but that could have unintended side effects by keeping the output alive longer than expected.

@chenyangyu1988, are you still working on this? Accumulating gradients with DDP is an important feature that can make a big speed difference when training with large batches.

@myleott @chenyangyu1988 I am picking this up.

We got multiple requests on this feature. Let's revive the discussion and try to get some consensus.

I agree with @apaszke that context manager is indeed more error prone.

And I also agree with @pietern that it is not favorable to let DDP keep the outputs around until exiting the context. Given that people requesting this feature because they need to deal with tight memory limit, keeping the output a little longer might even hit OOM errors.

So, how about let us just provide a DistributedDataParallel.no_sync() context manager? It does not sync at all within the entire context. And it is true that users need to call another model(input).backward() after exiting the context. But I think it is OK as one major concern in the previous discussion is to prevent users from running into errors without knowing it. This is indeed more verbose, but should reaffirm the expected behavior.

The application would then look like:

with ddp.no_sync():
  for input in inputs:
    ddp(input).backward()

ddp(one_more_input).backward()

@myleott @chenyangyu1988 @soumith Does this API look OK?

Yep that works for us.

Proposed API is implemented in #21736

This was superseded by #21736.