Support for PyTorch ZeroRedundancyOptimizer

[schmidt-jake]

🚀 Feature

Implement a DDP-ZeroRedundancyOptimizer strategy.

Motivation

I know that Zero is available in Fairscale, but since PyTorch now has its own implementation at torch.distributed.optim.ZeroRedundancyOptimizer, I think it would make sense to support it.

Pitch

A new strategy called e.g. DDPZero that would support DDP with torch's ZeroRedundancyOptimizer, along with support for overlapping ZeroRedundancyOptimizer with DDP ("overlap_with_ddp") and the related DDP comms hooks.

Alternatives

As mentioned above, Fairscale implements Zero. However, this requires taking on an additional dependency. And since pytorch has a native implementation, it could make sense to support that in lightning.

[rohitgr7]

the support was added in #14208
and will be available in next v1.8 release in a few days :)

you would still be able to do it:

class LitModel(LightningModule):
    def configure_optimizers(self):
        return ZeroRedundancyOptimizer(self.layer.parameters(), optimizer_class=torch.optim.Adam, lr=0.1)

model = LitModel()
trainer = Trainer(accelerator="gpu", devices=2, strategy='ddp')
trainer.fit(model)

[schmidt-jake]

Fantastic! Apologies I missed that.