There has been some work on DDP to reduce the memory overhead, perhaps @mrshenli may be able to shed some more light on that.

Thanks @osalpekar. The feature is gradient_as_bucket_view arg in DDP ctor. It will modify param.grad field and let it point to DDP communication bucket views. So that can save one copy of the model. This is still a prototype feature and subject to changes.

  
      pytorch.org
  
  
    

DistributedDataParallel — PyTorch 1.7.0 documentation 8



  
  
    
    
  
  


+1 to @osalpekar’s comment that torchelastic 8 is the recommended solution for DDP OOMs. The RuntimeError you saw is caused by desync, and the desync is caused by OOM in one process. Because DDP expects all processes to launch the same number of AllReduce comm ops in every iteration. If one process hit OOM and skip/redo some comm ops, it will break this assumption. TorchElastic handles this problem by checkpointing model states and let the entire DDP gang to revert to the previous checkpoint when it detects failures in any process, which seems to be very nice fit for this infrequent and unpredictable OOM.