流水并行训练

在 常见的分布式并行策略 一文中介绍了流水并行的特点。

在 OneFlow 的 一致性视角 下，通过简单的设置 Tensor 的 placement 属性，就可以实现流水并行。

以下代码是简单的示范，它将 快速上手 中的网络，以流水并行的方式运行。前几层的 Module nn.Flatten、nn.Linear(28*28, 512)、nn.ReLU() 在 GPU0 上运行；剩余的网络部分在 GPU1 上运行。

Code

  1. import oneflow as flow
  3. BATCH_SIZE = 16
  4. BROADCAST = [flow.sbp.broadcast]
  5. P0 = flow.placement("cuda", {0: [0]})
  6. P1 = flow.placement("cuda", {0: [1]})
  8. class Stage0Module(flow.nn.Module):
  9.     def __init__(self):
  10.         super().__init__()
  11.         self.flatten = flow.nn.Flatten()
  12.         self.linear0 = flow.nn.Linear(28*28, 512)
  13.         self.relu0 = flow.nn.ReLU()
  15.     def forward(self, x):
  16.         out = self.flatten(x)
  17.         out = self.linear0(out)
  18.         out = self.relu0(out)
  19.         return out
  21. class Stage1Module(flow.nn.Module):
  22.     def __init__(self):
  23.         super().__init__()
  24.         self.linear1 = flow.nn.Linear(512, 512)
  25.         self.relu1 = flow.nn.ReLU()
  26.         self.linear2 = flow.nn.Linear(512, 10)
  27.         self.relu2 = flow.nn.ReLU()
  29.     def forward(self, x):
  30.         out = self.linear1(x)
  31.         out = self.relu1(out)
  32.         out = self.linear2(out)
  33.         out = self.relu2(out)
  34.         return out
  36. class PipelineModule(flow.nn.Module):
  37.     def __init__(self):
  38.         super().__init__()
  39.         self.m_stage0 = Stage0Module()
  40.         self.m_stage1 = Stage1Module()
  42.         self.m_stage0.to_consistent(placement=P0, sbp=BROADCAST)
  43.         self.m_stage1.to_consistent(placement=P1, sbp=BROADCAST)
  45.     def forward(self, x):
  46.         out_stage0 = self.m_stage0(x)
  47.         in_stage1 = out_stage0.to_consistent(placement=P1, sbp=BROADCAST)
  48.         out_stage1 = self.m_stage1(in_stage1)
  49.         return out_stage1
  51. module_pipeline = PipelineModule()
  52. sgd = flow.optim.SGD(module_pipeline.parameters(), lr=0.001)
  54. class PipelineGraph(flow.nn.Graph):
  55.     def __init__(self):
  56.         super().__init__()
  57.         self.module_pipeline = module_pipeline
  58.         self.module_pipeline.m_stage0.config.stage_id = 0
  59.         self.module_pipeline.m_stage1.config.stage_id = 1
  60.         self.loss_fn = flow.nn.CrossEntropyLoss()
  61.         self.config.set_gradient_accumulation_steps(2)
  62.         self.add_optimizer(sgd)
  64.     def build(self, x, y):
  65.         out = self.module_pipeline(x)
  66.         loss = self.loss_fn(out, y)
  67.         loss.backward()
  68.         return loss
  70. graph_pipeline = PipelineGraph()
  72. x = flow.randn(BATCH_SIZE, 1, 28, 28)
  73. x = x.to_consistent(P0, BROADCAST)
  74. y = flow.randint(0, 10, (BATCH_SIZE,))
  75. y = y.to_consistent(P1, BROADCAST)
  77. for i in range(20):
  78.     loss = graph_pipeline(x, y)
  79.     print(loss.to_local())

以上代码，保存为脚本（如 pipeline.py）后，使用 launch 模块启动分布式训练：

  1. python3 -m oneflow.distributed.launch --nproc_per_node 2 ./pipeline.py

代码解读

设置 placement 与 sbp

将需要使用的 placement 与 sbp 设置提前准备好：

  1. BROADCAST = [flow.sbp.broadcast]
  2. P0 = flow.placement("cuda", {0: [0]})
  3. P1 = flow.placement("cuda", {0: [1]})

P0、P1 分别代表第0号机器上的第0个 GPU 和第1个 GPU。

通过调用 nn.Module.to_consistent 或 Tensor.to_consistent 就可以将模型或张量分配到指定的计算设备上运行，将一个网络拆分为多个流水阶段（stage）。

在此我们定义了一个 PipelineModule 专门设置各阶段的流水。

  1.     class PipelineModule(flow.nn.Module):
  2.         def __init__(self):
  3.             #...
  5.             self.m_stage0.to_consistent(placement=P0, sbp=BROADCAST)
  6.             self.m_stage1.to_consistent(placement=P1, sbp=BROADCAST)
  8.         def forward(self, x):
  9.             out_stage0 = self.m_stage0(x)
  10.             in_stage1 = out_stage0.to_consistent(placement=P1, sbp=BROADCAST)
  11.             out_stage1 = self.m_stage1(in_stage1)
  12.             return out_stage1

Local Tensor 与 Consistent Tensor 的转换

示例中使用了随机生成的数据作为输入。

  1.     x = flow.randn(BATCH_SIZE, 1, 28, 28)
  2.     x = x.to_consistent(P0, BROADCAST)

当使用 launch 模块启动训练时，因为命令行参数为 --nproc_per_node 2，launch 会启动 2 个进程。两个进程均为执行脚本中的代码。

其中 x = flow.randn(BATCH_SIZE, 1, 28, 28) 返回的是 Local Tensor（只在本进程中有效的本地数据），当运行 x = x.to_consistent(P0, BROADCAST) 时，OneFlow 会自动将所有进程中的 Local Tensor 整合为 Consistent Tensor。

在实际训练中，各个计算设备也可以加载属于各自的本地数据，然后通过 to_consistent 实现 Local Tensor 到 Consistent Tensor 的转化。

Stage ID 及 梯度累积设置

通过设置 Module 的 config.stage_id 属性，设置 Stage ID，Stage ID 从0开始编号，依次加1。 调用 self.config.set_gradient_accumulation_steps 方法，设置梯度累积的步长。 OneFlow 通过这两项配置，获取实现流水并行中的 micro batch 技术所需的信息。

  1.     self.module_pipeline.m_stage0.config.stage_id = 0
  2.     self.module_pipeline.m_stage1.config.stage_id = 1
  3.     self.config.set_gradient_accumulation_steps(2)

为正常使用来必力评论功能请激活JavaScript