Skip to content

大模型分布式推理中的并行化技术

引言:为什么需要并行化推理?

随着深度学习模型(尤其是 Transformer 类模型)的规模不断扩大,单块 GPU 已难以满足模型推理的内存和计算需求。例如,具有数百亿参数的 GPT-3 模型无法完整加载到一块 GPU 中,必须将模型拆分到多块 GPU 上执行推理。这就引出了并行化推理的各种技术手段,包括数据并行、模型并行(如张量并行、专家并行、序列并行)、流水并行,以及零冗余并行等。通过合理利用这些并行技术,我们可以将超大模型分布到多个设备,降低单设备的内存占用,同时充分利用多设备的算力,加速推理过程。

在本教程中,我们将深入介绍大规模模型推理中的六大关键并行概念:

  • 张量并行(Tensor Parallel, TP)
  • 数据并行(Data Parallel, DP)
  • 流水并行(Pipeline Parallel, PP)
  • 专家并行(Expert Parallel, EP)
  • 序列并行(Sequence Parallel, SP)
  • 零冗余并行(Zero Redundancy Parallel, ZP)

我们将讨论每种并行技术的定义、工作原理及适用场景,分析当前主流框架(如 DeepSpeed、Megatron-LM、FasterTransformer)如何在推理中利用这些技术,并探讨并行推理领域的发展趋势与挑战。最后,我们将展望一些新兴概念(如层次化并行、推理并行、量子并行),并尝试提出一个优化性能和效率的新的理论框架。

六大并行化技术概述

张量并行(Tensor Parallel, TP)

定义与原理:张量并行是一种模型并行(Model Parallel)策略,指在层内并行化(intra-layer parallelism),将神经网络某层的参数张量按一定维度切分到不同设备上。各设备各自对分片的参数执行运算,最后通过集合同步得到与单设备计算相同的结果。例如,对于矩阵乘法 $C = A \times B$,可以将矩阵 $B$ 按列分成 $[B_0, B_1, …, B_n]$ 并分配给不同 GPU。每个 GPU 计算 $A \times B_i$ 得到部分结果 $C_i$,然后通过 all-gather 操作将各部分结果收集拼接,从而得到完整的 $C$ 。在 Transformer 的张量并行实现中,通常对多头注意力和前馈网络的大矩阵进行切分,每个 GPU 负责部分头或部分神经元的计算,过程中仅需少量同步通信(如 All-Reduce 用于汇总结果) 。

适用场景:张量并行适用于超大模型的单层参数过大、单 GPU 内存无法容纳整个层参数的情况。通过将权重矩阵按列或行分割,每块 GPU 只需存储和计算一部分权重,从而降低单卡内存占用。例如,Megatron-LM 针对 Transformer 层采用列并行+行并行的分割策略,使得每个前馈层的 GEMM 计算只需在各 GPU 上进行局部矩阵乘法,并在末端通过两次 all-reduce 获得最终结果。这种方法开销小、扩展性好,在多 GPU 单机环境下尤为高效 。在一次 Transformer 前向计算中,典型的张量并行实现只需要约 2 次 all-reduce 通信(反向也类似) 。

主流实现:NVIDIA 的 Megatron-LM 是张量并行的早期实践者,其高效的层内并行算法使得 8.3 亿到数十亿参数模型的训练和推理成为可能 。Megatron-LM 将权重矩阵按最佳方式切分,并利用 NCCL 库进行 GPU 间通信 。Hugging Face Transformers 也支持将某些大型模型以张量并行方式加载到多 GPU 上,实现内置的张量并行推理 。此外,诸如 Colossal-AI、DeepSpeed 等框架也提供了张量并行的支持或与之集成,用以在推理时对模型权重切片后分发到多个设备执行计算。

数据并行(Data Parallel, DP)

定义与原理:数据并行是最常见、最简单的并行范式,即将输入数据分片给多个设备,每个设备各自拥有一份完整的模型副本,同时对不同的输入数据子集执行推理或训练 。在训练场景下,各设备会独立计算梯度并周期性进行 All-Reduce 梯度同步,以保持模型权重一致 。在推理场景下,数据并行意味着可以同时处理多条输入,实现批量并行推理:每张 GPU 加载相同的模型,但处理不同的输入,从而成倍提高吞吐量。

适用场景:数据并行最适用于模型能够放入单卡但需要更高吞吐的情况。对于中小规模模型,在推理服务中通过多卡数据并行可以线性扩展每秒处理请求数。同时,在训练场景,如果模型规模较小无法充分利用多 GPU 算力,那么通过增大数据并行度也能加速训练。然而,需要注意的是,数据并行对单次推理延迟并无帮助——对于单个输入,它仍然只能在一张 GPU 上完成推理。而且对于超大模型(单 GPU 放不下的),仅靠数据并行无法解决内存问题,因为每张卡仍需要完整模型副本 。因此,超大模型通常结合模型并行一起使用数据并行,以在多节点环境下进一步扩展规模 。例如,Megatron-LM 在 8.3B 参数模型的训练中结合了 8 路模型并行和 64 路数据并行,总共使用 512 张 GPU 来容纳模型并提高吞吐 。

主流实现:几乎所有深度学习框架都原生支持数据并行(例如 PyTorch 的 DistributedDataParallel、Horovod 等)。在推理方面,许多部署框架(如 TensorFlow Serving、ONNX Runtime、TensorRT 等)也支持通过复制模型来并行处理多请求。在训练方面,数据并行常与其他并行技术组合形成混合并行方案,例如 “3D 并行”(数据并行 + 张量并行 + 流水并行)被广泛应用于大型 Transformer 模型的训练 。

流水并行(Pipeline Parallel, PP)

定义与原理:流水并行是一种层间并行(inter-layer model parallelism)技术,即将模型按层序划分为若干阶段(stage),每个阶段包含连续的若干层,并部署在不同的设备上 。推理时,输入数据(或小批次 micro-batch)按顺序经过第 1 阶段设备的层计算,然后将中间结果传递给第 2 阶段设备,依次类推,仿佛“流水线”传送带一样流经各个分段的模型层。为了提高设备利用率,流水并行通常与微批处理结合:将输入批次再细分为多个微批次,流水线各阶段可以并发处理不同微批次,从而减少等待流水线气泡(pipeline bubble)的影响。例如,GPipe 提出将模型等分为多段,每段在一块 GPU 上,使用 8 个微批次循环填充流水线,实现多卡并行执行前向和后向 。

适用场景:流水并行擅长解决极深的模型(层数很多)或序列执行依赖强的模型在单卡内存放不下的问题。通过让不同设备各负责模型的一段,单卡只需保存本段的参数和中间激活,从而显著降低每卡的内存需求。对于超大模型(如包含数百层 Transformer block 的模型),将其划分到多卡能实现近乎线性扩展。但流水并行引入了额外的延迟:完整推理需串行经过各阶段设备,单输入的端到端延迟会增加。此外,只有当存在足够多的并行微批次时,才能填满流水线避免设备空闲,否则就会出现性能损失的气泡期。因此,流水并行更适合批量推理或需要高吞吐的场景,对于要求极低单样本延迟的在线推理要慎重考虑阶段划分带来的延迟开销。

主流实现:GPipe 和 PipeDream 是早期经典的流水并行方案,它们通过调度算法减少了流水延迟和内存开销。NVIDIA 的 Megatron-LM 同样在其 3D 并行中集成了流水并行,以支持超过单机多卡容量的模型规模 。Megatron-LM 还提出了一种交错(interleaved)流水调度,通过在每张设备上放置不连续的层块交错执行,进一步减少流水线空闲时间,提升了 10%左右的吞吐 。DeepSpeed 提供了 PipelineModule 等工具简化模型的手工分层,并支持与张量并行、ZeRO 等结合的混合并行。在推理框架中,NVIDIA Triton Inference Server 也允许用户将模型划分到不同设备并顺序执行,从而实现在服务端的流水并行推理部署。

专家并行(Expert Parallel, EP)

定义与原理:专家并行主要针对稀疏专家模型(MoE,Mixture of Experts)中的特殊并行方式。MoE 模型在某些层(通常是前馈层)并行地拥有许多“专家”子网络,每个输入 token 只激活其中一小部分专家。专家并行即将众多专家拆分到不同的 GPU 上,每个 GPU 负责一部分专家的参数和计算 。在前向推理时,根据门控路由器的判定,不同 token 的数据被发送到对应专家所在的 GPU 进行计算,然后再收集回主流程。这种并行方式使得可以大幅增加模型参数总量(通过增加专家数量)而不会线性增加每次推理的计算量,因为每个样本只经过少数几个专家。换言之,MoE 是一种激活稀疏策略,专家并行则是实现该稀疏策略的分布式手段。

适用场景:专家并行适用于超大规模稀疏模型,例如谷歌的 Switch Transformer 使用了 2048 个专家,使模型参数达 1.6 万亿,但每个输入只调用其中一个专家 。这显然无法在单机存储,必须将专家分散在大规模分布式集群中。通过专家并行,不同 GPU 各自存储一定数量的专家,每个专家仍可较独立地处理被分配的 token,实现近似线性扩展。同时,专家并行也需要应对负载均衡挑战:如果某些专家被过多 token 选择,而其他专家空闲,就会出现负载不均。为此,MoE 模型常采用随机噪声或限制每批次每专家容量等机制来平衡路由 。总的来说,当模型参数规模可以远超计算规模且训练/推理想提高参数利用效率时,MoE 结合专家并行是有利的。

主流实现:DeepSpeed-MoE 和 Fairseq-MoE 均实现了大规模 MoE 的训练与推理,将专家分片到 GPU 并提供通信收发路由的基础架构。NVIDIA NeMo 框架通过 Megatron-LM 的 MoE 实现提供了专家并行选项,只影响带专家层的部分,不干扰其余层的并行映射 。谷歌的 GShard 和 Switch Transformer 则在 TPU 上通过将专家映射到不同核心,实现了数千专家的并行训练。推理方面,这些框架通常要求和训练相同规模的并行环境,以加载所有专家参数。一些优化技术如专家剪枝(减少不常用专家)或分层 MoE(hierarchical MoE,将专家再用 MoE 细分 )也在探索中,以减少推理时的通信和不均衡开销。

序列并行(Sequence Parallel, SP)

定义与原理:序列并行是一种相对较新的并行技术,旨在沿着序列长度维度划分计算和内存负载。换言之,如果输入序列很长,可以将每个 Transformer 层的输入序列拆分成几段,由不同 GPU 分别处理不同位置的 token。例如,序列长度为 1024,使用 4 路序列并行时,每个 GPU 处理其中 256 个 token 的位置的计算。序列并行通常和张量并行配合使用:当张量并行度 > 1 时启用 SP,将跨所有张量并行 GPU 重复存储的激活拆分,降低内存占用 。在实现上,用 Reduce-Scatter 替代原本的 All-Reduce 聚合,用 All-Gather 在需要时收集完整序列激活 。由于 All-Reduce 本质上等价于先 reduce-scatter 再 all-gather,这意味着巧妙地调整通信顺序即可实现不增加通信量的序列分片 。因此,SP 能够在不增加通信开销的情况下大幅减少每张卡需要保留的激活数据副本 。

适用场景:序列并行主要用于长序列、大批量场景下的内存优化。在超长文本或大型图片序列模型训练中,反向传播需要存储大量中间激活,占用显存成为瓶颈。传统方案常用激活重计算(checkpointing)来折中内存和计算,而序列并行提供了一种减少激活占用的新途径。在推理阶段,如果单卡内存不足以放下整条长序列的中间表示,也可以考虑序列并行:将序列拆片到多卡处理,可以避免单卡存储全序列中间结果。例如在进行长上下文(上千 token)推理时,序列并行能让多卡各自处理一段,从而突破单卡内存限制。不过,序列并行目前更多是在训练中与张量并行结合减少重计算需求 。由于推理不涉及反向和梯度,相比训练场景,序列并行对推理的益处主要体现在内存足够大时可并行加速长序列计算,实际应用还相对有限。

主流实现:NVIDIA 的 Megatron-LM 和 NeMo 已实现了序列并行功能。在 Megatron-LM 中,SP 与 TP 搭配几乎消除了激活重算需求,在训练 5300 亿参数、序列长 2048 的 GPT-3 模型时,将由于重算导致的额外计算开销降低了 90%以上 。用户只需设置配置开关,NeMo/Megatron 会在 Transformer 层自动将激活按序列维度划分并处理。值得注意的是,序列并行要求张量并行度>1 才能生效,因为需要多个并行 GPU 协作存储不同序列片段 。目前除 Megatron 体系外,其他框架对序列并行的直接支持较少,但这一理念为今后长序列任务的并行化提供了新的思路。

零冗余并行(Zero Redundancy Parallel, ZP)

定义与原理:“零冗余”并行源自微软 DeepSpeed 的 ZeRO 技术,即通过跨设备划分存储模型各类状态,以消除数据并行中的冗余拷贝 。在训练中,ZeRO 分三个阶段逐步消除冗余:阶段 1 将优化器状态分区,阶段 2 将梯度分区,阶段 3 则将模型参数本身分区 。当达到 ZeRO-3 阶段时,每张 GPU 只保存模型参数的一部分(通常是按参数列表均分的一个分片 shard),从而使得总体模型参数不再每卡完整保存,实现 100%去重。在前向计算需要某层参数时,ZeRO-3 会将对应分片的参数从各 GPU 聚合(gather)或通信计算,使得该层的计算能够正确进行,然后再将结果分发 。在推理场景,没有梯度和优化器的额外开销,ZP 主要指将模型权重按层或张量完整切分到不同设备。这与张量并行的差别在于:ZeRO 的参数分区不要求按运算维度划分,可以直接把不同层的参数放在不同 GPU,推理时按需取用。这种方式本质上类似于一种自动化的模型划分(可视作流水并行的一种特例或者通用化)。

适用场景:零冗余并行适用于超大模型在内存上的极限场景,即模型参数总量远超单卡内存,但又希望在单一大型推理任务中利用多 GPU 共同承载模型。通过参数分片,每卡内存占用降低约 $1/N$(N 为 GPU 数),使得总参数量几乎线性扩展。例如,原本需要 800GB 显存的千亿级模型,通过 8 卡 ZP 每卡仅需 100GB 左右即可运行(仍需一定的通信开销)。和张量并行相比,ZeRO 式的分区更灵活,不要求模型结构高度均匀;和流水并行相比,它可以细粒度到参数张量级别而不仅仅是层级划分。这使其在需要最小化冗余时非常有价值。然而,ZP 也带来了巨大的通信量:几乎每层权重都需要在 GPU 间传输或广播,特别是在多节点环境下,通信延迟会成为主要瓶颈 。实际应用中,除非模型实在过大无法做其他并行,否则纯粹依赖 ZeRO-3 进行推理不是最高效的方案;往往会将 ZP 与张量并行和流水并行结合,在每个阶段内再做参数分片,以平衡内存和通信。

主流实现:DeepSpeed 是 ZeRO 技术的发明者和主要实现者。其 ZeRO-Inference 功能允许使用训练时未并行的模型权重,在推理初始化时切片到多 GPU 上执行 。也就是说,即便模型原始检查点是单卡的,DeepSpeed 也能自动将其按需要的并行度进行张量切分或参数分区。PyTorch 在 v1.11 引入了 Fully Sharded Data Parallel (FSDP),本质上等价于 ZeRO-3,将模型参数在多个 GPU 间分片存储和并行计算,这同样可用于推理阶段加载超大模型。Megatron-LM 等训练框架由于本身采用张量/流水并行进行训练,通常不需要 ZeRO 来省内存,但在模型部署时,如果需要在不同并行度运行模型,也可以借助 ZeRO 将不同切分的检查点合并/拆分 。总之,零冗余并行是确保最大模型规模适配给定硬件的利器,在 GPU 内存成为主要瓶颈时非常实用。

主流框架中的并行化推理实践

DeepSpeed 的并行推理优化

DeepSpeed 是微软推出的深度学习加速库,以其 ZeRO 优化器闻名于大规模模型训练。除了训练,它也针对推理进行了深度优化。DeepSpeed-Inference 模块引入了多项技术以高效服务 Transformer 模型 : • 模型并行支持:DeepSpeed-Inference 支持将未经模型并行的模型切片后部署到多 GPU 上运行,即它内置了张量并行/模型并行(MP)功能来适应超大模型的推理需求 。即使模型在训练时没有使用模型并行,部署时也可以通过 init_inference 接口指定一个 tp_size 将模型权重按需分片到多块 GPU。这使我们能够用多块 GPU 推理那些原本需要单卡的大模型,或者在模型较小时用多卡并行计算降低推理延迟 。 • 推理定制内核:为了进一步降低推理时的延迟和成本,DeepSpeed 提供了高度优化的内核实现(Kernel Fusion 等) 。例如针对 Transformer 层的 Attention 和 MLP 部分,DeepSpeed 用经过优化的 CUDA 内核替换 PyTorch 原始实现,并支持一键将 HuggingFace 模型中的子模块替换为优化 kernel。这些定制内核减少了推理过程中的算子数量和显存读写开销,在多 GPU 并行情况下也能更好地利用带宽。 • 混合精度与量化:DeepSpeed-Inference 引入了一种名为 MoQ (Mixed precision on Quantization) 的量化方法,可将模型权重量化以缩小体积,并结合混合精度计算以降低推理延迟 。通过量化,模型大小和内存带宽占用显著下降,这在多 GPU 分布式推理时可以缓解部分通信瓶颈(因为需要传输的数据变少)。同时 DeepSpeed 也支持将部分不太重要的层用低精度计算,以换取性能提升。 • 兼容多种并行组合:DeepSpeed 训练引擎本身支持数据并行+流水并行+ZeRO 的混合并行方案。在推理模式下,虽然不涉及梯度同步,但 DeepSpeed 仍然可以与 Megatron-LM 风格的张量并行或流水并行配合。例如,可以在多节点上使用数据并行来同时服务不同请求,而每节点内部用张量并行+ZeRO 来承载模型。DeepSpeed 可以与 HuggingFace Transformers 和 Megatron 提供的并行检查点无缝衔接,不需要手动转换模型格式 。这一兼容性使得用户可以方便地将训练好的大模型直接用 DeepSpeed 部署推理,而无需关心底层并行细节。

综上,DeepSpeed 在推理阶段充分利用了上述并行技术:通过 ZeRO 和张量并行承载大模型,通过数据并行拓展多实例吞吐,通过优化内核和量化减少单次推理延迟。其最新的推理引擎(如 DeepSpeed-FastGen)还在持续改进,以支持更大的模型和更新的硬件。

Megatron-LM 的并行策略

Megatron-LM 是 NVIDIA 开源的超大模型训练框架,以 3D 并行(数据并行 + 张量并行 + 流水并行相结合)著称 。Megatron-LM 的许多并行化策略同样适用于推理阶段: • 张量并行:Megatron-LM 最早实现了 Transformer 层内部的张量分割方法。通过只对现有 PyTorch 实现做少量修改并插入 all-reduce,同步局部计算结果,其模型并行算法能够在不依赖特殊编译器的情况下扩展到数十亿参数。在推理时,如果模型在训练中使用了 k 路张量并行,那么通常推理时沿用相同划分,各 GPU 各自计算部分矩阵乘并 all-reduce 获取完整输出。Megatron 的经验表明,张量并行更适合在单机内多卡进行,可最大化利用 NVLink 等高速通信,而尽量避免跨节点的频繁小通信 。 • 流水并行:对于超出单机总显存的大型模型,Megatron-LM 结合 Pipeline 并行按层划分模型到多台服务器。为了减少流水线开销,Megatron 提出将模型层分块交错分配,以及使用足够的小 micro-batch 来填充流水 。推理时,这意味着输入需要依次经过各服务器的模型段,但可以并行处理一批输入来提高利用率。Megatron 的研究指出,在跨节点扩展时,采用流水并行优于纯张量并行,因为跨服务器的通信延迟高,长流水的通信主要是大的激活传输,相对更高效。这也是其在 5300 亿参数 GPT-3 训练中选择 8 层流水并行(每 8 张 GPU 一组)的原因。 • 混合并行与调优:Megatron-LM 非常强调不同并行维度的组合调优。经验表明,如果张量并行度和流水并行级配置不当,即使有高速网络,吞吐也可能下降一倍以上 。Megatron 提供了配置选项允许在训练时灵活设置数据/张量/流水并行度。对于推理,用户需要在相同并行度环境加载模型权重(或使用 DeepSpeed 等工具转换并行度) 。Megatron 也集成了序列并行(SP)和选择性重算等高级特性用于优化显存 。这些特性在推理时主要帮助降低显存需求,从而能以更小的 batch 或更长的序列运行模型。

总的来说,Megatron-LM 代表了当前最成熟的大模型训练并行方案之一。它在推理阶段延续了训练时的并行配置,通过张量并行分担计算、流水并行分担内存,并辅以高带宽互联和通信优化,实现大模型推理的高效扩展。例如,其 5300 亿参数模型在 2240 张 A100 GPU 上推理时,由于几乎无需重算激活,模型 FLOPs 利用率达到 54.2%,比需要重算时快 29% 。这些优化理念也逐渐影响着其它推理框架对并行策略的采用。

FasterTransformer 的高性能推理

FasterTransformer 是 NVIDIA 提供的开源库,专门用于 Transformer 模型的高性能推理加速。它通过底层优化和并行计算相结合,在较少改动模型的情况下显著提高推理吞吐和降低延迟。 • 高效内核与图优化:FasterTransformer 实现了许多针对 NVIDIA GPU 的优化 kernel,比如 fused Multi-Head Attention、fused FFN、INT8 低精度 kernel 等。这些内核充分利用了 GPU 张量核心和共享内存,将原本需要多步的运算融合为一步,大幅提高了单卡性能。对于 BERT 类编码器模型和 GPT 类解码器模型,FT 都提供了特定优化(如自定义的 Beam Search 加速等),在相同硬件上比直接用深度学习框架推理要快很多。 • 张量并行支持:为了应对超大模型,FasterTransformer 从 4.0 版本开始支持多 GPU 和多节点的并行推理 。对 GPT-3 等模型,可以通过 tensor parallel 的方式在多 GPU 上拆分模型权重并行执行推理。FasterTransformer 与 NVIDIA Triton Inference Server 集成后,还支持在 Triton 中使用张量并行来服务大模型(如在 GPU 数量不足单节点放下时,跨节点共同承载模型) 。文档指出,FT 4.0 已支持 GPT 模型在单节点多 GPU 以及多节点多 GPU 的推理,包含对 PyTorch 和 TensorRT 的接口。这让开发者能够将例如 6 亿或几十亿参数的 GPT-J、T5 模型切分到两块或更多 GPU 上以满足内存和实时性要求 。 • 批量并行与流式生成:在部署方面,FasterTransformer 针对批量请求和流式输出做了优化。例如支持批量采样(batch sampling)和流式生成,可以一边生成一边输出而非等待整个序列完成。这对于长文本生成的在线服务很重要。在多请求场景下,FT 允许动态批处理、不同序列长度混合等,提高 GPU 利用率的同时确保延迟可控。 • 混合精度与内存优化:FT 提供 FP16、BF16 等混合精度支持,以及 INT8 权重量化选项用于加速。特别是在多 GPU 推理时,低精度模型能够显著降低各 GPU 间需要通信的数据量,提升并行效率。此外,FT 5.x 版本开始还支持了一些内存优化技巧(如 KV 缓存高效管理、big batch 内存池等)来减小在长序列生成时的显存占用。

总体而言,FasterTransformer 通过软硬件结合的优化将大模型推理性能推向极致:一方面利用并行技术(模型切分到多 GPU、批量并行)突破单 GPU 限制;另一方面在单 GPU 上榨取每滴性能,使并行部署的收益最大化。许多企业在部署 GPT、BERT 这类模型时,会采用 FasterTransformer(往往与 Triton 服务器集成)实现高吞吐和低延迟的推理服务 。

并行推理的趋势与挑战

随着模型规模持续攀升和应用需求增长,并行推理技术也在不断演进。下面总结当前的一些主要趋势和所面临的挑战: • 多维度混合并行成为常态:早期可能只采用单一的数据并行或模型并行,而如今针对超大模型的部署,往往需要综合多种并行技术。例如,DeepSpeed/Megatron 等使用数据并行+张量并行+流水并行的三维并行(3D Parallel)训练,被广泛证明可以扩展到上千亿参数 。在推理时,这些并行维度的组合调优也很关键:需要根据模型大小、GPU 拓扑结构来决定并行划分策略。例如,多 GPU 单机可以优先用张量并行减小通信延迟,而跨服务器则需要引入流水并行降低每步通信量 。未来趋势是出现自动并行划分调度工具,输入模型和硬件配置即可给出优化的混合并行方案。 • 通信优化成为瓶颈突破点:在分布式并行中,设备间的通信开销往往决定了扩展效率上限。随着 GPU 数量增加,总通信量也呈上升趋势 。尤其在多节点场景下,网络带宽和延迟可能拖慢整个推理速度。例如研究显示,在 1750 亿参数模型跨节点运行时,ZeRO-3 方案因需要频繁跨节点传输参数,性能比优化的张量+流水并行方案低约 70% 。为此,新技术在致力于减小通信开销,如:压缩通信(梯度压缩、激活压缩)、利用更快的互联(NVSwitch、InfiniBand、PCIe Gen5)、以及通信与计算重叠。这些优化在推理阶段同样关键,尤其是对低延迟要求的应用,需要最大限度地减少等待通信的时间。 • 内存与计算的权衡:大模型推理不仅消耗算力,更会消耗巨量内存(显存)。如何在内存受限的情况下仍高效推理是重要挑战。例如 GPU 可利用的显存有限时,需要权衡预先加载更多参数(减少运行时 PCIe 拷贝)和分块加载参数(降低峰值内存)。ZeRO-Infinity 这类技术甚至将参数和优化状态部分交换到 CPU 或 NVMe,以突破单机显存限制。但在推理时,将权重放在慢速介质可能导致严重延迟,除非能提前分批次异步预取。因此业界也在探索内存池、分段执行(切割计算图分段执行释放内存)、重计算等手段,来降低峰值内存占用,使得并行推理可以在更廉价的硬件上运行超大模型。 • 负载均衡与调度:对于某些并行策略,负载不平衡会显著影响效率。例如流水并行需要各阶段耗时相近,否则最快的设备也必须等待最慢的阶段完成,系统吞吐受限于最慢阶段。为解决这一问题,可以按计算量不等地划分层数(非均匀 pipeline),或引入冗余副本执行最慢阶段以扩展其吞吐。再如 MoE 专家并行中,不同专家的调用频率可能差异很大,需要路由器算法确保每个专家负载均匀,或者通过动态调整专家数目来匹配输入分布 。在推理服务场景,还涉及多请求的排队与调度:如何将不同请求动态地分配到 GPU 集群,尽量保持所有 GPU 忙碌且避免长尾延迟,是一项挑战。这需要智能的调度算法,可能考虑请求的 batching、模型副本选择、以及预留部分资源处理高优先级请求等。 • 软硬件协同设计:未来的并行推理将越来越依赖于硬件特性的进步,包括更快的互联、更大的显存以及专用加速芯片。例如,NVIDIA 的 Hopper GPU 引入了 Transformer Engine,对 Transformer 类模型的计算提供了硬件级 FP8 支持和更强大的 Tensor Core,这将降低每 FLOP 的能耗并允许更大 batch。而在互联方面,NVLink/NVSwitch 大规模全互联拓扑、PCIe 和以太网融合的通信架构都在演进,使上百 GPU 的并行训练/推理成为可能。另外,异构并行也是趋势,比如利用 CPU/内存作为辅助:部分非实时性的工作在 CPU 上进行(如长序列的前部分处理在 CPU,后面在 GPU),或用 FPGA/ASIC 加速某些模块。这些软硬件协同的发展将继续拓展并行推理的边界。

总结来说,并行推理已经成为部署大模型的必备手段,但也面临通信开销、内存限制、负载均衡等多方面瓶颈。研究者和工程师们正通过改进算法、调度和硬件来化解这些瓶颈。未来,我们有望看到更加智能的并行框架,自动根据模型和硬件特性选择最优的并行策略,让大模型推理既快又高效。

新概念探索:层次化并行、推理并行和量子并行

在现有并行技术基础上,研究者们也在探索新的并行化概念,以进一步提高大规模推理的性能和效率。本节将讨论三个新兴方向:层次化并行、推理并行和量子并行,以及它们可能的应用场景。

层次化并行

层次化并行是指在不同层级上组合并行策略的思想。例如,将整个模型的并行划分视为多层次:在集群级别,不同节点之间可以采用数据并行以服务独立请求;在节点内部,采用模型并行(张量并行)和流水并行划分模型结构;而在更细的层内部,再使用序列并行或专家并行处理特殊层块。这样形成一个并行的层次体系,每一层级解决不同维度的瓶颈。例如:集群层次解决吞吐扩展,节点层次解决模型内存放不下的问题,层内层次解决单运算超大和激活开销问题。

一个具体设想是混合分粒度并行:顶层使用粗粒度的流水并行将模型拆成模块分配给节点,中层对每个模块内部用张量并行加速矩阵计算,底层针对 MoE 等稀疏模块用专家并行进一步扩展容量,针对长序列部分用序列并行降低内存。层次化并行更像是对已有技术的组合编排,但核心在于统一的调度和通信管理:使得在不同层次并行之间切换时开销最小化,尽可能保持流水作业。例如,如果某阶段使用了专家并行,那么调度器需要确保路由通信与整体流水线协同,不成为新的瓶颈。这种多层次的并行需要更加智能的框架支持,但有望充分榨取每层级的性能潜力,让不同类型并行各尽其长。

推理并行

“推理并行”可以理解为专门针对推理工作负载的并行优化理念。与训练不同,推理阶段有其特殊性:通常是无状态的前向计算,需要优化延迟(Latency)和吞吐(Throughput)并存,还可能需要处理动态的输入输出(如文本生成过程中序列长度不断变化)。推理并行方向探索如何在推理算法本身引入并行。例如,近年来提出的**Speculative Decoding(推测解码)**可以被视为一种推理并行思想:通过并行地生成多个候选序列,提前猜测未来的 tokens,从中选择正确的继续生成,从而减少顺序等待时间。这相当于用额外并行计算来换取推理延迟降低。

另一个想法是在解码过程引入并行。例如对需要逐字生成的自回归模型,一般必须序列化输出,但可以通过批并行解码的方法:对同一个模型维护多份状态,让其并行生成不同序列(例如 beam search 的不同 beam 分配到不同 GPU)同时进行,从而利用多 GPU 提升单次请求的生成速度。这类似于数据并行但作用于同一请求的多个备选输出上。

推理并行还涉及推理服务的并发调度。不同于训练集中在一个任务,推理服务往往要并行处理很多独立请求。如何利用 GPU 完成多流并行(multi-stream)成为课题:例如将 GPU 的计算划分时隙,同时处理多个推理 kernel,以减少每个请求的等待时间。NVIDIA TensorRT 和 Triton Server 支持将一个 GPU 划分为多个隔离的执行上下文,从而并发执行多个小模型推理。类似地,对于一个大模型服务多个小请求,也可以通过微批次动态组合的方式同时推理。在未来,推理并行会进一步发展,如异步流水线(不同请求穿插经过模型不同部分),跨请求的算子融合等,使 GPU 资源利用率最大化且单请求延迟最小化。

量子并行

量子并行指利用量子计算的并行特性来加速模型推理的前瞻性方向。量子计算具有天然的并行性——通过量子叠加和纠缠,一个量子电路可以一次性并行处理指数级状态空间的信息。这为矩阵乘法等线性代数运算提供了可能的超高速实现。例如,Pisa 大学的研究团队提出过一种量子算法子程序,将两个矩阵相乘的结果编码在量子态中,利用量子并行和叠加原理,有望较经典方法实现指数级的速度提升 。

尽管当前的量子计算机尚无法直接处理大型神经网络的推理,但一些初步研究已经在探索量子机器学习(Quantum ML)。其中一个思路是将模型的某些计算(如大型矩阵乘、搜索最大概率的操作)转化为量子电路来完成,借助量子并行性显著减少计算步骤。此外,量子计算可以天然生成高维分布的采样,这或许可以用于加速生成模型的采样过程。然而,量子并行目前在实践中面临很多限制,包括量子比特数目有限、纠错困难以及数据读入读出瓶颈 。因此短期内,我们更多将量子并行视为一种理论上的极限:它提示我们硬件并行能力的天花板可能远高于经典计算,只要能加以利用。

在可见的将来,更实际的“量子并行”应用可能是类量子技术,如光学计算。光学计算利用光子的并行传播特性,可以在硬件上并行地执行矩阵运算和卷积操作 。一些研究已将光学技术与“量子般”的干涉原理结合,实现光学矩阵乘法加速 。如果这些技术成熟,或许能够在不涉及真正量子计算复杂度的情况下,实现对大模型推理某些部分的大幅加速,从而补充传统电子计算架构的并行能力。

总之,量子并行目前还主要停留在理论和早期实验阶段。但它代表了并行计算的终极形态之一——利用物理层面的性质突破经典并行架构的限制。随着量子技术的发展,未来某一天我们可能看到量子处理单元(QPU)与 GPU 协同工作:GPU 负责大部分经典并行任务,而 QPU 承担极其耗时的子模块(例如超大矩阵运算或组合优化搜索)的计算,为大模型推理提供一个全新的加速途径。

构建统一的并行推理理论框架

鉴于以上各种并行方法各有侧重和局限,一个令人期待的发展是统一的理论框架,能够指导我们在不同层面综合应用并行技术以优化性能和效率。这样的框架需要抽象出并行化的共性原理,并提供系统化的策略来选择和组合并行手段。下面我们提出一个初步的构想: 1. 并行维度的抽象:首先,将模型计算过程拆解为若干可并行的维度,例如:样本维度(对应数据并行)、层维度(对应流水并行)、张量维度(对应张量并行)、专家维度(对应 MoE 专家并行)、序列/特征维度(对应序列并行)等等。每个维度代表一种正交的并行切分视角。理论上,这些维度彼此独立,但在实际实现中会相互影响通信模式和效率 。统一框架需要刻画不同维度的开销模型,比如张量并行主要引入 All-Reduce 通信,流水并行引入顺序依赖和激活传输,数据并行引入参数梯度同步等。 2. 成本模型与优化目标:接着,为系统定义一个成本函数,综合考虑计算开销、通信开销和存储开销。例如,可以建立公式估计给定并行配置下的总延迟和总内存占用。类似的分析在 Megatron-LM 的论文中已有尝试,将通信量和计算时间与并行度关联。优化目标可能是最小化推理延迟(对单请求)或最大化吞吐/成本比(对批处理或多请求)。针对不同应用场景,目标函数可以调整。然后,将问题转化为在并行维度上分配资源(如多少 GPU 用于哪种并行)使成本函数最优的优化问题。 3. 约束条件:现实中还有各种硬约束需要纳入框架。例如硬件拓扑约束:GPU 之间的连接带宽不同,影响并行策略选择(如同一台服务器内的 GPU 适合张量并行,而跨服务器更适合流水并行 );内存容量约束:某些并行组合可能超出单 GPU 显存,需引入 ZeRO 分区;负载均衡约束:需要保证没有设备极端空闲或过载,这可以用约束或惩罚项形式加入优化模型。统一框架应当能接受这些约束输入,并在解空间中排除不可行或低效的方案。 4. 解空间搜索与自动并行:由于并行配置空间往往离散且组合复杂,一个实用的理论框架需要一个自动搜索或推理算法。这可以是启发式算法(基于经验规则缩小范围) ,也可以是准确算法(如对小规模问题用 ILP 求解),或是基于学习的方法(让模型自己学习在不同场景下如何选择并行策略)。早期工作如 FlexFlow 就尝试自动探索最佳并行方案并取得不错效果 。在统一框架中,我们可以借鉴类似思想,更全面地搜索包括数据/模型/流水/专家等各维度的组合,而不仅仅是训练中的数据+操作并行。强化学习也是一种可能途径,让代理反复试验不同并行组合,最终学到接近最优的配置。 5. 层次化执行计划:有了并行划分决策后,框架需要生成具体的执行计划,指导集群各设备如何协同完成推理。这类似于分布式编译器:划分计算图、插入必要的通信原语、安排流水调度顺序等。例如,可借鉴谷歌的 GShard 和 Mesh-TensorFlow,通过在计算图中插入分片算子和 AllReduce 算子来表示并行策略,然后由系统负责执行。统一框架应该提供抽象表示,比如在计算图层面标注“这个矩阵乘用 4 路张量并行”“这些层跨 2 机流水”等,然后下层 runtime 据此实现。 6. 动态调整和弹性:真正先进的框架还应考虑动态并行调整。推理负载可能是变化的,例如高峰期需要加更多 GPU 做数据并行提升吞吐,低谷期缩减资源;又或者针对不同长度的序列,要灵活决定是否启用序列并行以节省内存。未来的系统可以实现弹性调度,在运行时根据监测到的性能瓶颈调整并行度配置。这需要快速的模型并行切换技术(如 DeepSpeed 能够快速切分/合并权重 )以及作业调度层的协同。

通过上述要素,我们可以初步勾勒一个统一的并行推理理论框架。简而言之,它把各种并行方法看作在不同“维度”对模型计算图做划分,然后通过优化算法找到最优划分策略组合,最后将之转换为实际的分布式执行。这样的框架旨在让开发者不再需要手工微调并行参数,而是输入模型和硬件信息,系统自动产出高效的推理并行方案。例如,给定一个千亿参数的语言模型和若干带 NVLink 的 GPU 服务器,框架可能自动决定:每台服务器内部 4 路张量并行,模型按层分为 2 段放在两台服务器(流水并行),总共用 8 台服务器数据并行来服务批量请求,参数采用 ZeRO-2 优化梯度冗余而不对权重分片。这样的配置对于人类来说复杂繁琐,但框架可以基于代价模型快速评估出其延迟和吞吐,并确认其优于其他备选方案。

需要强调的是,统一并行框架的提出并不意味着替代具体的实现技术,恰恰相反,它建立在前述各种技术之上。框架提供的是指导原则和决策支持,而真正执行层面,仍将调用诸如 NCCL 通信库、DeepSpeed runtime、FasterTransformer 内核等实现具体并行计算。可以预见,在这个方向上学术界和工业界的合作将尤为重要:学术界提供模型、算法和理论分析,工业界提供大规模真实环境的测试和调优反馈。

结语

大规模模型的推理加速是一项系统性工程,需要综合考虑模型结构、硬件资源和并行算法。通过本教程的讲解,我们了解了张量并行、数据并行、流水并行、专家并行、序列并行和零冗余并行等关键概念,以及它们在不同场景下的作用和实现方式。在当前的实践中,DeepSpeed、Megatron-LM、FasterTransformer 等框架已经将这些技术应用于实际,推动了像 GPT-3、Switch Transformer 这类超大模型的训练和部署。与此同时,我们也看到并行推理领域存在的挑战,如通信瓶颈和负载均衡问题,但这些正在通过更好的硬件和更聪明的算法得到缓解。

展望未来,新的并行化思路(层次化并行、针对推理的并行优化、甚至量子并行)为我们打开了思路,有望在不远的将来应用于下一代 AI 系统。最终的目标是建立完善的理论与工具,让分布式并行推理变得像单机推理一样透明和高效。希望本教程能帮助对分布式推理感兴趣的工程师和研究者打下基础,在实际工作中灵活运用各种并行技术,并激发对未来并行化创新的思考。通过持续的技术积累和创新,我们终将突破规模与性能的壁垒,让人工智能模型的智慧更快速、更广泛地服务人类。

Powered by VitePress