Support uneven DDP inputs with pytorch model.join

[edenlightning]

See more details: pytorch/pytorch#38174

cc @Borda @tchaton @rohitgr7 @akihironitta @awaelchli

[stale]

This issue has been automatically marked as stale because it hasn't had any recent activity. This issue will be closed in 7 days if no further activity occurs. Thank you for your contributions, Pytorch Lightning Team!

[carmocca]

Interested in this issue! Hopefully some progress is done soon 👍

[xvr-hlt]

Interested in this also :)

[rohan-varma]

Is there any progress on this issue? Happy to help in any way.

[edenlightning]

@rohan-varma that would be great!! Want to try and submit a draft PR? And we can help from there?

[rohan-varma]

@edenlightning Sounds good, I also pinged the slack channel for any feedback/discussions.

[alanhdu]

We'd also be very interested in this feature. Let us know if there's anything I can do to help!

[rohan-varma]

The PR #5141 is ready for review, in case anyone wants to take a look.

[stale]

This issue has been automatically marked as stale because it hasn't had any recent activity. This issue will be closed in 7 days if no further activity occurs. Thank you for your contributions, Pytorch Lightning Team!

[ananthsub]

I discussed this more with @rohan-varma - DDP join docs: https://pytorch.org/docs/stable/_modules/torch/nn/parallel/distributed.html#DistributedDataParallel.join

This module currently does not support custom distributed collective operations in the forward pass, such as SyncBatchNorm or other custom defined collectives in the model’s forward pass.

As the LightningModule is wrapped in another module which is then wrapped with DDP, the LightningModule's training_step becomes the forward pass run by the DDP wrapped module: https://github.com/PyTorchLightning/pytorch-lightning/blob/d471fa30b3bf95cfe601014bac544754067241ca/pytorch_lightning/plugins/training_type/ddp.py#L223-L227

As a result, any collective call (such as metric syncing or all_gather) that happens during training step will cause this to not work. Therefore I lean towards closing this out given the caveats. @awaelchli @justusschock what do you think?

[awaelchli]

agree, I also don't see how this can be supported at the moment.

[carmocca]

To recap, the plan would be:

• Enable "join" as an optional feature of the DDP strategy: Trainer(strategy=DDPStrategy(uneven_input_support: bool). We could also add a registry string for it.
• Add support for "joining" the training_step.
  • Is there a benefit to doing it for validation_step and test_step? Probably not
  • Could validation_step and test_step use UnrepeatedDistributedSampler just as trainer.predict? Probably yes.
• When the feature is enabled, we don't automatically use the generic DistributedSampler as we wouldn't want to duplicate data to make inputs even.
  • Could we use UnrepeatedDistributedSampler? Related issues:
    • torch.utils.data.DistributedSampler allow uneven inputs pytorch/pytorch#49180
    • Incorrect Validation Accuracy Due to Distributed Sampler pytorch/pytorch#25162
• We print a big warning about how this feature is experimental and describe all its caveats.
• This would be Torch 1.10+ only

Some sources:
https://pytorch.org/docs/stable/distributed.algorithms.join.html#torch.distributed.algorithms.Join
https://pytorch.org/tutorials/advanced/generic_join.html

[awaelchli]

Is there a benefit to doing it for validation_step and test_step? Probably not

I assume there is, if collectives are being used. For example, sync_dist=True in self.log or similar. However, we don't wrap the model in ddp during val and test, so join won't be available anyways.

When the feature is enabled, we don't automatically use the generic DistributedSampler as we wouldn't want to duplicate data to make inputs even.

pytorch/pytorch#49180 is great! Hopefully this will clarify the drop_last argument which has a slightly misleading/incomplete description :) We would indeed need the UnrepeatedDistributedSampler.

[otaj]

Hi, everyone, I'm gathering information on what is needed in order to support this properly.

1. Use torch.distributed.algorithms.Join (https://pytorch.org/docs/stable/distributed.algorithms.join.html) as a context manager in which is the model run.
2. Use UnrepeatedDistributedSamplerWrapper
3. Check for all modules, that could use syncing (such as SyncBatchNorm) and have them as arguments to the Join context manager from 1.
4. Figure out what to with calls to self.log(..., sync_dist=True)

Is that it? cc @awaelchli, @carmocca.

[justusschock]

@otaj Almost.

Additionally, all metrics from torchmetrics would have to be considered as well as they are also capable of issuing syncs on their own. And in general, the user can run arbitrary syncing calls within each of the steps which have to be considered as well (which will be the trickiest part I guess)

[otaj]

oh, those torchmetrics are going to be fun... 😅 I think capturing user calls can be solved with yet another contextmanager (our, custom one), what do you think?

[justusschock]

if we can capture user calls with that, it might work similarly with torchmetrics. So let's ignore those metrics for now and if you got a working solution for everything else, I'm sure we'll manage to integrate metrics with that :D

[Borda]

let's check the option with LigthingLite first 🦦

[awaelchli]

Here is the corresponding issue as suggested in planning: #14635

[yygle]

hi, any updates of this issue?

[DucoG]

Is there a (hacky) method for implementing the PyTorch .join() context into your own PL setup? Or is currently the only solution to revert back to PyTorch?

[alberthli]

Any updates on this issue?