Разница между virtualGraph и pipeStage Graphcore в библиотеках PopART / Poplar

Question

Я пытаюсь реализовать основную c конвейерную модель, используя PopART Framework Graphcore (часть Poplar API), чтобы ускорить мою модель, которая разделена на несколько процессоров.

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

Небольшой фрагмент приведенного ниже примера:

#  Dense 1
W0 = builder.addInitializedInputTensor(
    init_weights(num_features, 512))
b0 = builder.addInitializedInputTensor(init_biases(512))

with builder.virtualGraph(0):
    x1 = builder.aiOnnx.gemm([x0, W0, b0], debugPrefix="gemm_x1")
    x2 = builder.aiOnnx.relu([x1], debugPrefix="relu_x2")

#  Dense 2
W1 = builder.addInitializedInputTensor(init_weights(512, num_classes))
b1 = builder.addInitializedInputTensor(init_biases(num_classes))

with builder.virtualGraph(1):
    x3 = builder.aiOnnx.gemm([x2, W1, b1], debugPrefix="gemm_x3")
    x4 = builder.aiOnnx.relu([x3], debugPrefix="relu_x4")

И наоборот, пример Берт, кажется, создает контекст, который объединяет virtualGraph() с pipelineStage():

self.stack.enter_context(self.builder.pipelineStage(self.pipelineStage))

Я не уверен, какой стиль должен быть предпочтительным. Есть ли какие-либо последствия только для использования virtualGraph()?

Answer 1

virtualGraph и pipelineStage - это две разные концепции в среде GraphCore PopART, хотя они также связаны между собой.

virtualGraph (см. Установка номера IPU для операций в Руководстве пользователя PopART и PopART C ++ API для справки) позволяет разбить график на несколько части, для запуска на нескольких IPU. Использование virtualGraph само по себе, как показано в примере кода , на который вы ссылались, означает последовательный запуск выделенных частей модели.

С другой стороны, pipelineStage позволяет вам разбить ваш график на несколько этапов, которые, по возможности, могут выполняться параллельно на разных IPU. У вас есть возможность выбрать, какие операции следует размещать на каждой стадии конвейера. Конвейерная обработка включается параметром opts.enablePipelining (см. PopART C ++ API ). По умолчанию конвейерная передача неявно создает один pipelineStage на virtualGraph.

Однако можно указать более одного pipelineStage для одного virtualGraph. Это может быть полезно, когда две части модели совместно используют один и тот же большой набор данных. Вот почему комбинация pipelineStage и virtualGraph используется в модели BERT : BERT имеет большую матрицу вложений, которая используется в начале и в конце модели. Можно разместить обе операции в одном и том же virtualGraph, который будет расположен на одном и том же IPU, чтобы общие данные не копировались на несколько IPU, а оставались доступными на разных этапах конвейера, поскольку они не выполняются последовательно.