在算力机房的深处，总厨阿明接到一张让人望而生畏的满汉全席菜单——一个千亿参数的模型。他最大的厨房只有一个小冷藏柜（16GB 显存），别说做菜，连菜单的一角都塞不下。阿明的对策很朴素：把菜单撕成几大摞，分给四个分厨房去分别备料，需要半成品时再通过传菜窗口递来递去。这套“分厨房、分食材、传半成品”的智慧，就是模型并行（Model Parallel）。 本文用这个“多厨房合办满汉全席”的故事，带你理解 JAX 中的模型并行。你会看到：食材怎么切（分片策略）、用什么工具切（pjit）、总厨怎么自动安排传菜（GSPMD 编译器），以及所有传菜窗口背后的通信原语。读完你会发现，大模型训练并没有想象中神秘——它不过是一场精心编排的群体烹饪。

岗位关联

掌握模型并行，意味着你能真正驾驭多卡训练大模型的底层逻辑；对 AI 基础设施工程师而言，这是千亿大模型落地时绕不开的“必考手艺”。

目录
  1. 可视化演示
  2. 模型并行概念
  3. 分片策略
  4. JAX编程模型
  5. pjit分片注解
  6. GSPMD自动推导
  7. 通信原语
  8. 混合并行策略
  9. 实战调优
  10. AI 应用场景
  11. Python 演示
  12. 常见误区

可视化演示

动态演示Model Parallel的核心机制，通过全局/局部/动态/静态四种模式展示关键概念与交互流程。

本节配有交互动效（进入视口后自动加载，约 3KB）。

Model Parallel：GSPM GSPMD自动化模型引导 pjit函数与Partitio 核心机制可视化 Model Parallel整体架构与关键组件

模型并行概念

阿明的冷柜只有 16GB，可菜单上单是“烤全羊”的备料就要 50GB。他灵机一动：把羊腿放 1 号厨房，羊排放 2 号厨房，羊腩放 3 号厨房……每个厨房只保存一小部分食材，做菜时再把半成品传来传去。这套“食材分库存放”的思路，正是模型并行最朴素的定义：把模型参数切分到多个设备上，让每个设备只负责一部分计算。

在深度学习训练中，显存消耗远不止模型参数本身。反向传播需要保存每一层的中间激活值，优化器还要保存动量与方差等状态。一个 175B 参数模型仅参数就需要 700GB（FP16），加上激活和优化器状态，单卡根本无法容纳。模型并行的动机由此而来：当模型太大装不进一张卡时，就把参数甚至激活切到多张卡上共同承载。

阿明把食材账本翻开，发现“冷柜装不下”的根本原因是四类东西都在抢空间。模型并行没有魔法，它只是让每间厨房各背一小部分食材，把大问题拆成小问题。

Mtotal=4P+4A+4G+4OM_{\text{total}} = 4P + 4A + 4G + 4O
上式估算训练时单卡显存需求：参数（P）、激活（A）、梯度（G）与优化器状态（O）各按 4 字节（FP32）计算。切分后，每张卡只需存总显存的若干分之一。
  1. 参数项 P模型所有权重矩阵的体积，单位是参数量。切分后，每张卡只持有参数总量 \(P\) 的一部分。
    PP
  2. 激活项 A前向传播中每层输出的中间张量。保存下来是为了反向传播求梯度，分片后可分摊到各设备。
    AA
  3. 梯度项 G 与优化器状态 O梯度张量与 Adam 等优化器额外维护的动量和方差。采用混合精度时，这部分仍然占据可观显存。
    G,OG, O

分片策略

食材都分出去了，怎么分才不浪费灶台？阿明有三种刀法。第一种是“硬菜分部位”：同一道菜的食材横切成几段，不同厨房各拿一段，最后拼盘——这叫张量并行（Tensor Parallel）。第二种是“按菜序分工”：1 号厨房做前菜，2 号做主菜，3 号做甜品，一道流程走下来——这叫流水线并行（Pipeline Parallel）。第三种是“复制菜单分客桌”：每个厨房都有一份完整食材，但接待不同的客人——这是数据并行（Data Parallel）。

张量并行核心是把矩阵运算本身切开。例如一个线性层 \(Y = XW\)，把权重 \(W\) 按列切成多块，每个设备只算一部分输出列，最后拼接。这一过程等价于把大的矩阵乘法拆成若干小的独立乘法再拼装，中间需要一次通信来交换边界数据。

阿明看着菜单点了点头：硬菜横切（张量并行）适合大矩阵层，菜序分工（流水线并行）适合层数多的网络，复制菜单（数据并行）则适合数据量大的场景。真正的满汉全席，通常三种刀法一起上。

Y=[XW1⊕XW2⊕⋯⊕XWk]Y = [XW_1 \oplus XW_2 \oplus \cdots \oplus XW_k]
其中 \(\oplus\) 表示按输出特征维拼接。\(W\) 被竖切成 \(k\) 份，每个设备计算 \(XW_i\) 后拼接出完整输出。
  1. 竖切权重 W把权重矩阵 \(W\) 沿列方向切成 \(W_1, W_2, \dots, W_k\)。每个设备只持有其中一份。
    W=[W1∣W2∣⋯∣Wk]W = [W_1 \mid W_2 \mid \cdots \mid W_k]
  2. 局部矩阵乘每个设备用自己的分片 \(W_i\) 与完整输入 \(X\) 做乘法，得到部分输出。
    XWiXW_i
  3. 拼接输出把所有局部输出按列拼回完整结果 \(Y\)，此时需要一次 AllGather 通信。
    Y=XW1⊕XW2⊕⋯⊕XWkY = XW_1 \oplus XW_2 \oplus \cdots \oplus XW_k

JAX编程模型

阿明发现，每次给四个厨房分配任务都要写一大堆传菜指令，太累了。于是他找到一本“魔法委托书”——JAX。只要把做菜函数 \(f\) 交给它，它就能自动改写成适合多厨房执行的版本。

JAX 的核心思想是“函数变换”：一切并行方式都是对纯函数的改写。pmap 是最早的并行工具，它把函数映射到多个设备上，但隐含“每个设备计算同一份模型、不同数据”的数据并行语义。pjit 则更进一步：它允许你为函数的每个参数指定分片方式，让设备分别持有模型的不同部分，实现真正的模型并行。jit 则是单设备上的编译加速，三者本质都是“函数式转换”。

总厨阿明不再亲自传菜，他只把菜谱往魔法委托书上一放，剩下的切分与调度都由 JAX 接管。这份“只写一份代码、任意方式并行”的体验，正是 JAX 在科研社区快速流行的原因。

fθ(x)  →  pjit    {fθ1(x1),  fθ2(x2),  … }f_{\theta}(\mathbf{x}) \;\xrightarrow{\;\text{pjit}\;}\; \{ f_{\theta_1}(\mathbf{x}_1),\; f_{\theta_2}(\mathbf{x}_2),\; \dots \}
表示同一个函数 \(f\) 经过 pjit 转换后，成为一组各持参数分片 \(\theta_i\) 与数据分片 \(\mathbf{x}_i\) 的局部执行体。
  1. 原始函数 f一个普通的、纯函数的训练步骤，输入参数 \(\theta\) 与数据 \(\mathbf{x}\)。
    fθ(x)f_{\theta}(\mathbf{x})
  2. pjit 转换JAX 编译器读取分片注解，把函数改写成多设备并行版本。
    →  pjit  \xrightarrow{\;\text{pjit}\;}
  3. 局部执行体集合每个设备持有自己的参数分片 \(\theta_i\) 与数据分片 \(\mathbf{x}_i\)，执行同样的局部逻辑。
    {fθ1(x1),fθ2(x2),… }\{ f_{\theta_1}(\mathbf{x}_1), f_{\theta_2}(\mathbf{x}_2), \dots \}

pjit分片注解

魔法委托书虽然强大，但总厨得告诉它食材往哪里放。阿明在菜谱上贴起标签：“羊腿 → 1 号厨房”“羊排 → 2 号厨房”。在 JAX 中，这些标签就是 PartitionSpec 分片注解。

在 pjit 中，每个张量可以配一个 PartitionSpec，它描述张量每个维度如何映射到设备网格的某一轴。例如 PartitionSpec('batch', None) 表示：第一个维度（batch 维）按设备网格的 'batch' 轴切分，第二个维度完全不切。None 是“不切分”的标记。设备网格则由 Mesh 定义，比如一个 2×2 的网格有 ('row', 'col') 两个逻辑轴。

阿明的标签贴得越清楚，分厨房就越有章法。None 就是“这道食材每个厨房都要有一份”的备注，而标了名字的维度则是“大家分着拿”。一个准确的注解，能避免一半以上的并行 bug。

PartitionSpec("batch",None):RB×F  ⟶  RBP×F\text{PartitionSpec}(\text{"batch"}, \text{None}):\quad \mathbb{R}^{B \times F} \;\longrightarrow\; \mathbb{R}^{\frac{B}{P} \times F}
表示一个形状为 \(B \times F\) 的张量，在 batch 维被 \(P\) 个设备均分，每个设备只持有 \(\frac{B}{P}\) 行，特征维保持完整。
  1. 逻辑轴切分PartitionSpec 中第一个元素 "batch" 告诉编译器，张量的 batch 维和设备网格的 batch 轴绑定。
    ""batch""→B 维按 P 切分\text{""batch""} \rightarrow B \text{ 维按 } P \text{ 切分}
  2. None 不切分第二个元素 None 表示特征维 \(F\) 在每个设备上都保留完整副本。
    None→F 维不切\text{None} \rightarrow F \text{ 维不切}
  3. 形状变化切分后，每个设备上的张量形状由 \(B \times F\) 变为 \(\frac{B}{P} \times F\)。
    RB×F→RBP×F\mathbb{R}^{B \times F} \to \mathbb{R}^{\frac{B}{P} \times F}

GSPMD自动推导

传菜最头疼的不是备料，而是“谁在什么时候需要什么”。阿明惊喜地发现，只要有少数关键标签，他的副手——GSPMD 编译器——就会自动补全所有传菜路线。

GSPMD（General and Scalable Parallelization of ML Computation Graphs）是 XLA 编译器中的自动并行化模块。它接收用户写下的少量分片注解，将其视为“软约束”，然后自动传播到整个计算图中：每个中间张量该分还是不该分、沿哪个轴分，都由编译器统一推导。当某些设备需要其他设备的数据时，GSPMD 自动插入集合通信原语，开发者几乎看不到通信代码。它还支持“激活重计算”：不保存前向的中间激活，反向时重新计算一遍，以算力换显存。

阿明的副手把每个人的需求都安排得明明白白，传菜员在最恰当的时机出现在最合适的窗口。这就是 GSPMD 的价值：开发者只管贴标签，最优传菜路线由编译器自己算出来。

Mpeak=Mparam+Mgrad+∑l∈storedAl+∑l∈recompute0M_{\text{peak}} = M_{\text{param}} + M_{\text{grad}} + \sum_{l \in \text{stored}} A_l + \sum_{l \in \text{recompute}} 0
激活重计算让部分层的激活值不再驻留显存，运行时重新计算，峰值显存从“所有激活”降至“仅保留层”的激活。
  1. 参数与梯度\(M_{\text{param}}\) 与 \(M_{\text{grad}}\) 始终占用显存，无法省略。
    Mparam+MgradM_{\text{param}} + M_{\text{grad}}
  2. 存储的激活被标记为保留的层，其激活 \(A_l\) 仍驻留显存。
    ∑l∈storedAl\sum_{l \in \text{stored}} A_l
  3. 重计算的激活被标记为重计算的层，显存占用记 0，改为反向时重新执行前向。
    ∑l∈recompute0\sum_{l \in \text{recompute}} 0

通信原语

厨房之间必须有一批标准化的传菜窗口。这些窗口的名字，就是集合通信原语。

最常用的是 AllReduce：所有厨房把各自的局部结果拿出来，归约成一个完整的全局结果，再让每个厨房都拿到这份结果。反向传播中，梯度需要跨设备求平均，这正是 AllReduce 的典型应用。与之对应的还有 AllGather：把每个设备的部分数据拼成完整数据，让所有设备共享——张量并行中拼接输出时就会用到。

传菜窗口越多，菜越容易凉——通信也会拖慢训练。阿明意识到，硬件带宽是厨房之间唯一的传菜通道，如果频繁传菜，时间就浪费在等待上。因此，减少通信次数、提高通信与计算的重叠度，是并行调优的关键。

xˉ=1P∑p=1Px(p)\bar{x} = \frac{1}{P} \sum_{p=1}^{P} x^{(p)}
AllReduce 的数学本质：\(P\) 个设备各有一个局部值 \(x^{(p)}\)，归约后每台设备都得到所有设备局部值的均值 \(\bar{x}\)。
  1. 局部值每个设备持有自己的局部结果 \(x^{(p)}\)，例如一个梯度切片。
    x(p)x^{(p)}
  2. 求和通信链路把所有局部值相加，得到全局和。
    ∑p=1Px(p)\sum_{p=1}^{P} x^{(p)}
  3. 求平均并广播除以设备数 \(P\) 后，把 \(\bar{x}\) 分发给每个设备，确保更新一致。
    xˉ=1P∑p=1Px(p)\bar{x} = \frac{1}{P} \sum_{p=1}^{P} x^{(p)}

混合并行策略

满汉全席的规模，让单一策略捉襟见肘。阿明决定把三种刀法同时用上：同一条流水线上，每个阶段内再做张量并行，每个设备上再做数据并行。

业界的两大流派：Megatron 风格强调“层内张量并行”，把 Transformer 的每个大矩阵层切开，通过 AllReduce 同步；DeepSpeed 风格强调“层间流水线并行”，把网络按层切成多个 stage，设备间只传递每层边界处的激活。JAX 的 pjit + GSPMD 可以同时表达两种风格：开发者给出混合分片注解，编译器自动插通信，并在通信与计算重叠的调度上做优化。

传菜和做菜要重叠起来：A 厨房等菜时，B 厨房正好在切配。阿明让“传菜员”和“大厨”各自忙各自的，谁也不用干等，这就是混合并行策略的核心：消除空闲，压榨每一秒算力。

Eff=TcalcTcalc+Tcomm\text{Eff} = \frac{T_{\text{calc}}}{T_{\text{calc}} + T_{\text{comm}}}
并行效率是纯计算时间占（计算 + 通信）总时间的比例。通信占比越低，多卡扩展越接近线性。
  1. 计算时间\(T_{\text{calc}}\) 是设备真正做矩阵乘法的时间。分片后单设备计算量下降。
    TcalcT_{\text{calc}}
  2. 通信时间\(T_{\text{comm}}\) 是传菜等待的时间，由通信原语类型和数据量决定。
    TcommT_{\text{comm}}
  3. 效率比值当通信时间趋近于 0，效率趋近 100%；通信越多，效率越低。
    Eff=TcalcTcalc+Tcomm\text{Eff} = \frac{T_{\text{calc}}}{T_{\text{calc}} + T_{\text{comm}}}

实战调优

做完一桌菜，阿明还要试菜、调整火候。模型并行上线前，同样需要一套“试菜”流程。

第一步是测量显存：用 jax.device_memory_profile() 或 XLA 的 HLO dump，查看每个设备上参数、激活、梯度的真实占比。第二步是定位通信热点：开启 profiler 看 AllReduce 耗时是否超过计算耗时。第三步是按“由粗到细”的原则调分片：先决定要不要流水线，再决定张量并行的切法，最后才调 batch 维的数据并行度。

阿明在试菜笔记上写下心得：先确认每间厨房的承载量，再规划传菜路线，最后才追求菜品的华丽。模型并行调优同样如此——内存、通信、吞吐量，三张表看完，哪里堵哪里等，一目了然。

Throughput=B×STper-step\text{Throughput} = \frac{B \times S}{T_{\text{per-step}}}
吞吐量的直观定义：每个训练步处理 \(B\) 个样本、每个样本长度 \(S\)，除以单步耗时 \(T_{\text{per-step}}\)。调优的目标就是让吞吐量最大化。
  1. 样本数与序列长\(B\) 是 batch size，\(S\) 是序列长度，两者相乘是一个训练步处理的 token 总数。
    B×SB \times S
  2. 单步耗时\(T_{\text{per-step}}\) 是该训练步从开始到结束的总墙钟时间，包含计算与通信。
    Tper-stepT_{\text{per-step}}
  3. 吞吐量比值越大，单位时间内训练的数据越多，算力利用越充分。
    Throughput=B×STper-step\text{Throughput} = \frac{B \times S}{T_{\text{per-step}}}

AI 应用场景

千亿级大语言模型训练

阶段：大规模分布式训练

解决的问题：GPT 类模型参数量高达千亿，单卡显存完全无法容纳；传统做法是简单数据并行，但模型本身放不进一张卡，训练直接失败。模型并行把参数切分到数十甚至数千张 GPU 上，配合流水线并行让训练跑起来。

具体案例：某团队用 512 块 A100 训练 175B 参数模型，采用“8 路张量并行 + 8 段流水线并行”的混合策略。JAX 的 pjit 注解标注每个 Transformer 层的大矩阵如何切分，GSPMD 自动插入 AllReduce 通信，团队只需关注数据加载和 loss 收敛，大幅降低并行编程门槛。

推荐系统巨型 Embedding 分片

阶段：特征工程与模型训练

解决的问题：推荐系统的 embedding 表常达数十亿行，单卡显存动辄数百 GB；传统方案需要手工实现“参数服务器”，运维复杂。模型并行将 embedding 表按行切分到多个设备，用 pjit 的 PartitionSpec('row', None) 标注即可自动完成分片与查表通信。

具体案例：某电商推荐团队将 50 亿行的商品 embedding 分片到 16 张 GPU 上，通过 JAX 的 pjit 对 embedding 查表函数注解切分维度，GSPMD 自动处理查表时的 AllGather 通信，线上 CTR 模型训练吞吐提升了 3 倍。

多模态模型跨设备推理

阶段：模型推理

解决的问题：多模态大模型（如视觉-语言模型）单卡放不下，推理时又无法使用数据并行。模型并行在推理阶段把不同模态的编码器切到不同设备，让单次前向的峰值显存下降，同时保证低延迟。

具体案例：一个 80B 参数的视觉-语言模型在 4 张 A100 上做生成推理，开发者用 pjit 将视觉编码器和 LLM 解码器分别标注分片策略，GSPMD 自动推导跨模态 attention 的通信方案，首 token 延迟从 15 秒降到 4 秒。

Python 演示

python65 行
"""手动模拟模型并行：把两层MLP切分到4个虚拟设备，理解pjit幕后的分片逻辑"""import numpy as np  class FakeDevice:    """模拟一个只拥有模型分片的计算设备"""     def __init__(self, rank):        self.rank = rank        self.w1 = None  # 第一层局部权重 (8, 4)        self.w2 = None  # 第二层局部权重 (4, 4)  def shard_weights(W, n):    """把矩阵按列切成 n 块，对应 PartitionSpec 的列分片"""    cols_per_shard = W.shape[1] // n    return [W[:, i * cols_per_shard:(i + 1) * cols_per_shard] for i in range(n)]  def main():    rng = np.random.default_rng(42)    x = rng.normal(size=(2, 8))   # 输入批次 (batch=2, features=8)    y = rng.normal(size=(2, 4))   # 目标输出     # 全量权重（想象中“只存在于单机”的大模型参数）    w1_full = rng.normal(size=(8, 16))    w2_full = rng.normal(size=(16, 4))     # 分片：w1 按列分给4个设备，w2 按行分给4个设备    devices = [FakeDevice(i) for i in range(4)]    for i, dev in enumerate(devices):        dev.w1 = w1_full[:, i * 4:(i + 1) * 4]        dev.w2 = w2_full[i * 4:(i + 1) * 4, :]     # 前向第一层：每个设备只算自己那部分输出（列分片 -> 需要拼接）    h_parts = [x @ dev.w1 for dev in devices]    h = np.concatenate(h_parts, axis=1)   # AllGather 效果：拼回 (2, 16)     # 前向第二层：每个设备算自己那部分贡献（行分片 -> 需要求和）    out_parts = [h[:, i * 4:(i + 1) * 4] @ dev.w2 for i, dev in enumerate(devices)]    pred = sum(out_parts)                  # AllReduce 效果：求和得到 (2, 4)     # 反向：残差    err = pred - y                          # (2, 4)     # 第二层分片梯度：dW2_i = h_i^T @ err    d_w2_list = [h[:, i * 4:(i + 1) * 4].T @ err for i in range(4)]     # 第一层分片梯度：dW1_i = x^T @ dh_i    dh_full = err @ w2_full.T               # (2, 16)    d_w1_list = [x.T @ dh_full[:, i * 4:(i + 1) * 4] for i in range(4)]     # 梯度更新：每个设备只更新自己的局部权重    for i, dev in enumerate(devices):        dev.w1 -= 0.01 * d_w1_list[i]        dev.w2 -= 0.01 * d_w2_list[i]     print("前向 pred 形状:", pred.shape)    print("各设备 w1 分片形状:", [dev.w1.shape for dev in devices])    print("各设备 w2 分片形状:", [dev.w2.shape for dev in devices])    print("分片参数总量:", sum(d.w1.size + d.w2.size for d in devices))  if __name__ == '__main__':    main()
目的：用纯 NumPy 模拟模型并行的“切分—局部计算—通信聚合”三阶段，直观展示 JAX pjit 在底层替你做的事。
  1. shard_weights(W, n)把权重矩阵按列切成 n 块，等价于 PartitionSpec('row', 'col') 中列维被切分的效果。
  2. h_parts = [x @ dev.w1 for dev in devices]每个设备只做局部矩阵乘法，模拟张量并行的分片计算。
  3. h = np.concatenate(h_parts, axis=1)对应 AllGather：把各设备的局部输出拼回完整激活。
  4. out_parts = [h[:, i * 4:(i + 1) * 4] @ dev.w2 ...]第二层按行分片，每个设备计算局部贡献。
  5. pred = sum(out_parts)对应 AllReduce：把局部贡献求和，得到完整输出。
  6. dev.w1 -= 0.01 * d_w1_list[i]每个设备只更新自己持有的分片权重，这正是模型并行与数据并行的本质差异。

常见误区

误区我觉得模型并行就是把一个模型复制到四张 GPU 上，每张卡都跑一份完整模型，最后取平均。

正解模型并行是“切分”而非“复制”。每张卡只持有模型的一部分参数，前向与反向时通过通信把各部分拼接、聚合起来。数据并行才是复制模型，模型并行是分尸存放与协作加工。

为何常见：初学者最先接触的是数据并行，容易把“多卡训练”与“每卡一份模型”画等号，忽略了显存放不下的场景。

误区四卡能分四份，八卡就能分八份，分片越多并行度越高，速度一定越快。

正解增加分片会引入更多通信。AllReduce 的通信量随设备数对数增长，频繁传菜会让设备大部分时间在干等。实际调优要找到计算—通信的平衡点，分片过多反而可能让吞吐量下降。

为何常见：直觉上“并行=越多越好”，但忽略了总厨传菜窗口的物理带宽上限。

误区GSPMD 是全自动的，我随便写个函数传给 pjit，它就能帮我找到最优分片方案。

正解GSPMD 只能在用户给出的分片注解约束下做传播优化。若用户不提供任何标注，GSPMD 会退化成最保守的方案，甚至无法正确切分。理解 PartitionSpec 的语义，才能写出有意义的“软约束”。

为何常见：GSPMD 的自动推导确实强，容易让人误以为它是万能的，忽略了“自动优化必须基于正确意图”的前提。

误区既然模型并行能解决显存问题，那以后所有训练都用模型并行就好了。

正解模型并行引入了昂贵的通信开销，而数据并行通信频率更低、扩展性更好。业界标配是“数据并行 + 模型并行”混合：能塞进一张卡的模型先用数据并行，卡放不下了再叠加张量或流水线并行。

为何常见：初学分布式时只看到模型并行的显存优势，没看到通信成本与实现复杂度。

误区pmap 能多卡跑，pjit 也能多卡跑，它们只是名字不同。

正解pmap 隐含“数据并行”语义——所有设备上执行同一份模型、处理不同数据；而 pjit 允许任意分片模式，包括张量并行、流水线并行、混合并行，且由 GSPMD 自动推导通信。

为何常见：两者在 API 形式上相似，容易忽略它们在并行语义上的本质差异。