想象你经营着一家口碑极好的连锁餐厅，总店有几家分店，每家分店的大厨都拿着完全相同的菜谱。客人（数据样本）源源不断地涌来，总店把订单按批次分给各家分店：一号分店做东北风味的订单，二号分店做川味的订单……每位大厨用同一本菜谱做出菜，并记录顾客的反馈（损失）。一天营业结束后，总店把所有分店的改进建议（梯度）收集起来，取一个平均值，重新印刷成新菜谱，第二天再发给所有分店。这样，每家分店的菜谱时刻保持一致，但服务的客人总量却翻了好几倍——这就是多GPU数据并行。 在JAX的世界里，pmap就是那位传递命令、分派订单的“总店调度员”。它把你的训练函数复制到每个GPU上，让每个GPU处理不同的数据切片，并通过一次集体通信（all-reduce）让梯度在所有设备间达成共识。你不用手写复杂的并行逻辑，只需掌握几个核心概念，就能让训练飞起来。本文将带你从餐厅经营的视角，一步步搭建完整的多GPU数据并行训练流水线，并学会排查错误、优化性能。

岗位关联

学习本节后，你将理解JAX数据并行的核心原理，能独立编写pmap训练代码，并应对真实的分布式训练开发与调优工作——这是大型模型训练工程师的必备技能。

目录
  1. 可视化演示
  2. 数据并行原理
  3. pmap入门
  4. 前向与反向
  5. 完整训练流程
  6. 性能分析
  7. 常见错误
  8. 调优技巧
  9. AI 应用场景
  10. Python 演示
  11. 常见误区
  12. 面试题

可视化演示

展示系统处理流水线、实时性能指标曲线及多方案对比分析，包含数据流转、吞吐延迟监控与方案评估。

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

系统可视化 处理流程：输入→处理→模型→输出 性能监控：实时指标 P=L×T

数据并行原理

餐厅决定扩建分店：总店有4家分店，每家分店的大厨技术完全一样。现在有800桌客人，总店把订单平均分成4份，每家分店服务200桌。每桌客人就是一条训练样本，每种菜式对应模型的一个参数，大厨做菜的过程就是模型的推理与学习。关键点是：4家分店同时开工，每家的工作量只有原来的四分之一，整体出餐速度提升却接近4倍——这就是数据并行最直观的收益。

数据并行（Data Parallelism）的核心思想是：在多个计算设备（GPU）上保存一份完全相同的模型参数，将训练数据集切成若干份，分别交给不同设备计算，最后将各设备算出的梯度汇总取平均，用平均梯度更新所有设备上的参数。这里“完全相同”不是天然成立的——如果没有同步步骤，各家分店会拿着不同的菜谱越走越远。因此，梯度同步是数据并行的灵魂。

在JAX中，这种并行模式通过pmap自动实现。假设共有 \(K\) 个设备，一批大小为 \(N\) 的数据被切成 \(K\) 份，每份含 \(N/K\) 个样本。每个设备用同一份参数计算局部损失 \(L_i\)，总损失为所有局部损失的平均：

回到餐厅的比喻：每家分店记录下的“顾客改进建议”各不相同，有的嫌菜咸，有的嫌菜淡，总店把所有建议逐条平均后，得到的是一份平衡的“修改意见”，照此更新菜谱不会矫枉过正。数据并行正是这样——看似各家各做各的，但最终用同一个声音说话。

Ltotal=1K∑i=1KLi,∇θLtotal=1K∑i=1K∇θLi\displaystyle L_{\text{total}} = \frac{1}{K}\sum_{i=1}^{K} L_i,\qquad\nabla_{\theta} L_{\text{total}} = \frac{1}{K}\sum_{i=1}^{K} \nabla_{\theta} L_i
公式含义：全局损失和梯度都是各设备结果的算术平均。这一步保证了所有设备朝着同一个“最终改进方向”前进。
  1. 符号约定\(K\) 表示GPU设备数，\(L_i\) 是第 \(i\) 个设备上的局部损失值。每个设备只看到自己分到的数据子集，因此局部损失各不相同。
    Li=1N/K∑j=1N/Kℓ(fθ(xi,j),yi,j)\displaystyle L_i = \frac{1}{N/K}\sum_{j=1}^{N/K} \ell(f_{\theta}(x_{i,j}), y_{i,j})
  2. 求平均值将所有设备的局部损失相加再除以 \(K\)，得到全局损失 \(L_{\text{total}}\)。分母 \(K\) 是设备数，等价于把每个设备的损失视为同等重要。
    Ltotal=1K∑i=1KLi\displaystyle L_{\text{total}} = \frac{1}{K}\sum_{i=1}^{K} L_i
  3. 梯度平均损失对参数 \(\theta\) 求导，由于求导是线性运算，全局梯度等于各局部梯度的平均。后续更新参数时，所有设备都使用这个平均梯度。
    ∇θLtotal=1K∑i=1K∇θLi\displaystyle \nabla_{\theta} L_{\text{total}} = \frac{1}{K}\sum_{i=1}^{K} \nabla_{\theta} L_i

pmap入门

总店要发布“明天同时按新菜谱开工”的命令，不能挨家打电话，需要一套广播机制。JAX里的pmap就是这套广播系统：你写好一个函数（比如训练步骤），pmap会把它复制到所有GPU上，并自动处理数据分发、执行和结果回收。你不需要关心GPU之间如何通信，只需告诉它“这个函数要做并行版本”。

pmap是JAX中“单程序多数据”（SPMD）的入口。它会将输入数组沿第一个轴切分，分给每个设备一份。例如，输入形状为 (8, 224, 224, 3) 的批次，在4个GPU上运行时，每个设备会得到形状 (2, 224, 224, 3) 的子批次。返回的结果则会多出一个设备轴，形状变成 (4, 2, 224, 224, 3)。这个“多出来的轴”代表设备维度。

在JAX中，每个GPU对应一个Device对象，可通过jax.local_devices()获取。用pmap转换训练函数时，需要指定axis_name，这个名字后面会被lax.pmean等集体操作引用。较新版本中，jax.experimental.global_device_array（GDA）可以将全局数据分布到多设备，并支持跨主机扩展。基本使用形式如下：

pmap就像餐厅里的“同步开工广播”：一个口令下去，所有分店同时点火、同时翻勺、同时装盘。虽然每位大厨面前只有一部分订单，但整体厨房的吞吐量成倍增加。理解“设备轴”是看懂JAX分布式代码的关键，后续所有并行操作都围绕这根轴展开。

pmap(f,axis_name=′devices′)(x)=[ f(xi) ]i=1K\text{pmap}(f, \text{axis\_name}='\text{devices}')(x)= \left[\, f(x_i)\,\right]_{i=1}^{K}
公式含义：pmap(f) 会把输入 \(x\) 按第一维切成 \(K\) 份，每个设备执行一份，最后返回一个带有设备轴的输出。
  1. 切分输入\(x\) 的形状可写为 \((B, \dots)\)，其中 \(B\) 是第一维大小。pmap会将其均分为 \(K\) 份，每份形状 \((B/K, \dots)\)。这里 \(B\) 必须是 \(K\) 的整数倍。
    x↦x1,x2,…,xK(xi 形状为 (B/K,… ))\displaystyle x \mapsto x_1, x_2, \ldots, x_K \quad (x_i \text{ 形状为 } (B/K, \dots))
  2. 各设备执行每个设备 \(i\) 独立运行函数 \(f\)，输入为 \(x_i\)。由于是同一份代码副本，\(f\) 的参数和逻辑在所有设备上完全相同。
    yi=f(xi)\displaystyle y_i = f(x_i)
  3. 收集结果所有设备的结果按设备顺序拼接，形成新的第一维，最终输出 \(y\) 的形状为 \((K, B/K, \dots)\)。这个多出的 \(K\) 维就是“设备轴”。
    y=[y1,y2,…,yK]\displaystyle y = [y_1, y_2, \ldots, y_K]

前向与反向

每家分店做完菜后，需要记录顾客反馈：这道菜太咸还是太淡？大厨根据反馈调整下一道菜的做法，这就是“反向传播”。在数据并行中，每家分店（GPU）都会根据自己的客户群体得出一个调整建议（梯度），但这些建议可能互相冲突。因此总店需要收集所有分店的建议，逐项取平均值，给出一个大多数人都能接受的“新菜谱”。

技术层面，前向传播计算模型在本地batch上的损失，反向传播用jax.grad或jax.value_and_grad得到梯度。但grad计算出的梯度是“局部梯度”，只反映本设备数据的影响。为了同步，JAX提供了lax.pmean（跨设备求平均）或lax.psum（求和）。在pmap函数内，用lax.pmean(grads, axis_name='devices')即可完成AllReduce操作，所有设备会自动交换梯度并得到相同结果。

参数更新公式是所有设备一致执行的：

分店大厨们每晚把“少放盐”“多放糖”等建议汇总成一个平均方案，第二天大家按新方案烹饪。因为每家的调整幅度一致，菜谱永远不会分叉。JAX中的pmean就是总店里的“意见平均器”，它让并行训练既快又稳。

θ(t+1)=θ(t)−η⋅1K∑i=1K∇θLi(θ(t))\displaystyle \theta^{(t+1)} = \theta^{(t)} - \eta \cdot \frac{1}{K}\sum_{i=1}^{K} \nabla_{\theta} L_i\left(\theta^{(t)}\right)
公式含义：新参数 = 旧参数 - 学习率 × 平均梯度。平均梯度由all-reduce负责计算，因此每个设备的 \(\theta\) 更新后完全一致。
  1. 局部梯度每个设备 \(i\) 使用自己的数据子集计算损失对参数的梯度 \(\nabla_{\theta} L_i\)。这个梯度是单独前向反向的产物。
    gi=∇θLi(θ(t))\displaystyle g_i = \nabla_{\theta} L_i(\theta^{(t)})
  2. AllReduce平均通过 pmean 将所有 \(g_i\) 求平均，得到 \( \bar{g} \)。这个操作涉及设备间通信，是数据并行的成本中心。
    gˉ=1K∑i=1Kgi\displaystyle \bar{g} = \frac{1}{K}\sum_{i=1}^{K} g_i
  3. 更新参数所有设备用同一个 \(\bar{g}\) 结合学习率 \(\eta\) 更新参数，更新后的 \(\theta^{(t+1)}\) 在各设备上完全相同。
    θ(t+1)=θ(t)−ηgˉ\displaystyle \theta^{(t+1)} = \theta^{(t)} - \eta \bar{g}

完整训练流程

餐厅一天营业的完整流程：开门前备好原料（数据加载），营业中分派订单（batch分片），打烊后汇总反馈（梯度同步），更新菜谱（优化器更新），然后进入下一天。多GPU训练也是这样一个周而复始的循环。

在JAX中，完整的训练流程通常包括：构造数据迭代器，定义损失函数和更新函数，用pmap包装更新函数，然后在循环中不断调用。一个容易被忽视的点是：数据必须先切成设备维+本地batch维。例如4个GPU、每设备32条样本，则一个全局batch的形状应为 (4, 32, ...)，其中4是设备轴，32是每个GPU实际看到的batch大小。

优化器更新一般借助optax或手写。在使用pmap时，优化器的状态（如动量缓冲）也需要是“每设备一份”且保持同步。若使用optax，只需把优化器状态放入params的树状结构中一起pmap即可。训练更新的数学形式如下：

整个训练循环就像餐厅每天重复的“开工—打烊—更新菜谱”。每天的新菜谱是前一天所有反馈的结晶，日复一日，菜品越来越受欢迎，模型损失越来越低。只要保证同步逻辑正确，循环无需人工干预即可稳定运行。

for t=1,…,T:θ(t+1)=θ(t)−ηt⋅1K∑i=1K∇θLi(θ(t))\displaystyle \text{for } t=1,\ldots,T:\quad\theta^{(t+1)} = \theta^{(t)} - \eta_t \cdot \frac{1}{K}\sum_{i=1}^{K} \nabla_{\theta} \mathcal{L}_i(\theta^{(t)})
公式含义：这是整个训练循环的浓缩。每个时间步都重复“算梯度→平均→更新”的流程，\(T\) 是总迭代次数。
  1. 循环开始\(t\) 从1到 \(T\) 循环，代表训练轮数。每轮包含一次完整的数据并行同步。
    t=1,2,…,T\displaystyle t = 1, 2, \ldots, T
  2. 计算平均梯度每个设备计算局部损失 \(\mathcal{L}_i\) 的梯度，再由 pmean 获得平均梯度 \(\bar{g}_t\)。
    gˉt=1K∑i=1K∇θLi(θ(t))\displaystyle \bar{g}_t = \frac{1}{K}\sum_{i=1}^{K} \nabla_{\theta} \mathcal{L}_i(\theta^{(t)})
  3. 更新参数利用当前学习率 \(\eta_t\)（可能随调度器改变）更新参数，得到下一步的参数 \(\theta^{(t+1)}\)。
    θ(t+1)=θ(t)−ηt⋅gˉt\displaystyle \theta^{(t+1)} = \theta^{(t)} - \eta_t \cdot \bar{g}_t

性能分析

餐厅开了4家分店，但总店和分店之间的“意见传递”需要时间。如果每天打烊后汇总反馈要花3小时，而4家分店一天总共节省的工作时间只有2小时，那“开分店”反而亏了。性能分析要回答的核心问题是：多GPU到底带来了多大收益？通信成本是否吞掉了计算收益？

常用指标有两个：吞吐量（每秒处理的样本数）和扩展效率（加速比与设备数的比值）。设单GPU吞吐为 \(S_1\)，\(K\) 个GPU的总吞吐为 \(S_K\)，则扩展效率为：

通信开销与模型参数数量 \(M\) 近似成正比，因为每次all-reduce都要传输 \(M\) 个浮点数。若每设备batch为 \(B\)，则计算量与 \(B\) 成正比。当 \(B\) 太小、通信占比过高时，扩展效率会大幅下降。适当调大每设备batch、采用梯度累积或混合精度，都可以提高计算/通信比，让分店之间的“意见传递”不再成为瓶颈。

EK=SKK×S1\displaystyle E_K = \frac{S_K}{K \times S_1}
公式含义：如果 \(E_K\) 接近1，说明设备增加后吞吐几乎线性增长；如果明显小于1，说明通信或负载不均拖了后腿。
  1. 单GPU基线\(S_1\) 是使用1个GPU时每秒钟训练样本数。先测量基准，才能对比加速效果。
    S1=样本数训练时间\displaystyle S_1 = \frac{\text{样本数}}{\text{训练时间}}
  2. 多GPU吞吐\(S_K\) 是使用 \(K\) 个GPU时的实际总吞吐。由于通信开销，\(S_K\) 通常小于 \(K \times S_1\)。
    SK=样本数并行训练时间\displaystyle S_K = \frac{\text{样本数}}{\text{并行训练时间}}
  3. 扩展效率\(S_K\) 除以 \(K \times S_1\)，得到一个0到1之间的比值。比值越接近1，扩展性越好。
    EK=SKK×S1\displaystyle E_K = \frac{S_K}{K \times S_1}

常见错误

分店协作时最容易出什么岔子？订单总数不能被分店数整除、某家分店偷偷用了旧菜谱、或者厨房案板太小放不下这么多食材。多GPU训练中的错误也集中在这些地方：数据切分、维度、同步状态和显存。

第一个常见错误是切分失败。当batch大小不能被GPU数整除时，pmap会直接报错。例如4个GPU，batch=30，就无法均分。第二个是维度错误：忘记把batch reshape成(n_devices, per_device_batch, ...)，导致pmap认为第一维不是设备轴。第三个是随机状态未同步：如果每个设备使用不同的随机数生成器状态，dropout等操作会让模型产生分叉，破坏数据并行的一致性。第四个是显存不足（OOM）：每设备batch设置过大，超出单卡显存。这些问题可以总结为条件约束：

排查顺序也很重要：先检查batch能否整除，再打印张量形状确认设备轴，接着检查各设备随机key是否一致，最后用nvidia-smi观察显存占用。就像餐厅开分店前先数客人、算桌数、看厨房面积，每一步都验一遍才不会在营业时手忙脚乱。

Bglobal mod K=0且BglobalK×激活内存+参数内存<单卡显存\displaystyle B_{\text{global}} \bmod K = 0 \quad \text{且} \quad \frac{B_{\text{global}}}{K} \times \text{激活内存} + \text{参数内存} < \text{单卡显存}
公式含义：全局batch必须能被设备数整除；每设备batch占用的总显存不应超过单卡容量。
  1. 整除约束全局batch \(B_{\text{global}}\) 除以设备数 \(K\) 必须余0，否则无法均分数据。
    Bglobal≡0(modK)\displaystyle B_{\text{global}} \equiv 0 \pmod{K}
  2. 显存约束每设备batch为 \(B_{\text{global}}/K\)，其产生的激活值内存加上模型参数和优化器状态，必须小于单卡显存总量。
    BglobalK×memact+memparam<VRAM\displaystyle \frac{B_{\text{global}}}{K} \times \text{mem}_{\text{act}} + \text{mem}_{\text{param}} < \text{VRAM}

调优技巧

当你熟悉基本流程后，就该琢磨如何让餐厅更赚钱了。经验丰富的经营者会采用几个绝招：一次多接几桌客人（增大batch）、把难以消化的订单攒到一起反馈（梯度累积）、用平价食材替代昂贵食材（混合精度）、提前备好半成品（数据预取）。

增大每设备batch是提高吞吐最直接的方法。通信量固定为 \(O(M)\)，计算量随batch增大而增大，计算/通信比上升，GPU利用率更高。但batch过大会影响收敛，所以常配合学习率调整。梯度累积是在显存不足时，把多个小batch的梯度累加后再更新，模拟大batch效果。混合精度（如bfloat16）能减少显存占用和通信量，但要注意数值稳定性。数据预取使用tf.data或jax的异步加载，让GPU不因等数据而空闲。梯度累积的数学表达为：

这些技巧就像餐厅经营中的“备菜”和“批量烹饪”：既提高了出餐速度，又不牺牲菜品质量。需要注意的是，梯度累积会延长参数更新的间隔，因此学习率和batch大小的匹配要仔细验证。多做实验、观察loss曲线，往往比盲调参数更有效。

θ(t+1)=θ(t)−η⋅1n∑j=1ngˉ(j)\displaystyle \theta^{(t+1)} = \theta^{(t)} - \eta \cdot \frac{1}{n}\sum_{j=1}^{n} \bar{g}^{(j)}
公式含义：将 \(n\) 个小batch的平均梯度先存在本地，加起来后再统一更新，相当于用多个小batch拼出一个大batch的梯度。
  1. 累积变量在循环外初始化梯度累积量 \(G_{\text{acc}} = 0\)。
    Gacc←0\displaystyle G_{\text{acc}} \leftarrow 0
  2. 逐步累加每个小batch计算出的平均梯度 \(\bar{g}^{(j)}\) 累加到 \(G_{\text{acc}}\) 中。这里 \(\bar{g}^{(j)}\) 已经是跨设备平均后的梯度。
    Gacc←Gacc+gˉ(j)\displaystyle G_{\text{acc}} \leftarrow G_{\text{acc}} + \bar{g}^{(j)}
  3. 平均后更新累积 \(n\) 次后，用 \(G_{\text{acc}} / n\) 作为最终梯度更新参数，并清空累积器。
    θ(t+1)=θ(t)−η⋅Gaccn\displaystyle \theta^{(t+1)} = \theta^{(t)} - \eta \cdot \frac{G_{\text{acc}}}{n}

AI 应用场景

大规模图像分类训练

阶段：模型训练

解决的问题：使用单张GPU训练ResNet-50处理百万级图像，一轮训练需要数天甚至数周。通过多GPU数据并行，可将训练时间缩短数倍，让团队能更快迭代模型。

具体案例：某电商团队在8张A100 GPU上使用JAX的pmap训练商品分类模型。他们将1024的图像batch按8个设备切成每设备128张，前后向计算后通过pmean同步梯度，训练速度提升6.8倍，扩展效率达85%。团队还配合混合精度，将显存占用降低40%，batch可进一步增大。

自然语言处理预训练

阶段：预训练与超参数调优

解决的问题：GPT类模型的预训练成本极高，需要大量数据并行加速。但模型参数和优化器状态占据大量显存，需要结合梯度累积、混合精度与数据并行。

具体案例：某研究团队用JAX在多节点（每节点8 GPU）上预训练一个1.3B参数的Transformer。他们使用pmap将batch沿设备维切分，用lax.pmean同步梯度，并使用梯度累积实现超大全局batch（4M tokens）。通过性能分析发现，通信占比约12%，扩展效率达到90%。

强化学习仿真数据并行

阶段：数据生成与训练

解决的问题：强化学习需要大量仿真环境交互数据，数据并行可以让多个GPU同时收集不同环境的经验，再统一更新策略网络。

具体案例：一个机器人控制项目使用JAX pmap在4个GPU上并行运行64个仿真环境，每个GPU负责16个环境。收集到的轨迹拼接成一个大的batch，通过数据并行训练策略网络，训练速度提升了3.5倍。同时，随机种子在每个设备上独立设置，保证环境探索的多样性。

Python 演示

python65 行
import jaximport jax.numpy as jnpfrom jax import random, value_and_grad, laximport numpy as np # 获取当前所有可用设备devices = jax.local_devices()num_devices = len(devices)print("可用设备数量:", num_devices) # 定义线性回归模型def model(params, x):    return x * params['w'] + params['b'] def loss_fn(params, x, y):    pred = model(params, x)    return jnp.mean((pred - y) ** 2) # 定义单步训练函数，包含跨设备梯度平均def train_step(params, batch, lr=0.1):    x, y = batch  # 每个设备上的本地batch    loss, grads = value_and_grad(loss_fn)(params, x, y)    # 关键：所有设备梯度求平均    grads = lax.pmean(grads, axis_name='devices')    new_params = jax.tree_util.tree_map(        lambda p, g: p - lr * g, params, grads    )    # 损失也做平均，用于监控    loss_avg = lax.pmean(loss, axis_name='devices')    return new_params, loss_avg # 用pmap包装训练函数，指定设备轴名称pmap_train = jax.pmap(train_step, axis_name='devices') # 制造模拟数据：y = 3x + 1 + 噪声key = random.PRNGKey(42)per_device_batch = 8data_size = per_device_batch * num_devicesx = random.normal(key, (data_size, 1))y = 3.0 * x + 1.0 + 0.1 * random.normal(key, (data_size, 1)) # 数据reshape为 (设备数, 每设备batch数, 特征数)batch = (    x.reshape(num_devices, per_device_batch, 1),    y.reshape(num_devices, per_device_batch, 1)) # 初始参数，并复制到所有设备（replicated）params = {'w': jnp.zeros((1,)), 'b': jnp.zeros(())}params = jax.device_put_replicated(params, devices) # 训练20轮for epoch in range(20):    params, loss = pmap_train(params, batch)    loss = jnp.mean(loss)  # 跨设备再求平均    # 从每设备的参数中取第一个设备的副本（值相同），再复制回所有设备    params = jax.tree_util.tree_map(lambda p: p[0], params)    params = jax.device_put_replicated(params, devices)    if epoch % 5 == 0:        print(f"epoch {epoch}, loss = {loss:.4f}") # 输出拟合结果final_w = float(params['w'][0])final_b = float(params['b'])print(f"拟合结果: w={final_w:.3f}, b={final_b:.3f}")
目的：用JAX的pmap实现多GPU数据并行线性回归训练，演示数据分片、梯度平均和参数同步的完整流程。
  1. devices = jax.local_devices()获取机器上所有GPU设备，决定数据切分数。
  2. grads = lax.pmean(grads, axis_name='devices')对所有设备的梯度求平均，这是数据并行同步的核心。
  3. pmap_train = jax.pmap(train_step, axis_name='devices')将单设备训练函数转换成多设备并行版本。
  4. x.reshape(num_devices, per_device_batch, 1)构造“设备轴+本地batch”的标准张量形状。
  5. params = jax.device_put_replicated(params, devices)将同一份参数复制到所有设备，保证起始状态一致。
  6. params = jax.tree_util.tree_map(lambda p: p[0], params)从设备分片参数中取出第一个设备的值（因为已同步，所以任取一个即可）。

常见误区

误区我以为只要用了pmap，模型参数、优化器状态、batch norm的moving mean都会自动在所有设备间保持一致，不需要额外操作。

正解pmap只负责“把函数复制到多设备”和“按轴切分数据”。函数内部如果用lax.pmean计算梯度，那梯度会同步；但其他变量（如batch norm的统计量、优化器的动量）如果不显式做跨设备操作，它们会在各设备独立更新，逐渐发散。就像总店只统一了菜谱，但每家分店自己记录的“老顾客口味档案”不会自动同步，需要额外流程汇总。

为何常见：因为pmap的“自动切分”让人误以为它也自动做了所有集体通信，实际通信必须通过pmean/psum等原语显式触发。

误区设在pmap中传入了batch大小为1024，我以为每个GPU都会处理1024条样本，总batch就是1024。

正解在pmap中，输入张量的第一维被视为设备轴，会自动切分给所有GPU。传入形状为(1024, 32, 32, 3)的张量到4个GPU，每个GPU拿到的是(256, 32, 32, 3)。因此，你真正需要看的是“每设备batch”，总batch = 每设备batch × 设备数。比如想实现全局batch=1024，在4卡时应传入(4, 256, ...)形状的数据。

为何常见：习惯了单卡代码，忘记pmap会沿第一维切分，导致实际batch比预期小K倍，影响学习率设置和训练效果。

误区加一倍GPU，训练速度应该快一倍，扩展效率接近100%。

正解多GPU带来的吞吐提升受通信开销制约。每轮都要做all-reduce，通信量随模型参数数量线性增长。如果每设备batch很小，计算时间短，通信占比就高，扩展效率可能只有60%。提高计算/通信比（如调大batch、用梯度累积）才能接近线性加速，但由于总线带宽等因素，扩展效率极少达到100%。

为何常见：只看到计算并行，忽略了设备间同步也是耗时环节。

误区让每个GPU使用不同的随机种子，dropout和初始化都不同，会不会提升模型多样性、效果更好？

正解数据并行要求所有设备上的模型完全一致，随机状态不一致会导致每设备前向结果不同，梯度平均后再更新，参数虽然同步了，但每个设备的dropout mask不同会让梯度估计更加噪声化，反而降低训练稳定性。正确做法是：模型参数、随机数生成器状态都要复制到所有设备保持一致，而数据增强的随机性应通过数据生成时给每个设备分配不同的数据子集来实现。

为何常见：把“数据多样性”和“模型随机性”混为一谈，忽略了数据并行的基本假设。

面试题

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

单选题 · easy

  • A.将函数映射到多个设备上并行执行← 正确答案
  • B.自动将模型转换为混合精度
  • C.对输入数据执行数据增强
  • D.自动生成训练集

答案：A。pmap是JAX的单程序多数据并行转换，将函数复制到所有设备并沿着第一个轴切分输入，实现数据并行。它不负责混合精度或数据增强。

Q2. 在多GPU数据并行中，每个GPU处理的数据是？

单选题 · easy

  • A.相同的数据副本
  • B.整个batch中不同的子集← 正确答案
  • C.模型所有层参数
  • D.随机抽取的单个样本

答案：B。数据并行的核心是把全局batch切成多份，每个GPU只处理其中一份子集。所有GPU共享同一份模型参数，但数据不同。

Q3. 以下哪个JAX函数用于对跨设备的梯度求平均？

单选题 · easy

  • A.jax.nn.relu
  • B.jax.random.split
  • C.lax.pmean← 正确答案
  • D.jax.numpy.reshape

答案：C。lax.pmean执行的是集体通信操作，沿着指定的axis_name对所有设备上的值求平均，常用于梯度同步。其他函数与设备间通信无关。

Q4. pmap默认会将输入张量沿哪个轴切分？

单选题 · mid

  • A.最后一个轴
  • B.第一个轴← 正确答案
  • C.中间任意轴
  • D.不切分，直接复制

答案：B。pmap按照约定切分第一个轴（axis 0），例如形状(global_batch, ...)中global_batch就是设备轴。这也是为什么需要将数据reshape为(n_devices, per_device_batch, ...)的原因。

Q5. 以下哪些策略可以提升多GPU数据并行的扩展效率？

多选题 · mid

  • A.增大每设备batch大小← 正确答案
  • B.计算与通信重叠← 正确答案
  • C.减小每设备batch大小
  • D.使用混合精度减少通信量← 正确答案

答案：ABD。增大每设备batch能提高计算/通信比；计算与通信重叠可隐藏通信延迟；混合精度减少浮点传输量。减小batch会让通信占比更高，降低效率。

Q6. 梯度同步完成后，各设备上的模型参数应当？

单选题 · mid

  • A.各不相同
  • B.完全一致← 正确答案
  • C.部分一致
  • D.由第一个设备独有

答案：B。数据并行的关键约束是所有设备始终保存同一份参数。通过pmean得到平均梯度后，所有设备用相同梯度更新，参数保持一致。

Q7. 下列哪些属于JAX数据并行中的常见错误？

多选题 · mid

  • A.全局batch不能被设备数整除← 正确答案
  • B.忘记同步各设备的随机数key← 正确答案
  • C.每设备batch过大导致显存溢出← 正确答案
  • D.使用CPU训练但调用了pmap

答案：ABC。A导致切分失败；B导致dropout等状态不一致；C导致OOM。D中pmap可以在CPU上运行（多CPU核），不算错误。

Q8. 扩展效率 \(E_K = S_K / (K \times S_1)\) 中，\(S_K\) 表示什么？

单选题 · hard

  • A.\(K\)个GPU的总吞吐量← 正确答案
  • B.单GPU的训练速度
  • C.\(K\)个GPU的总显存
  • D.通信延迟

答案：A。\(S_K\) 是使用K个GPU时每秒处理的样本数。扩展效率衡量理想线性加速与实际加速的比值，越接近1越好。

Q9. 关于pmap，以下说法正确的是？

多选题 · hard

  • A.pmap的输入会沿第一维切分← 正确答案
  • B.pmap内部不会自动同步所有变量← 正确答案
  • C.使用pmap时必须手动调用all-reduce同步梯度
  • D.输出结果会多出一个设备轴← 正确答案

答案：ABD。C不准确，可以在函数内用pmean/psum同步，但不强制手动调用；所有未同步的变量不保证一致。A、D是pmap的基本语义。

Q10. 在数据并行中，通信开销大致与什么成正比？

单选题 · extreme

  • A.模型参数数量← 正确答案
  • B.训练轮数
  • C.每设备batch大小
  • D.学习率大小

答案：A。AllReduce需要传输每个参数的梯度，通信量 \(O(M)\)，\(M\) 为模型参数量。训练轮数只影响总通信次数，不影响单次通信量；每设备batch影响计算量；学习率仅影响更新大小。