老张在城里开了一家“嘉禾餐厅”，这天接到一张大单：一场1000人的婚宴，每桌都点招牌菜“红烧数据”。后厨只有一位大厨掌勺时，从备菜、下锅到装盘，一桌一桌地来，客人怕是等到花儿都谢了。老张灵机一动：厨房里不是还有三台空灶吗？把1000份订单平分成4份，让四位大厨各领250份同时开工，出菜速度立刻变成原来的四倍！ JAX里的 pmap，干的就是这件事——它把“做菜”这个函数，连同切好的食材，分发给多台设备（GPU/TPU），让它们各做各的一份，再汇总上桌。与之相对的 vmap，则像一位大厨手脚麻利地按更快的节拍连做1000份。一个是“加人手”，一个是“提手速”。理解了这对兄弟的差别，你就拿下了JAX并行编程的半壁江山。从大规模矩阵运算到百亿参数模型训练，pmap 处处是主角。

岗位关联

掌握 pmap 是 JAX 工程师迈向分布式训练与大规模并行计算的必经之路。它能让你把单卡跑不动的任务分发到多卡上，把训练时间从“天”压缩到“小时”，是算法工程师从单机走向集群的敲门砖。

目录
  1. 可视化演示
  2. 认识pmap
  3. 设备管理
  4. axis_size
  5. pmap与vmap
  6. 嵌套组合
  7. 并行矩阵运算
  8. 模型并行实战
  9. 调试与降级
  10. AI 应用场景
  11. Python 演示
  12. 常见误区
  13. 面试题

可视化演示

动态演示JAX pmap并行编程模式的核心机制，通过入门/进阶/实战/优化四种模式展示关键概念与交互流程。

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

JAX pmap并行编程模式：一pma 一、pmap并行编程基础概念 二、axis_size与dev 核心机制可视化 JAX pmap并行编程模式整体架构与关键组件

认识pmap

嘉禾餐厅的后厨里，老张接了个大单：一场1000人的婚宴，每桌都点招牌菜“红烧数据”。只有一位大厨掌勺时，从备菜、下锅到装盘，一桌一桌地来，客人等到花儿都谢了。老张灵机一动：后厨不是还有三台空灶吗？把1000份订单平分成4份，四位大厨各领250份同时开工，出菜速度立刻变成原来的四倍。JAX里的 pmap，干的就是这件事——它把“做菜”这个函数，连同切好的食材，复印到多台设备上，让它们各做各的一份。

在JAX的语境里，“菜谱”是一个纯函数 f，“灶台”是一台台设备（GPU/TPU），“订单”是输入数据 X。平时我们写 f(X) 是在一台设备上把整个数组算完；而 pmap(f)(X) 会自动把 X 沿着第0个维度——也就是 batch（批量）维——切成与设备数相同的份数，每台设备只拿到其中一份。对每台设备来说，它眼中的 X 从 (1000, ...) 变成了 (250, ...)，就像每位大厨只看到自己那份待处理的食材。

执行完毕后，JAX 会把每台设备上的输出重新沿 batch 维拼接起来，所以在函数外部看来，结果形状仍然是 (1000, ...)，就像四位大厨做的半成品最终又拼成一桌完整的宴席。在函数内部，可以通过 jax.lax.psum、pmean 这类集合通信操作让设备间互相交流数据。pmap 的返回值默认只完整地出现在0号设备上，其他设备保留自己那份。

老张看着后厨四灶齐开、热气腾腾，心里明白：这道菜之所以能快四倍，不是哪位大厨手速变快了，而是“人多力量大”——这正是 pmap 的本质：加人手，而不是提手速。

y=pmap(f)(X),X∈RB×D,B=ndev×blocaly = \mathrm{pmap}(f)(X), \quad X \in \mathbb{R}^{B \times D}, \quad B = n_{\text{dev}} \times b_{\text{local}}
这个公式说明：总批量 B 被拆成 n_dev 份，每台设备拿到 b_local 份；pmap 把函数 f 分发到每台设备，对各自的子集并行计算。
  1. 多设备化pmap(f) 把普通函数 f 包装成一个“会分发的函数”，返回新函数对象。
    pmap(f)\mathrm{pmap}(f)
  2. 输入切分调用时，X 的形状 (B, D) 被按第0维切成 n_dev 份，每台设备得到 (b_local, D)。
    X→[X1,X2,…,Xn]X \rightarrow [X_1, X_2, \ldots, X_n]
  3. 输出聚合各设备输出沿第0维拼接，外部看到完整形状 (B, D')。
    y=[y1,y2,…,yn]y = [y_1, y_2, \ldots, y_n]

设备管理

老张要把1000份订单分给四位大厨，不能凭空分发——他得先清点后厨有几位大厨、谁的灶还空着，然后把一份份食材端到对应的备菜台。在 JAX 里，jax.devices() 就是老张的清点动作：它返回当前机器上所有可见设备的列表，比如四张 GPU 就是 [gpu:0, gpu:1, gpu:2, gpu:3]。

默认情况下，pmap 会“征用”全部设备，就像老张让四位大厨全员上阵。但有时你只想用其中两台——比如把 gpu:2 留给其他任务。这时可以给 pmap 传 devices 参数指定子集，例如 pmap(f, devices=jax.devices()[:2])。把数据放到指定设备的操作叫 jax.device_put(x, device)，相当于把一篮食材端到某位大厨的灶台上，之后再传给 pmap 就免去了隐式搬运。

设备间传输不是免费的。pmap 执行中存在一个隐形“同步屏障”：所有设备必须等最慢的那台做完，才能进行下一次通信——这就像婚宴必须等所有桌的菜都齐了才开席。设备间的数据交流通过集合通信原语完成：jax.lax.psum 做全局求和，pmean 做全局平均，它们会在各设备之间形成一次“交头接耳”的同步点。用得越多，同步开销越大，因此在设计并行策略时要尽量减少跨设备通信次数。

老张把食材整整齐齐码在四位大厨的备菜台上，谁也不用去隔壁灶台“借酱油”——数据在本地、计算在本地，这就是并行高效的秘密。

Xdev(i)=device_put(Xi,di),i=1,…,ndevX_{\text{dev}}^{(i)} = \mathrm{device\_put}\left(X_i, d_i\right), \quad i = 1, \ldots, n_{\text{dev}}
意思是：把切分后的第 i 份数据 X_i 显式放到设备 d_i 上，避免隐式传输带来的额外拷贝开销。
  1. 切分数据X 被沿第0维切成 n_dev 份。
    X→{X1,X2,…,Xn}X \rightarrow \{X_1, X_2, \ldots, X_{n}\}
  2. 指定设备每一份 X_i 被放到对应的设备 d_i 上。
    device_put(Xi,di)\mathrm{device\_put}(X_i, d_i)
  3. 分布完成所有设备各自持有本地数据，等待并行执行。
    {d1:X1,d2:X2,…}\{d_1: X_1, d_2: X_2, \ldots\}

axis_size

老张说“把订单分成4份”，每位大厨负责250份。这里的“4“就是 axis_size——切分的份数，它必须恰好等于灶台（设备）的数量。如果老张脑筋一热说”分成8份“，可后厨只有4位大厨，那活就没法干。在 pmap 中，axis_size 就是”设备映射轴”的大小，默认等于设备总数。

这个数字决定了输入数据怎么切：总批量 batch_size 必须能被 axis_size 整除。比如 batch=1024、设备=4，那么每台设备拿到 256 个样本，完美；如果 batch=1000、设备=4，JAX 会直接报错——“1000 不能被 4 整除”，就像250份卤肉饭没法均分给4个灶台一样尴尬。更严格的写法是：batch_size = axis_size × batch_per_device，其中 axis_size 就是设备数，batch_per_device 是每台设备分配到的样本数。

axis_size 的值也决定了 pmap 返回结果第0维的长度。把 pmap(f)(X) 看作一次整体运算，它的输出形状里第0维正好等于 axis_size。当只有一台设备时，axis_size=1，pmap 退化为普通函数调用——但注意它仍会引入设备通信的固定开销，所以单卡场景用普通函数或 vmap 更划算。

老张的婚宴菜单能顺利开席，全靠“订单数能被厨师数整除”这条铁律。下次 JAX 报错说 shape 不匹配时，先拿这条算一算。

batch_size=axis_size×batch_per_device\text{batch\_size} = \text{axis\_size} \times \text{batch\_per\_device}
这一等式是 pmap 数据切分的基本盘：左端是总订单量，右端是“灶台数×每灶菜量”，两边不相等就无法开工。
  1. 总批量一次喂给模型的所有样本数。
    batch_size\text{batch\_size}
  2. 切分份数axis_size 等于设备数，决定切几份。
    axis_size=ndev\text{axis\_size} = n_{\text{dev}}
  3. 每份大小总批量除以设备数，得到每台设备的本地批量。
    batch_per_device=batch_sizeaxis_size\text{batch\_per\_device} = \frac{\text{batch\_size}}{\text{axis\_size}}

pmap与vmap

后厨里有两类“快”。一类是手速快：一位大厨练就了“一手颠锅一手翻炒”的绝活，单锅能同时处理10份菜，这叫 vmap——在同一台设备上把一个操作批量应用到很多样本上。另一类是人力多：四位大厨各自掌一灶，同时各做各的250份，这叫 pmap——在多台设备上把一个函数复制到每台设备并行执行。一个是“提手速”，一个是“加人手”。

从代码签名看，vmap(f)(X) 和 pmap(f)(X) 长得几乎一模一样，返回结果的形状都是 (B, D')。但底层执行完全不同：vmap 是在单台设备内部做 SIMD 式向量化，它把“对每个样本做操作”编成高效的向量指令，数据不需要离开这张卡；pmap 则是 SPMD（单程序多数据）模型，它把数据切块分发到多张卡，每张卡上运行一份完整的 f 副本。

维度语义上有一个关键差异：vmap 不关心数据量能否被整除，它可以逐个处理任何大小的 batch；pmap 则严格要求 batch_size 是设备数的整数倍。另一个现象是：当设备数只有1时，pmap 和 vmap 的结果几乎一致，但 pmap 会因为通信机制多出固定开销。因此工程上常有一个技巧——先用 vmap 验证函数逻辑，再换成 pmap 上多卡，逻辑不变但性能倍增。

老张总结道：vmap 是“一位快手大厨”，pmap 是“一支后厨团队”。要办千人大宴，光靠手快不够，还得人手多——但调菜单时，先用快手大厨试菜，总比让整个团队白忙活强。

vmap(f):RB×D→RB×D′,pmap(f):R(ndev⋅b)×D→R(ndev⋅b)×D′\mathrm{vmap}(f): \mathbb{R}^{B \times D} \rightarrow \mathbb{R}^{B \times D'}, \quad\mathrm{pmap}(f): \mathbb{R}^{(n_{\text{dev}} \cdot b) \times D} \rightarrow \mathbb{R}^{(n_{\text{dev}} \cdot b) \times D'}
两者输入输出形状看起来一样，但 vmap 只在单设备上循环向量化，pmap 把第0维拆份分发到 n_dev 台设备。
  1. vmap 映射在单设备内，把 f 应用到每个样本上。
    vmap(f):B 个样本依次向量化\mathrm{vmap}(f): B \text{ 个样本依次向量化}
  2. pmap 映射把 B 拆成 n_dev 份，每台设备一份。
    pmap(f):B=ndev⋅blocal\mathrm{pmap}(f): B = n_{\text{dev}} \cdot b_{\text{local}}
  3. 形状一致两者外部签名相同，但执行硬件不同。
    out_shape=(B,D′)=(ndev⋅blocal,D′)\text{out\_shape} = (B, D') = (n_{\text{dev}} \cdot b_{\text{local}}, D')

嵌套组合

老张很快发现，光有分工还不够：四位大厨虽然各炒各的区域，但每位大厨自己颠勺时如果能“一锅炒10份”（vmap），整场婚宴的出菜速度还能再翻几倍。于是他把两者叠起来：外层 pmap 负责把订单分给四位大厨，内层 vmap 负责让每位大厨批量颠勺。这就是 pmap(vmap(f))——设备间并行 × 设备内批量，双管齐下。

反过来也有 vmap(pmap(f)) 的写法：外层先把样本逐个展开，内层把每个样本的操作映射到设备上。这种模式适合每设备只处理一个样本、但要跑很多个样本的实验。实际训练中最常见的组合其实是 pmap(grad(loss))：jax.grad 先算出损失对参数的梯度函数，再用 pmap 让每台设备基于自己的数据块计算梯度，最后用 pmean 把各设备的梯度平均起来，更新一份共享参数。

嵌套时有个大坑：随机数。如果你在 pmap 内部使用 jax.random.normal，所有设备会拿到相同的随机种子，产生一模一样的随机数——就像四位大厨用了同一罐用到底的盐，味道完全一样。正确做法是用 jax.random.split 把主密钥拆成 n_dev+1 个子密钥，每个设备各拿一份，才能保证每份数据对应的随机扰动各不相同。

老张的后厨最终形成了这样的流水线：先分单（pmap），再批量颠勺（vmap），最后汇总调味（pmean 梯度聚合）。每一层各司其职，婚宴的菜出得又快又齐。

pmap(vmap(f))(X):(ndev,blocal,D)→(ndev,blocal,D′)\mathrm{pmap}(\mathrm{vmap}(f))(X) \quad : \quad (n_{\text{dev}}, b_{\text{local}}, D) \rightarrow (n_{\text{dev}}, b_{\text{local}}, D')
嵌套后第0维是设备维，第1维是每台设备内的批量维，维度语义清晰分层。
  1. 外层设备维形状第0维 n_dev 对应设备数量。
    (ndev,…)(n_{\text{dev}}, \ldots)
  2. 内层批量维每个设备内部处理 b_local 个样本。
    (ndev,blocal,…)(n_{\text{dev}}, b_{\text{local}}, \ldots)
  3. 特征维剩余维度是样本自身的特征。
    (…,D)→(…,D′)(\ldots, D) \rightarrow (\ldots, D')

并行矩阵运算

婚宴上有一道“矩阵炖汤”：一锅 4096×4096 的浓汤，要让所有客人都尝到。老张的办法是让四位大厨各守一角，每人在自己负责的行区域里搅动。矩阵乘法 A @ B 天然适合这种分块：把 A 按行切成4块，每台设备拿到一块 A_i 和完整的 B，计算 A_i @ B，最后把四块结果垂直拼回去，就得到了完整的 A @ B。

用 pmap 实现只需一行核心逻辑：pmap(lambda A_block, B: A_block @ B)(A, B)。JAX 会自动把 A 的第0维切成4份分发给设备，而 B 因为第0维不是“设备映射维”，会被完整广播到每台设备上——就像每位大厨案板上都摆着一整本菜单，但只负责做自己那几页的菜。

性能上有个重要权衡：矩阵乘法的计算量是 O(n³)，而设备间通信量是 O(n²)。这意味着矩阵越大，并行收益越显著——通信成本被海量计算摊薄；反过来说，如果矩阵只有 32×32，通信开销可能比计算本身还大，pmap 反而更慢。工程上判断“值不值得并行”的简单标准是：单块计算时间是否远大于跨设备通信时间。

四位大厨每人一勺、各搅一角，看似各做各的，但合起来就是一锅均匀入味的浓汤。计算量大到一定程度，“加人手”的收益就远超“提手速”的极限。

C=[A1A2A3A4]B=[A1BA2BA3BA4B]C = \begin{bmatrix} A_1 \\ A_2 \\ A_3 \\ A_4 \end{bmatrix} B = \begin{bmatrix} A_1 B \\ A_2 B \\ A_3 B \\ A_4 B \end{bmatrix}
分块矩阵乘法：A 按行切成四块，每块与 B 相乘，结果垂直拼接得到 C。
  1. 按行分块把 A 切为四个行块。
    A→[A1;A2;A3;A4]A \rightarrow [A_1; A_2; A_3; A_4]
  2. 本地乘法每个设备独立完成 A_i 与 B 的乘法。
    Ci=AiBC_i = A_i B
  3. 垂直拼接四块结果按行合并为完整 C。
    C=[C1;C2;C3;C4]C = [C_1; C_2; C_3; C_4]

模型并行实战

老张的生意越做越大，开始开连锁店：每家店独立接待客人（前向计算），晚上把各店营业额汇总到总部（梯度聚合），总部统一调整菜单配方（参数更新）。这就是“数据并行”训练：把一个大 batch 分散到多台设备上，每台设备持有一份完整模型副本和一份独特数据，算自己的梯度，再通过 AllReduce 把梯度聚合成全体的平均梯度，用这个平均梯度更新模型。

在 JAX 中，数据并行的流程被 pmap 包装得非常简洁：假设你已经写好了单设备的 loss_fn(params, batch)，那么 pmap(jax.grad(loss_fn))(params, big_batch) 会在每台设备上算出梯度 g_i。这些梯度在每台设备上是“不完整”的——就像每家店只知道自家客人的口味。要用 jax.lax.pmean(g, axis_name='batch') 把所有店的梯度取平均，这样每台设备都拿到全局一致的梯度，之后才能安全地更新参数。

除了数据并行，还有“模型并行”：当模型大到单卡显存放不下时（比如百亿参数大模型），可以把不同层的参数放在不同设备上，数据按顺序流经各设备。这种模式下 pmap 不再是主角，jax.jit 配合 shard_map 或手动 device_put 更常用。梯度聚合则是模型并行的关键——各设备只算了局部梯度，必须通过集合通信汇总成完整梯度。理解 pmean 这类原语，是打通数据并行和模型并行的共同基础。

连锁店的精髓在于“千店一面”：各家口味必须一致，靠的是每晚统一汇总流水、调整配方。数据并行训练也一样——各设备的模型保持一致，靠的是梯度聚合这关键一步。

gˉ=1N∑i=1Ngi,θt+1=θt−ηgˉ\bar{g} = \frac{1}{N} \sum_{i=1}^{N} g_i, \quad \theta_{t+1} = \theta_t - \eta \bar{g}
参数更新公式：N 台设备各算梯度 g_i，取平均得到全局梯度，再沿负梯度方向更新参数。
  1. 局部梯度第 i 台设备用自己的数据块算梯度。
    gi=∇θLi(θ)g_i = \nabla_\theta \mathcal{L}_i(\theta)
  2. 梯度平均AllReduce 聚合所有局部梯度取平均。
    gˉ=1N∑i=1Ngi\bar{g} = \frac{1}{N} \sum_{i=1}^{N} g_i
  3. 参数更新用平均梯度更新共享参数。
    θ←θ−ηgˉ\theta \leftarrow \theta - \eta \bar{g}

调试与降级

后厨出问题的时候，老张有个习惯：先不让四位大厨同时上阵，而是让一位大厨把整道菜从头到尾做一遍，确认流程无误，再正式分给四灶齐开。这个“一人演练”的思路，正是 pmap 调试的第一原则：先用 vmap 或普通函数把逻辑跑通，再换成 pmap 上多卡。因为 pmap 的报错信息往往涉及设备维度、通信失败，排查起来就像在四灶轰鸣的后厨里听出一粒盐放多了——几乎不可能。

另一个实用技巧是“单设备降级”：把 n_devices 设为1重新运行。虽然 pmap 传1台设备会保留通信框架，但报错信息会简化很多，方便定位是数据切分问题还是函数本身问题。JAX 也提供了调试打印能力：在 pmap 内部用 jax.debug.print 可以按设备打印值，看到每台设备上的实际数据分布——就像是给每位大厨配了个对讲机，随时汇报手上的活儿。

pmap 最常见的三类错误：一是 batch 不能整除设备数，报错会指出 shape 和 axis_size 不匹配；二是在 pmap 内部用了 Python 的 if 控制流——pmap 需要的是可以被 XLA 编译的纯函数，条件分支必须写成 jax.lax.cond 或使用函数式技巧；三是忘记处理随机数种子，所有设备生成相同随机序列。记住口诀：先 vmap 验逻辑，再 pmap 验性能；先单卡跑通，再多卡上阵。

老张的厨房管理哲学，也是并行程序员的哲学：团队协作前，先保证每个成员单打独斗都能过关。先降级验证，再并行上阵，才能在后厨和 GPU 集群里都游刃有余。

Nrequired=⌈Btotalbcap⌉N_{\text{required}} = \left\lceil \frac{B_{\text{total}}}{b_{\text{cap}}} \right\rceil
设备需求估算：总样本量除以单卡能容纳的最大批量，向上取整，就是做这个任务最少需要几张卡。
  1. 总样本量待处理的数据总量。
    BtotalB_{\text{total}}
  2. 单卡容量一张卡显存能装下的最大批量。
    bcapb_{\text{cap}}
  3. 向上取整算出最少需要的设备数。
    ⌈Btotalbcap⌉\left\lceil \frac{B_{\text{total}}}{b_{\text{cap}}} \right\rceil

AI 应用场景

数据并行训练大规模视觉模型

阶段：模型训练

解决的问题：单卡训练 ResNet-50 处理 ImageNet 级别的数据需要数周，传统串行训练无法在合理时间窗口内迭代模型，数据并行让多卡同时算梯度、聚合更新，训练时间可缩短到数天甚至数小时。

具体案例：某算法团队用 8 张 A100 训练一个图像分类模型。他们把 batch size 设为 4096，调用 pmap 将数据切成 8 份分发到每张卡，每卡算完梯度后用 lax.pmean 聚合，再统一更新参数。原来单卡 20 天的训练周期压缩到不到3天，且模型精度几乎无损——就像嘉禾餐厅用四位大厨同时上灶接下了千人大单。

大规模矩阵运算加速科学计算

阶段：数据预计算 / 数值模拟

解决的问题：物理模拟中需要反复计算超大矩阵乘法（如 8192×8192 的刚度矩阵），单卡计算资源紧张且耗时长。pmap 按行分块并行计算，把 O(n³) 的乘法拆分到多设备，显著缩短仿真迭代时间。

具体案例：某计算物理课题组在模拟流体力学时，每步迭代要算 A @ B，其中 A 为 8192×8192。使用 pmap(lambda A_block, B: A_block @ B)(A, B) 后，4 卡并行把单步耗时从 320ms 降至 90ms，且结果与串行版完全一致——正是“矩阵炖汤各搅一角”的思路。

大规模语言模型的推理分片

阶段：模型部署 / 推理

解决的问题：百亿参数模型单卡显存装不下，推理时无法启动；模型并行把参数分片放在多卡，配合设备间通信完成逐层前向计算，让超大模型在受限硬件上跑起来。

具体案例：某 AI 公司部署 130 亿参数的对话模型时，用 4 张 40GB A100 按层分片：第0层到第10层在卡0，第11层到第20层在卡1，逐层接力计算。推理延迟 1.2 秒，满足实时对话需求——这好比后厨炉灶不够大，就把一整锅汤分到四个灶上接力炖。

Python 演示

python50 行
import jaximport jax.numpy as jnpfrom jax import pmap, vmap  def main():    # 1. 查看当前机器可用的设备    devices = jax.devices()    n_dev = len(devices)    print('可用设备数量:', n_dev)     # 2. pmap 并行矩阵乘法：A 按行分块，B 广播到各设备    A = jnp.zeros((4096, 64))    B = jnp.zeros((64, 4096))     def matmul_row(A_block, B):        return A_block @ B     C = pmap(matmul_row)(A, B)    print('并行矩阵乘法结果形状:', C.shape)     # 3. vmap 套 pmap：设备间并行 + 设备内批量    def block_fn(x):        return x * 2.0     X = jnp.arange(32.0).reshape(8, 4)    Y = pmap(vmap(block_fn))(X)    print('vmap+pmap 组合结果形状:', Y.shape)     # 4. 用 pmap 模拟分布式梯度下降    key = jax.random.PRNGKey(0)    x_data = jax.random.normal(key, (1024, 4))    w = jnp.zeros((4, 1))     def loss_fn(w, x):        return jnp.mean((x @ w) ** 2)     def train_step(w, x_batch):        x_batch = jnp.reshape(x_batch, (n_dev, -1, x_batch.shape[-1]))        grads = pmap(jax.grad(loss_fn))(w, x_batch)        grad_mean = jnp.mean(grads, axis=0)        w = w - 0.01 * grad_mean        return w     w = train_step(w, x_data)    print('训练后参数形状:', w.shape)  if __name__ == '__main__':    main()
目的：演示 JAX pmap 在矩阵运算和梯度计算中的典型用法，运行环境需安装 jax 和 jaxlib（CPU 版本也可体验单设备 pmap）。
  1. devices = jax.devices()获取可用设备列表，确定 n_dev，这是 pmap 切分数据的依据。
  2. C = pmap(matmul_row)(A, B)核心并行调用，A 的第0维自动按设备数切分，B 广播到每台设备。
  3. Y = pmap(vmap(block_fn))(X)外层 pmap 分发到设备，内层 vmap 在每个设备内批量处理。
  4. x_batch = jnp.reshape(x_batch, (n_dev, -1, ...))手动把 batch 重塑成“设备数 × 每设备样本数”的结构，明确切分语义。
  5. grads = pmap(jax.grad(loss_fn))(w, x_batch)每台设备用本地数据算出梯度，返回形状为 (n_dev, ...) 的梯度数组。
  6. grad_mean = jnp.mean(grads, axis=0)模拟 AllReduce：对设备维取平均，得到全局一致梯度。

常见误区

误区“我刚入门的时候觉得 pmap 不就是多卡版 vmap 吗？反正都是把函数应用到数组上，写哪个都一样。”

正解两者底层执行模型完全不同。vmap 是单设备内 SIMD 向量化，不改变设备数量，batch 大小任意；pmap 是 SPMD 多设备并行，要求 batch 能被设备数整除，而且涉及设备间通信。最直观的区分：vmap 是“一位快手大厨”，pmap 是“四人团队分工”。只有逻辑调试时可以用 vmap 模拟 pmap 的行为，性能上完全不可互相替代。

为何常见：因为 vmap(f)(X) 和 pmap(f)(X) 的代码签名和返回形状几乎一模一样，初学者很容易从“长得像”推出“差不多”的错误结论。

误区“我只要把训练代码套上 pmap，多卡一起跑，肯定比单卡快好几倍。”

正解并行加速有前提：计算量要远大于通信开销。pmap 每次执行都有数据切分、设备间同步、结果聚合的成本。小矩阵、小 batch 场景下，通信成本可能超过计算收益，pmap 反而更慢。老张不会为了炒一盘花生米动用四位大厨——只有像千人大宴这种大批量任务，“加人手”才划算。

为何常见：因为人们本能地认为“人多一定快”，忽略了并行系统里沟通协调本身要花时间，这个误区在分布式计算领域非常普遍。

误区“我以为 axis_size 只是个超参数，想设多大就设多大，甚至比设备数大也没关系。”

正解axis_size 是 pmap 的“设备映射轴”大小，默认必须等于设备数。如果 batch_size 不能整除 axis_size，JAX 会直接报错；如果显式设置的 axis_size 与设备数不匹配，同样会报错。这就是老张的“订单数必须能整除厨师数”铁律。唯一例外是配合 SPMD 分区等高级特性时才有更灵活的分片语义。

为何常见：因为 axis_size 这个名字听起来像“可以让用户指定的轴长度”，容易与 batch size 或特征维度混淆，忽略了它和设备数量的强绑定关系。

误区“我在 pmap 函数里直接调 jax.random.normal，跑出来的结果看起来也正常，应该没问题吧？”

正解问题很大。默认情况下，每台设备会使用相同的随机种子，生成完全相同的随机数，相当于四位大厨用了同一罐盐，味道一点差异都没有，等于白并行。必须用 jax.random.split(key, n_dev + 1) 拆出独立子密钥，每个设备各取一份，才能保证各设备的数据增强、Dropout 掩码互不相同。

为何常见：因为单设备编程中随机种子只是“设一次就行”的直觉，移植到多设备时没意识到随机状态也需要按设备分发。

面试题

Q1. pmap 在 JAX 中的核心作用是什么？

单选题 · easy

  • A.在单设备上做向量化批量计算
  • B.将函数复制到多台设备并行执行← 正确答案
  • C.自动求导的逆过程
  • D.把 JAX 数组转为 NumPy 数组

答案：B。pmap 的本质是“多设备并行”，类似老张让四位大厨同时开工。A 是 vmap 的作用，C 是 reverse-mode autodiff 的职责，D 是 np.asarray 的功能，都与 pmap 的核心定位无关。

Q2. pmap 沿输入数据的哪个维度进行切分分发？

单选题 · easy

  • A.最后一个维度
  • B.第0个维度（batch 维）← 正确答案
  • C.中间任意维度
  • D.不需要切分直接复制

答案：B。pmap 默认沿第0维把数据切成设备数份，像把订单按人头平分。其余维度作为特征维完整保留在同一设备上。如果从中间切分，设备的“部分数据”就不再是同分布的样本，会破坏 batch 语义。

Q3. 关于 axis_size，下列说法正确的是：

单选题 · mid

  • A.表示每台设备上的数据量
  • B.表示设备数量且结果第0维长度一致← 正确答案
  • C.表示模型参数量
  • D.表示梯度累积步数

答案：B。axis_size 是“设备映射轴”的大小，等于设备数，也等于 pmap 返回结果第0维的长度——就像婚宴分成几桌，答案的维度就是几。A 描述的是 batch_per_device，C 是模型规模，D 是优化器超参。

Q4. vmap 与 pmap 最本质的区别是什么？

单选题 · mid

  • A.vmap 单设备内向量化，pmap 多设备间并行← 正确答案
  • B.vmap 多设备并行，pmap 单设备向量化
  • C.两者完全相同
  • D.vmap 用于推理，pmap 用于训练

答案：A。vmap 是“一位快手大厨提高手速”，pmap 是“多位大厨同时开工”，一个提升单设备计算密度，一个增加并行设备数。B 说反了，C 忽略底层执行差异，D 是没有依据的用途划分。

Q5. 在 pmap 中组合使用其他 JAX 特性时，哪些做法是正确的？

多选题 · mid

  • A.用 vmap 嵌套在 pmap 内实现设备内批量← 正确答案
  • B.用 jax.random.split 为每台设备分配独立随机种子← 正确答案
  • C.用 jax.lax.pmean 在设备间聚合梯度← 正确答案
  • D.直接在 pmap 函数内部使用 Python 的 if 语句

答案：ABC。A 是经典嵌套组合，B 保证各设备随机数不重复，C 是梯度聚合的标准手段。D 错误——pmap 要求纯函数，Python 控制流无法被 XLA 编译，必须用 jax.lax.cond 或函数式写法。

Q6. 数据并行训练中，各设备算完梯度后通常做什么？

单选题 · mid

  • A.直接用各自的梯度更新本地参数
  • B.将梯度求平均后统一更新共享参数← 正确答案
  • C.丢弃梯度不做聚合
  • D.只保留0号设备的梯度

答案：B。数据并行的关键是“千店一面”：各设备必须保持一致模型。做法是 AllReduce 求平均梯度再更新，对应 pmap 里的 lax.pmean。A 会导致模型分叉，C、D 会丢失信息。

Q7. 用 4 台设备并行计算 A(4096×4096) @ B(4096×64)，各设备承担什么工作？

单选题 · hard

  • A.每台设备算一个标量累加
  • B.每台设备算一块行切片 A_i 与 B 的乘积← 正确答案
  • C.每台设备复制完整 A，各算一遍再比较
  • D.每台设备只存 B 不参与计算

答案：B。这是“矩阵炖汤各搅一角”：A 按行切成4块，每设备持有一块 A_i 和完整 B，算 A_i @ B，最后垂直拼接。A 是归约操作而非矩阵乘法，C 是重复计算没有并行意义，D 完全没参与计算。

Q8. 当 pmap 在训练中报错时，合理的排查手段有哪些？

多选题 · hard

  • A.将 pmap 换成 vmap 先验证函数逻辑← 正确答案
  • B.检查 batch_size 能否整除设备数← 正确答案
  • C.把设备数设为1做单设备降级测试← 正确答案
  • D.删除所有随机代码避免干扰

答案：ABC。A、C 都是“先一人演练再团队上阵”的降级思路，B 检查 axis_size 整除规则这一最常见的报错源。D 没有意义——随机代码不是报错根源，删掉反而破坏功能。

Q9. batch_size = 1024、设备数 = 4 时，pmap 每台设备实际处理的数据量是：

单选题 · hard

  • A.实际分到 1024 个样本自动分配
  • B.每台设备得到 256 个样本← 正确答案
  • C.每台设备得到 1024 个样本完全相同
  • D.数据不切分，只在瓶颈设备计算

答案：B。batch_size 必须能被设备数整除，每台设备分到的本地批量 = 1024 / 4 = 256。这是 pmap 切分的基本运算规则，也是 axis_size × batch_per_device 公式的直接应用。

Q10. pmap 返回结果的形状与设备数的关系，下列描述正确的是：

多选题 · extreme

  • A.结果第0维的长度等于设备数← 正确答案
  • B.结果会自动拼接回完整 batch 形状← 正确答案
  • C.结果只分布在0号设备上不拼接
  • D.结果第0维与批次大小无关

答案：AB。pmap 输出的第0维具有“设备语义”：其长度等于设备数（axis_size），并在外部自动聚合为完整 batch 形状——就像四份半成品拼回整桌宴席。C 说反了（聚合后完整结果在0号设备可见），D 忽略了第0维正是 batch 维。