Обучение модели TensorFlow с последовательным распределением операций на нескольких графических процессорах

Question

Я хочу реализовать модель, операции которой выполняются на разных графических процессорах (и чьи соответствующие веса размещаются на разных графических процессорах). Я знаю, что это не оптимально, и мне, вероятно, следует вместо этого заняться полноценным параллелизмом моделей, но это всего лишь первый шаг.

Когда я попробовал разместить устройства вручную, мне показалось, что мне удалось выполнить прямой проход на нескольких графических процессорах, но не обратное распространение.Как я могу убедиться, что обратное распространение также происходит на нескольких графических процессорах (если даже возможно) в TensorFlow при выполнении параллелизма моделей?

Минимальный воспроизводимый пример моего кода для двух графических процессоров:

import tensorflow as tf
from tensorflow.keras.layers import Conv3D
from tensorflow.keras.models import Model

tf.debugging.set_log_device_placement(True)

class MyBigModel(Model):
    def __init__(self, n_convs=2, n_filters=32, **kwargs):
        super().__init__(**kwargs)
        self.convs = []
        self.n_convs = n_convs
        for i in range(n_convs - 1):
            with tf.device(f'/gpu:{int(i<self.n_convs//2)}'):
                conv = Conv3D(
                    n_filters,
                    3,
                    padding='same',
                    activation='relu',
                )
            self.convs.append(conv)
        with tf.device('/gpu:1'):
            last_conv = Conv3D(
                1,
                3,
                padding='same',
            )
        self.convs.append(last_conv)
        
    def call(self, inputs):
        outputs = inputs
        for i, conv in enumerate(self.convs):
            with tf.device(f'/gpu:{int(i<self.n_convs//2)}'):
                outputs = conv(outputs)
        return outputs

model = MyBigModel(n_convs=10, n_filters=256)
size = 128
inputs = tf.random.normal([1, size, size, size, 1])
outputs = tf.random.normal([1, size, size, size, 1])
res = model(inputs)  # just a forward pass
model.compile(loss='mse', optimizer='sgd')
model.train_on_batch(inputs, outputs)  # forward and backward pass

Я пытался убедиться, что обратное распространение также происходит на нескольких графических процессорах, путем мониторинга использования графических процессоров с помощью watch -n 1 nvidia-smiна глаз. Я заметил ступеньки во время выполнения train_on_batch метод:

• GPU0 на 100%
• GPU1 на 100%
• GPU0 находится на 100% в течение значительного количества времени

Из этого я сделал вывод, что прямой проход действительно происходит на двух графических процессорах, но не обратное распространение, которое происходит только на первом. Возможно, я неправильно диагностирую это, и в этом случае, пожалуйста, скажите мне, как это сделать.

Чтобы прояснить, меня не интересует параллелизм данных (разделение пакетов между разными графическими процессорами), а на самом деле последовательное выполнение разных компонентов моей модели на разных графических процессорах. Конечно mesh-tensorflow- это альтернатива, которую я рассматриваю для параллелизма реальных моделей, но сначала я хотел попробовать все последовательно.