Federated Tensorflow memory-overloads a single GPU

Admin 1
Viewed 222

I'm trying to do federated learning using tensorflow. I've created a model based on this tutorial for MNIST. Now I change it: I supply it with images of size 112x112 with 3 channels each (I.e. the size of the input layer is 112x112x3). When I try to use multiple models (around 50 with about 100 images each), I get "out of memory" exception. Looking at output of nvidia-smi, it makes sense: for some reason my program takes memory only from a single GPU. How can I avoid it?

Some comments:

  • I use allow_growth, so I'm pretty sure that the shown memory usage is the actual memory usage. And I do can see it growing over time.
  • TensorFlow recognizes other GPUs: before I run the program, they have 0MB occupied. Also there is line Adding visible gpu devices: 0, 1, 2, 3, 4, 5, 6, 7 in logs.

enter image description here

tensorflow-federated tensorflow out-of-memory python
1 Answers

You have hit an interesting and thorny edge in TFF.

TLDR: TFF uses TF in an unorthodox manner; this hits a bug in TF, in which device placement fails to propagate through tf.data.Dataset.reduce correctly, and TFF's ops only get placed on a single GPU.

A TFF contributor is currently working on adding a workaround to TFF's default configuration, which essentially involves more-or-less "tricking" TensorFlow's AutoGraph to generate a tf.while loop instead of a dataset reduction. If you are building from source, you can swap the dataset.reduce for a specific for-loop:

num_examples_sum = 0
for batch in iter(dataset):
  num_examples_sum = reduce_fn(num_examples_sum, batch)

This should cause all the GPUs to be utilized.

If they are still ooming, try adjusting the clients_per_thread parameter (the inverse of the former num_client_executors parameter) upwards; TFF will attempt to run all clients in parallel by default.

user_75079320
Related