virtualGraph和pipelineStage Graphcore的PopART/Poplar库的差异

Question

我正在尝试使用Graphcore的PopART framework ( Poplar API的一部分)来实现一个基本的流水线模型，以加快我的模型的速度，该模型被分到多个处理器上。

我遵循了他们的example code，但我注意到示例中没有使用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")

相反，Bert示例似乎创建了一个结合了virtualGraph()和pipelineStage()的上下文

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

我不确定哪一个应该是首选的风格。只使用virtualGraph()有什么含义吗？

Answer 1

在Graphcore PopART框架中，virtualGraph和pipelineStage是两个不同的概念，虽然它们也是相关的。

virtualGraph (请参阅PopART用户指南的Setting the IPU number for operations部分和PopART C++ API中的参考)支持将图形拆分为多个部分，以便在多个IPU上运行。如您提到的code example中所示，单独使用virtualGraph意味着在模型的已分配部分上顺序运行。

另一方面，pipelineStage允许您将图形分成几个阶段，这些阶段可以在可能的情况下在不同的IPU上并行运行。您可以灵活地选择应将哪些操作放置在每个管道阶段中。通过选项opts.enablePipelining (请参阅PopART C++ API)启用流水线。默认情况下，管道为每个virtualGraph隐式创建一个pipelineStage。

但是，可以为单个virtualGraph指定多个pipelineStage。当模型的两个部分共享相同的大型数据集时，这可能很有用。这就是为什么在BERT模型中使用pipelineStage和virtualGraph的原因: BERT有一个很大的嵌入矩阵，它在模型的开始和结束时使用。可以将这两个操作放在位于同一IPU上的同一virtualGraph中，这样共享的数据就不会被复制到多个IPU，而是在不同的流水线阶段保持可用-因为它们不是按顺序执行的。