模型并行

模型并行是什么

模型并行（model parallelism）是把一个神经网络本身切开、分摊到多块加速器（accelerator）上训练的办法：切的是参数和用到这些参数的计算，而不是训练数据。更常见的数据并行（data parallelism）里，每块设备都保存一份完整模型，各自处理不同的数据批次；模型并行里，每块设备只保存模型的一部分，一层的计算由多块设备共同完成。它要解决的是容量问题：有些模型单块设备根本放不下，切分是让训练得以开始的前提。^[1]

卡住大模型的是显存，不是算力

Megatron-LM 论文把显存，而不是原始的算术吞吐，看作超大模型上的硬约束。^[1] 这句话说的是瓶颈所在的位置。参数、它们的梯度、优化器状态，以及前向过程中产生的中间激活（activation），必须同时驻留在设备上；一台加速器每秒能做的浮点运算再多，装不下这些张量也无济于事。权重一旦放不下，训练就无法进行；唯一的出路是把模型的不同部分、以及层内产生的激活，分散保存在不同设备上。^[1]

层内怎么切

英伟达（NVIDIA）2019 年的 Megatron-LM 论文给出的是一种层内（intra-layer）实现：把 transformer 层内部的矩阵乘法切分到多块 GPU 上。^[1] 矩阵乘法是最自然的切分单位。可以想象一个权重矩阵按块拆开，每块设备拿自己那一块去乘输入向量，得到的部分结果再合并起来——切的是这一层的运算本身，而不是把整层搬到别的设备上。

这套实现引人注意的地方在工程代价：它不需要新的编译器，也不需要新的库，只要在原生 PyTorch 代码里插入少量通信操作就能写出来。^[1] 张量程序还是原来的样子，并行结构表现为几次收发，而不是一套被重写的框架。

论文还记录了一个随规模变大而愈发重要的数值细节：对 BERT 这类模型来说，层归一化（layer normalization）放在什么位置会影响被切分层的数值行为。^[1]

与流水线并行的关系

模型并行不止一种形态。流水线模型并行（pipeline model parallelism）把不同的层放在不同设备上，激活值像流水线一样从一块设备传到下一块。论文把自己的层内方法描述为与流水线模型并行正交且互补，因此两者可以组合使用。^[1] 层内切分和跨层切分回答的是两个不同的问题——一层的运算怎么放得下、整摞层怎么放得下——训练系统可以同时做这两件事。

报告的结果

该实现用 512 块 GPU 训练并收敛了参数规模最高 83 亿（8.3 billion）的 transformer 模型。^[1] 整个训练应用持续维持了 15.1 PetaFLOPs 的算力，相对一个很强的单 GPU 基线，缩放效率（scaling efficiency）为 76%。^[1] 缩放效率衡量的是多设备运行实际拿到的那部分理想加速：设备之间交换结果要花时间，所以实测吞吐会低于把单卡性能直接乘以卡数。

局限

模型并行的代价写在它的定义里：既然一层由多块设备共同计算，它们就必须不断交换中间结果，通信成了绕不开的开销。它回答的是「放不放得下」，而不是「算不算得快」——切分并不会让运算量减少，只是把运算分散开。论文关于层归一化位置的提醒也指向另一项不太显眼的成本：一层被切开之后，它的数值行为需要单独照看。^[1]