分布式管道并行简介¶
作者:Howard Huang
注意
在 github 上查看和编辑本教程。
本教程使用 gpt 风格的转换器模型来演示使用 torch.distributed.pipelining API 实现分布式管道并行。
如何使用
torch.distributed.pipelining
API如何将管道并行应用于转换器模型
如何在一组微批次上使用不同的调度
熟悉 PyTorch 中的 基本分布式训练
设置¶
使用 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).float() 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
,在本例中,每个阶段只有一个等级,因此索引等同于等级;以及 num_stages
,它等同于 world_size。
num_stages
用于设置管道并行调度中将使用的阶段数量。例如,对于 num_stages=4
,微批次需要经过 4 次前向传播和 4 次反向传播才能完成。stage_index
对于框架了解如何进行阶段间通信是必要的。例如,对于第一个阶段 (stage_index=0
),它将使用来自数据加载器的数据,不需要从任何之前的对等方接收数据来执行计算。
步骤 1:划分 Transformer 模型¶
有两种不同的模型划分方法。
第一种是手动模式,我们可以通过删除模型属性的某些部分来手动创建模型的两个实例。在本例中,对于 2 个阶段(2 个等级),模型被分成两半。
def manual_model_split(model, example_input_microbatch, model_args) -> 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
stage_input_microbatch = example_input_microbatch
elif stage_index == 1:
# prepare the second stage model
for i in range(4):
del model.layers[str(i)]
model.tok_embeddings = None
stage_input_microbatch = torch.randn(example_input_microbatch.shape[0], example_input_microbatch.shape[1], model_args.dim)
stage = PipelineStage(
model,
stage_index,
num_stages,
device,
input_args=stage_input_microbatch,
)
return stage
如您所见,第一个阶段没有层归一化或输出层,它只包含前四个 Transformer 块。第二个阶段没有输入嵌入层,但包含输出层和最后四个 Transformer 块。然后,该函数返回当前等级的 PipelineStage
。
第二种方法是基于跟踪器的模式,它根据 split_spec
参数自动划分模型。使用管道规范,我们可以指示 torch.distributed.pipelining
在哪里划分模型。在下面的代码块中,我们在第 4 个 Transformer 解码器层之前进行划分,这与上面描述的手动划分相同。同样,在完成划分后,我们可以通过调用 build_stage
来检索 PipelineStage
。
步骤 2:定义主执行¶
在主函数中,我们将创建一个特定的管道调度,阶段应该遵循该调度。torch.distributed.pipelining
支持多种调度,包括每等级单阶段调度 GPipe
和 1F1B
,以及每等级多阶段调度,例如 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, example_input_microbatch, model_args)
# Option 2: Tracer model splitting
# stage = tracer_model_split(model, example_input_microbatch)
x = x.to(device)
y = y.to(device)
def tokenwise_loss_fn(outputs, targets):
loss_fn = nn.CrossEntropyLoss()
outputs = outputs.view(-1, model_args.vocab_size)
targets = targets.view(-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)
dist.destroy_process_group()
在上面的示例中,我们使用手动方法划分模型,但可以取消注释代码以尝试基于跟踪器的模型划分函数。在我们的调度中,我们需要传入微批次的数量和用于评估目标的损失函数。
.step()
函数处理整个小批量,并根据之前传入的 n_microbatches
自动将其划分为微批次。然后,根据调度类对微批次进行操作。在上面的示例中,我们使用 GPipe,它遵循一个简单的前向传播和后向传播调度。等级 1 返回的输出将与模型在单个 GPU 上运行并使用整个批次运行时相同。同样,我们可以传入一个 losses
容器来存储每个微批次相应的损失。
步骤 3:启动分布式进程¶
最后,我们准备运行脚本。我们将使用 torchrun
创建一个单主机、2 个进程的作业。我们的脚本以等级 0 执行管道阶段 0 所需的逻辑,等级 1 执行管道阶段 1 的逻辑的方式编写。
torchrun --nnodes 1 --nproc_per_node 2 pipelining_tutorial.py
结论¶
在本教程中,我们学习了如何使用 PyTorch 的 torch.distributed.pipelining
API 实现分布式管道并行。我们探讨了如何设置环境、定义 Transformer 模型并将其划分以进行分布式训练。我们讨论了两种模型划分方法,手动和基于跟踪器的方法,并演示了如何在不同阶段的微批次上调度计算。最后,我们介绍了管道调度的执行以及使用 torchrun
启动分布式进程。
有关管道并行的生产就绪用法以及与其他分布式技术的组合,请参阅 TorchTitan 3D 并行的端到端示例。