virtualGraph 和 pipelineStage Graphcore 的 PopART/Poplar 库的区别

Question

我正在尝试使用 Graphcore 的PopART 框架（Poplar API 的一部分）来实现一个基本的流水线模型，以加速我的模型，该模型被拆分到多个处理器上。

我正在关注他们的示例代码，但我注意到该示例没有使用在pipelineStage()他们的其他一些应用程序（即Bert）中使用的调用，而是使用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

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

virtualGraph（请参阅设置操作的 IPU 编号部分 PopART 用户指南和PopART C++ API以供参考）可以将图形拆分为多个部分，以在多个 IPU 上运行。使用virtualGraph其自身，如图中的代码示例你简称，装置在模型上的所分配的部分顺序地运行。

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

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