SPMD编程模型与JAX分布式训练实践

可视化演示

动态演示SPMD编程模型与JAX分布的核心机制，通过源头/管道/节点/终端四种模式展示关键概念与交互流程。

SPMD核心思想

深夜食堂后厨的比喻再往前推一步：主厨发现客人太多、一口锅忙不过来，于是把同一本菜谱复印四份，四口锅同时开火。每个厨子的动作完全一致，但案板上的食材不同——厨子一处理牛肉，厨子二处理豆腐，互不干扰。这正是 SPMD 的图景：一份程序（菜谱），多份数据（食材），多个执行单元（灶台）。

在并行计算的历史里，SPMD 是形态最朴素、也最实用的一种并行范式。它不同于 SIMD——SIMD 是处理器内部的“单指令多数据”，所有计算单元被同一个时钟驱动，严格同步；SPMD 则是进程/设备级别的并行，每台设备有自己的指令流、自己的内存，只是执行同一份源码。它也不同于“多程序多数据”（MPMD）——后者让不同设备跑不同程序，而 SPMD 的团队协作方式高度统一，代码只有一份，大幅降低了维护和调试成本。

SPMD 给分布式训练带来了一个关键优势：可扩展性。只要数据量足够大，把数据切成 4 份分给 4 台设备，理想情况下训练时间能缩短近 4 倍。但现实世界总有无法并行的“串行尾巴”——数据加载、日志上报、设备同步等。Amdahl 定律定量刻画了这种上限：

回到厨房：炒菜本身可以并行，但“切菜备料”如果只有主厨一个人做，那么即便灶台翻倍，整体出菜速度也会被切菜拖住。这也是为什么真实的分布式训练系统要尽量把数据加载、预处理等也并行化，让 \(f\) 尽可能接近 1。

JAX并行API

主厨既想多开灶，又不想让每位厨子分头研究菜谱如何改编，于是雇佣了一位“分厨长”。分厨长精通整本菜谱，负责把菜谱翻译成每口灶台上可直接执行的分步骤指令。在 JAX 的世界里，这位分厨长就是并行 API：pmap、pjit，以及如今更统一的 jax.jit 加 in_shardings 组合。

pmap 是 JAX 最早的并行原语之一：它把一个函数编译成单程序多数据的形态，并且自动把输入张量的第 0 维均匀切分到所有设备上。它的学习曲线平缓，但有局限——如果你想切分第二维、第三维，或者想表达更复杂的模型并行，pmap 就力不从心了。

于是 pjit（parallel jit）出现了：它允许用户通过 PartitionSpec 指定“张量的哪个维度映射到逻辑网格的哪个轴”，从而同时覆盖数据并行和模型并行。随着 JAX 演进，pjit 的能力被并入了 jax.jit——只需传入 in_shardings 和 out_shardings，编译器就会在编译期为每个设备生成对应的 SPMD 程序。这种统一的接口，让并行策略变成了函数签名的一部分：

分厨长的好处是，厨师们不需要自己分配食材——只需在门口写清楚“第几位厨子拿第几筐菜”，实际的分菜动作由 JAX 在编译和运行时自动完成。同理，使用 jax.jit 声明分片后，开发者不需要手写设备间的数据搬运代码，程序会自己生成通信指令。

Sharding分片机制

分厨长在发食材前，要先在案板上画好切线：哪几根萝卜归一号灶，哪几根藕归二号灶。这张“分切图纸”在 JAX 中称为 Sharding，它描述了一个张量如何沿各个维度被切分并映射到设备网格上。

具体的“图纸语言”由三件套构成：Mesh 定义逻辑网格与物理设备的绑定，PartitionSpec 描述张量维度与网格轴的对应关系，NamedSharding 把两者合成为一个完整的 sharding 声明。例如 PartitionSpec('data', None) 表示：张量的第 0 维沿名为 'data' 的网格轴切分，第 1 维不切分、在每个设备上完整保留。

更厉害的是，sharding 信息不是孤立的：JAX 编译器会根据输入的分片规则，在函数内部自动“传播”分片——某个中间结果该用哪种分片方式，编译器会基于计算图推断出来，用户只需在函数边界（输入和输出）说清楚即可。

$$\displaystyle \text{shard}(X, d) = X\left[ l_d : r_d, : \right], \quad l_d = \left\lfloor \frac{d \cdot D_0}{N} \right\rfloor, \; r_d = \left\lfloor \frac{(d+1) \cdot D_0}{N} \right\rfloor$$

1. 分片区间：每台设备拿到的不是完整张量，而是从行号 \(l_d\) 到 \(r_d\) 的一个切片。

$$\displaystyle X[l_d : r_d, :]$$

2. 起始行号：对设备 \(d\)，起始行是总行数乘以设备序号再除以设备总数后向下取整。

$$\displaystyle l_d = \left\lfloor \frac{d \cdot D_0}{N} \right\rfloor$$

3. 结束行号：结束行是下一个设备起始行的位置，从而保证所有设备的切片连续且不重叠。

$$\displaystyle r_d = \left\lfloor \frac{(d+1) \cdot D_0}{N} \right\rfloor$$

画切线越清楚，后厨的分工就越顺畅。同理，分片声明写得越细致，JAX 越能精准地生成通信方案，避免不必要的设备间数据搬运。这也是为什么在 JAX 中调分布式性能，本质上是在调 sharding 策略。

分布式执行原理

厨房不是一堆灶台随意摆放，而是被分成了备菜区、炒菜区、传菜口；每位厨子在哪口锅、旁边是哪位厨子，都事先排定。JAX 把这种物理空间的编排抽象为“设备网格”（Mesh）。

设备网格的构建方式很直白：把当前可见的设备排列成逻辑形状，并给每个维度命名。例如 8 台设备可以排成 Mesh(devices, ('data', 'model'))，形成 4×2 的逻辑网格。之后 PartitionSpec('data', None) 中的轴名，会被网格解析成物理设备的编号。执行时，XLA 编译器拿到完整的计算图和分片约束，会把图“切”成与设备数相同的若干子图，每个子图跑在对应设备上；图与图之间需要交换数据的地方，自动插入 all-reduce、all-gather、scatter 等集合通信原语。

分布式执行的性能可以用一个朴素的时间模型理解：

$$\displaystyle T_{\text{并行}} \approx \frac{T_{\text{串行}}}{N} + T_{\text{通信}}$$

1. 计算分摊：把一份工作在 \(N\) 台设备上均分，理想情况下每台设备只需付出原来的 \(1/N\)。

$$\displaystyle \frac{T_{\text{串行}}}{N}$$

2. 通信代价：设备间同步梯度、交换数据需要时间，记作 \(T_{\text{通信}}\)。

$$\displaystyle T_{\text{通信}}$$

3. 并行总耗时：真正的耗时是计算与通信之和，这解释了为什么小规模并行容易见效、大规模并行收益递减。

$$\displaystyle T_{\text{并行}} \approx \frac{T_{\text{串行}}}{N} + T_{\text{通信}}$$

在真实的厨房里，传菜通道只有那么宽——菜做得再快，传不过来也是白搭。分布式系统同理：设备越多，通信拓扑越复杂，通信与计算的比值（computation-to-communication ratio）就成了衡量扩展效率的核心指标。

数据并行策略

最经典的大锅饭做法是：食材按桌数平均分配到每口锅，每个厨子手里都有一份完整菜谱，炒完后把各自调好的“酱汁”汇到一起拌匀。这就是数据并行：每台设备保存一份完整的模型参数，各自基于本地的 batch 分片计算梯度，最后通过 AllReduce 通信把梯度汇总、求平均，再同步更新参数。

数据并行的优势在于实现简单。在 JAX 中，只需要把 batch 维度映射到 'data' 轴，模型参数用 PartitionSpec(None, None) 保持完整复制即可。训练开始后，每台设备独立完成前向和反向计算，得到局部梯度：

$$\displaystyle \bar{g} = \frac{1}{N} \sum_{i=1}^{N} g_i, \qquad \theta \leftarrow \theta - \eta \cdot \bar{g}$$

1. 局部梯度：第 \(i\) 台设备在自己的数据分片上反向传播，得到局部梯度 \(g_i\)。

$$\displaystyle g_i = \nabla_{\theta} L_i(\theta)$$

2. 梯度平均：所有设备通过 AllReduce 通信把各自梯度求和再除以设备数，得到全局梯度。

$$\displaystyle \bar{g} = \frac{1}{N} \sum_{i=1}^{N} g_i$$

3. 参数更新：每台设备用相同的全局梯度更新本地参数副本，保证模型一致。

$$\displaystyle \theta \leftarrow \theta - \eta \cdot \bar{g}$$

大锅饭的缺点也很明显：如果菜谱太大，一口锅装不下所有配料，数据并行就无能为力——因为每台设备都必须持有完整模型副本。这也是为什么当模型规模超过单卡显存时，人们开始转向模型并行。

模型并行策略

当“一口锅装不下整份菜谱”时，后厨会换一种协作方式：不再每口锅复制同样的食材，而是把整桌宴席拆成几道大菜，不同灶台各负责一道。在分布式训练里，这就是模型并行：把模型本身切碎，分配到不同设备上。

常见的模型并行有三类：参数分片（如 ZeRO，把优化器状态、梯度、参数分布到各设备）、张量并行（一个 Transformer 层的矩阵乘被切成多块，不同设备各算一部分）、流水线并行（不同层放在不同设备，前向传播像流水线一样按层传递）。三者可以组合，构成二维甚至三维的并行策略。张量并行最简单直观的例子是切分权重矩阵：

$$\displaystyle W = [\,W_1 \mid W_2\,], \qquad y = xW = [\,xW_1,\; xW_2\,]$$

1. 按列切分权重：参数矩阵 \(W\) 的列被分成两部分 \(W_1\) 和 \(W_2\)，属于张量并行。

$$\displaystyle W = [\,W_1 \mid W_2\,]$$

2. 分布计算：设备 1 计算 \(xW_1\)，设备 2 计算 \(xW_2\)，输入 \(x\) 可以复制或分片。

$$\displaystyle xW_1, \quad xW_2$$

3. 拼接结果：把两个部分结果沿特征维拼接，就得到与原始矩阵乘相同的结果。

$$\displaystyle y = [\,xW_1,\; xW_2\,]$$

这就像把一道复杂的佛跳墙拆成“各自炖汤、最后合盅”。模型并行让超大模型的训练成为可能，但每一层都涉及设备间的通信——切得越细，通信越频繁，因此实际中常把张量并行与现代通信技术（如 NVLink、ICI）配合使用。

AI 应用场景

Python 演示

import jaximport jax.numpy as jnpfrom jax.sharding import Mesh, PartitionSpec, NamedSharding # 1. 查看当前可见的物理设备（如 8 张 GPU/TPU core）devices = jax.devices()print("物理设备数：", len(devices)) # 2. 建立逻辑设备网格：一维网格，轴名为 "data"mesh = Mesh(devices, ("data",)) # 3. 同一份程序：一个两层 MLPdef mlp(x, w1, b1, w2, b2):    return jnp.tanh(x @ w1 + b1) @ w2 + b2 # 4. 构造输入与参数x = jnp.ones((64, 128))w1 = jnp.ones((128, 64))b1 = jnp.zeros((64,))w2 = jnp.ones((64, 16))b2 = jnp.zeros((16,)) # 5. 声明分片：x 沿 data 轴切分，模型参数完整复制data_sharding = NamedSharding(mesh, PartitionSpec("data", None))model_sharding = NamedSharding(mesh, PartitionSpec(None, None)) # 6. 编译为 SPMD 程序mlp_spmd = jax.jit(    mlp,    in_shardings=(        data_sharding,        model_sharding,        model_sharding,        model_sharding,        model_sharding,    ),    out_shardings=data_sharding,) # 7. 执行：每个设备持有 64/设备数 行，参数在每台设备上完整保留y = mlp_spmd(x, w1, b1, w2, b2)print("输出形状：", y.shape)print("输出分片设备：", sorted({d.id for d in y.devices}))

1. mesh = Mesh(devices, ("data",))定义逻辑网格，把物理设备绑定到名为 data 的逻辑轴上
2. PartitionSpec("data", None)声明“第 0 维被切分，第 1 维不切分”
3. NamedSharding(mesh, ...)把网格和分片规范合成一个完整的 sharding 对象
4. x = jnp.ones((64, 128))初始张量在 host 上创建，传给带 sharding 约束的 jit 后会被自动分发
5. in_shardings=(...)描述每个入参的分片方式，函数内部自动推导中间结果的分片
6. out_shardings=data_sharding描述输出也是按 batch 维分片的，避免不必要的收集
7. mlp_spmd(...)调用时 JAX 只编译一次，生成多设备上的 SPMD 可执行程序
8. y.devices查看输出张量实际分布在哪些物理设备上，验证分片是否生效

常见误区

面试题