跳到主要内容
Version: v0.1.9

并行配置

作者: Shenggui Li, Siqi Mai

预备知识:

简介

我们在 Colossal-AI 中支持多种并行技术。代码库中的混合并行是指您可以轻松地结合数据并行、流水线并行和张量并行(1D、2D、2.5D、3D)的优势共同来进行并行训练。

每种并行方式需要不同的网络拓扑结构,因此要初始化不同的进程组。您可以通过在配置文件中设置 parallel 来初始化相应的进程组。 parallel 的配置必须遵从以下格式。数据并行度的大小将被根据您对流水线并行和张量并行的输入自动推断。colossalai.launch 将根据您的配置自动初始化这些分布式进程组。

我们为您提供了一些配置的例子以供参考。

# sampler format
parallel = dict(
pipeline=dict("size": int),
tensor=dict("size": int, "mode": '1d' or '2d' or '2.5d' or '3d', "kwargs": Any)
)

# this is ok
parallel = dict(
pipeline=dict(size=2),
tensor=dict(size=4, mode='2d')
)

# this is ok
parallel = dict(
pipeline=2,
tensor=dict(size=4, mode='2d')
)

# this is not ok
# as you need to specify the mode for tensor parallelism
parallel = dict(
pipeline=2,
tensor=4
)

# this is ok as well as tensor will be default to size 1
# and mode None
parallel = dict(
pipeline=2
)

# this is ok as well as pipeline will default to size 1
parallel = dict(
tensor=dict(size=4, mode='2d')
)

关键字 size 指的是并行维度的并行大小。 例如,流水线大小为2意味着有 将有2个流水线阶段。张量并行配置中的关键字 mode 意味着相应的张量并行技术 将被初始化,如1D、2D、2.5D、3D。

您也可以选择不在您的配置中使用 "并行",此时流水线和张量的并行度都将默认为大小1。

GPU的总数量必须等于 数据并行大小 x 张量并行大小 x 流水线并行大小

数据并行

数据并行是最常见的分布式训练方式。它将数据分割成几个碎片分别在每个设备上进行训练。数据并行的配置会自动检测并为您设置。您不需要在您的配置中明确地设置它们。在Colossal-AI 中,有两种方法来处理数据并行的 all-reduce。

  1. 如果您设置了梯度handler,梯度handler将会all-reduce梯度。
  2. 若没有指定相应的配置,Colossal-AI 将会使用 PyTorch 的 DistributedDataParallel。

在大多数情况下,若您对梯度没有复杂的处理的需求,您将会使用第二种模式。

1D, 2D, 2.5D 和 3D 并行

为了实现混合并行,我们提供了一系列张量并行方法。您可以阅读相应的学术论文进行深入的了解。这些并行模式需要和 Colossal-AI 提供的分布式层一同工作。

# 1D parallel
parallel = dict(
tensor=dict(size=4, mode='1d')
)

# 2D parallel
parallel = dict(
tensor=dict(size=4, mode='2d')
)

# 2.5D parallel
parallel = dict(
tensor=dict(size=8, mode='2.5d', depth=2)
)

# 3D parallel
parallel = dict(
tensor=dict(size=8, mode='3d')
)

当您在配置中指定了张量并行模式,您就可以使用其相应的分布式算子。例如,若您设置模式为 2d,那么在模型构建中就能使用 colossalai.nn.Linear2D 了。

流水线并行

流水线并行是将模型按层分成几个部分。例如,假设我们有一个简单的模型,它由两个线性层组成。我们有两个 GPU,我们可以将第一个线性层分配给第一个 GPU 而第二层则分配给第二个 GPU。

您可以在您的配置文件中设置流水线并行度的大小。当流水线并行度大于1,Colossal-AI 将会自动地创建流水线并行的 schedule,这将会为您定义好模型训练的 forwardbackward

parallel = dict(
pipeline=dict(size=4), # number of pipeline stages
)

序列并行

针对处理大图片、视频、长文本、长时间医疗监控等数据的需要,Colossal-AI 还提供了序列并行的方法。该方法是在论文Sequence Parallelism: Making 4D Parallelism Possible中提出的。您可以指定模式为 sequence 来初始化进程组。

parallel = dict(
tensor=dict(size=4, mode='sequence')
)