想象你是一家大型餐厅的后厨主管。后厨里 8 位厨师站在 2 行 4 列的灶台前，有人旁边是烤箱，有人旁边是水池——这张灶台布局图就是“设备网格”。你收到一份写给“一位厨师”的菜谱：洗菜、切丝、下锅、装盘，清清楚楚写完了整道菜。可问题来了：8 位厨师总不能干等着。于是你这位主厨亲自上阵，把菜谱默读一遍，心中算出“三百斤土豆分 8 份，每份 37.5 斤；牛肉先切后炖、炖肉的同时炒配菜”，然后按布局图把任务一一派下去，连“谁切好的菜该递给谁”都安排妥当。写菜谱的人一个字都没提“几个人做”“怎么分”，成品却和一个人做的一模一样。这套打法，就是 JAX 里 Mesh 与 GSPMD 干的事：前者描述灶台怎么摆，后者把“单厨师菜谱”自动变成“八厨师作战方案”。本文就用这个后厨故事，把分布式训练的底层逻辑一层层剥开。

岗位关联

学会 Mesh 与 GSPMD，你就能把“单设备代码”一键放大到上百台加速器，这是大模型时代 AI 工程师必备的分布式基本功。

目录
  1. 可视化演示
  2. Mesh设备网格：后厨布局图
  3. 单设备编程：一人菜谱
  4. GSPMD自动切分：主厨分派
  5. 显式切分控制：菜谱批注
  6. 并行策略：分工模式
  7. 对比手动并行：主厨 vs 帮厨
  8. 调试与优化：试菜环节
  9. AI 应用场景
  10. Python 演示
  11. 常见误区
  12. 面试题

可视化演示

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

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

JAX中Mesh与GSPMD：Mesh Mesh设备网格拓扑定义与配置 GSPMD单设备代码自动并行化 核心机制可视化 JAX中Mesh与GSPMD整体架构与关键组件

Mesh设备网格：后厨布局图

一家餐厅要接 500 人的宴席，光靠一个灶台肯定不行。主厨做的第一件事，是把 8 个灶台按 2 行 4 列摆好，画成一张布局图：图上每个格子代表一个灶台，横向叫“data 轴”、纵向叫“model 轴”。有了这张图，主厨才知道该把哪道菜的食材往哪个方向传。

JAX 里的 Mesh 就是这张布局图。它由两个部分组成：设备数组（哪些设备参与）和轴名（每条坐标轴叫什么）。创建 Mesh 时，你告诉 JAX“我有 8 台设备，排成 2×4，横轴叫 data、纵轴叫 model”，之后所有关于“往哪台设备上放数据”的决策，都以这张图为基准。实际代码只有两行：

设备数组的排列不是随意的——相邻设备之间的通信带宽通常更高，因此把频繁通信的两个并行维度放在相邻位置，能显著降低通信开销。轴名则是给每个维度起名字，方便之后用 PartitionSpec 引用。这张图能容纳的设备总数由所有轴的尺寸相乘决定：

布局图画得好，后厨就顺畅：切好的土豆从左边传到右边，中间的“传菜走廊”（设备间通信）几乎不堵车；画得不好，食材会在过道里来回穿梭，整个厨房的效率就垮了。所以第一步，永远是把 Mesh 这张图画明白。

∣M∣=∏i=1ndi=ddata×dmodel\displaystyle |\mathcal{M}| = \prod_{i=1}^{n} d_i = d_{\text{data}} \times d_{\text{model}}
网格中的设备总数，等于每条轴上设备数量的乘积。
  1. 网格符号\displaystyle \mathcal{M} 指代整个设备网格，如同图纸上所有灶台的集合
    M\displaystyle \mathcal{M}
  2. 设备总数\displaystyle |\mathcal{M}| 表示网格内设备的总数量，即后厨总灶台数
    ∣M∣\displaystyle |\mathcal{M}|
  3. 累乘运算\displaystyle \prod_{i=1}^{n} d_i 把所有轴上的设备数目乘起来，得到全部设备数
    ∏i=1ndi\displaystyle \prod_{i=1}^{n} d_i
  4. 具体实例\displaystyle d_{\text{data}} \times d_{\text{model}} 例如 2×4=8 台设备，两个轴名各管一个维度
    ddata×dmodel\displaystyle d_{\text{data}} \times d_{\text{model}}
python4 行
from jax.experimental.mesh_utils import create_device_meshfrom jax.sharding import Mesh mesh = Mesh(create_device_mesh((2, 4)), ('data', 'model'))

单设备编程：一人菜谱

主厨定了一条规矩：所有厨师交上来的菜谱，只准按“一个人做”来写。不许写“张三把菜递给李四”，也不许写“王五负责看火”。为什么？因为派活的活儿归主厨，你们只要把菜做好吃。

JAX 的编程模型正是如此。你用普通 JAX 函数写出前向和反向逻辑——一个输入进来，一个输出出去，中间不碰全局状态、不改外部变量。这种“纯函数”就是一人菜谱。它不需要任何并行标注，因为 GSPMD 会在编译期负责“分派”。

用 jax.jit 包装函数后，JAX 会把函数编译成一张计算图。GSPMD 拿到这张图，像主厨默读菜谱一样，分析每个张量的形状、运算和依赖关系，然后决定如何把它拆成多设备版本。拆完之后，每台设备上跑的是同一份程序的“分片版本”，这就是 SPMD（单程序多数据）的精髓：

你交上去的菜谱依旧只有六步：洗、切、焯、炒、勾芡、装盘。至于哪一步在哪个灶台做、做完递给谁，主厨心里有数。这就是“单设备编程”的意义——把并行复杂度从写菜谱的人身上，转移到了主厨（编译器）身上。

∀ i∈{1,…,d}:yi=fθ(xi),x=[x1,x2,…,xd]\displaystyle \forall\, i \in \{1, \dots, d\}: \quad y_i = f_{\theta}(x_i), \qquad x = [x_1, x_2, \dots, x_d]
每台设备执行同一个函数，但各自处理输入的一个分片。
  1. 设备编号\displaystyle i \in \{1, \dots, d\} 遍历 d 台设备，每台都有唯一的编号
    i∈{1,…,d}\displaystyle i \in \{1, \dots, d\}
  2. 同一程序\displaystyle y_i = f_{\theta}(x_i) 每台设备执行同一个函数，只是输入不同
    yi=fθ(xi)\displaystyle y_i = f_{\theta}(x_i)
  3. 切片输入\displaystyle x_i 第 i 台设备拿到完整输入 x 的第 i 片
    xi\displaystyle x_i
  4. 拼回整体\displaystyle x = [x_1, x_2, \dots, x_d] 所有分片拼接起来仍是原来的张量
    x=[x1,x2,…,xd]\displaystyle x = [x_1, x_2, \dots, x_d]

GSPMD自动切分：主厨分派

宴席开动前夜，主厨把菜谱翻来覆去看了三遍。他得决定：三百斤土豆是平均分 8 份，还是留两份做土豆泥？牛肉是整块炖还是切块炖？这些决定直接关系到每位厨师手里有多少活、灶台之间要传递多少半成品。

GSPMD 就是这位主厨。它在编译期对计算图做三件事：第一，分析张量之间的数据依赖；第二，为每个张量决定切分方案；第三，在切分边界插入通信操作，比如 all-reduce、all-gather。这三步做完，单设备程序就变成了多设备程序，而用户代码一行都没动。

切分方案的核心是“每个张量沿哪些维度切成几片”。对常见的矩阵乘法，GSPMD 会权衡计算开销和通信开销，选择总体代价最小的方案。给定张量第 j 维的大小和切分因子，每片的大小为：

主厨切土豆的标准不是“切得越小越好”，而是“每口锅都能同时开火、传递次数最少”。GSPMD 的切分决策同样是通信与计算的平衡——这是理解它一切行为的关键。

shard_sizej=Tjsj,sj∈{1,2,4,8,… }\displaystyle \text{shard\_size}_j = \frac{T_j}{s_j}, \qquad s_j \in \{1, 2, 4, 8, \dots\}
某一维切分后，每片长度等于原长度除以切分份数。
  1. 维度总长\displaystyle T_j 是张量在第 j 维的完整长度，相当于整块食材的尺寸
    Tj\displaystyle T_j
  2. 切分因子\displaystyle s_j 表示第 j 维被切成的份数，必须能整除原长度
    sj∈{1,2,4,8,… }\displaystyle s_j \in \{1, 2, 4, 8, \dots\}
  3. 每片长度\displaystyle \frac{T_j}{s_j} 用总长除以份数，得到每个设备上那一小段的长度
    Tjsj\displaystyle \frac{T_j}{s_j}
  4. 不切情形\displaystyle s_j = 1 表示整条维度不切，每台设备都保留完整的一份
    sj=1\displaystyle s_j = 1

显式切分控制：菜谱批注

大多数时候，主厨自己定切法。但有一道祖传招牌菜，主厨会在菜谱上加一行红字批注：“这块和牛必须送到最左边那排灶台，不能切。”批注只说明“放哪、切不切”，具体怎么做还是厨师自己定。

JAX 提供了同样的“批注”机制：NamedSharding 配合 PartitionSpec。PartitionSpec 描述张量的每个维度对应 Mesh 的哪条轴，或者对应 None（表示不切）；NamedSharding 则把这份描述和具体的 Mesh 绑定，形成一个完整的布局声明。

比如一个形状为 (B, F) 的输入张量，想让它在 data 轴上切分、特征维保持完整，就可以写 PartitionSpec('data', None)。声明之后，GSPMD 在编译期会把它当作初始分片线索，并沿着这条线索传播到后续所有算子。显式切分并不等于手动编写并行逻辑——它只是给编译器一张“起始提示”：

批注写得越清楚，主厨越省心；但批注不等于“主厨每次都要重新排班”。绝大多数张量你根本不用管，只要给关键数据（输入、权重）标好位置，GSPMD 会帮你把整条链路串起来。

PartitionSpec(’data’,None):x↦[xi∈RBddata×F]i=1ddata\displaystyle \text{PartitionSpec}(\text{'data'}, \text{None}): \quad x \mapsto \left[ x_i \in \mathbb{R}^{\frac{B}{d_{\text{data}}} \times F} \right]_{i=1}^{d_{\text{data}}}
按 data 轴切分批次维，特征维保持完整，每个设备拿到一个子张量。
  1. 轴名映射\displaystyle \text{'data'} 表示张量第 0 维按照 Mesh 的 data 轴来切
    ’data’\displaystyle \text{'data'}
  2. None 保留\displaystyle \text{None} 表示第 1 维不切，每台设备保留完整的 F 列
    None\displaystyle \text{None}
  3. 分片形状\displaystyle \mathbb{R}^{\frac{B}{d_{\text{data}}} \times F} 批次被平均分到 d_data 台设备上
    RBddata×F\displaystyle \mathbb{R}^{\frac{B}{d_{\text{data}}} \times F}
  4. 分片集合\displaystyle [x_i]_{i=1}^{d_{\text{data}}} 所有设备上的分片合在一起，还原成原张量
    [xi]i=1ddata\displaystyle \left[ x_i \right]_{i=1}^{d_{\text{data}}}

并行策略：分工模式

后厨有三种典型的派活方式：一种是把同样的菜做 8 份，每份让一个厨师从头做到尾；一种是一道佛跳墙，8 位厨师每人负责几样食材，最后汇到一个锅里；还有一种像流水线，切菜、炖煮、摆盘分段接力。

这三种方式对应分布式训练的三种并行策略：数据并行、张量并行与流水并行。数据并行最简单，每台设备持有完整模型、处理不同批次数据；张量并行把单个大矩阵切成块分给多台设备，解决单卡放不下大模型的问题；流水并行则把网络层切段，设备之间接力计算。

数据并行中，每台设备算完自己的梯度后，需要用一次 all-reduce 把梯度求和再平均，保证所有设备上的参数一致。这一步是数据并行的主要通信成本：

主厨选哪种派活方式，取决于菜的复杂度、灶台数量和食材量。GSPMD 的价值在于：这些策略可以混合使用，而菜谱只需要写一次——主厨自动安排，厨师只需照做。

gˉ=1d∑i=1dgi,θt+1=θt−η gˉ\displaystyle \bar{g} = \frac{1}{d} \sum_{i=1}^{d} g_i, \qquad \theta_{t+1} = \theta_t - \eta\, \bar{g}
所有设备的梯度求平均后，统一更新参数。
  1. 局部梯度\displaystyle g_i 是第 i 台设备用自己的批次算出的梯度，相当于每位厨师的试吃反馈
    gi\displaystyle g_i
  2. 梯度平均\displaystyle \bar{g} = \frac{1}{d} \sum_{i=1}^{d} g_i 把所有局部梯度求和后除以设备数 d
    gˉ=1d∑i=1dgi\displaystyle \bar{g} = \frac{1}{d} \sum_{i=1}^{d} g_i
  3. 参数更新\displaystyle \theta_{t+1} = \theta_t - \eta\, \bar{g} 用平均梯度更新模型参数
    θt+1=θt−η gˉ\displaystyle \theta_{t+1} = \theta_t - \eta\, \bar{g}
  4. 学习率\displaystyle \eta 控制每一步更新的步子大小，太小收敛慢、太大不稳定
    η\displaystyle \eta

对比手动并行：主厨 vs 帮厨

老式后厨有个“帮厨”制度：每道菜都要在菜谱里写明“张三洗菜、李四切菜、王五掌勺”，连“做完第一步，把盆递给隔壁灶台”这种话都要写进去。菜谱写出来厚得像本书，换个人还看不懂。

JAX 的 pmap 就像这位帮厨——它要求你写清楚：设备怎么映射、数据怎么分发、结果怎么收集。代码里到处是设备下标和通信函数，逻辑与并行细节完全耦合。GSPMD 则是把主厨请了回来：你交单设备菜谱，它负责排班。

手动并行难以组合：一个用 pmap 写的子函数，很难嵌进另一个并行策略里；而 GSPMD 的切分声明可以放在任何函数边界上，编译器自动推导。从性能模型看，任何并行方案的总耗时都由计算和通信两部分构成，手动并行容易引入多余通信：

pmap 是“人肉排班”，GSPMD 是“智能排班”。前者面对 8 台设备还能应付，面对千卡集群时就崩溃了；后者无论厨房多大，菜谱都保持不变。

Ttotal=Tcompute+Tcomm,Tcomm∝∑e∈EcutVe\displaystyle T_{\text{total}} = T_{\text{compute}} + T_{\text{comm}}, \qquad T_{\text{comm}} \propto \sum_{e \in E_{\text{cut}}} V_e
总时间 = 计算时间 + 通信时间，通信量正比于被切开的连接上的数据量。
  1. 总耗时\displaystyle T_{\text{total}} 一轮训练从开始到结束的总时间
    Ttotal\displaystyle T_{\text{total}}
  2. 计算耗时\displaystyle T_{\text{compute}} 设备做矩阵运算等真正“干活”的时间
    Tcompute\displaystyle T_{\text{compute}}
  3. 通信耗时\displaystyle T_{\text{comm}} 设备之间传数据的时间，相当于后厨里传菜的时间
    Tcomm\displaystyle T_{\text{comm}}
  4. 通信总量\displaystyle \sum_{e \in E_{\text{cut}}} V_e 所有被切开的边上传输的数据量之和
    ∑e∈EcutVe\displaystyle \sum_{e \in E_{\text{cut}}} V_e

调试与优化：试菜环节

宴席开席前，主厨一定会试菜。试菜不只是尝味道，还要看：每口锅是不是都在同时烧？是不是有厨师闲着手等别人的半成品？如果一盘菜等了十分钟才上桌，问题出在切菜还是传菜？

分布式训练同样需要“试菜”。JAX 提供了分片可视化工具——比如 jax.debug.visualize_sharding，可以把张量实际的分片布局画成一幅网格图，让你一眼看出数据被切成了几片、落在哪些设备上。配合性能剖析工具，可以定位哪台设备在等数据。

衡量并行方案好坏的最直观指标是设备利用率：真正花在计算上的时间占比。利用率越低，说明通信占比越高或负载越不均：

试菜发现问题后，主厨会调整切分方式、重新排班。调试分布式训练也是这样：先看分片图，再量利用率，最后调整 Mesh 布局或 PartitionSpec 批注——循环往复，直到整条流水线顺畅运转。

U=TcomputeTcompute+Tcomm×100%\displaystyle U = \frac{T_{\text{compute}}}{T_{\text{compute}} + T_{\text{comm}}} \times 100\%
设备利用率越高，说明并行方案越高效。
  1. 利用率定义\displaystyle U 是设备真实计算时间在总时间中的占比
    U\displaystyle U
  2. 分子\displaystyle T_{\text{compute}} 纯计算时间，利用率想高，分子要尽量大
    Tcompute\displaystyle T_{\text{compute}}
  3. 分母\displaystyle T_{\text{compute}} + T_{\text{comm}} 计算与通信的总时间，是分母
    Tcompute+Tcomm\displaystyle T_{\text{compute}} + T_{\text{comm}}
  4. 优化目标\displaystyle U \to 100\% 通信占比越小，利用率越接近百分之百
    U→100%\displaystyle U \to 100\%

AI 应用场景

大语言模型预训练

阶段：模型训练

解决的问题：数百亿参数无法放进单卡，传统手动并行代码量大、难维护、难复现。GSPMD 允许研究员只写单设备的前向/反向代码，用 Mesh + NamedSharding 声明数据与模型布局，编译器自动生成数据并行 + 张量并行的混合方案，就像主厨自动给 64 位厨师排好班。

具体案例：某团队在 64 块 TPU v4 上训练 130 亿参数的 Transformer。研究员定义 8×8 的 Mesh（data 轴 8、model 轴 8），用 PartitionSpec 把注意力权重按 model 轴切分、batch 按 data 轴切分，GSPMD 自动完成 all-reduce 与 all-gather。训练脚本和单卡实验几乎一样简洁，收敛速度却快了近 50 倍。

推荐系统大规模 Embedding

阶段：特征工程与训练

解决的问题：推荐系统的 embedding 表往往有数百 GB，单机装不下，且不同特征访问频率差异很大。用 Mesh 定义设备组，PartitionSpec 指定 embedding 表按行切分，GSPMD 自动处理索引查询的跨设备通信——相当于把一本巨大的食材账本拆成 16 本小账本，分给 16 个备菜台。

具体案例：某广告推荐团队把 500GB 的 user embedding 表切到 16 张 A100 上，线上 batch 按 data 轴切分。原本需要手写 all-to-all 的 embedding 查询逻辑，改用 NamedSharding 后代码量减少 70%，吞吐提升 2.3 倍。

多模态模型分布式推理

阶段：推理部署

解决的问题：超大模型推理时单设备放不下参数，需要把参数切分到多设备，同时保持低延迟。GSPMD 在编译期把矩阵乘法切块并插入通信，推理代码保持单设备形态，部署时只需声明一次布局。

具体案例：某视觉语言模型（80B 参数）部署在 8 卡 A100 上，使用 NamedSharding 把 transformer 层权重按 model 轴切分、输入按 data 轴切分，GSPMD 生成张量并行推理图，p99 延迟从 800ms 降到 350ms。

Python 演示

python40 行
"""JAX Mesh 与 NamedSharding 自动切分演示。"""import jaximport jax.numpy as jnpfrom jax.sharding import Mesh, NamedSharding, PartitionSpec as Pfrom jax.experimental import mesh_utils  def two_layer_net(x, w1, b1, w2, b2):    """两层全连接网络：一条写给单一厨师的菜谱。"""    h = jax.nn.relu(x @ w1 + b1)    return h @ w2 + b2  def main():    # 1. 画好灶台布局图：把本机所有设备排成 (n_dev, 1) 网格    n_dev = jax.local_device_count()    mesh = Mesh(mesh_utils.create_device_mesh((n_dev, 1)), ('data', 'model'))    # 2. 约定数据布局：batch 维按 data 轴切，特征维不切    sharding = NamedSharding(mesh, P('data', None))     rng = jax.random.PRNGKey(0)    x = jax.random.normal(rng, (64 * n_dev, 128))    w1 = jax.random.normal(rng, (128, 256))    b1 = jnp.zeros(256)    w2 = jax.random.normal(rng, (256, 64))    b2 = jnp.zeros(64)     # 3. 声明输入输出布局，其余交给 GSPMD 编译期自动处理    train_step = jax.jit(        fun=two_layer_net,        in_shardings=(sharding, None, None, None, None),        out_shardings=sharding,    )    y = train_step(x, w1, b1, w2, b2)    print('输出形状:', y.shape)    print('输出分片:', y.sharding)  if __name__ == '__main__':    main()
目的：演示如何用 Mesh 定义设备网格，用 NamedSharding + PartitionSpec 声明张量布局，再交给 GSPMD 自动完成分布式切分与编译。
  1. mesh = Mesh(mesh_utils.create_device_mesh((n_dev, 1)), ('data', 'model'))画好“灶台布局图”，把本机设备排成一行网格，并给坐标轴起名。
  2. sharding = NamedSharding(mesh, P('data', None))给张量贴标签：批次维按 data 轴切，特征维保持完整。
  3. in_shardings=(sharding, None, None, None, None)只声明输入 x 的布局，四个权重参数全部复制到每台设备。
  4. out_shardings=sharding告诉编译器输出也要按同样的方式切分，保持上下游一致。
  5. y.sharding打印输出张量的实际分片布局，验证 GSPMD 是否按预期工作。

常见误区

误区我买了 8 台加速器，把它们都塞进 Mesh，训练速度应该提升 8 倍吧？加设备永远是划算的。

正解设备增加带来收益的前提是计算量足够、通信占比小。当通信成为瓶颈时，加设备反而可能更慢。就像后厨从 2 人加到 8 人，如果传菜通道只有一条，切好的菜全堵在过道里，效率反而下降。训练速度的提升取决于计算与通信的平衡，而不是设备数量的线性增长。

为何常见：线性外推是人的直觉，忽略了通信开销和 Amdahl 定律——短板决定总收益。

误区我用 jax.jit 包装了函数，GSPMD 应该自动把所有计算并行了吧？不需要做任何其他配置。

正解jit 只负责编译，自动并行需要你先初始化 Mesh 并把数据用 NamedSharding 或可推断的方式布置。没有 Mesh，GSPMD 无从谈起；没有分片线索，编译器默认按单设备执行。主厨再厉害，也得先有一张灶台布局图才能派活。

为何常见：“写单设备代码”的宣传让很多人以为零配置，忽略了 Mesh 是 GSPMD 的前提。

误区我用了 NamedSharding 指定切分，这不就是手动并行吗？跟 pmap 是一回事。

正解NamedSharding 只声明“数据长什么样”（布局），何时计算、如何通信、怎么同步仍由编译器和运行时决定；pmap 则直接规定“程序跑在哪些设备上、以什么方式映射”。一个是图纸上的标注，一个是执行时的命令。主厨批注“这块肉放左边灶台”，和亲手把肉端过去，是两件完全不同的事。

为何常见：两者都涉及设备与张量的关系，概念边界容易混淆。

误区GSPMD 是全自动的，我随便写个形状的张量，它都会切得最优吧？布局跟我没关系。

正解自动切分受初始布局线索影响很大。如果输入全采用复制布局，GSPMD 可能会选择保守方案，产生大量额外通信。理解 Mesh 轴与数据维度的对应关系，是写出高性能代码的前提。主厨再聪明，如果菜谱上一笔批注都没有，他也只能按最稳妥的方式做。

为何常见：“自动”二字被绝对化，忽略了编译器也需要好的线索才能做出好的决策。

误区大模型训练，把参数切得越碎、分到越多设备上，才是专业的做法。

正解张量并行会引入大量通信。当显存足够时，数据并行往往更划算；实际大模型训练常是数据并行 + 张量并行 + 流水并行的混合方案。GSPMD 的价值恰恰在于让你可以用同一份代码自由组合这些策略。不是每道菜都要 8 个厨师同时碰，分工要匹配菜的复杂度。

为何常见：对大模型的敬畏导致“参数切得越碎越专业”的错觉，忽略了通信成本。

面试题

Q1. 在 JAX 中，Mesh 的主要作用是什么？

单选题 · easy

  • A.连接设备之间的网络线路
  • B.定义设备排列与坐标轴名← 正确答案
  • C.自动生成设备间通信指令
  • D.替代 jit 完成代码编译

答案：B。Mesh 描述设备的排列方式以及每条坐标轴的名字，就像后厨的灶台布局图。通信指令由 GSPMD 在编译期生成，不是 Mesh 的职责，因此 A、C 错误；Mesh 也不参与编译，D 错误。

Q2. GSPMD 让用户写什么样的代码？

单选题 · easy

  • A.必须显式指定每台设备任务
  • B.只写单设备语义的代码← 正确答案
  • C.用 MPI 接口编写通信
  • D.手动切分所有输入张量

答案：B。GSPMD 的核心理念是用户只写“一个厨师”能看懂的单设备代码，编译器自动负责切分与通信。A、C、D 都是手动并行的做法，与 GSPMD 的设计相反。

Q3. PartitionSpec 中的 None 表示什么？

单选题 · easy

  • A.该维度按 model 轴切分
  • B.该维度保持完整不切分← 正确答案
  • C.该维度直接被丢弃不用
  • D.该维度自动选择切分

答案：B。None 表示“不切”，对应维度在每台设备上都保留完整副本，就像整块牛排原样送到每桌。按轴切分要写轴名，丢弃维度不是合法的布局声明。

Q4. GSPMD 在哪个阶段完成自动切分决策？

单选题 · mid

  • A.运行期实时推断
  • B.数据加载阶段
  • C.编译期静态分析← 正确答案
  • D.梯度更新阶段

答案：C。GSPMD 在 jit 编译期分析计算图与数据依赖，决定每个张量的分片方案并插入通信，随后生成可执行代码。它是“主厨提前排班”，不是开席后现想。

Q5. 数据并行训练中，梯度同步通常使用哪种通信操作？

单选题 · mid

  • A.all-to-all 全量交换
  • B.all-reduce 梯度归约← 正确答案
  • C.broadcast 参数广播
  • D.scatter 数据散射

答案：B。每台设备算出局部梯度后，需要 all-reduce 把梯度求和再平均，保证所有设备参数一致。像八位厨师把各自调好的味道汇总成同一锅底料，再用同一锅底料做下一道菜。

Q6. 相比 pmap，GSPMD 自动并行的最大优势是什么？

单选题 · mid

  • A.完全不产生任何通信
  • B.可组合稀疏注意力机制
  • C.自动插入通信并保持语义一致← 正确答案
  • D.不需要初始化任何设备

答案：C。pmap 要求手写设备映射和通信，GSPMD 则在编译期自动插入通信指令，并保证结果与单设备计算一致。A 不可能实现，D 忽略了 Mesh 是前提。

Q7. NamedSharding 与 pmap 的本质区别是什么？

单选题 · hard

  • A.前者声明布局，后者编排执行← 正确答案
  • B.前者只能运行在 GPU
  • C.后者不产生任何通信
  • D.两者本质上没有区别

答案：A。NamedSharding 只声明“张量按什么轴切”，不规定执行顺序；pmap 则是显式的并行执行原语。前者是图纸上的布局标注，后者是直接下达的派工命令。

Q8. 张量并行主要解决什么问题？

单选题 · hard

  • A.单设备显存放不下整个模型← 正确答案
  • B.数据量太大无法加载
  • C.梯度计算太不稳定
  • D.学习率无法被选择

答案：A。张量并行把单个大矩阵按行或列切到多台设备上共同计算，解决模型参数超出单设备显存的问题。它像是几口锅同时炖一锅大汤，而不是每口锅做一道菜。

Q9. 下列哪些属于 GSPMD 在编译期要做的事？

多选题 · mid

  • A.分析数据依赖关系← 正确答案
  • B.决定张量分片方案← 正确答案
  • C.自动插入通信指令← 正确答案
  • D.重写用户的训练算法

答案：ABC。GSPMD 做依赖分析、分片决策、通信插入三件事，但不会改写你的数学逻辑。你的训练算法就像菜谱本身，主厨只负责排班，不改菜谱内容。

Q10. 调试分布式训练时，哪些指标能帮助定位性能瓶颈？

多选题 · hard

  • A.设备利用率 U← 正确答案
  • B.通信时间的占比← 正确答案
  • C.分片后的形状分布← 正确答案
  • D.代码风格与行数

答案：ABC。设备利用率、通信占比、分片形状都能反映并行方案的好坏：利用率低说明在等数据，通信占比高说明切分导致传菜太多，分片形状异常说明布局声明有误。代码行数与性能无关——菜谱页数不代表菜好吃。