• Tutorials >
  • Introduction to Distributed Pipeline Parallelism
Shortcuts

分布式管道并行性简介

创建于:2024年7月9日 | 最后更新:2024年12月12日 | 最后验证:2024年11月5日

作者: Howard Huang

注意

editgithub上查看和编辑本教程。

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

What you will learn
  • 如何使用torch.distributed.pipelining API

  • 如何将管道并行应用于变压器模型

  • 如何在一组微批次上利用不同的调度

Prerequisites

设置

使用torch.distributed.pipelining,我们将对模型的执行进行分区,并在微批次上调度计算。我们将使用一个简化版本的变压器解码器模型。该模型架构用于教育目的,并包含多个变压器解码器层,因为我们希望展示如何将模型分割成不同的块。首先,让我们定义模型:

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

rankworld_sizeinit_process_group() 代码对你来说应该很熟悉,因为这些在所有分布式程序中都很常用。管道并行性特有的全局变量包括 pp_group,它是用于发送/接收通信的进程组,stage_index,在这个例子中,每个阶段只有一个 rank,所以索引等同于 rank,以及 num_stages,它等同于 world_size。

num_stages 用于设置管道并行调度中将使用的阶段数。例如,对于 num_stages=4,一个微批次需要经过4次前向和4次反向才能完成。stage_index 是框架知道如何在阶段之间进行通信的必要条件。例如,对于第一个阶段(stage_index=0),它将使用来自数据加载器的数据,并且不需要从任何前一个对等节点接收数据来执行其计算。

第一步:分割Transformer模型

有两种不同的模型分区方式:

首先是手动模式,在这种模式下,我们可以通过删除模型的部分属性来手动创建两个模型实例。在这个例子中,对于两个阶段(2个等级),模型被切成两半。

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

正如我们所看到的,第一阶段没有层归一化或输出层,它只包括前四个变压器块。 第二阶段没有输入嵌入层,但包括输出层和最后四个变压器块。该函数 然后返回当前等级的PipelineStage

第二种方法是基于追踪器的模式,它根据split_spec参数自动分割模型。使用管道规范,我们可以指示torch.distributed.pipelining在何处分割模型。在下面的代码块中,我们在第四个变压器解码器层之前进行分割,与上述手动分割相对应。同样地,在完成分割后,我们可以通过调用build_stage来获取PipelineStage

步骤2:定义主执行

在主函数中,我们将创建一个特定的管道调度,各个阶段应遵循该调度。torch.distributed.pipelining 支持多种调度,包括支持多种调度,如单阶段每秩调度GPipe1F1B, 以及多阶段每秩调度,如Interleaved1F1BLoopedBFS

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是一个简洁、最小化的代码库,用于使用原生PyTorch进行大规模LLM训练。有关生产环境中使用管道并行以及与其他分布式技术组合的示例,请参阅TorchTitan的3D并行端到端示例