Is it possible to access the global rank and world size outside of LightningModule?

[EricWiener]

Hi,

I'm using an IterableDataset and need to manually split up the data depending on the GPU (I'm using DDP). The lightning docs say that replace_sampler_ddp doesn't automatically handle IterableDatasets, so I need to do this myself.

Is it possible to access the global rank or world size in a dataset (aka are there any utility functions available for this purpose)? I saw _get_rank here, but ideally I wouldn't use a private helper function.

I tried using:

file_paths_filtered = file_paths[
    torch.distributed.get_rank() :: torch.distributed.get_world_size()
]

inside the IterableDataset, but it fails because the default process group has not been setup yet (which makes sense):

RuntimeError: Default process group has not been initialized, please make sure to call init_process_group.

Is there any way to access the global rank/world size? If not, how would you recommend I check which GPU the dataset is on without causing an issue with init_process_group not being called yet (should I maybe call this manually)?

[ananthsub]

Where are you calling this? If this is inside a DataModule, you could implement this splitting inside your train_dataloader() funtion, which is called after the distributed training setup and process group creation. in this case, accessing the torch.distributed ranks should work

[EricWiener]

Thanks for the reply! I was using this in __next__ of a torch.utils.data.IterableDataset (that was contained within a LightningDataModule's train_dataloader(). The issue turned out to be torch.distributed.get_rank() and torch.distributed.get_world_size() aren't available when on a single GPU (locally) which I was testing on before sending to multi-GPU server

[EricWiener]

Issue ended up being torch.distributed.get_rank() and torch.distributed.get_world_size() aren't available when on a single GPU, which I was testing on locally before sending to multi-GPU server.

The following works:

if torch.distributed.is_available() and torch.distributed.is_initialized():
    file_paths = file_paths[
        torch.distributed.get_rank() :: torch.distributed.get_world_size()
    ]

[tchaton]

If you are using accelerator=ddp and without any launcher, you could use os.getenv("LOCAL_RANK", '0').

[Yusifu]

hi, I have encountered the same problem. I solved this problem by using trainer = pl.Trainer(..., accelerator='ddp', ...) in main.py and local_rank = os.getenv("LOCAL_RANK") in IterableDataset's iter(self) function.