Каков правильный способ параллелизма модели в тензорном потоке?

Question

У меня есть несколько 4 Гб узлов GPU, поэтому я хочу, чтобы они запускали огромную модель параллельно. Надеюсь, что просто разделение слоев на несколько частей с соответствующими областями устройств просто позволяет моделировать параллелизм, но оказывается, что он не уменьшает площадь памяти для главного узла (задача 0). (Конфигурация 10 узлов - master: 20g, последователи: 2g, 1 конфигурация узла - master: 6 ~ 7g)

Подозрительным является то, что градиенты не распределены, потому что я не настроил для них правильную область устройств. .

моя модель доступна на GitHub (https://github.com/nakosung/tensorflow-wavenet/tree/model_parallel_2)

устройство журнала размещения здесь: https://gist.github.com/nakosung/a38d4610fff09992f7e5569f19eefa57

Answer 1

Так что хорошая новость заключается в том, что вы используете colocate_gradients_with_ops, что означает, что вы обеспечиваете, что градиенты быть вычисленный на том же устройстве, что и операционные устройства. (https://github.com/nakosung/tensorflow-wavenet/blob/model_parallel_2/train.py#L242))

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

Во-вторых, вы можете попытаться увидеть, как размеры ваших операций сопоставляются с устройствами - возможно, что самые большие уровни (наибольшие активации или наибольшие веса) могут быть непропорционально размещены на некоторых узлах, чем другие. Вы можете попытаться использовать https://github.com/tensorflow/tensorflow/blob/6b1d4fd8090d44d20fdadabf06f1a9b178c3d80c/tensorflow/python/tools/graph_metrics.py, чтобы проанализировать график, чтобы получить лучшее представление о том, где нужны ресурсы на вашем графике.

Более длительный срок мы хотели бы попытаться решить некоторые из этих проблем с размещением автоматически, но пока модель параллелизма требует некоторой осторожности, чтобы точно разместить вещи.