In no way to steal from the focus of this super important issue you raise, I am curious how your code would fair with the DeepSpeed integration - since it manages temp memory allocations itself. And in my few experiments it needs 3-5x less gpu RAM with its ZeRO algorithms. that is if you’re using the HF trainer. If you don’t use the latter, you can see from the code how to activate deepspeed in your own trainer.
We are getting close to completing the integration but you can already try my branch:
  
      github.com/huggingface/transformers
  
  
    
  
    
  

  
    
      [trainer] deepspeed integration 10
    

    
      huggingface:master ← stas00:ds
    

    
      
        opened 
        
          
        
        Dec 19, 2020
      
      

      
        
          
          stas00
        
      

      
        
          +384
          -23
        
      
    

  


  
  
    
    
  
  


There is a doc in it that explains how to activate it.
You’d need to install DeepSpeed from its master.
If you have any follow up questions please ask in that PR thread so that we don’t derail this thread.
We also recently added support for Sharded DDP via fairscale 7. if you use HF trainer, just install fairscale and add --sharded_ddp to the training args and you should also see a huge improvement in memory utilization. Again if you have questions let’s continue this discussion on HF forums.
Back to the topic of this thread.