Model parallelism in Keras does not seem possible as layers cannot be assigned to devices. Working with large data grids and large complex models eg using the Keras functional API,  means that we run out of memory on GPUs very quickly so an approach to model parallelism is essential particularly for the application of AI to science problems. It would be very helpful to understand what is happening here?
If you are interested here is paper that talks about these issues related to weather forecasting:-

  
      

      researchsquare.com
  

  
    

Data driven global weather predictions at high resolutions 4

  Research Square is a preprint platform that makes research communication faster, fairer, and more useful.


  

  
    
    
  

  


I first opened this as a bug under TensorFlow but have not had a response…

  

      github.com/tensorflow/tensorflow
  

  
    
  
	  
  

  
    
      tf.device scope not working correctly 7
    

    
      
        opened 
        
          
        
        Nov 2, 2020
      
      


      
        
          
          JohnTaylor2000
        
      
    

    
        stat:awaiting tensorflower
        type:support
        comp:keras
        comp:gpu
        TF 2.2
    
  


  
    <em>Please make sure that this is a bug. As per our
[GitHub Policy](https://git…hub.com/tensorflow/tensorflow/blob/master/ISSUES.md),
we only address code/doc bugs, performance issues, feature requests and
build/installation issues on GitHub. tag:bug_template</em>

**System information**
- Have I written custom code (as opposed to using a stock example script provided in TensorFlow): Y es
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): ppc64le-linux
- Mobile device (e.g. iPhone 8, Pixel 2, Samsung Galaxy) if the issue happens on mobile device:
- TensorFlow installed from (source or binary)
- TensorFlow version (use command below):
- Python version: python3.7
- Bazel version (if compiling from source):
- GCC/Compiler version (if compiling from source):
- CUDA/cuDNN version:10.1.243
- GPU model and memory: V100 16GB

You can collect some of this information using our environment capture
[script](https://github.com/tensorflow/tensorflow/tree/master/tools/tf_env_collect.sh)
You can also obtain the TensorFlow version with:
1. TF 1.0: `python -c "import tensorflow as tf; print(tf.GIT_VERSION, tf.VERSION)"`
2. TF 2.0: `python -c "import tensorflow as tf; print(tf.version.GIT_VERSION, tf.version.VERSION)"`


**Describe the current behavior**

tf.device command does not correctly assign a GPU device to tf.keras layers on node with 4 GPUs so cannot implement model parallelism. All layers appear on device GPU:0 with the exception of some IO based on output of tf.debugging.set_log_device_placement(True)

**Describe the expected behavior**
tf.keras layers are correctly assigned to a device.

**Standalone code to reproduce the issue**

import tensorflow as tf
from tensorflow import keras

tf.debugging.set_log_device_placement(True)

print("On GPU:1")
inputs = keras.Input(shape=(784,))
with tf.device("/device:GPU:1"): # Or GPU:1 for the 2nd GPU, GPU:2 for the 3rd etc.
   x = keras.layers.Dense(256, activation="relu")(inputs)
   print(x)
   assert x.device.endswith("/GPU:1")

Provide a reproducible test case that is the bare minimum necessary to generate
the problem. If possible, please share a link to Colab/Jupyter/any notebook.

I have a larger test problem that will run on a 4 GPU node. If you turn off the assert statement, then using nvidia-smi you can see that all memory and computational work is happening on GPU:0 and almost none is assigned to other GPUs. Happy to supply this code if needed.

**Other info / logs** Include any logs or source code that would be helpful to
diagnose the problem. If including tracebacks, please include the full
traceback. Large logs and files should be attached.

Tensor("dense/Identity:0", shape=(None, 256), dtype=float32)
Traceback (most recent call last):
  File "py_test.py", line 11, in <module>
    assert x.device.endswith("/GPU:1")
AssertionError