Model parallelism in Keras

JOHN_TAYLOR · August 18, 2021, 2:27am

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:-

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

opened 04:15AM - 02 Nov 20 UTC

closed 02:22PM - 26 Apr 22 UTC

JohnTaylor2000

stat:awaiting response type:support stalled 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

Stonepia · August 18, 2021, 5:20am

As far as I know, even if the model you wrote in the issue could work, it will encounter the curse of dimensions because of the transferring time, so you would expect a super low scaling efficiency for your simple implementation.

I would recommend you only do model parallelism unless you couldn’t even fit one reasonably small batch size. Because the model parallel could always expect smaller scaling efficiency than data parallel currently (from my knowledge).

Considering model parallelism, there is a work called mesh-tensorflow, you may refer to that for your purpose. Secondly, you could also write a custom distributed strategy if you wish (similar to TPUStrategy, but that’s too tedious for just one model testing and I wouldn’t recommend you to do so).

Stonepia · August 18, 2021, 5:45am

Besides, there are two issues that I think why your model could not work (not tested, just thoughts):

You split models on different devices, but the inputs is not explicitly transferred to the device. Thus maybe the computing graph may experience unfound variables(possibly).
You may think the model parallelism as a serialized model: y1 = model_part0 (input) , y2 = model_part1(y1). Thus, you could not write something like

with tf.device("/device:GPU:0"):
  x = ...
with tf.device("/device:GPU:1"):
  x = ...

with tf.device("/device:GPU:2"):
  x = f(x) # <- This will get confused

This may make the graph have duplicated variable. That’s why the assertion assert x.device.endswith("/GPU:1") failed.