分布式流水线并行简介¶

Created On: Jul 09, 2024 | Last Updated: Dec 12, 2024 | Last Verified: Nov 05, 2024

作者：Howard Huang

备注

在 github 中查看并编辑本教程。

本教程使用 gpt 风格的 transformer 模型来演示如何使用 torch.distributed.pipelining API 实现分布式流水线并行。

What you will learn

如何使用 torch.distributed.pipelining API
如何将流水线并行应用于 transformer 模型
如何在一组微小批次上使用不同的调度方式

Prerequisites

熟悉 PyTorch 中基本分布式训练

设置¶

使用 torch.distributed.pipelining 时，我们将对模型的执行进行分片，并对微型批次进行计算调度。我们将使用 transformer 解码器模型的简化版本。模型架构是出于教育目的，并且涵盖了多个 transformer 解码器层，因为我们希望演示如何将模型分成不同的块。首先，让我们定义模型：

import torch
import torch.nn as nn
from dataclasses import dataclass

@dataclass
class ModelArgs:
   dim: int = 512
   n_layers: int = 8
   n_heads: int = 8
   vocab_size: int = 10000

class Transformer(nn.Module):
   def __init__(self, model_args: ModelArgs):
      super().__init__()

      self.tok_embeddings = nn.Embedding(model_args.vocab_size, model_args.dim)

      # Using a ModuleDict lets us delete layers witout affecting names,
      # ensuring checkpoints will correctly save and load.
      self.layers = torch.nn.ModuleDict()
      for layer_id in range(model_args.n_layers):
            self.layers[str(layer_id)] = nn.TransformerDecoderLayer(model_args.dim, model_args.n_heads)

      self.norm = nn.LayerNorm(model_args.dim)
      self.output = nn.Linear(model_args.dim, model_args.vocab_size)

   def forward(self, tokens: torch.Tensor):
      # Handling layers being 'None' at runtime enables easy pipeline splitting
      h = self.tok_embeddings(tokens) if self.tok_embeddings else tokens

      for layer in self.layers.values():
            h = layer(h, h)

      h = self.norm(h) if self.norm else h
      output = self.output(h).clone() if self.output else h
      return output

接下来，我们需要在脚本中导入必要库并初始化分布式训练过程。在这种情况下，我们定义了一些全局变量以供以后在脚本中使用：

import os
import torch.distributed as dist
from torch.distributed.pipelining import pipeline, SplitPoint, PipelineStage, ScheduleGPipe

global rank, device, pp_group, stage_index, num_stages
def init_distributed():
   global rank, device, pp_group, stage_index, num_stages
   rank = int(os.environ["LOCAL_RANK"])
   world_size = int(os.environ["WORLD_SIZE"])
   device = torch.device(f"cuda:{rank}") if torch.cuda.is_available() else torch.device("cpu")
   dist.init_process_group()

   # This group can be a sub-group in the N-D parallel case
   pp_group = dist.new_group()
   stage_index = rank
   num_stages = world_size

rank、world_size 和 init_process_group() 的代码应该对你来说很熟悉，因为这些代码是所有分布式程序中常用的。与流水线并行相关的全局变量包括 pp_group，即将用于发送/接收通信的进程组；stage_index，在这个示例中，每个阶段一个 rank，因此索引等同于 rank；还有 num_stages，等同于 world_size。

num_stages 用于设置流水线并行计划中将使用的阶段数。例如，对于 num_stages=4，一个微小批次需要经过 4 次前向传播和 4 次反向传播才能完成。stage_index 是框架知道如何在阶段之间进行通信所必需的。例如，对于第一个阶段（stage_index=0），它将使用数据加载器中的数据，并且不需要从任何其他进程接收数据即可执行其计算。

步骤 1：分割 Transformer 模型¶

有两种不同的模型分割方式：

第一种是手动模式，我们可以通过删除模型属性的一部分来手动创建两个模型实例。在这个示例中，对于两个阶段（2 个 rank），模型被分割为两半。

def manual_model_split(model) -> PipelineStage:
   if stage_index == 0:
      # prepare the first stage model
      for i in range(4, 8):
            del model.layers[str(i)]
      model.norm = None
      model.output = None

   elif stage_index == 1:
      # prepare the second stage model
      for i in range(4):
            del model.layers[str(i)]
      model.tok_embeddings = None

   stage = PipelineStage(
      model,
      stage_index,
      num_stages,
      device,
   )
   return stage

我们可以看到，第一个阶段没有层归一化或输出层，它仅包含前四个 transformer 块。第二个阶段没有输入嵌入层，但包含输出层和最后四个 transformer 块。然后函数返回当前 rank 的 PipelineStage。

第二种方法是基于追踪器的模式，它基于 split_spec 参数自动分割模型。使用流水线规范，我们可以指示 torch.distributed.pipelining 在模型中分割的位置。在下面的代码块中，我们在第 4 个 transformer 解码器层之前进行分割，与上述的手动分割类似。同样，我们可以在分割完成后通过调用 build_stage 获取 PipelineStage。

步骤 2：定义主要执行流程¶

在主函数中，我们将创建一个特定的流水线计划供各个阶段遵循。torch.distributed.pipelining 支持多种计划，包括单阶段单 rank 的 GPipe 和 1F1B，以及多阶段单 rank 的 Interleaved1F1B 和 LoopedBFS。

if __name__ == "__main__":
   init_distributed()
   num_microbatches = 4
   model_args = ModelArgs()
   model = Transformer(model_args)

   # Dummy data
   x = torch.ones(32, 500, dtype=torch.long)
   y = torch.randint(0, model_args.vocab_size, (32, 500), dtype=torch.long)
   example_input_microbatch = x.chunk(num_microbatches)[0]

   # Option 1: Manual model splitting
   stage = manual_model_split(model)

   # Option 2: Tracer model splitting
   # stage = tracer_model_split(model, example_input_microbatch)

   model.to(device)
   x = x.to(device)
   y = y.to(device)

   def tokenwise_loss_fn(outputs, targets):
      loss_fn = nn.CrossEntropyLoss()
      outputs = outputs.reshape(-1, model_args.vocab_size)
      targets = targets.reshape(-1)
      return loss_fn(outputs, targets)

   schedule = ScheduleGPipe(stage, n_microbatches=num_microbatches, loss_fn=tokenwise_loss_fn)

   if rank == 0:
      schedule.step(x)
   elif rank == 1:
      losses = []
      output = schedule.step(target=y, losses=losses)
      print(f"losses: {losses}")
   dist.destroy_process_group()

在上面的示例中，我们使用手动方法分割模型，但代码可以通过取消注释尝试基于追踪器的模型分割功能。在我们的计划中，我们需要传入微小批次数量以及用于评估目标的损失函数。

.step() 函数处理整个小批量，并根据之前传入的 n_microbatches 自动将其分割为微小批次。然后根据计划类对微小批次进行操作。上述示例中，我们使用 GPipe，它遵循一个简单的全前向然后全反向的计划。返回的 rank 1 的输出将与模型在单 GPU 上使用整个批量运行的结果相同。同样，我们可以传入一个 losses 容器来存储每个微小批次的相应损失。

步骤 3：启动分布式进程¶

最后，我们准备运行脚本。我们将使用 torchrun 创建一个单主机、2 个进程的任务。我们的脚本已经写好，rank 0 执行流水线阶段 0 所需的逻辑，而 rank 1 执行流水线阶段 1 的逻辑。

torchrun --nnodes 1 --nproc_per_node 2 pipelining_tutorial.py

结论¶

在本教程中，我们学习了如何使用 PyTorch 的 torch.distributed.pipelining API 实现分布式流水线并行性。我们探索了如何设置环境、定义 Transformer 模型并将其分割用于分布式训练。我们讨论了模型分割的两种方法：手动分割和基于追踪器的分割，并演示了如何跨不同阶段对微批次进行计算计划。最后，我们介绍了流水线计划的执行以及使用 torchrun 启动分布式进程。

附加资源¶

我们已经成功将 torch.distributed.pipelining 集成到 torchtitan 仓库中。TorchTitan 是一个简洁、最小化的大规模 LLM 训练代码库，使用原生 PyTorch。要了解生产级流水线并行性以及与其他分布式技术的组合，请参阅 TorchTitan 的 3D 并行性端到端示例。