Expectations for custom data parallel implementations

[ananthsub]

I have a need to use a custom [DistributedDataParallel](https://pytorch.org/docs/stable/notes/ddp.html) implementation. I'd like to do this with the existing DDPBackend today.

With Lightning, the API and docs are unclear as to whether I need to extend LightningDistributedDataParallel, or if I can directly extend torch DistributedDataParallel.

The LightningModule API docs suggest configure_ddp should work with torch DistributedDataParallel: https://pytorch-lightning.readthedocs.io/en/latest/lightning-module.html#configure-ddp

However, there are spots in Lightning which rely on checking isinstance of the custom Lightning overrides: https://github.com/PyTorchLightning/pytorch-lightning/blob/master/pytorch_lightning/trainer/model_connector.py#L31-L34 1

The Lightning DDP also forwards calls to train/val/test step. Is this a requirement for custom DDP implementations when used with Lightning?

TLDR: should I subclass LightningDistributedDataParallel or DistributedDataParallel when I implement the model hook for configure_ddp?

Also asked here: https://forums.pytorchlightning.ai/t/expectations-for-custom-data-parallel-implementations/162

[awaelchli]

Almost certainly you need to inherit LightningDistributedDataParallel for the reasons you have already mentioned. As @williamFalcon already explained, this is probably not what should be done today. I have come across this as well. There are libraries like torchgeometry that have a custom DataParallel class and I believe it is currently not possible to just drop that into Lightning, because of the reasons you have already mentioned.

[ananthsub]

@awaelchli it seems like the workaround for now is to extend LightningDistributedDataParallel while the accelerator backend refactors are happening. We should update the code and docs to reflect that: https://pytorch-lightning.readthedocs.io/en/latest/lightning_module.html#configure-ddp . Right now they're misleading and will break because of those checks. Though I'm pretty sure you can check if isinstance(model, DistributedDataParallel) to grab the module reference and it'll work. What happens if a LightningModule doesn't implement forward? Would this raise an exception?

[edenlightning]

@williamFalcon PTAL

[awaelchli]

@ananthsub Not sure if you have seen it already, there is now a DDP Plugin in which you can override the DistributedDataParallel. It can be passed into the trainer via plugins list.
https://pytorch-lightning.readthedocs.io/en/latest/plugins.html?highlight=DDPPlugin

[ananthsub]

yup I saw. that resolves this issue 👍