DataModule PrepareData Only Called on Node 0

[Queuecumber]

🐛 Bug

The LightningDataModule prepare_data function should be able to be called on the local rank 0 process on any node if the user chooses, and there are APIs around controlling this, but they don't work. I've reproduced this bug and traced it to one (rather obvious) line:

https://github.com/PyTorchLightning/pytorch-lightning/blob/master/pytorch_lightning/core/datamodule.py#L384

This line wraps the prepare_data function in rank_zero_only, a function which means global zero not local zero. There is already logic in the data connector around making sure that the function is only run when it is supposed to. Apparently the multi-node setup was never tested?

This fix is simple: just delete the rank_zero_only part.

Additional context

This is a quick one-line fix that I'm happy to PR, but someone needs to engage with me on it so that it actually gets merged.

[awaelchli]

Hi, great observation. I was not aware that there is a decoration happening.

I can summarize:

prepare_data is handled in DataConnector.prepare_data and called by the trainer.
https://github.com/PyTorchLightning/pytorch-lightning/blob/c1eac483e93494a173f5db7527df0008f551c734/pytorch_lightning/trainer/connectors/data_connector.py#L61

Two cases:

• only when node_rank=0 and local_rank=0 (default case): This is useful e.g. when the nodes share a filesystem
• only on local_rank=0 (so every node): this is the case when you set Trainer(prepare_data_per_node=True).

I agree with you that there should be no decorator on that function. The hook is called by Lightning, not the user.

[awaelchli]

If you could send a PR, that would be great! We can look for existing tests for prepare_data and extend them to check the calls on different settings for local and node rank.

[Queuecumber]

Sure I should be able to get to it tomorrow and I'll ping you when I send it in