transfer_batch_to_device doesn't work under DP

[ZhaofengWu]

🐛 Bug

This is discussed under #1756 and I'm opening a separate issue here for visibility.

In the training loop, for DP/DDP/DDP2, we do not move the data to devices ourselves, but instead use the default scatter to transfer data. This results in transfer_batch_to_device not being called.

https://github.com/PyTorchLightning/pytorch-lightning/blob/16a7326e5259a3cdd20a508c34a0f84806d88f8e/pytorch_lightning/trainer/training_loop.py#L736-L737

Expected behavior

Ideally, we want transfer_batch_to_device to work in all settings. If it's not possible at all to override this behavior, at least a run-time warning and/or some warning in the doc should be given.

[williamFalcon]

Ummm... yeah good point. i'm not sure we can add a hook here. Maybe @awaelchli can look into this

[stale]

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.

[williamFalcon]

@edenafek @awaelchli did we add a hook for this now?

[awaelchli]

No it is not there yet. This would require a custom scatter/gather in LightningDataParallel/LightningDistributedDataParallel that the user defines. Here I am not sure what the recommended way is in Lightning. Should the user subclass these classes and init them in configure_ddp hook?

[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!

[ZhaofengWu]

Not sure if the earlier label removal counts towards a new "activity" by the stale bot, so commenting here to indicate that this is not stale and still needs to be addressed.

[edenlightning]

@awaelchli is this still relevant?

[rohitgr7]

@edenlightning yes. Still not supported for DP AFAIK.

[edenlightning]

@awaelchli is it possible to support?

[awaelchli]

Unlikely it can be supported in the near future. The blocker here is torch.nn.DataParallel which just inherently through its design cannot scatter and gather arbitrary objects. At this time, I don't see a non-hacky solution.

Users can still provide their own DataParallel wrapper. For example, torch-geometric people have to import the custom DP module from here: https://github.com/rusty1s/pytorch_geometric/blob/master/torch_geometric/nn/data_parallel.py and use that in Lightning (they would also have to do this without Lightning).

[rubencart]

I am not sure if these belongs here, but it seems related. I can also open another issue if you prefer.

For me, when I use DDP, my (custom) LightningDataModule.transfer_batch_to_device fails because its batch argument is now a tuple that contains not only the batch but also the batch index and the optimizer index. When I don't use an accelerator, the batch argument only contains a batch, as I would expect.

With normal execution on 1 GPU and without accelerator, the _step method in gpu_accelerator.py (line 59) calls args[0] = self.to_device(args[0]), where args contains the batch, its index and the optim index. transfer_batch_to_device is then correctly called with only the batch (and the device) as argument.

However, with DDP, the _step method in ddp_accelerator.py (line 173) calls self.ddp_plugin.on_before_forward(self.trainer.get_model(), *args), where args is again the batch, idx and optim idx. The ddp_plugin passes all arguments to transfer_batch_to_device without doing anything with them. It seems that this could thus maybe be fixed by changing line 173 into self.ddp_plugin.on_before_forward(self.trainer.get_model(), args[0])?

Or is this the thing you are talking about, which cannot be fixed in the near future? Then what would you recommend us to do to solve this? Add a check for which accelerator is used to transfer_batch_to_device and handle the batch accordingly?

[awaelchli]

@rubencart This seems to be unrelated, and should not be a problem for DDP. This part of the code underwent a lot of changes lately. If you don't mind, would you send us a repro example in a new issue? If you ping me there and it is fixable, I will fix it.

[rohitgr7]

However, with DDP, the _step method in ddp_accelerator.py (line 173) calls self.ddp_plugin.on_before_forward(self.trainer.get_model(), *args), where args is again the batch, idx and optim idx. The ddp_plugin passes all arguments to transfer_batch_to_device without doing anything with them. It seems that this could thus maybe be fixed by changing line 173 into self.ddp_plugin.on_before_forward(self.trainer.get_model(), args[0])?

this has been updated in the recent refactors, so won't be a problem now. you can try master. the new version will be officially released next week I guess.

[rubencart]

@rohitgr7 Okay thanks!

[edenlightning]

@rohitgr7 is issue still in master?

[rohitgr7]

@edenlightning yes it is just with DP since pytorch DataParallel handles the splitting and transferring of batches to the device internally.

[awaelchli]

@edenlightning Even though this issue was originally reported as a bug, it is really not a bug. We document that DP is not supported for this hook.

As mentioned already #2350 (comment)
the users can provide their own DP wrapper as they would anyway have to in regular pytorch code. We could add documentation for how to do that in a faq section somewhere? What do you think?

[edenlightning]

@awaelchli maybe in the multi_gpu guide under DP?

[edenlightning]

@SeanNaren maybe add to your PR?