记录-DeepSpeed框架分布式训练(一)

URL

type

status

date

slug

summary

📝 3D并行

什么是3D并行？

3D并行中的3指的是3中并行方法，D指的是维度（以维度理解是为了便于大家计算GPU数量）。

名称	含义
PP	流水线并行，用于将模型划分成多个阶段。
TP	张量并行，用于将模型参数拆分。
DP	数据并行，用于将数据拆分，可以理解为将一个batch_size拆分为多个mirco_batch_size。

上述的三种方法分别和Pytorch中的PP、TP和DP理论相似，这里博主不在详细记录。伙伴们可以查看Pytorch框架分布式训练系列的博客。

Tips

3D并行中的D指的是维度。因此，在计算需要多少张显卡时，不能随意指定。

eg:假设我们现在想要实现流水线并行。Deepspeed的流水线并行需要将模型分为不同阶段（简单理解为分层），将不同阶段分发到不同的GPU上（假设模型分为4个阶段，需要4张GPU）。在训练时，通过通信协议进行数据传播。

如果我们的GPU资源充足，想要在流水线并行的基础上增加数据并行。假设我们想要2份数据进行数据并行，那么数据并行将batch_size拆分为batch_size/2。数据并行的前提是模型的完整性，所以2份数据并行，需要2份完整的模型才能进行。如下图，数据并行时，每一份数据都会对应一份完整的模型。

所以，最终我们需要4*2=8张GPU进行。

在使用流水线并行时，拆分阶段，请尽量设置成2的倍数。

在处理流水线并行时，需要注意以下问题：

流水线并行拆分的各个阶段(Stage)的输入和输出只能是单个张量或者一个张量元组。

在实现流水线并行时，我们需要注意各个GPU负载均衡。不可以某个GPU负责很多层，某个GPU只负责一层等。

模型有些层是共享的，在构建流水线模型时，如何处理这些共享层。

CPU内存优化问题：在使用pytorch的Sequential进行序列化后实现流水线并行时，每个GPU节点会在CPU备份一个完整的模型。这种情况会导致CPU内存占用暴增。

DeepSpeed针对上诉问题做了一些优化。接下来，我们一起了解DeepSpeed的优化方法。

流水线并行优化

这里博主使用DeepSpeed官方提供的案例进行记录。

输入和输出

模型经过序列化后，输入和输出只能是单个张量或者一个张量元组。以transformer的attention_block为例(简写)：

GPU负载均衡

DeepSpeed的流水线对象提供了设置负载均衡的partition_method参数-[parameters,uniform,type:[attention_block]]。其中，parameters按照参数量划分阶段，uniform按照层均匀划分阶段，type:[attention_block]按照层名划分阶段。

CPU内存优化

博主将DeepSpeed优化的方法和普通方法进行对比，观察CPU内存占用率，发现对应的线程并没有太大的差距。以下是博主的代码，如有问题请及时指正！

模型层共享

PP、TP和DP如何结合起来一起使用？

PP和DP

因为博主只有两块GPU，在测试PP和DP时，博主将PP的num_stages设置为1（即每个GPU节点拷贝一个完整的模型，但是可以设置micro_batch_size）。

模型推理输出结果：

PP和TP

DeepSpeed中的TP和PP不能一起使用，会报错。AssertionError: mpu must be None with pipeline parallelism。mpu是设置张量并行的参数。所以，我们需要自己实现相应的张量并行(TP)，以及GPU间的通信。同时，我们利用DeepSpeed的并行拓扑结构进行管理。代码如下(因为博主只有两个GPU，所以流水线并行阶段设置为1，张量并行设置为2，数据并行设置为1)：

LEVSONGSW/DeepLearnLog: 深度学习日志记录