We have built uneven input support in DDP for this purpose. It is designed to allow users to train DDP models when there are a different # of inputs across ranks (if model has no communication other than those done by DDP), as well as throw an exception that can be caught and recovered from if the model does have custom user communication.
Here are the docs: torch.nn.parallel.distributed — PyTorch 1.9.0 documentation
and tutorial Distributed Training with Uneven Inputs Using the Join Context Manager — PyTorch Tutorials 1.9.0+cu102 documentation 3. Feel free to follow up with anymore questions!