JAX pmap 多设备并行计算

可视化演示

展示JAX pmap如何实现跨多GPU/TPU设备的数据并行计算，包括设备映射、数据分片、all-reduce/all-gather通信原语与多设备训练流程。

并行基础：订单太大，分给帮厨做

周五晚上，餐厅接到一张 1000 份宫保鸡丁的团购订单。主厨一个人就算手脚再快，也要忙到天亮。最简单的解法，是把订单按份数拆成四堆，让四位帮厨各拿一堆，在自己的灶台上同时开工——四倍的灶台，四倍的出锅速度。这种“把一份大工作拆成几份小工作、同时干”的思路，放在 AI 训练里就叫数据并行（data parallelism）。JAX 里负责实现这一步的 API，就是今天的主角：pmap。

在 JAX 的世界里，程序永远分为两种角色。一种是 host（宿主），通常指 CPU 和主内存，它像餐厅老板：负责接单、派单、收账，但绝不亲自下厨。另一种是 device（设备），也就是 GPU 或 TPU，它们像后厨的帮厨：重活全在它们手里。pmap 这个名字拆开是 parallel + map，意思是“把一个函数并行地映射到多台设备上执行”。你写一个函数 f，pmap 会把这个函数复制到每一台设备上，大家步调一致地跑同一段代码。

不过帮厨开工前，得先把食材分好。pmap 默认沿着输入数组的第一个维度——也就是样本（batch）维度——把数据切成 D 份，每台设备只拿到一份。切分之后每台设备处理多少样本，由一个非常简单的除法决定：

订单拆成几堆，帮厨谁也不用等谁，各自锅里的菜互不干扰。这就是数据并行最舒服的地方：切分是自动的，计算是独立的。唯一需要大家坐下来商量的事，到第三节再讲。先把“分堆”这个动作刻在脑子里——pmap 的一切，都从这一刀开始。

设备映射：帮厨各就各位

开饭前，老板要做的第一件事是数人头：今天几位帮厨上班？各站哪口灶台？在 JAX 里，对应两行代码：jax.devices() 查看当前可见的全部设备，jax.local_device_count() 数出本机一共有几台。它们返回的信息，就是 pmap 分配任务时手里那张“排班表”。

有了排班表，pmap 开始干活。它做了一件很有意思的事：把“设备”这个概念，变成数组的一个新维度。当函数被 pmap 包装后，输入数组的第一个维度被解释为设备轴（parallel axis）：下标为 d 的那一块数据，交给第 d 台设备去算。你可以把这个轴想象成帮厨工位上贴的编号——1 号灶台拿第 1 堆食材，2 号灶台拿第 2 堆。

计算结束后，pmap 会把所有设备的输出按设备编号堆叠起来，返回一个前导维度大小为 D 的数组。这个数组长得和普通 JAX 数组一模一样，却暗藏玄机：每一层切片都来自不同的设备。

从外部看，你只是调用了 pmap(f)；但从设备的角度看，每个帮厨都守着自己那口灶，锅里的菜只属于自己。设备映射的本质，就是让“多台设备”这件事在数组形状上显形——看一眼 shape 的第一个维度，就知道老板今天派了几位帮厨。

通信原语：那口混酱汁的盆

帮厨们各炒各的，可宫保鸡丁讲究“整桌一个味儿”。于是主厨定下规矩：所有人把调好的酱汁倒进一个大盆，搅匀，再各舀一勺回去。这个“倒进去—搅匀—舀回来”的动作，在并行计算里叫 all-reduce（全归约）：所有设备先把自己的数值贡献出来，汇总成一个结果，再广播给每一台设备。

JAX 里最常用的通信原语有两个。lax.pmean（mean 是平均）做的事正如其名：跨设备求平均值再广播回所有设备，是 all-reduce 家族里专做“平均”的那一款。lax.all_gather 则像拼盘环节：每个帮厨把自己做的特色配菜交出来互相分享，最后每台设备手里都有一份拼好的完整数据。前者让数值变得一致，后者让数据变得完整。

为什么训练离不开 pmean？因为梯度更新要求所有设备使用同一个平均梯度。假设设备 d 算出了梯度 g_d，全局平均梯度就是：

那口盆就是 pmean，规矩就是 axis_name——所有帮厨得先知道“往哪口盆里倒”，动作才能统一。没有这口盆，每个帮厨按自己的口味调酱，整桌菜就成了大杂烩。

手动并行：老板亲自下场分堆

大多数时候 pmap 会自动切分数据，但有些时候帮厨太忙，需要老板先把食材按堆分好再递过去。这种“亲自下场分堆”的做法叫手动并行。它通常只有一步：把输入 reshape 成（设备数，每台样本数）的形状，像把 1000 份食材装进 4 个贴着编号的筐里。

手动并行还要求你理解一个隐藏的规矩：通信即同步。当某台设备执行到 lax.pmean 或 lax.all_gather 时，它必须停下来，等其他所有设备都执行到同样的位置，交换完数据，才能继续往前走。这个“必须一起到齐”的点叫同步屏障（synchronization barrier）——就像混酱汁时，最后一个帮厨没到，谁也不能先开饭。

同步是要花时间的。并行训练一个 step 的总时间可以拆成两块：

手动并行的精髓，是让老板对“分多少、传什么、等多久”心里有数。数据分得越细，计算越轻；设备越多，通信越重。懂得在两者之间找平衡，才算真正会派单。

jit组合：主厨的预制菜谱

假设帮厨做菜全凭记忆，每炒一份都要停下来回忆下一步——太慢了。高明的做法是，主厨把菜谱提前印好，贴在每口灶台前，帮厨照方抓药，流畅得像肌肉记忆。这个“提前印好的菜谱”，在 JAX 里就是 jit 编译。

jit 会把函数编译成底层的 XLA 程序，跳过 Python 的解释开销。和 pmap 组合时，标准写法是把 jit 放在里层：pmap(jit(f))。意思是：先把菜谱印好（jit 编译一份 XLA 程序），再分发给所有设备（pmap 分发到 D 个灶台）。每台设备执行的，都是同一份编译好的程序——完全符合“先印菜谱，再派给帮厨”的顺序。

组合之后，自动微分依然畅通。要在 pmap(jit(f)) 里用 grad，JAX 会在每台设备上独立计算损失对参数的梯度，再靠 pmean 跨设备取平均。要保证这一切成立，函数必须是纯函数：不能读写外部变量，就像帮厨只看自己灶台上的食材，不能偷看别人的锅——一切输入输出都通过参数和返回值传递。

预制菜谱加帮厨分发，再加酱汁共识，这套组合拳让多设备训练兼顾了编译速度和并行效率。jit 负责“无脑快”，pmap 负责“人多力量大”，pmean 负责“心往一处想”。

训练实战：一整桌宴席的演练

最后，把所有环节串成一次完整的宴席。后厨收到订单（一批数据），第一步永远是确认灶台数量，然后把订单按灶台数分堆：x.reshape(设备数, 每台样本数)。准备妥当，pmap 包装的训练函数就开火了。

每轮训练 step 里发生的事，与前面几节完全吻合：每台设备用自己的子批次做前向计算得到损失，反向传播算出梯度；pmean 把梯度平均并广播；所有设备用同一个平均梯度，更新自己手里那份参数副本。因为每台设备都持有参数副本，且每步都用相同的梯度更新，所以参数始终保持同步。

最后一件事是验收。衡量并行训练值不值，看两个指标：吞吐量（每秒处理多少样本）和扩展比（D 台设备相对 1 台设备的加速倍数）。理想扩展比是线性，实际因为通信开销总会差一点。更新参数的公式长这样：

从派单、分堆、炒菜、混酱，到按同一份菜谱改良口味，一整桌宴席就这样有条不紊地出锅了。pmap 把“人多”这件事变得安全可预期，而你要做的，只是记住那口盆（pmean）和那张菜谱（jit），其余交给 JAX。

AI 应用场景

Python 演示

import jaximport jax.numpy as jnpfrom jax import grad, pmap, jitfrom jax import lax  def main():    device_count = jax.local_device_count()    print("设备数量:", device_count)    print("设备清单:", jax.devices())     # 1. 模拟一批 8 条样本的特征数据    x = jnp.arange(8.0)    y = 3.0 * x + 1.0    print("原始批次 x:", x)     # 2. pmap: 沿第一个维度把批次切分到各设备    def compute(v):        return v * 2 + 1     mapped = pmap(compute)(x)    print("各设备计算结果:", mapped)     # 3. all-reduce: 跨设备求平均并广播    def compute_with_mean(v):        return lax.pmean(v, axis_name="data")     averaged = pmap(compute_with_mean, axis_name="data")(x)    print("all-reduce 全局平均:", averaged)     # 4. all-gather: 收齐所有分片并拼接成完整批次    def collect(v):        return lax.all_gather(v, axis_name="data", tiled=True)     gathered = pmap(collect, axis_name="data")(x)    print("all-gather 拼接形状:", gathered.shape)     # 5. 数据并行训练: 每台设备只看自己的子批次    per = x.shape[0] // device_count    x_shards = x[:device_count * per].reshape(device_count, per)    y_shards = y[:device_count * per].reshape(device_count, per)    w = jnp.array([0.0])     def loss_fn(w, x_local, y_local):        return jnp.mean((w * x_local - y_local) ** 2)     def train_step(w, x_local, y_local):        g = grad(loss_fn)(w, x_local, y_local)        avg_g = lax.pmean(g, axis_name="data")        return w - 0.01 * avg_g     step = pmap(train_step, axis_name="data")    for _ in range(300):        w = step(w, x_shards, y_shards)    print("训练后的参数 w:", w)  if __name__ == "__main__":    main()

1. device_count = jax.local_device_count()数一数后厨有几口灶：查询本机可用的加速设备数量。
2. pmap(compute)(x)把 compute 函数映射到所有设备上，并沿第一个维度把 x 切成设备数份。
3. lax.pmean(v, axis_name="data")混酱汁：跨设备求平均，再把平均值广播回每台设备。
4. lax.all_gather(v, axis_name="data", tiled=True)拼盘：把各设备的分片拼回完整批次，让每台设备都持有全量数据。
5. x_shards = x[:device_count * per].reshape(device_count, per)手动分堆：把批次切成“设备数 × 每台样本数”的块。
6. g = grad(loss_fn)(w, x_local, y_local)每台设备只用自己的子批次算出梯度，互不偷看别人的数据。
7. avg_g = lax.pmean(g, axis_name="data")梯度跨设备平均，保证所有设备用同一个“共识梯度”更新参数。

常见误区

面试题