Keeping DDP override in sync with upstream torch

[edenlightning]

From @ananthsub:
how should Lightning keep its DDP override in sync with the upstream torch DistributedDataParallel? these implementations have now diverged. I think this leads to performance degradations with Lightning + gradient accumulations, since the require_backward_grad_sync attribute isn't checked before the backwards pass

[edenlightning]

@awaelchli @tchaton

[ananthsub]

@pritamdamania87 suggested this workaround: Instead of overriding DDP, we can wrap the LightningModule in another nn.Module. This wrapper module will define its forward function to call the LightningModule's *_step function, depending on the training/testing flags.

Then when we wrap this module in DDP, we can rely on the wrapper's forward to do the right thing, and we don't have to worry about keeping the overrides in sync

[pritamdamania87]

As @ananthsub mentioned, I'd suggest always calling DDP.forward and not relying on the internals of DDP. The current implementation in LightningDistributedDataParallel could break if we make changes in DDP and the reducer. As an example, mmcv was doing something similar and broke when we refactored some code related to DDP and reducer: open-mmlab/mmcv#636.

[awaelchli]

@ananthsub I think this could work, nice idea! And I guess we should unwrap the model when passing it to the callbacks etc. so the user never sees the wrapper.

[stale]

This issue has been automatically marked as stale because it hasn't had any recent activity. This issue will be closed in 7 days if no further activity occurs. Thank you for your contributions, Pytorch Lightning Team!