How to do checkpointing when optimizer can only use distributed state dict

[RuABraun]

I'm doing multinode, multi GPU training and am using this optimizer, which requires using distributed_state_dict() and load_distributed_state_dict().

I think the loading I could do by overwriting load_from_checkpoint(). But for saving it seems I might have to subclass the Strategy class? Is that correct or is there an easier way.