Do you think that your case is still in the perimeter of:

  

      github.com/tensorflow/tensorflow
  

  
    
  
	  
  

  
    
      GradientTape: Allow to execute backward functions on same device as forward functions 2
    

    
      
        opened 
        
          
        
        Oct 24, 2019
      
      


      
        
          
          olesalscheider
        
      
    

    
        
          stat:awaiting tensorflower
        
        
          type:feature
        
        
          comp:ops
        
    
  


  
    **System information**
- TensorFlow version (you are using): 2.0
- Are you wil…ling to contribute it (Yes/No): Yes



**Describe the feature and the current behavior/state.**
Currently GradientTape.gradient() is executed on the device of the scope it is called in. Have a look at the following code:

```
with tf.GradientTape() as tape:
    with tf.device('/gpu:1'):
        x = f1(input)
    with tf.device('/gpu:2'):
        x = f2(x)
    with tf.device('/gpu:0'):
        g = tape.gradient(x, f_vars)
```

Here all gradient calculations will be carried out by GPU:0 and all variables needed for the gradient calculation will also be allocated on GPU:0. This is a problem if these temporary variables are too large to fit into the VRAM of GPU:0.

Please provide a way to execute the backward functions on the device of the corresponding forward function and allocate temporary variables for gradient calculation there. This allows to split a large model and distribute it among as many GPUs as necessary.

**Will this change the current api? How?**
It will add a parameter to tf.GradientTape that controls if the user wants the current or the suggested behavior.

**Who will benefit with this feature?**
Anyone who wants to train large models that do not fit into the VRAM of a single GPU.