Have you tried with these envs on TF 2.7:

  

      github.com/tensorflow/tensorflow
  

  
    
  
    
  

  
    
      [PJRT] Allow GPU memory oversubscription when unified memory is enabled. 7
    

    
      
        committed 
        
          
        
        Jul 9, 2021
      
      

      
        
          
          tensorflower-gardener
        
      

      
        
          +5
          -2
        
      
    
  


  
    With this CL, we can enable GPU memory oversubscription via env flags.
For examp…le, `TF_FORCE_UNIFIED_MEMORY=1 XLA_PYTHON_CLIENT_MEM_FRACTION=8.0` provides 8x the GPU memory to the program. The 'extra' memory is physically located on the other GPU devices and the host's RAM, with swapping done transparently by CUDA.

PiperOrigin-RevId: 383819164
Change-Id: Id139d3184d3a62983c1e86bf95ca4078a08db4f4