More explanation needed in documentation on DP vs. DDP in multi-GPU training
❓ Questions and Help
The documentation on multi-GPU training says DP is suitable for softmax since it can combine all pieces in the end. I do not understand why it is the case. The gradients calculated from a batch of samples are equal to the sum (of average) of gradients calculated from individual samples. Therefore I feel that DP is also averaging the gradients from each sample.
I acknowledge the difference from the following example provided in the documentation, since training_step_end concatenates along dim=1. However, in most cases, batches should be concatenated along dim=0.
def training_step(self, batch, batch_idx):
    x, y = batch
    y_hat = self(batch)

    # on dp or ddp2 if we did softmax now it would be wrong
    # because batch is actually a piece of the full batch
    return y_hat

def training_step_end(self, batch_parts_outputs):
    # batch_parts_outputs has outputs of each part of the batch

    # do softmax here
    outputs = torch.cat(outputs, dim=1)
    softmax = softmax(outputs, dim=1)
    out = softmax.mean()

    return out

I appreciate any comment on this issue.