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(请参阅《PopART用户指南》和PopART C++API的"设置操作的IPU编号"部分以供参考(可将图形拆分为多个部分，以便在多个IPU上运行。如您提到的代码示例所示，单独使用virtualGraph意味着在模型的已分配部分上按顺序运行。

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

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