模型并行是一种分布式计算技术，用于训练或推理无法容纳在单个设备（如GPU或TPU）内存中的大型机器学习和深度学习神经网络。与数据并行不同，数据并行中每个设备持有模型的完整副本并处理不同的数据批次，而模型并行则对模型本身进行划分，通常将不同层或层子集分配给不同设备。这种方法对于扩展基于Transformer (architecture)的大型语言模型等模型至关重要，这些模型可能拥有数十亿或数万亿个参数。
概述
在模型并行中，神经网络沿其深度（按层）进行划分，或在更高级的形式中，沿其他维度（如宽度甚至单个张量操作）进行划分。最直接的形式是逐层划分，即将连续的层放置在不同的设备上。在前向传播过程中，激活值从一个设备流向下一个设备；在反向传播过程中，梯度则反向流动。这形成了一种类似流水线的执行方式，如果调度不当，可能会引入空闲时间（气泡）。模型并行通常与数据并行形成对比，但两者可以结合在混合方法中，例如训练大型模型时使用的3D并行。
动机
现代神经网络，尤其是大型语言模型，其规模呈指数级增长。例如，GPT-3（2020年）拥有1750亿个参数，仅以32位精度存储权重就需要超过350 GB的内存，远超A100等高端GPU的80 GB容量。即使采用混合精度和梯度检查点等内存优化技术，单设备训练也是不可行的。模型并行允许研究人员通过将模型分布到由高速互连连接的一组设备上，来训练比单个设备所能容纳的大数个数量级的模型。
模型并行的类型
逐层（流水线）并行
这是最简单的形式，将模型拆分为顺序阶段，每个阶段分配给不同的设备。例如，一个100层的网络可以拆分为10个阶段，每个阶段10层，每个阶段放在单独的GPU上。在前向传播过程中，数据依次流经阶段1、阶段2等。主要挑战是负载均衡和减少流水线气泡。GPipe和PipeDream等技术引入了微批处理和调度以提高利用率。
张量并行
张量并行将单个操作（如矩阵乘法）拆分到多个设备上。例如，在transformer的注意力机制中，查询、键和值投影可以跨GPU进行划分，并通过全归约操作合并结果。这种方法用于Megatron-LM，能有效减少每个设备的内存占用，同时保持较高的计算效率，但需要高带宽通信。
专家并行
用于混合专家（MoE）模型，专家并行将不同的专家网络放置在不同的设备上，而路由器（门控网络）将token分发到相应的专家。这是一种模型并行形式，能够在不按比例增加计算量的情况下实现庞大的参数数量，因为每个token只激活一部分专家。
实现挑战
模型并行引入了若干挑战：
- 通信开销：设备之间必须交换激活值和梯度，这可能成为瓶颈，尤其是在互连速度较慢的情况下。
- 负载不均衡：层大小或计算需求不均可能导致某些设备空闲，而其他设备仍在工作。
- 流水线气泡：在流水线并行中，设备可能等待来自前序阶段的数据，从而降低利用率。
- 内存碎片化：模型划分可能导致内存使用不均，需要仔细的放置和调度。
- 容错性：设备数量增多时，故障概率增加，需要检查点和恢复机制。
与数据并行的比较
在数据并行中，每个设备持有模型的完整副本，并处理不同的数据小批量。每一步之后，梯度在设备间进行平均。这种方法简单，且对于适合单个设备内存的模型扩展性良好。而模型并行则是在模型本身过大时必需的。然而，模型并行通常具有更高的通信成本和更低的效率，因为存在串行依赖。在实践中，大规模训练会同时使用两者：节点间采用数据并行，节点内采用模型并行。
应用
模型并行对于训练GPT-4、PaLM和LLaMA等大型语言模型以及在生产环境中部署它们至关重要。OpenAI、Google（通过Google Cloud）和Anthropic等公司依赖模型并行来训练和部署其模型。它也被用于其他领域，如计算机视觉（例如3D CNN）和科学计算，在这些领域中，模型对于单个加速器来说过大。
软件支持
多个框架提供了对模型并行的内置支持：
- PyTorch：提供
torch.distributed，支持流水线并行（例如torch.distributed.pipeline.sync.Pipe）以及通过tensor_parallel库实现的张量并行。 - TensorFlow：提供
tf.distribute，包含TPUStrategy和MultiWorkerMirroredStrategy等策略，可整合模型并行。 - Megatron-LM：NVIDIA的专门库，用于张量和流水线并行。
- DeepSpeed：来自Microsoft，提供ZeRO（零冗余优化器），这是一种模型并行形式，对优化器状态、梯度和参数进行划分。
- JAX：通过
jax.sharding，允许显式设备放置和并行执行。
硬件考量
模型并行高度依赖设备间的高速通信。在单个节点内，NVLink和PCIe提供快速连接；跨节点时，使用InfiniBand或高速以太网。NVIDIA GPU和Google TPU等专用硬件通常包含专用互连（例如NVSwitch、TPU链接）以促进模型并行。Cerebras和SambaNova等公司设计了具有大容量片上内存的系统以减少对模型并行的需求，但对于极端规模，分布式方法仍然必不可少。
近期进展
近期研究聚焦于提高模型并行的效率。序列并行（拆分序列维度）和上下文并行（用于长序列）等技术已被开发。Alpa和FlexFlow等自动模型划分算法优化了跨设备的操作放置。此外，结合数据、张量和流水线并行的混合并行现已成为大规模训练的标准做法，例如NVIDIA和Microsoft训练的Megatron-Turing NLG模型（530B参数）。
局限性与未来方向
尽管模型并行功能强大，但也存在局限性。对于小模型或慢速网络，通信开销可能占主导地位。随着模型规模和集群异构性的增加，划分和调度的复杂性也随之增加。未来方向包括更智能的自动并行化、更好的通信压缩，以及本质上更易于并行化的新型架构。随着生成式AI的兴起和模型规模的不断增大，模型并行将继续作为可扩展深度学习的基础。
参见
参考文献
- Narayanan, D., et al. (2019). PipeDream: Generalized Pipeline Parallelism for DNN Training.
- Shoeybi, M., et al. (2019). Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism.
- Rajbhandari, S., et al. (2020). ZeRO: Memory Optimizations Toward Training Trillion Parameter Models.
- Huang, Y., et al. (2019). GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism.