Hi,
I’m trying to figure out how to use Torch.utils.checkpoint.checkpoint to save memory consumption with multi-gpu training (nn.DataParallel or nn.parallel.DistributedDataParallel)
But it seem that using checkpointing and multi-gpu training simultaneously greatly slows down training speed. See:

  
    
    
    Torch.utils.checkpoint.checkpoint with multi-gpus autograd
  
  
    Hi, 
The program is blocked when I use Torch.utils.checkpoint.checkpoint with multi-GPUs. And when I change to single GPU training, the program works! I also try to remove the checkpoint function and the program also works with multi-GPUs. Does checkpoint function not support multi-GPUs training? 
Thanks!
  


and

  
      github.com/gpleiss/efficient_densenet_pytorch
  
  
    
  
	  
  

  
    
      MultiGPU efficient densenets are slow 7
    

    
      
        opened 
        
          
        
        Apr 30, 2018
      
      


      
        
          
          wandering007
        
      
    
  



  I just want to benchmark the new implementation of efficient densenet with the code here. However, it seems that the used...





  
  
    
    
  
  


I’ve found an unofficial implementation of gradient checkpointing for DenseNet which works fine for multi-gpu: https://github.com/csrhddlam/pytorch-checkpoint 12. But I wonder is there any efficient way to use checkpointing and multi-gpu training simultaneously in the official pytorch implementation.
Thanks!