How to define my own sampler in ddp?

[zhhao1]

Normally, i will defind my own sampler in dataloader. When using ddp as the accelerator, i don't know where to place it.
If i define a class inherited from pl.LightningDataModule and have my own sampler class inherited from DistributedSampler, which function is right to place my own sampler, def prepare(self) or def setup(self, stage) or def train_dataloader(self) ?

[ydcjeff]

You can add your sampler to sampler or batch_sampler arguments of DataLoader and set False to replace_ddp_sampler of Trainer flag.

[pgmikhael]

Following up on this, custom ddp samplers take rank as an argument and use that to partition the data (e.g., DistributedSampler). In lightning, we would need to pass the global_rank argument to the sampler. However, it seems that global_rank is set after trainer.fit() and within ddp_train. The dataloaders need to be defined before trainer.fit() is called since they are inputs to the function. How can we create our custom sampler and provide it with the correct rank?

[zhhao1]

Following up on this, custom ddp samplers take rank as an argument and use that to partition the data (e.g., DistributedSampler). In lightning, we would need to pass the global_rank argument to the sampler. However, it seems that global_rank is set after trainer.fit() and within ddp_train. The dataloaders need to be defined before trainer.fit() is called since they are inputs to the function. How can we create our custom sampler and provide it with the correct rank?

You can define it in a class inherited from pl.LightningDataModule. global_rank is set before this. I have try it and no error produce.

[pgmikhael]

Following up on this, custom ddp samplers take rank as an argument and use that to partition the data (e.g., DistributedSampler). In lightning, we would need to pass the global_rank argument to the sampler. However, it seems that global_rank is set after trainer.fit() and within ddp_train. The dataloaders need to be defined before trainer.fit() is called since they are inputs to the function. How can we create our custom sampler and provide it with the correct rank?

You can define it in a class inherited from pl.LightningDataModule. global_rank is set before this. I have try it and no error produce.

Could you provide a sample code for this to see if we can reproduce results with no error, specifically including from where global_rank is being obtained to be passed into the sampler?

From what I could tell, global_rank is first set as 0, then changed to the correct rank after the call to ddp_train. I may be missing something -- could you point me to the logic in the code where this is not the case or if LightningDataModule wouldn't suffer from this issue?

[zhhao1]

import pytorch_lightning as pl
from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler
class data_prep(pl.LightningDataModule):
    def __init__(self, args):
        #your code
    def train_dataloader(self):
        train_sampler = DistributedSampler(your_data)
        train_dataloader = DataLoader(your_data, batch_sampler=train_sampler)
    return train_loader
data = data_prep()
trainer = pl.trainer(gpu=4, accelerator='ddp', replace_sampler_ddp=False)
trainer.fit(your_model, data)

Global_rank is set in distributedSampler.

[zhhao1]

Following up on this, custom ddp samplers take rank as an argument and use that to partition the data (e.g., DistributedSampler). In lightning, we would need to pass the global_rank argument to the sampler. However, it seems that global_rank is set after trainer.fit() and within ddp_train. The dataloaders need to be defined before trainer.fit() is called since they are inputs to the function. How can we create our custom sampler and provide it with the correct rank?

You can define it in a class inherited from pl.LightningDataModule. global_rank is set before this. I have try it and no error produce.

Could you provide a sample code for this to see if we can reproduce results with no error, specifically including from where global_rank is being obtained to be passed into the sampler?

From what I could tell, global_rank is first set as 0, then changed to the correct rank after the call to ddp_train. I may be missing something -- could you point me to the logic in the code where this is not the case or if LightningDataModule wouldn't suffer from this issue?

You can see this.

[zhhao1]

Following up on this, custom ddp samplers take rank as an argument and use that to partition the data (e.g., DistributedSampler). In lightning, we would need to pass the global_rank argument to the sampler. However, it seems that global_rank is set after trainer.fit() and within ddp_train. The dataloaders need to be defined before trainer.fit() is called since they are inputs to the function. How can we create our custom sampler and provide it with the correct rank?

You can define it in a class inherited from pl.LightningDataModule. global_rank is set before this. I have try it and no error produce.

Could you provide a sample code for this to see if we can reproduce results with no error, specifically including from where global_rank is being obtained to be passed into the sampler?

From what I could tell, global_rank is first set as 0, then changed to the correct rank after the call to ddp_train. I may be missing something -- could you point me to the logic in the code where this is not the case or if LightningDataModule wouldn't suffer from this issue?

DistributedSampler is triggered at the same time by different. They have different rank(refer to self.rank in DistributedSampler). According this rank, different gpu samples different data.