想象你是一位古老港口的搬运总管。这天，富商运来一头十吨重的石象，要把它安然渡过河安置在对岸庭院里。可码头上只有几条载重两吨的小木船——整只象抬上船，船会当场沉没，就像把大模型直接加载到一张显卡上，立刻迎来 “CUDA out of memory“ 的刺眼报错。聪明的做法，是请石匠把石象精确切成几块，分装到不同的船上；航行途中船与船用旗语互报方位、协同划桨；到岸后再用榫卯把石象拼回原样。这，就是”模型并行“的朴素画面：把一个大模型拆成多个部分，放到多张显卡上协作完成训练或推理。JAX 里的 pjit 就是那位老练的船长——你只需告诉他”船队多宽、货物多重“，他便自动安排切法、装船顺序和船队通信。本文要把这套”拆象运过河”的学问，从头到尾讲给你听。

岗位关联

学会模型并行，你就能在单卡装不下大模型时优雅地解决显存瓶颈——这是训练千亿参数大模型、走向分布式系统方向的必备核心技能。

目录
  1. 可视化演示
  2. 模型并行概述
  3. 分片策略
  4. 设备间通信
  5. pjit机制
  6. 实现步骤
  7. 性能优化
  8. 应用实践
  9. AI 应用场景
  10. Python 演示
  11. 常见误区
  12. 面试题

可视化演示

\u5c55\u793a\u5927\u6a21\u578b\u57288\u5c42\u5206\u7247\u52304\u5f20GPU\u7684\u8fc7\u7a0b\uff0c\u5bf9\u6bd4\u5355\u5361OOM\u6ea2\u51fa\u4e0e\u591a\u5361\u5e76\u884c\u7684\u663e\u5b58\u5747\u8861\u5206\u914d\u4e0e\u6d41\u6c34\u7ebf\u901a\u4fe1\u3002

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

\u6a21\u578b\u5e76\u884c\u5206\u7247\u793a\u610f \u6a21\u578b\u5c42 \u6570\u636e\u6d41 8\u5c42>4GPU \u6bcf\u53612\u5c42\u5747\u8861\u5206\u914d

模型并行概述

你站在码头，面对的是一头完整石象。强行把它塞上载重两吨的小船，船沉、象落水——这就像把百亿参数模型直接加载到一张 80GB 显卡上，训练脚本立刻抛出 “CUDA out of memory“。这不是你不够努力，而是”货“天生超出了”船”的承载极限。大模型的参数规模，决定了单卡显存总有见底的一天。

这头“石象”就是模型参数。一个百亿参数的 Transformer，光参数在半精度下就占约 20GB，再算上梯度、Adam 优化器的双份动量、前向激活值，整体显存需求轻松突破 150GB。显然，单张显卡装不下完整模型。模型并行（Model Parallelism）的思路，就是把模型本身切开，让不同设备各管一块。显存总需求可以写成公式：

当你把石象切分到多只船上，每只船都有富余的载重空间——这就是模型并行最直接的红利：显存压力成倍释放。当然，货物被拆开只是第一步，接下来的问题是怎么拆。

Mtotal=P×b+P×b+2P×b+A\displaystyle M_{total} = P\times b + P\times b + 2P\times b + A
公式解释：P 是模型参数量，b 是每个参数占用的字节数（半精度下 b=2），A 是激活值占用。模型并行的目标就是把这个总量按设备数摊薄。
  1. 参数项模型本身有 P 个参数，每个参数以 b 字节保存，这是最基础的显存占用。
    P×b\displaystyle P\times b
  2. 梯度项反向传播需要为每个参数计算并保存一个梯度，大小同为参数体的体积。
    P×b\displaystyle P\times b
  3. 优化器状态Adam 等优化器还要额外保存一阶、二阶动量，通常是参数大小的两倍。
    2P×b\displaystyle 2P\times b
  4. 激活值项前向传播的中间结果 A 同样占用显存，它与批量大小、序列长度成正比。
    A\displaystyle A

分片策略

石匠面对石象要反复思量切法：按部位切，象牙归象牙、躯干归躯干，每一块都是完整的可搬运结构；躯干太大，还得横向再切成两层；最终每只船的载重必须大致相等，否则船队会偏航。这三种考量，恰好对应模型并行中的三种分片视角：层级、张量与布局。

神经网络天然按“层”堆积，把不同层分配给不同设备，让数据依次流经各设备，这叫层级分片，也叫流水线并行。但 Transformer 的单个线性层动辄几万乘几万，一张卡同样装不下，所以还要把单层内的大矩阵切开：按行切、按列切、按二维网格切，这就是张量分片。分片后，每台设备需要保存的显存变为：

分片布局还要兼顾通信负载。在 JAX 中，分片方案用 PartitionSpec 描述：P('dev', None) 表示把行维度沿命名轴 'dev' 切分；P(None, 'dev') 表示切分列维度；P(None, None) 表示不切分、完整复制。石匠切得越均衡，船只负载越均匀；但切得越碎，船与船之间的“拼接”就越频繁——分片策略的智慧，就是在负载与通信之间找平衡。

Mper=MtotalN\displaystyle M_{per} = \frac{M_{total}}{N}
公式解释：N 是参与并行的设备总数。理想均衡分片下，每台设备只承担总显存的 N 分之一。
  1. 总显存需求M_total 是模型不切分时单卡所需的总显存，包含参数、梯度、优化器状态等。
    Mtotal\displaystyle M_{total}
  2. 设备总数N 表示参与并行计算的设备数量，船队越多，单船负担越轻。
    N\displaystyle N
  3. 单卡分量理想均衡分片后，每台设备只需承载总量的 N 分之一，显存瓶颈就此化解。
    MtotalN\displaystyle \frac{M_{total}}{N}

设备间通信

石象切成八块，装上了八条船。可它终究是一头完整的象：船队想同步靠岸，必须靠旗语和缆绳互通信息。第二船要先知道第一船的速度才能调整桨频；到岸前大家还要互相确认每块象石的位置，保证拼接时严丝合缝。在模型并行中，被切开的权重并不是孤立的。

例如矩阵乘法把权重按行切到两台设备后，每台设备只算出“部分和”，必须把两者的部分和相加才能得到完整结果，于是有了 All-reduce（全规约）这一系列集体通信操作。反向传播时，各设备也要先规约梯度再更新参数。All-reduce 一次同步的通信量约为：

除了 All-reduce，还有 All-gather（把各分片拼接成完整张量）、Reduce-Scatter（先求和再分块）和点对点传输。设备之间的物理连接，例如 NVLink、InfiniBand，决定了这些旗语传递的速度。旗语虽帮了船队大忙，但嗓门再大也要花时间；如果八条船都在拼命喊话、没时间划桨，运输速度照样上不去——通信开销，是模型并行绕不开的代价。

Ccomm=2M×N−1N\displaystyle C_{comm} = 2M \times \frac{N-1}{N}
公式解释：M 是模型数据量，N 是参与通信的设备数。设备越多，需要汇总交换的数据总量越大。
  1. 模型数据量M 表示需要参与规约的模型或梯度大小，例如 10GB 的权重矩阵。
    M\displaystyle M
  2. 设备数量N 是参与 All-reduce 的设备总数。
    N\displaystyle N
  3. 两趟传输All-reduce 包含“规约”与“广播”两个阶段，所以通信量乘以 2。
    2M\displaystyle 2M
  4. 随 N 缩放设备越多，需要汇总的数据越多，通信量按 (N-1)/N 的比例增长。
    N−1N\displaystyle \frac{N-1}{N}

pjit机制

港口来了一位经验丰富的老船长。你不必告诉他“象牙放三号船、躯干切两半”这些细节，只需说：“船队八条船，目标对岸，货物是这头石象。”他自会指挥石匠切分、安排装船、约定旗语节奏。这位老船长，就是 JAX 中的 pjit——partitioned jit，分区即时编译。

pjit 把“写分布式程序”变成“标注分片意图”。你用 PartitionSpec 告诉它每个张量希望沿哪个维度切给哪些设备，剩下的设备网格布局、通信算子插入、内存管理，全部由编译器在编译期自动编排。逻辑上，一次 pjit 调用可以用下面的映射描述：

值得注意的是，在新版 JAX 中，pjit 的能力已整合进 jax.jit(..., in_shardings=..., out_shardings=...)，写法有所演进，但“声明式分片”的内核不变。有了老船长，港口不再需要每一趟都重新制定方案；你只需把规则写在运输单上，船队便自动各就各位。

y=Pjit(f)(x;mesh,spec)\displaystyle y = \mathrm{Pjit}(f)(x; \mathrm{mesh}, \mathrm{spec})
公式解释：f 是要并行执行的函数，mesh 是设备网格，spec 是分片规格。编译后，所有设备执行同一份代码、处理自己的分片。
  1. 目标函数f 是要并行化的前向函数或训练步骤。
    f\displaystyle f
  2. 设备网格mesh 定义设备如何排列成多维命名网格，例如一维的 ('dev',) 或二维的 ('dp', 'fsdp')。
    mesh\displaystyle \mathrm{mesh}
  3. 分片规格spec 指明每个张量维度沿哪个命名轴切分，相当于货物上的搬运标签。
    spec\displaystyle \mathrm{spec}
  4. 输出结果y 与输入的逻辑形状一致，但物理上分布在多个设备，设备间通过通信协作完成计算。
    y\displaystyle y

实现步骤

你要做的，是学会填写老船长的“运输单”：船队排成一行还是两行？每只船上应该装货物的哪一部分？在 JAX 中，这张运输单由三部分组成：设备网格、命名分片、分片函数。

第一步，构建设备网格 Mesh(jax.devices(), ('dev',))，把本机所有 GPU 排成一维命名轴；第二步，定义分片规格 NamedSharding(mesh, P('dev', None))，相当于在货物上贴标签；第三步，用 jax.device_put 把标签应用到数据上，或用 jax.jit(..., in_shardings=...) 把标签应用到函数上。多轴网格本质上是多个维度的笛卡尔积：

运输单写得规范，老船长的命令才有据可依。同理，PartitionSpec 里的维度必须与张量形状对得上，否则编译器会直接报维度错误。在后面的 Python 演示中，你会看到一张完整“运输单”的模样。

Mesh=D1×D2×⋯×Dk\displaystyle \mathrm{Mesh} = D_1 \times D_2 \times \cdots \times D_k
公式解释：D_i 是第 i 个命名轴上的设备数量，总设备数等于各轴设备数相乘。
  1. 设备命名轴每个 D_i 表示第 i 个命名轴上的设备数量，例如 ('dp', 'fsdp') 中的数据并行轴和分片轴。
    Di\displaystyle D_i
  2. 网格总规模所有命名轴的设备数相乘，得到网格可用的总设备数，二维网格常用于真实的大模型训练。
    D1×D2×⋯×Dk\displaystyle D_1\times D_2\times\cdots\times D_k

性能优化

船队往返几次后你发现，每次靠旗语协调都要停下来，白白浪费时间。一位船手提出：“我可以在划桨的同时喊话！”于是，报方位的时间就被划桨动作“藏”起来了。这个朴素技巧，就是计算与通信重叠（overlap）。

模型并行中，设备不仅要计算，还要与其他设备通信。如果先算完、再通信、再计算，整条链路会有大量空闲等待；若把一部分通信安排到计算的同时进行，单步训练时间会明显下降。重叠后单步耗时的理想模型为：

工程上，通过增大单次计算的粒度、使用微批次流水线、选择通信量更低的分片方案，可以让 T_calc 尽量接近甚至超过 T_comm。聪明的船队不是不喊话，而是把喊话揉进桨声里——模型并行优化的核心，就是让通信时间尽可能被计算掩盖。

Tstep=max⁡(Tcalc,Tcomm)\displaystyle T_{step} = \max(T_{calc}, T_{comm})
公式解释：T_calc 是计算耗时，T_comm 是通信耗时。通信与计算并行时，总耗时取两者较大值，而不是相加。
  1. 计算耗时T_calc 是设备完成本分片矩阵乘、激活等计算所需的时间。
    Tcalc\displaystyle T_{calc}
  2. 通信耗时T_comm 是交换分片结果、规约梯度所需的时间。
    Tcomm\displaystyle T_{comm}
  3. 重叠耗时两者并行执行时，单步时间近似取较大值而非求和。
    max⁡(Tcalc,Tcomm)\displaystyle \max(T_{calc}, T_{comm})

应用实践

石象运输从传说变成了码头日常：大象搬家、寺庙立柱、整座牌坊迁移都用上了同一套“切分-共运-拼装”的智慧。而在今天，任何千亿参数模型的训练，也都离不开模型并行这块基石。

GPT 级模型参数量动辄千亿，训练集群动辄几千张卡。实际工程中，通常把张量并行、流水线并行、数据并行组合成“3D 混合并行”：数据并行负责复制模型副本、切分数据，模型并行把单个副本切到多卡，流水线并行再把层分批。模型并行能换来的收益，可以用加速比刻画：

因此，工程选型时要结合显存约束与吞吐指标做综合判断：如果你的模型单卡能放下，优先考虑数据并行；只有权重或激活值确实放不下时，才需要让模型并行登场。从一头石象到一支船队，再到整个港口的常态作业，模型并行的核心思想不过十六个字：切得开、装得下、传得快、拼得回。

S=TserialTparallel\displaystyle S = \frac{T_{serial}}{T_{parallel}}
公式解释：T_serial 是单设备串行完成的时间，T_parallel 是模型并行后的时间。理想加速比接近设备数 N，但受通信限制往往小于 N。
  1. 串行时间T_serial 是参数全部放在单设备时的端到端训练耗时。
    Tserial\displaystyle T_{serial}
  2. 并行时间T_parallel 是采用模型并行后，多设备协作的端到端耗时。
    Tparallel\displaystyle T_{parallel}
  3. 加速比S 大于 1 说明并行有效，但受通信和同步开销限制，实际 S 小于设备数 N。
    TserialTparallel\displaystyle \frac{T_{serial}}{T_{parallel}}

AI 应用场景

千亿参数大模型预训练

阶段：分布式训练

解决的问题：GPT-4 级别的万亿参数模型，半精度权重就要占用 2TB 以上显存，单卡 80GB 完全不可能容纳。传统做法只能靠数据并行，但一个模型副本都放不下时，数据并行也无济于事。模型并行把 Transformer 每层的大矩阵切成多块放到不同芯片，配合流水线并行，让超大规模训练成为可能。

具体案例：PaLM 540B 使用 6144 块 TPU，通过 JAX 的 pjit 机制完成模型并行与数据并行的混合训练。开发者只需用 PartitionSpec 描述每个张量的分片方式，编译器自动生成分布式执行计划。

大模型多卡推理服务

阶段：模型推理部署

解决的问题：70B 参数的大语言模型，FP16 权重约 140GB，一张 A100 80G 放不下。若不做模型并行，推理服务只能退而求其次使用量化或小模型。张量并行把权重按行或列切到 2-4 张卡，让多卡同时参与解码计算，既解决显存问题，又降低单请求延迟。

具体案例：vLLM、TensorRT-LLM 等推理引擎支持多卡张量并行部署，单个请求的生成吞吐可以提升数倍。JAX 生态中，pjit 也可用于把训练好的模型参数按分片方案加载到多设备，直接提供服务。

超长序列训练与混合并行选型

阶段：工程架构设计

解决的问题：长序列输入会导致激活值爆炸式增长，即使权重放得下，单卡也会因中间结果而 OOM。模型并行不仅能切分权重，还能把激活值、序列维度分片到多设备。

具体案例：Gemini 1.5 支持百万 token 超长上下文，训练时采用序列并行与模型并行组合，把注意力矩阵的中间结果分散到多张卡上，避免单点显存峰值。工程师需要根据显存画像决定采用哪种并行组合。

Python 演示

python35 行
import jaximport jax.numpy as jnpfrom jax.sharding import Mesh, NamedSharding, PartitionSpec as P  def main():    # 1. 构建设备网格: 把本机所有设备排成一维命名轴 'dev'    mesh = Mesh(jax.devices(), ('dev',))    key = jax.random.PRNGKey(0)    big = jax.random.normal(key, (1024, 1024))     # 2. 不切分直接放置: 每台设备保存完整矩阵副本    unsharded = NamedSharding(mesh, P(None, None))    whole = jax.device_put(big, unsharded)    print('复制模式: 全局形状', whole.shape,          '单卡本地形状', whole.addressable_data(0).shape)     # 3. 模型并行分片: 按行维度切到所有设备    sharded_spec = NamedSharding(mesh, P('dev', None))    part = jax.device_put(big, sharded_spec)    local = part.addressable_data(0)    print('模型并行: 全局形状', part.shape,          '单卡本地形状', local.shape)     # 4. 模拟 All-gather 聚合还原并校验    gathered = jnp.concatenate(        [part.addressable_data(i) for i in range(len(jax.devices()))],        axis=0,    )    assert jnp.allclose(gathered, big)    print('聚合校验通过: 各分片拼回后与原矩阵完全一致')  if __name__ == '__main__':    main()
目的：演示 JAX 中“复制模式”与“模型并行分片模式”的区别，展示如何用 Mesh、NamedSharding 和 P 把一个大型权重矩阵按行切到多张卡上，并验证分片后能完整拼回。
  1. Mesh(jax.devices(), ('dev',))把多张显卡排成一条一维“船队”，命名轴叫 'dev'。
  2. NamedSharding(mesh, P(None, None))所有维度都不切分，矩阵被完整复制到每台设备。
  3. NamedSharding(mesh, P('dev', None))行维度沿 'dev' 轴切分，列维度保持完整，这是零一维的模型并行分片。
  4. jax.device_put(big, sharded_spec)按分片规则把参数分布到各设备，就像老船长按运输单装船。
  5. part.addressable_data(0)查看 0 号设备上实际持有的本地分片形状。
  6. jnp.concatenate([...], axis=0)模拟 All-gather 把各设备分片拼回完整矩阵。
  7. assert jnp.allclose(gathered, big)校验分片前后数值完全一致，确保拆得开也拼得回。

常见误区

误区模型并行很简单，把 Transformer 的 48 层平均分到 8 张卡上，每张卡放 6 层，不就行了吗？层与层本来就是顺序计算的，这样切应该没什么通信量吧。

正解你说的这种切法是“流水线并行”，但它只切了层与层的边界。每一层内部的超大权重矩阵仍然完整保存在单卡上，矩阵乘法照样会撑爆显存。真正的模型并行还包括把单个大矩阵切开的多卡协作，即张量并行。而且流水线并行中，前一层算完才能算后一层，设备之间会出现“气泡”空闲，需要配合微批次调度才能提高利用率。

为何常见：初看网络结构图时，层与层是串行的方块，很容易把“模型”误解为一串层序列，忽略了层内部巨大的矩阵维度。

误区既然 pjit 带“自动”两个字，我只要给模型函数加上 pjit 装饰器，百亿参数就会自动塞进多张卡，通信开销它也会帮我优化好。

正解pjit 是“按你给出的 PartitionSpec 自动实现分片”，不是“自动替你决定最优切法”。你仍然要自己选择设备网格、分片维度和通信策略。分片规格选得不好，可能导致设备间频繁通信，甚至维度不匹配直接编译失败。它优化的是“执行”，而不是“决策”。

为何常见：把“自动编译”误解为“自动优化一切”，忽略了 pjit 需要使用者提供正确的分片标注作为输入。

误区显存不够就多插几张卡，模型并行会把模型切成越来越小的块，显存永远够用，训练速度也会随卡数线性上升。

正解显存确实可以随设备数扩展，但模型并行引入了通信开销和流水线气泡。设备增多时，通信总量也在增加，加速比受 Amdahl 定律制约，不会线性增长。实际工程中常常出现“卡越多、每卡算得越少、通信占比越高”的情况，必须依靠计算通信重叠等手段调优。

为何常见：只看到显存数字被摊薄，忽略了设备间协作本身也要付出时间成本。

误区我写了 `P('dev', None)` 之后运行报错，提示维度切不开，那我再多买几张卡不就行了吗？

正解PartitionSpec 中带命名轴的维度，其大小必须能被该命名轴上的设备数整除。例如 1024 行矩阵切到 3 张卡就无法均分，编译直接失败。此外 'dev' 不是内置魔法字符串，必须先在 Mesh 中定义命名轴才能引用它。分片标注必须与张量形状、网格结构严格对应。

为何常见：把 PartitionSpec 当成随手写的魔法字符串，忽略了它与张量维度、设备网格之间的整数约束关系。

误区模型并行是切模型，数据并行是切数据，两种思路完全不同。所以一个训练任务要么用模型并行，要么用数据并行，不能混着来。

正解现代大模型训练普遍采用“混合并行”：数据并行负责复制模型副本、切分训练数据；模型并行负责把单个副本切到更多设备；流水线并行再把层分批。三者组合使用，才能在大规模集群上充分压榨算力。

为何常见：教程通常单独讲解模型并行和数据并行，初学者容易把它们当成互斥的选项，没有意识到它们解决的是不同维度的瓶颈。

面试题

Q1. 单卡显存溢出的直接原因通常是下列哪一项？

单选题 · easy

  • A.训练数据量实在太大
  • B.模型规模超出显存容量← 正确答案
  • C.显卡散热能力出现不足
  • D.学习率设置得过高了

答案：B。显存放不下是因为模型参数、梯度、优化器状态和激活值合计超过了显卡容量，这就像小船载不动十吨重的石象。训练数据大是数据并行要解决的问题，散热和学习率与显存溢出没有直接关系。

Q2. 模型并行与数据并行的主要区别是什么？

单选题 · easy

  • A.模型并行切分参数，数据并行复制参数← 正确答案
  • B.数据并行切分模型，模型并行切分数据
  • C.模型并行只能用于推理阶段场景
  • D.数据并行只能用于训练阶段场景

答案：A。模型并行把模型参数切到不同设备，数据并行则让每台设备保存完整模型副本、只切分数据。船队运石象时，模型并行是“大象切块分船装”，数据并行是“每条船上装一头完整的小象”。

Q3. 下列哪些场景需要使用模型并行方案？

多选题 · mid

  • A.单卡显存无法容纳整个模型权重← 正确答案
  • B.训练数据总量超过主内存的大小
  • C.单个巨大矩阵需要拆开并行计算← 正确答案
  • D.想要提升一个小模型的推理速度

答案：AC。权重放不下需要切分模型，超大矩阵一层装不下也需要拆开计算，这两者都属于模型并行。训练数据超过内存应使用数据并行或数据流式加载；小模型推理优先考虑量化或批处理，不需要模型并行。

Q4. 把单个大权重矩阵按行切到多张卡上，属于哪种并行？

单选题 · mid

  • A.数据并行中的梯度平均
  • B.流水线并行中的层间切分
  • C.张量并行中的层内切分← 正确答案
  • D.模型量化中的精度压缩

答案：C。按行切分单个矩阵属于层内分片，即张量并行。流水线并行切的是层与层的边界；数据并行复制完整模型；量化是压缩精度而非切分设备。

Q5. All-reduce 在模型并行中主要用于完成什么？

单选题 · mid

  • A.把不同设备的梯度求平均后再分发← 正确答案
  • B.把不同设备的参数各自独立地保存
  • C.把输入数据切分成多个小批量样本
  • D.把损失函数值打印到所有设备日志

答案：A。All-reduce 对跨设备的部分结果进行规约求和，再把结果广播给所有参与者。矩阵切分后各设备只算出部分和，需要 All-reduce 汇总，就像船队靠旗语互相对齐位置一样。

Q6. PartitionSpec P('dev', None) 表示的含义是什么？

单选题 · mid

  • A.张量的所有维度都复制到每台设备上
  • B.行维度按设备切分，列维度保持完整← 正确答案
  • C.列维度按设备切分，行维度保持完整
  • D.行和列两维度都按设备数量切分

答案：B。P('dev', None) 中第一个位置对应行轴，带 'dev' 表示沿设备命名轴切分；第二个位置 None 表示列轴不切分、保持完整。P(None, None) 才是完全复制。

Q7. 在 JAX 中查看分片后单设备的本地形状，应使用？

单选题 · hard

  • A.调用 device_get 获取全局形状
  • B.调用 addressable_data 获取本地分片← 正确答案
  • C.调用 reshape 重新塑造出内存布局
  • D.调用 numpy_shape 查看分片后形状

答案：B。addressable_data(0) 返回 0 号设备上实际持有的本地分片数组；device_get 返回的是全局逻辑视图，看不到单卡本地大小。这个操作就像检查某条船上实际装了多少块象石。

Q8. 模型并行中需要设备间通信的根本原因是什么？

单选题 · hard

  • A.不同设备需要同步各自的风扇转速
  • B.分片后的张量仍存在逻辑上的数据依赖← 正确答案
  • C.主机内存太小，设备间必须互相转发
  • D.显卡驱动要求定期发送心跳保活消息

答案：B。矩阵切分后，每一片只是完整计算结果的一部分，后续计算需要把部分结果合起来，这种逻辑依赖迫使设备通信。石象虽然被拆开装在多条船上，但它终究是一头完整的象，到岸前必须交换信息才能拼装。

Q9. 下列哪些手段可以有效地降低模型并行通信开销？

多选题 · hard

  • A.增大单次计算的粒度以掩盖通信← 正确答案
  • B.选择通信量更小的分片维度方案← 正确答案
  • C.把权重完整复制到每一张设备上
  • D.使用更高带宽的设备互联技术← 正确答案

答案：ABD。增大计算粒度让通信时间被计算掩盖，选择合适分片维度能减少跨设备数据交换量，高带宽互联直接压缩通信耗时。把权重复制到每张卡属于数据并行思路，虽然消除了通信，但也失去了模型并行省显存的意义。

Q10. 关于 JAX pjit 的分区机制，以下说法正确的有哪些？

多选题 · extreme

  • A.通过 PartitionSpec 标注张量分片维度← 正确答案
  • B.设备网格需先定义命名轴才能引用← 正确答案
  • C.编译器会自动插入必要的通信算子← 正确答案
  • D.分片维度大小与设备数可以不必匹配

答案：ABC。pjit 依赖 PartitionSpec 声明分片，Mesh 必须预先定义命名轴，编译时会自动插入 All-reduce 等通信算子。分片维度必须能被设备数整除，否则编译报错，因此 D 错误。