深夜十一点，一家爆火餐厅的后厨仍灯火通明。四口智能灶台并排而立，每一口都能同时喷出上百束火焰；主厨举着喇叭一声令下，四位厨师同时抄起锅铲，各自翻炒面前那一筐配料，最后四份炒饭的调味被统一校准，分毫不差地端上桌。这家餐厅的秘密，就是“智能灶台 + 统一指挥 + 分工协作”三件套——放到 AI 世界里，它们分别叫 TPU、JAX 和分布式训练。 TPU 是专为深度学习打造的芯片，算力爆表；JAX 是连接 TPU 的最佳编程框架，能像指挥棒一样调度所有设备；分布式训练则是把一份大活儿拆给多台设备一起干的艺术。这篇文章会沿用这个后厨故事，带你从零看懂：TPU 凭什么快、JAX 的魔法变换是什么、数据并行和模型并行怎么实现，以及实战中如何定位瓶颈。读完你就能明白，那些动辄数百台设备的“大模型训练工厂”到底是如何运转的。

岗位关联

对 AI 从业者来说，TPU + JAX 是通往大规模分布式训练的捷径。掌握这套组合，你就能用更低的成本训练更大规模的模型，这是大厂 AI 工程岗位的核心硬技能，也是从“调包侠”走向“系统工程师”的关键一步。

目录
  1. 可视化演示
  2. TPU硬件基础：认识算力爆表的“智能灶台”
  3. JAX核心机制：主厨的“函数变换”魔法
  4. 分布式编程模型：一声口令，全员开炒
  5. 数据并行实现：同一道菜，四口锅同时炒
  6. 模型并行实现：佛跳墙的流水线拆解
  7. 实战与调优：出餐速度的全局复盘
  8. AI 应用场景
  9. Python 演示
  10. 常见误区
  11. 面试题

可视化演示

展示TPU数据并行分片训练、模型并行层切分与JAX pmap训练流水线的核心工作流程。

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

数据并行训练 数据分片到各TPU核心每个核心独立计算梯度后同步 梯度通过all-reduce聚合后更新模型参数 T=N/B 吞吐等于核数除以批大小

TPU硬件基础：认识算力爆表的“智能灶台”

午高峰的后厨是真正的战场。普通燃气灶一次只能照顾一口锅，火焰时大时小，全凭厨师手感。而这家餐厅新换的“智能灶台”把上百个火嘴排成密集矩阵，一秒钟能对几十个位置同时精准喷火——别的灶台还在等水烧开，它已经炒完三盘菜了。这台猛灶，就是 TPU。

TPU 最核心的武器叫 MXU，即矩阵乘法单元。深度学习模型九成以上的运算量都是矩阵乘法：词向量乘权重、激活值乘梯度，本质上都是把大块大块的数据成批相乘相加。MXU 把成千上万次乘加运算压进一个时钟周期，像智能灶台同时喷出百束火焰，专治这种“大块头、重复多”的运算。

光有猛火还不够，端菜还得快。TPU 的第二件法宝是高内存带宽，用超宽通道持续喂料；第三件是低精度设计，BF16 半精度让同样的芯片面积塞进更多计算单元，火力密度翻倍。但智能灶台也有脾气：在算术强度低、算子又碎又小的任务里，它的猛火根本喷不出来，反而可能比 GPU 更慢。

选灶台先看菜，选硬件先算算术强度——这个直觉，就是每年千万级算力预算背后的第一道决策公式。

I=CM\displaystyle I = \frac{C}{M}
算术强度等于计算量除以内存访问量。比值越高，越适合 TPU 这种“算得快”的硬件；比值太低，说明瓶颈在搬数据，再猛的灶台也只能干瞪眼。
  1. I（算术强度）一句话概括，就是“每搬一次食材能炒多少下锅”，它是判断任务吃计算还是吃带宽的核心指标。
    I\displaystyle I
  2. C（计算量）这道菜总共需要颠多少次勺，单位是 FLOPs（浮点运算次数），对应模型里的矩阵乘法总次数。
    C\displaystyle C
  3. M（内存访问量）从冰箱搬到灶台边的食材总重量，单位是字节，对应模型训练时要读写的数据量。
    M\displaystyle M
  4. I = C / M（比值判据）数值高说明是“久炒少搬”的硬菜，正好喂给智能灶台；数值低说明一直在搬货，灶台再猛也白搭。
    I=CM\displaystyle I = \frac{C}{M}

JAX核心机制：主厨的“函数变换”魔法

后厨新来的主厨有个怪癖：从不现场翻食谱，而是把每道菜的步骤提前背成肌肉反应，开火瞬间一气呵成。更神奇的是，他能把“炒一人份”的菜谱自动改写成“同时炒一百人份”的批量版本，还能在出锅前准确告诉你哪一味调料该减半。这位主厨，就是 JAX。

JAX 表面上是一套长得像 numpy 的数组库，灵魂却是“函数变换”：它可以对你写好的普通 Python 函数施法。jax.jit 把整个函数编译成 XLA 底层的融合指令，不再逐行解释，像主厨提前背熟菜谱，开火后绝不卡壳；jax.grad 对函数自动求梯度，你只管写前向传播，反向传播的偏导数它替你算，像主厨尝一口就报出该减哪味料。

第三个变换是 jax.vmap，它把“对单条数据计算”自动升级成“对整批数据并行计算”，一行帮你向量化，免去手写 for 循环。这三个魔法是分布式训练的基石：jit 提供编译后的高性能内核，grad 提供训练必需的梯度，vmap 则把单样本逻辑批量放大，为后面把数据切给多台设备铺路。

主厨不必亲自尝遍一百道菜，JAX 也不必手写反向传播——有了 grad 魔法，剩下的就是让差评率一路下滑。

∇θL(θ)=∂L∂θ\displaystyle \nabla_\theta L(\theta) = \frac{\partial L}{\partial \theta}
JAX 的 grad 变换就是自动求出损失函数对每个参数的偏导数，告诉你每一步该往哪个方向调整参数，损失才会下降。
  1. L(θ)（损失函数）客人的差评率，它衡量模型当前表现有多差；θ 是模型的所有参数，也就是整本菜谱的调料配比表。
    L(θ)\displaystyle L(\theta)
  2. θ（参数向量）模型里所有可调整的权重和偏置，相当于菜谱里每一项调料的用量。
    θ\displaystyle \theta
  3. ∇（梯度算子）主厨的“试菜敏锐度”，它指向损失上升最快的方向，取反就是下降最快的方向。
    ∇\displaystyle \nabla
  4. ∇θ L(θ)（梯度向量）把损失对每个参数求偏导，得到的向量告诉主厨：每一项调料该往哪个方向调、调多少。
    ∇θL(θ)=∂L∂θ\displaystyle \nabla_\theta L(\theta) = \frac{\partial L}{\partial \theta}

分布式编程模型：一声口令，全员开炒

四位灶台厨师各守一口锅，主厨举起喇叭喊：“听我口令，一起开炒！”四人同时抬手、同一套动作、同一份菜谱，只是面前各有一筐不同的配料。没人交头接耳，没人问下一句该干嘛——因为大家都知道，喊完这一声，所有人只做同一件事。这种“同一个程序，喂不同的数据”的协作模式，在分布式计算里叫 SPMD。

JAX 里对应的那句话，叫 pmap。它的作用是把纯函数“广播”到所有设备上：调用 pmap(f) 会返回一个新函数，这个函数被调用时，N 台设备同时执行 f，每台设备拿到输入数据的一个切片。你在单机上写的 f 完全是原来的样子，pmap 只是给它加了一层“分身术”。

pmap 返回的不再是普通数组，而是 ShardedDeviceArray——一个“设备数组”。它像一个遥控器：逻辑上你看到的还是完整形状的数据，物理上它早已被拆成 N 份，躺在 N 台设备的内存里。更妙的是，pmap 可以和 grad、jit 随意组合：pmap(grad(f)) 就是让每台设备各自完成前向和反向计算。这套语法，正是数据并行的地基。

一声口令、一份菜谱、N 份食材——SPMD 的精髓在于“人同法，菜不同”，这也是所有并行训练共同的心法。

X=[x1,x2,…,xN],yi=f(xi)\displaystyle X = [x_1, x_2, \ldots, x_N], \quad y_i = f(x_i)
SPMD 把完整数据集 X 切成 N 份，第 i 台设备只对自己那份 x_i 执行同一个函数 f，各自得到 y_i，互不干扰。
  1. X（完整数据集）这一整筐要炒的青菜，从外面看就是满满一筐，对应训练时完整的一个大 batch。
    X\displaystyle X
  2. x_i（第 i 份切片）主厨把一筐菜均匀分成 N 份，第 i 位厨师领走第 i 份，这就是数据切分。
    xi(i=1,2,…,N)\displaystyle x_i \quad (i = 1, 2, \ldots, N)
  3. f（同一个函数）那张贴在墙上的标准菜谱，所有厨师都照它炒，对应 pmap 广播给所有设备的模型函数。
    f\displaystyle f
  4. y_i = f(x_i)（各算各的）每位厨师炒自己那份菜，出锅的 N 份炒饭拼在一起，就是完整一桌的产出。
    yi=f(xi)\displaystyle y_i = f(x_i)

数据并行实现：同一道菜，四口锅同时炒

招牌炒饭是店里的爆款，但一口锅再快，午高峰也顶不住几百份订单。主厨的办法很朴素：菜谱只有一份，食材管够。把一大锅米饭分成四份，四口灶台同时开炒同一道炒饭，出锅前再由主厨把四份调味校准成完全一致。这就是数据并行——模型只有一份，数据切成多份。

落到代码上，数据并行的流程分三步：第一步，把一个大 batch 按设备数 N 切成 N 个小 batch；第二步，每台设备用完全相同的模型参数副本，对自己那份数据做前向和反向，算出一份“局部梯度”；第三步，四份梯度方向略有差异，不能各改各的参数，必须聚合成一份，用平均梯度更新那个唯一的模型。

第三步正是数据并行的命门。通信上的标准操作叫 All-Reduce，即全归约：把 N 台设备上的值两两合并、取平均，再把结果广播给所有人。JAX 里可以用 jax.lax.psum 显式完成求和，或者靠 optax 优化器里的梯度同步逻辑自动完成。数据并行的扩展性非常诱人：设备翻倍、batch 翻倍，吞吐量几乎线性上涨——只要通信量没有吃掉收益。

四口锅炒同一道菜，出锅前必须校准调味——少了 All-Reduce 这一步，四份模型迟早各炒各的，变成四道完全不同的菜。

gˉ=1N∑i=1Ngi\displaystyle \bar{g} = \frac{1}{N} \sum_{i=1}^{N} g_i
全归约把所有设备的梯度逐元素相加再除以设备数。这个平均梯度就是“主厨校准后的统一调味”，用来更新唯一的模型参数。
  1. g_i（第 i 台设备的局部梯度）第 i 位厨师尝过自己那锅炒饭后写下的调味心得。
    gi=∇Li(θ)\displaystyle g_i = \nabla L_i(\theta)
  2. Σg_i（梯度求和）主厨把四份心得汇总到同一张纸上，对应 All-Reduce 中的归约求和过程。
    ∑i=1Ngi\displaystyle \sum_{i=1}^{N} g_i
  3. 1/N（取平均）除以灶台数，去掉“某口锅特别咸”的偶然偏差，得到大家公认的调整方向。
    1N\displaystyle \frac{1}{N}
  4. ḡ（平均梯度）最后的统一意见——所有设备都按这个平均梯度更新同一份参数，模型才不会分叉。
    gˉ=1N∑i=1Ngi\displaystyle \bar{g} = \frac{1}{N} \sum_{i=1}^{N} g_i

模型并行实现：佛跳墙的流水线拆解

招牌佛跳墙工序复杂：泡发海参、煸炒底料、文火慢炖，一个人从头做到尾，又慢又占灶台。后厨于是把工序拆开——有人专管泡发，有人负责煸炒，有人守着砂锅。每个灶台只干一件事，但合起来就是一整锅佛跳墙。模型并行也是这个思路：模型太大，一口锅塞不下，就把它拆成几段，每台设备负责一段。

最常见的拆法有两种。第一种叫张量并行：一层 Transformer 的权重矩阵特别大，那就把这个矩阵按列切开，N 台设备各持一块，前向时各自算出一部分结果，再通过通信把它们拼回完整的输出。它像一张大菜谱被拆成 N 册，每册记一部分做法，做菜时还要互相传纸条对进度。

第二种叫流水线并行：按层把模型切成几段，设备 1 算完前几层，把中间结果传给设备 2 继续算，像工厂流水线。它的代价是“启动气泡”——上游还没把半成品传下来时，下游设备只能干等。JAX 里更现代的做法，是用 jax.sharding 配合 pjit 做声明式分片：你只告诉编译器这个张量要切几份、放哪些设备，通信方案由编译器自动生成；再配合 BF16 混合精度和梯度重计算，把显存压到极限。

佛跳墙一个人做不完，大模型一口锅也放不下——把菜谱拆开、把工序排开，这就是模型并行的生存智慧。

W=[W1,W2,…,WN]\displaystyle W = [W_1, W_2, \ldots, W_N]
张量并行把大权重矩阵 W 按列切成 N 块，各设备持一块，前向计算时拼合局部结果，等效于用完整的 W 做一次矩阵乘法。
  1. W（完整权重矩阵）一整本佛跳墙菜谱，一个灶台的手写板根本记不下，必须拆开。
    W\displaystyle W
  2. W_i（第 i 份权重切片）把菜谱拆成 N 册，第 i 台设备只持有第 i 册，记住属于自己的那部分做法。
    Wi\displaystyle W_i
  3. W_i x（局部计算）每台设备用自己的切片和输入 x 做矩阵乘法，得到一个部分结果，相当于各自先处理自己负责的配料。
    Wix\displaystyle W_i x
  4. 拼接完整输出通过通信把所有部分结果拼接成完整输出，从外面看就像用整块 W 算了一次。
    Y=concat(W1x,W2x,…,WNx)\displaystyle Y = \mathrm{concat}(W_1 x, W_2 x, \ldots, W_N x)

实战与调优：出餐速度的全局复盘

设备到位、菜谱改好，真正开店后，主厨盯的不再是某口锅有多猛，而是整条出餐线的吞吐率：哪口锅在等食材？哪位厨师闲得发呆？排队越来越长的客人，说明瓶颈卡在了某个暗处。分布式训练调优，干的也是这种全局复盘。

上手的起点是环境配置。在 TPU 虚拟机里，第一件事是确认能看到几颗芯片：jax.devices() 会列出所有可用设备，比如 TPU v3 的 8 核版就有 8 个 device。然后按设备数切数据，把大 batch 用 jnp.array_split 切成 N 份，再用 pmap（或 pjit）分发。记住一条铁律：先用单设备跑通逻辑，再加并行——否则报错时你分不清是算法错了还是并行写错了。

调优时优先盯三处瓶颈：第一是通信占比过高，梯度同步太频繁，All-Reduce 的时间超过了计算时间，这时要加大 batch 或降低同步频率；第二是负载不均，数据洗牌没做好，某台设备吃得多、其他设备在等，这时要给数据加随机打乱；第三是编译时间过长，动态 shape 会让 JAX 反复重编译，尽量固定 batch 大小。定位工具用 jax.profiler，抓一条 trace 就能看到设备利用率和通信开销。

后厨的真相是：瓶颈永远藏在最慢的那一环。找到它、消除它，比盲目增加灶台更管用。

S=1(1−p)+pN\displaystyle S = \frac{1}{(1-p) + \frac{p}{N}}
阿姆达尔定律告诉我们：加速比 S 由串行比例 (1-p) 钳制。当 p=0.9、N=100 时，极限加速也只有 10 倍——只有把串行部分压到最低，加设备才划算。
  1. p（可并行比例）这家店流程中能分给多口灶台同时干的比例，对应训练里能并行化的部分。
    p\displaystyle p
  2. 1-p（串行比例）洗锅、摆盘等只能一个人做的环节，对应梯度同步、数据加载等无法省略的开销。
    1−p\displaystyle 1-p
  3. N（设备数）灶台数量，对应 TPU 芯片数。
    N\displaystyle N
  4. S（加速比）加灶台后整体出餐速度的提升倍数。串行比例越小，加设备越划算。
    S=1(1−p)+pN\displaystyle S = \frac{1}{(1-p) + \frac{p}{N}}

AI 应用场景

大语言模型预训练

阶段：模型训练

解决的问题：千亿参数模型单设备放不下、单机训不动。传统方案要么漫长单机训练，要么手写多卡通信逻辑，开发周期长且易出错。用 JAX 在 TPU Pod 上做混合并行（数据并行 + 张量并行 + 流水线并行），代码量小且通信方案由编译器自动生成，能把数百台设备的利用率拉到可观水平。

具体案例：某开源团队在 TPU v4 上训练 130 亿参数模型：先用单设备调试模型逻辑，再用 jax.sharding 声明权重分片，用 pmap 把数据切到 256 份，训练吞吐相比单设备提升约 180 倍。

推荐系统模型持续更新

阶段：持续训练

解决的问题：推荐模型每天面对数亿新样本，单机训练一轮要几十分钟，模型新鲜度差。用 JAX 数据并行把 batch 切到 8 颗 TPU 核上，All-Reduce 聚合梯度，训练吞吐提升数倍，让模型能小时级滚动更新。

具体案例：某广告 CTR 预估模型原来单卡每轮 40 分钟，改用 JAX + 8 核 TPU 后每轮 6 分钟，日更新频率从 1 次提升到 6 次，线上点击率提升约 1.2%。

科学计算批量推理

阶段：批处理推理

解决的问题：蛋白质结构预测等任务需要对海量序列打分，单设备逐条推理太慢。用 jax.vmap 把推理函数批量向量化，再用 pmap 分发到多块 TPU 上并行打分。

具体案例：AlphaFold 类模型批量推理，单条序列打分耗时从 2 分钟降到约 5 秒，支撑起全基因组规模的筛选任务。

Python 演示

python30 行
"""用纯 Python 模拟 JAX 在 TPU 上的数据并行训练.四台虚拟 TPU 各吃一个数据切片, 独立算梯度, 最后 All-Reduce 取平均."""import numpy as np NUM_DEVICES = 4  # 模拟 4 台 TPU 芯片 def local_gradient(x_slice, w):    """单设备计算: 对简化损失 (w*x - y)^2 求导, 返回局部梯度"""    return 2.0 * w + x_slice.mean() def all_reduce_mean(grads):    """全归约: 汇总所有设备梯度并取平均, 广播给所有设备"""    return sum(grads) / len(grads) def train_step(data, w, lr=0.1):    """一次迭代: 切分数据 -> 并行算梯度 -> 聚合 -> 更新参数"""    shards = np.array_split(data, NUM_DEVICES)    local_grads = [local_gradient(s, w) for s in shards]    avg_grad = all_reduce_mean(local_grads)    return w - lr * avg_grad, local_grads, avg_grad if __name__ == "__main__":    np.random.seed(0)    data = np.random.randn(1024)   # 一整筐 1024 份食材    w = 3.0                        # 初始"调味参数"    for epoch in range(1, 6):        w, local_grads, avg_grad = train_step(data, w)        print("epoch {}: 四设备梯度={}, 平均梯度={:.4f}, w={:.4f}".format(            epoch, [round(g, 4) for g in local_grads], avg_grad, w))
目的：用纯 Python 完整模拟数据并行训练的核心流程——数据切分、各设备独立算梯度、All-Reduce 取平均、参数更新，帮你在没有 TPU 的环境下也能理解 JAX 分布式训练的每一步。
  1. shards = np.array_split(data, NUM_DEVICES)把一整筐数据均匀切成 4 份，对应数据并行第一步“数据切分”。
  2. local_grads = [local_gradient(s, w) for s in shards]模拟四台设备“同时”各自算梯度，对应 pmap 广播函数后各设备独立计算。
  3. avg_grad = all_reduce_mean(local_grads)全归约取平均，对应 JAX 里 jax.lax.psum 的梯度聚合。
  4. return w - lr * avg_grad, local_grads, avg_grad用唯一的平均梯度更新唯一参数 w，模型才不会分叉。
  5. np.random.seed(0)固定随机种子，保证每次运行结果一致，方便对比调试。

常见误区

误区听说 TPU 是专门为深度学习设计的芯片，那我把代码丢到 TPU 上，肯定比 GPU 快很多，训练时间直接砍半。

正解TPU 的强项是矩阵乘法密集、数据搬运少的 workload。算术强度高的任务它才能碾压 GPU；如果模型全是小算子、动态 shape、频繁分支，TPU 的高算力根本用不上，甚至因为编译和调度开销反而更慢。选硬件前先算一算自己的算术强度，别盲目换灶台。

为何常见：厂商宣传和公开 Benchmark 都拿大矩阵任务说事，给人“TPU 万能快”的错觉，忽略了 workload 匹配度。

误区我给训练函数套上 pmap 后，四台设备都在跑，觉得梯度聚合、参数更新、All-Reduce 这些事 JAX 肯定已经悄悄帮我完成了。

正解pmap 的语义只是“每台设备执行同一份函数、处理各自的数据切片”，它不做任何跨设备的梯度通信。要聚合梯度，必须在函数里显式写 jax.lax.psum，或用 optax 的梯度同步逻辑。忘了同步，四台设备就会各自为政，模型参数分叉，loss 永远不收敛。

为何常见：pmap 名字听起来像“并行全搞定了”，加上 JAX 常被形容为“自动编译、自动并行”，让人误以为所有分布式细节都被自动处理。

误区既然有两种并行方式，那我得选一种：要么每台设备复制整个模型然后切数据，要么把模型拆开。选数据并行就不能用模型并行。

正解真实的大模型训练几乎都是混合并行：模型大到单设备放不下，就用张量并行把一层切开；层数多，就用流水线并行串起来；数据量大，再叠加数据并行。JAX 的 pjit/sharding 可以给每个张量独立声明分片策略，同时混合多种并行方式。

为何常见：教材按“两种并行”分开讲，初学者容易把它们当成互斥方案，忘了现实世界是组合拳。

误区JAX 不是号称 numpy 的加速版吗？我直接把 `import numpy as np` 换成 `import jax.numpy as jnp`，训练应该就自动变快了。

正解没写 jax.jit 的话，JAX 还是逐行解释执行，甚至因为 Python 层开销比 numpy 更慢。真正的加速来自 jit 把整段函数编译成 XLA 的融合算子，以及 grad/vmap/pmap 这些变换的组合。换库名只是第一步，关键要吃透“函数式 + 变换”的编程范式。

为何常见：JAX 的 API 长得和 numpy 几乎一模一样，语法上确实能无缝替换，让人以为性能也能自动继承。

面试题

Q1. TPU 里负责把大量矩阵乘加运算压进一个时钟周期完成的核心单元是？

单选题 · easy

  • A.矩阵乘法单元MXU← 正确答案
  • B.统一着色器阵列
  • C.专用寄存器堆阵列
  • D.标量运算执行单元

答案：A。TPU 的核心武器是 MXU，它把成百上千次乘加运算并行执行，像智能灶台同时喷出百束火焰。着色器是 GPU 的旧概念，寄存器堆和标量单元都不是 TPU 的杀手锏。

Q2. 在 JAX 中，把整个 Python 函数编译成 XLA 底层指令、避免逐行解释的变换是？

单选题 · easy

  • A.jax.grad
  • B.jax.jit← 正确答案
  • C.jax.vmap
  • D.jax.pmap

答案：B。jax.jit 把函数整体编译成融合的 XLA 指令，像主厨提前背熟菜谱、开火后绝不卡壳。grad 是求梯度，vmap 是向量化，pmap 是设备并行，功能各不相同。

Q3. pmap 的并行语义属于下面哪一种模型？

单选题 · mid

  • A.单程序单数据
  • B.多程序多数据
  • C.单程序多数据← 正确答案
  • D.多程序单数据

答案：C。pmap 是典型的 SPMD：所有设备执行同一个函数（单程序），各自处理自己的数据切片（多数据）。这正是一声口令、全员开炒的后厨协作方式。

Q4. 数据并行训练中，把各设备梯度汇总取平均再广播回去的通信操作是？

单选题 · mid

  • A.All-Gather
  • B.All-Reduce← 正确答案
  • C.Broadcast
  • D.Scatter

答案：B。All-Reduce（全归约）把 N 台设备的梯度归约成一份再广播给所有人，对应主厨校准四份调味。All-Gather 只拼接不全约，Broadcast 只广播不归约，Scatter 是切分发散的逆操作。

Q5. 算术强度 I = C / M 描述的是？

单选题 · easy

  • A.代表计算量大小
  • B.每字节搬运对应算力← 正确答案
  • C.代表内存容量大小
  • D.代表设备数量多少

答案：B。I 是计算量 C 与内存访问量 M 的比值，衡量每搬一字节数据能做多少次运算。它越高越适合 TPU，越低说明瓶颈在搬数据，再猛的火力也白搭。

Q6. 模型并行常见的实现形式包括？

多选题 · mid

  • A.张量并行← 正确答案
  • B.流水线并行← 正确答案
  • C.随机失活
  • D.梯度裁剪

答案：AB。张量并行把权重矩阵切开分给多设备，流水线并行按层分段接力计算，二者都是模型并行。随机失活是正则化手段，梯度裁剪是稳定训练技巧，都不是并行方式。

Q7. 关于 pmap 返回的设备数组（ShardedDeviceArray），下列说法正确的是？

多选题 · mid

  • A.逻辑上是完整数组← 正确答案
  • B.物理分布多台设备← 正确答案
  • C.数据只能放在CPU
  • D.能够自动聚合梯度

答案：AB。设备数组逻辑上仍是完整形状，物理上被切分放在多台 TPU 上。它并不会自动做梯度聚合——聚合需要 psum 或优化器显式完成，这是新手最常见的误解。

Q8. 流水线并行中，上下游设备因等待中间结果而空闲的时间被称为？

单选题 · hard

  • A.通信拥堵字节数
  • B.设备等待的空闲时间← 正确答案
  • C.显存不足异常报错
  • D.梯度爆炸对应步数

答案：B。流水线启动阶段，下游设备要等上游算完前几层把半成品传下来，这段时间无事可做，称为 bubble（气泡）。它是流水线并行的固有开销，可通过微批次切分来缓解。

Q9. JAX 中把“对单条数据计算”自动升级为“对整批数据计算”的变换是？

单选题 · easy

  • A.自动求梯度方向
  • B.批量自动向量化← 正确答案
  • C.自动混合精度转换
  • D.自动并行执行调度

答案：B。jax.vmap 自动向量化，把单样本的 f(x) 变成能处理整批 X 的函数，免去手写 for 循环——相当于主厨把“炒一人份”的菜谱改写成“炒百人份”的批量版。

Q10. 关于 JAX 分布式训练的说法，正确的是？

多选题 · hard

  • A.pmap 内可继续使用 grad← 正确答案
  • B.pjit 支持声明式分片← 正确答案
  • C.全归约可用 psum 实现← 正确答案
  • D.TPU 不需要切分数据

答案：ABC。pmap(grad(f)) 完全合法，每台设备各自反向传播；pjit 配合 jax.sharding 声明张量怎么切；jax.lax.psum 正是全归约求和原语。TPU 必须显式切分数据，否则数据只落在单颗芯片上。