数据并行通过分散数据有效地扩展了训练,但它依赖一个主要假设:模型本身可以放入单个加速器的内存中。对于先进模型,特别是在自然语言处理和计算机视觉方面,这个假设不再成立。当模型的参数、梯度和优化器状态超出单个GPU的可用显存,就需要一种不同的方法。将模型本身分散到多个设备上的方案能够解决这个问题。层间模型并行 (张量并行)处理单个无法放入一个GPU的庞大层最直接的方法是拆分层的计算,这种方法通常称为张量并行。考虑一个大型线性层,由方程 $Y = XA$ 定义,这里 $X$ 是输入激活,$A$ 是权重矩阵。如果矩阵 $A$ 对单个设备来说太大,我们可以将其划分到多个GPU上。例如,将线性层的权重矩阵 $A$ 按列拆分到两个GPU上 ($A = [A_1, A_2]$),我们就可以并行计算输出的部分。$$ Y = X[A_1, A_2] = [XA_1, XA_2] $$GPU 1 计算 $Y_1 = XA_1$,GPU 2 计算 $Y_2 = XA_2$。输入 $X$ 被广播到两个GPU。并行计算完成后,结果 $Y_1$ 和 $Y_2$ 被聚合以形成最终输出 $Y$。这个过程需要较多的通信:输入必须发送到所有设备,并且部分结果必须组合。digraph G { rankdir=TB; splines=false; node [shape=record, style="rounded,filled", fontname="Arial", fontsize=10]; edge [fontname="Arial", fontsize=9]; bgcolor="transparent"; newrank=true; subgraph cluster_0 { style=invis; label=""; {rank=same; "Input_X" [label="输入 (X)", shape=box, style="filled", fillcolor="#e9ecef"];} } subgraph cluster_1 { style=invis; label=""; {rank=same; "GPU_0" [fillcolor="#a5d8ff", label="{<f0> GPU 0 | 存储 A_1 | 计算 Y_1 = XA_1}"]; "GPU_1" [fillcolor="#96f2d7", label="{<f0> GPU 1 | 存储 A_2 | 计算 Y_2 = XA_2}"]; } } subgraph cluster_2 { style=invis; label=""; {rank=same; "AllGather" [label="All-Gather\n(连接)", shape=cds, style="filled", fillcolor="#ffc9c9"];} } subgraph cluster_3 { style=invis; label=""; {rank=same; "Output_Y" [label="输出 (Y)", shape=box, style="filled", fillcolor="#e9ecef"];} } "Input_X" -> "GPU_0" [label="广播 X"]; "Input_X" -> "GPU_1" [label="广播 X"]; "GPU_0" -> "AllGather" [label="Y_1"]; "GPU_1" -> "AllGather" [label="Y_2"]; "AllGather" -> "Output_Y" [label="[Y_1, Y_2]"]; }一张张量并行图,其中单个层的权重矩阵被拆分到两个GPU上。输入被广播,并且部分结果通过All-Gather过程组合。反向传播遵循类似的反向方式。关于输出的梯度 $\nabla_Y$ 被拆分并发送到各自的GPU。每个GPU计算其局部权重梯度,并且关于输入的梯度在使用Reduce-Scatter过程组合后,才传递给上一层。这种细粒度并行带来了高频通信,使其对互连带宽高度敏感。它在具有NVLink等高速连接的多GPU服务器内部最为有效。流水线并行对于非常深的模型,另一种方法是在层之间而不是之内划分模型。这被称为流水线并行。在这种方法中,连续的层块(称为阶段)被放置在不同的加速器上。例如,在一个运行在4个GPU上的32层模型中:GPU 0 负责处理第1-8层。GPU 1 负责处理第9-16层。GPU 2 负责处理第17-24层。GPU 3 负责处理第25-32层。一个数据批次被送入GPU 0。一旦GPU 0完成其第1-8层的前向传播,它就将产生的激活传递给GPU 1,GPU 1随后开始其计算。这个过程沿着流水线顺序继续。digraph G { rankdir=TB; splines=true; node [shape=box, style="rounded,filled", fontname="Arial", fontsize=10]; edge [fontname="Arial", fontsize=9]; bgcolor="transparent"; "Input" [style="filled", fillcolor="#e9ecef", label="输入"]; "Output" [style="filled", fillcolor="#e9ecef", label="输出"]; "GPU0" [label="阶段 1\n(层 1-8)", fillcolor="#a5d8ff", group=g1]; "GPU1" [label="阶段 2\n(层 9-16)", fillcolor="#96f2d7", group=g1]; "GPU2" [label="阶段 3\n(层 17-24)", fillcolor="#b2f2bb", group=g1]; "GPU3" [label="阶段 4\n(层 25-32)", fillcolor="#d8f5a2", group=g1]; "Input" -> "GPU0" [label="批次"]; "GPU0" -> "GPU1" [label="激活"]; "GPU1" -> "GPU2" [label="激活"]; "GPU2" -> "GPU3" [label="激活"]; "GPU3" -> "Output"; }朴素流水线并行中的顺序流。每个GPU阶段都等待前一个完成,造成较多的空闲时间。解决流水线气泡问题上面所示的简单顺序方法效率低下。当GPU 0处理一个批次时,GPU 1、2和3完全处于空闲状态。这个空闲时间,被称为“流水线气泡”,严重影响了整体硬件利用率。为了减轻这种情况,输入数据批次被拆分成更小的微批次。一旦GPU 0完成处理第一个微批次并将其传递给GPU 1,它就可以立即开始处理第二个微批次。这使得GPU可以并行处理不同的微批次,创建了一个真正的流水线。{"data":[{"name":"GPU 0 (阶段 1)","x":[1,2,3,4,4,3,2,1],"y":["MB 1","MB 2","MB 3","MB 4","MB 4","MB 3","MB 2","MB 1"],"type":"timeline","orientation":"h","marker":{"color":"#a5d8ff"},"text":["前向","前向","前向","前向","反向","反向","反向","反向"],"hoverinfo":"text"},{"name":"GPU 1 (阶段 2)","x":[1,2,3,4,4,3,2,1],"y":["MB 1","MB 2","MB 3","MB 4","MB 4","MB 3","MB 2","MB 1"],"type":"timeline","orientation":"h","marker":{"color":"#96f2d7"},"base":[1,2,3,4,5,6,7,8],"text":["前向","前向","前向","前向","反向","反向","反向","反向"],"hoverinfo":"text"},{"name":"GPU 2 (阶段 3)","x":[1,2,3,4,4,3,2,1],"y":["MB 1","MB 2","MB 3","MB 4","MB 4","MB 3","MB 2","MB 1"],"type":"timeline","orientation":"h","marker":{"color":"#b2f2bb"},"base":[2,3,4,5,6,7,8,9],"text":["前向","前向","前向","前向","反向","反向","反向","反向"],"hoverinfo":"text"},{"name":"GPU 3 (阶段 4)","x":[1,2,3,4,4,3,2,1],"y":["MB 1","MB 2","MB 3","MB 4","MB 4","MB 3","MB 2","MB 1"],"type":"timeline","orientation":"h","marker":{"color":"#d8f5a2"},"base":[3,4,5,6,7,8,9,10],"text":["前向","前向","前向","前向","反向","反向","反向","反向"],"hoverinfo":"text"}],"layout":{"title":{"text":"流水线和微批次处理下的GPU利用率"},"xaxis":{"title":"时间步","zeroline":false},"yaxis":{"title":"微批次 (MB)","autorange":"reversed"},"barmode":"stack","showlegend":true,"height":400,"legend":{"traceorder":"normal"}}}4个微批次的4阶段流水线中GPU随时间变化的利用率。“Fwd”表示前向传播,“Bwd”表示反向传播。初始启动和最终清空阶段(即“气泡”)呈现为灰色区域,但中间阶段体现出高且重叠的利用率。这个调度方案,通常被称为GPipe,提高了利用率但不尽理想。在处理完整批次的开始(启动阶段)和结束(下降阶段)时,气泡仍然存在。更先进的调度器,例如DeepSpeed中使用的,可以通过改变反向传播顺序来进一步改进此问题,以填充更多气泡。选择合适的方法模型并行和流水线并行可以同时使用;它们解决了同一个问题的不同方面,并经常一起使用。选择取决于模型架构和硬件限制。模型并行何时使用: 当单个层太大,无法放入单个GPU内存时。这在大型语言模型中的注意力头和MLP块中很常见。通信方式: 单个层的前向/反向传播中需要高频、低延迟的通信(例如,all-reduce)。它最适合单个节点内紧密连接的GPU(例如,通过NVLink)。流水线并行何时使用: 对于可以逻辑上划分为顺序阶段的非常深的模型。它很适合通过较慢的网络互连(如InfiniBand或以太网)在多个节点之间进行扩展。通信方式: 较低频率的通信,包括每个微批次一次性在阶段之间传递较大的激活张量。在许多生产情况下,混合方法是最有效的解决方案。例如,一个大型模型可能使用流水线并行跨不同节点拆分成多个阶段。在每个节点内部,分配给单个阶段的GPU可能使用张量并行来管理该阶段层的内存,并使用数据并行来更快地处理微批次。这种组合使得模型能够扩展到大型加速器集群中的超大模型尺寸。