how to use Apex DistributedDataParallel with Lightining?
I was wondering if there's a way to use apex.parallel.DistributedDataParallel instead of pytorch native DistributedDataParallel. (I am trying to reproduce a paper that used Apex DDP and apex mixed precision and i am getting lower results using pytorch native one)
Here is a quick draft of what you could try:
from pytorch_lightning.plugins.training_type import DDPPlugin
from apex.parallel import DistributedDataParallel
class ApexDDPPlugin(DDPPlugin):

    def _setup_model(self, model: Module):
        return  DistributedDataParallel(module=model, device_ids=self.determine_ddp_device_ids(), **self._ddp_kwargs)

    @property
    def lightning_module(self):
        return self.module.module
I'm not sure if apex DistributedDataParallel supports device ids (it seems not??), you may need to remove it.
Use it in the trainer:
trainer = Trainer(gpus=2, strategy=ApexDDPPlugin(), precision=...)
trainer.fit(model)