Setup is running more than one time on ddp in rank 0. That is not happening on ddp_spawn.

[dhorka]

Hi,

I am a newbie using pl and I am trying to understand how it works with the DDP mode. To do that I create this silly code to try to understand the behaviour (it is a bit messy is just to understand how pl works):

import pytorch_lightning as pl    
import torch    
      
                                                                                                                                                                                                  
class TestModule(pl.LightningModule):                                                                                                                                                           
	def __init__(self):                                                                                                                                                                         
	  super(TestModule, self).__init__()                                                                                                                                                      
	  self.nn = torch.nn.Linear(3, 500)                                                                                                                                                       
	                                                                                                                                                                                          
	def training_step(self, batch, batch_indx):                                                                                                                                                 
	  print('rank: {} x: {}'.format(self.global_rank, batch[0]))                                                                                                                              
	                                                                                                                                                                                          
	def configure_optimizers(self):                                                                                                                                                                                                                                                                                                                            
	return torch.optim.Adam(self.parameters(), lr=0.001)                                                                                                                                    
	                                                                                                                                                                                          
	def prepare_data(self):                                                                                                                                                                     
	  print('Preparing train_data on rank {}'.format(self.global_rank))                                                                                                                                                                                                                                     
	                                                                                                                                                                                          
	def setup(self, stage):                                                                                                                                                                     
	  print('setup on rank: {}'.format(self.global_rank))                                                                                                                                     
	  self.train_data = torch.utils.data.TensorDataset(torch.randn(4, 1))                                                                                                                     
	                                                                                                                                                                                          
	def train_dataloader(self):                                                                                                                                                                 
          return torch.utils.data.DataLoader(self.train_data, batch_size=1, num_workers=2)                                                                                                        
                                                                                                                                                                                          
                                                                                                                                                                                          
if __name__ == '__main__':                                                                                                                                                                      
                                                                                                                                                                                          
  pl.seed_everything(0)                                                                                                                                                                       
                                                                                                                                                                                            
  net = TestModule()                                                                                                                                                                          
                                                                                                                                                                                           
  trainer = pl.Trainer(gpus=-1, max_epochs=2, auto_select_gpus=True, accelerator='ddp')                                                                                                       
  trainer.fit(net) 

The output of this program is the following one:

Global seed set to 0                                                                                                                                                                               
GPU available: True, used: True                                                                                                                                                                    
TPU available: None, using: 0 TPU cores                                                                                                                                                            
Preparing train_data on rank 0                                                                                                                                                                     
setup on rank: 0
Global seed set to 0
Preparing train_data on rank 0
setup on rank: 0
Global seed set to 0
initializing ddp: GLOBAL_RANK: 1, MEMBER: 2/4
Global seed set to 0
Preparing train_data on rank 0
setup on rank: 0
Global seed set to 0
initializing ddp: GLOBAL_RANK: 2, MEMBER: 3/4
Global seed set to 0
Preparing train_data on rank 0
setup on rank: 0
Global seed set to 0
initializing ddp: GLOBAL_RANK: 3, MEMBER: 4/4
Global seed set to 0
initializing ddp: GLOBAL_RANK: 0, MEMBER: 1/4

  | Name | Type   | Params
--------------------------------
0 | nn   | Linear | 2.0 K
--------------------------------
2.0 K     Trainable params
0         Non-trainable params
2.0 K     Total params
0.008     Total estimated model params size (MB)
Epoch 0:   0%|                                                                                                                                                               | 0/1 [00:00<?, ?it/s]
rank: 3 x: tensor([[-1.2829]], device='cuda:3')
rank: 1 x: tensor([[1.3468]], device='cuda:1')
rank: 2 x: tensor([[0.0642]], device='cuda:2')
rank: 0 x: tensor([[0.2676]], device='cuda:0')

Epoch 1:   0%|                                                                                                                                           | 0/1 [00:00<?, ?it/s, loss=nan, v_num=28]
rank: 2 x: tensor([[-1.2829]], device='cuda:2')
rank: 3 x: tensor([[0.2676]], device='cuda:3')
rank: 1 x: tensor([[0.0642]], device='cuda:1')
rank: 0 x: tensor([[1.3468]], device='cuda:0')

Reading the documentation about setup and prepare_data I was expecting a different output.
I was expecting to have just one print of the prepare_data hook in the main process. And multiple prints of the prepare hook one for each rank process.

Am I doing something wrong or I misunderstood something?

Thanks for your time,

[dhorka]

Hi, after doing some more tests I still do not understand why prepare_data runs two times in rank 0. In fact, after doing some more tests I observed that configuring optimizers is also running two times in rank 0. In fact, If increase the number of GPUs the numbers of runs of prepare_data and configuring_optimizers in rank 0 increases.

Code

import pytorch_lightning as pl    
import torch    
      
      
  class TestModule(pl.LightningModule):    
      def __init__(self):    
          super(TestModule, self).__init__()    
          self.nn = torch.nn.Linear(3, 500)    
      
      def training_step(self, batch, batch_indx):    
          pass    
      
      def configure_optimizers(self):    
          print('Configuring optimizers on rank {}'.format(self.global_rank))    
          return torch.optim.Adam(self.parameters(), lr=0.001)                                                                                                                                    
                                                                                                                                                                                                  
                                                                                                                                                                                                  
       def setup(self, stage):                                                                                                                                                                     
          print('Running: Setup on rank {}'.format(self.global_rank))                                                                                                                             
          x = torch.randn(2, 1)                                                                                                                                                                   
          self.train_data = torch.utils.data.TensorDataset(x)                                                                                                                                     
                                                                                                                                                                                                  
      def train_dataloader(self):                                                                                                                                                                 
         return torch.utils.data.DataLoader(self.train_data, batch_size=1, num_workers=2)    
      
      
  if __name__ == '__main__':                                                                                                                                                                      
                                                                                                                                                                                                  
      pl.seed_everything(0)                                                                                                                                                                       
                                                                                                                                                                                                  
      net = TestModule()                                                                                                                                                                          
                                                                                                                                                                                                  
      trainer = pl.Trainer(gpus=8, max_epochs=0, auto_select_gpus=True, accelerator='ddp')                                                                                                        
      trainer.fit(net)   

Output

Global seed set to 0                                                                                                                                                                               
GPU available: True, used: True                                                                                                                                                                    
TPU available: None, using: 0 TPU cores                                                                                                                                                            
Running: Setup on rank 0
Global seed set to 0
Running: Setup on rank 0
Configuring optimizers on rank 0
Global seed set to 0
initializing ddp: GLOBAL_RANK: 1, MEMBER: 2/8
Global seed set to 0
Running: Setup on rank 0
Configuring optimizers on rank 0
Global seed set to 0
initializing ddp: GLOBAL_RANK: 2, MEMBER: 3/8
Global seed set to 0
Running: Setup on rank 0
Configuring optimizers on rank 0
Global seed set to 0
initializing ddp: GLOBAL_RANK: 3, MEMBER: 4/8
Global seed set to 0
Running: Setup on rank 0
Configuring optimizers on rank 0
Global seed set to 0
initializing ddp: GLOBAL_RANK: 4, MEMBER: 5/8
Global seed set to 0
Running: Setup on rank 0
Configuring optimizers on rank 0
Global seed set to 0
initializing ddp: GLOBAL_RANK: 5, MEMBER: 6/8
Global seed set to 0
Running: Setup on rank 0
Configuring optimizers on rank 0
Global seed set to 0
initializing ddp: GLOBAL_RANK: 6, MEMBER: 7/8
Configuring optimizers on rank 0
Global seed set to 0
initializing ddp: GLOBAL_RANK: 0, MEMBER: 1/8
Global seed set to 0
Running: Setup on rank 0
Configuring optimizers on rank 0
Global seed set to 0
initializing ddp: GLOBAL_RANK: 7, MEMBER: 8/8

As far as I understood the setup method is used to download just one time a dataset when you are doing ddp. Am I right? In the case of the optimizer, I was expecting to obtain one for each gpu, however, I see that just one on rank 0 is obtained (but is obtained multiple times for the same rank). Moreover, I observed that this is not happening on ddp_spawn.

[mosheliv]

I think global_rank does not represent you gpu rank correctly.... I think its always 0. I am also looking for a way to know on what gpu I am running.