AI大语言模型之分布式训练策略

一、前言

流水线并行是模型并行的一种，核心按照模型的层切分到不同设备进行并行计算。

二、流水线并行

流水线并行是一种并行计算策略，将模型的各个层分段处理，并将每个段分布在不同的计算设备上，使得前后阶段能够流水式、分批进行工作。流水线行通常应用于大语言模型的并行系统中，以有效解决单个计算设备内存不足的问题。

给出了一个由四个计算设备组成的流水线并行系统，包含了前向计算和后向计算。其中F1、F2、F3、F4 分别代表四个前向路径，位于不同的设备上；而B4、B3、B2、B1 则代表逆序的后向路径，也分别位于四个不同的设备上。

从图中可以看出，计算图中的下游设备（Downstream Device）需要长时间持续处于空闲状态，等待上游设备（Upstream Device）的计算完成，才能开始计算自身的任务。这种情况导致了设备的平均使用率大幅降低，形成了模型并行气泡（Model Parallelism Bubble），也称为流水线气泡（Pipeline Bubble）。

2.1、GPipe流水线方法

GPipe方法将小批次（Mini-batch）进一步划分成更小的微批次（Micro-batch），利用流水线并行方案，每次处理一个微批次的数据。在当前阶段计算完成得到结果后，将该微批次的结果发送给下游设备，同时开始处理后一个微批次的数据，这样可以在一定程度上减少并行气泡。

前向F1计算被拆解为了F11、F12、F13、F14，在计算设备1 中计算完成F11 后，会在计算设备2 中进行F21 计算，同时在计算设备1 中并行计算F12。相比于最原始的流水线并行方法，GPipe 流水线方法可以有效降低并行气泡。

2.2 1F1B 流水线策略

1F1B 流水线策略，是一个前向通道和一个后向通道。1F1B 流水线策略引入了任务调度机制，使得下游设备能够在等待上游计算的同时执行其他可并行的任务，从而提高设备的利用率。1F1B 给出了非交错式和交错式两种调度模式。

2.2.1 1F1B 非交错式调度模式

1F1B 非交错式调度模式可分为三个阶段

• 首先是热身阶段，在计算设备中进行不同数量的前向计算。
• 接下来的阶段是前向-后向阶段，计算设备按顺序执行一次前向计算，然后进行一次后向计算。
• 最后一个阶段是后向阶段，计算设备再完成最后一次后向计算。 相比于GPipe 策略，非交错式调度模式在节省内存方面表现得更好。然而，它需要与GPipe 策略一样的时间来完成一轮计算。

2.2.2 1F1B 交错式调度模式

1F1B 交错式调度模式要求Micro-batch 的数量是流水线阶段的整数倍。每个设备不仅负责连续多个层的计算，还可以处理多个层的子集，这些子集被称为模型块。设备1 可以处理层1、2、9、10，设备2 处理层3、4、11、12，依此类推。在这种模式下，每个设备在流水线中被分配到多个阶段。