想象一家高端餐厅的晚市高峰。订单像雪花一样涌进来，但后厨里只有一张主灶台、几位备菜台和一排冷藏柜——空间有限，厨师却要同时处理几十道菜。聪明的总厨会提前把菜单检查一遍：哪些菜可以一起做、谁先谁后、哪个备菜台放哪类食材最省走动、什么时候需要把半成品传给下一道工序。这套“谁在哪张台子上做什么、食材放哪里、怎么互相传菜”的安排，就是厨房的调度方案。 在AI世界里，这张“厨房”就是一台训练服务器——里面的多位厨师就是多个GPU设备（统称设备），食材就是你喂给模型的张量数据，而总厨则是JAX的编译引擎XLA。设备放置（Device Placement），正是XLA来决定“哪个张量放到哪张台子上、哪个计算派给哪位厨师、什么时候在设备之间传数据”的制度。XLA这位总厨默认会自己计算最优安排，但当你遇到特殊菜式时，也可以用jit和pmap两大“指挥口令”亲自干预。安排得好，出餐飞快；安排得差，厨师们互相抢台子、传菜堵成一团。

岗位关联

设备放置是多GPU训练提速与省显存的核心功底。掌握它，你才能真正搭建起高性能、可复用的分布式训练系统，这是AI工程进阶的必修课。

目录
  1. 可视化演示
  2. 设备放置基础
  3. XLA自动决策
  4. jit手动控制
  5. pmap分片
  6. 通信开销优化
  7. 显存利用率
  8. 数据与模型分片
  9. 扩展实践
  10. AI 应用场景
  11. Python 演示
  12. 常见误区
  13. 面试题

可视化演示

动态演示Device Placeme的核心机制，通过总览/传输/转换/输出四种模式展示关键概念与交互流程。

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

Device Placeme：设备放置 设备放置基本概念与原理 XLA 自动设备放置推断 核心机制可视化 Device Placeme整体架构与关键组件

设备放置基础

让我们回到那间厨房。餐厅里有好几位厨师，每一位都有自己的灶台、刀具和冷藏柜。总厨每天早上会做一件事：现在这些订单（计算任务）来了，每一道菜该由谁来做、放在哪张台子上备料。这个决策，就是设备放置。在JAX中，“厨师”是物理设备——常见的是GPU，也可以是TPU甚至CPU。而“菜”则是对张量的一系列数学运算。

在JAX的视角里，你写的代码最终会被翻译成一张计算图——你可以把它理解为一份完整的“菜肴流程图”：先切洋葱（矩阵乘法），再下锅（激活函数），最后装盘（损失计算）。图上的每个节点是一个操作，每条边是一条数据流。设备放置要解决的，就是给这张图的每个节点分配一位“厨师”（设备），并决定数据在设备间如何流动。

这不是一件可以随意拍板的事。厨房里的灶台布局（设备拓扑）决定了厨师之间传菜要走多远：距离近的传菜快，距离远的传菜慢。GPU之间的数据传输带宽和延迟各不相同，有的在同机箱内通过NVLink高速互联，有的则要跨机器走网络。放置策略的核心权衡是：算得快的方案未必通信少，通信少的方案未必显存均衡。总厨必须找到一个综合最优解。

所以当你打开JAX，总厨XLA做的第一件事，就是拿着这张计算图，在心中推演N种分工方案——这是每一个多设备AI训练故事的起点。

Goal:min⁡P(∑v∈Vcomputev+∑(u,v)∈EcommP(u),P(v))\text{Goal:} \quad \min_{P} \Big( \sum_{v \in V} \text{compute}_v + \sum_{(u, v) \in E} \text{comm}_{P(u), P(v)} \Big)
这个公式的意思是：设备放置的终极目标，是把“所有设备的计算时间”加上“所有设备间的数据传递时间”加起来，让总和最小。
  1. 虚拟操作集合$V$ 表示计算图中的全部操作节点，也就是需要安排的“菜品工序”。
    V={v1,v2,…,vk}V = \{v_1, v_2, \ldots, v_k\}
  2. 可用物理设备$D$ 表示你手头所有的设备，也就是厨房里的“厨师们”。
    D={d1,d2,…,dN}D = \{d_1, d_2, \ldots, d_N\}
  3. 放置映射$P: V \to D$ 把每个操作指派给一个设备，这就是“设备放置”本身。
    P(vi)=djP(v_i) = d_j
  4. 通信开销累加$E$ 是计算图中的边集合，任何跨设备的边都会产生传菜成本。
    commP(u),P(v)=数据量  /  设备间带宽\text{comm}_{P(u), P(v)} = \text{数据量} \;/\; \text{设备间带宽}

XLA自动决策

既然设备放置这么纠结，能不能让总厨自己拍板？XLA就是那位总厨。它的工作方式非常有策略：首先把整张计算图摊开在案板上通读一遍，识别出哪些操作会生成大量中间数据（比如大矩阵乘法），哪些操作轻量到可以顺手做掉（比如ReLU激活）。这种计算图分析是它的第一步功夫。

第二步是通信模式识别。XLA会观察数据流向：某个张量如果被多个下游操作使用，它就像一锅高汤，应该在哪个位置熬好、再分装端到不同厨师台前。XLA会计算“把这张张量复制到4台设备”与“每次用时单独传”哪个更划算。换句话说，它在寻找一种放置，使得跨设备的“传菜”次数和传菜量降到最低。

第三步是最优放置搜索。XLA内部会生成若干候选放置方案，用代价模型估算每个方案的总耗时，然后挑选成本最低的那一个。这像是在心里做了一场“沙盘推演”。之所以能做到这一点，是因为JAX是静态编译的：你调用jit后，整个计算在运行前就变成了一个固定结构的、可被全面分析的对象。XLA可以大胆地重组、移动、融合计算，因为它已经提前看完了整本菜谱，而不是像解释执行那样端一道菜想一道菜。

XLA的自动决策并不神秘，它只是把总厨的“直觉”变成了一套可枚举、可计算的搜索过程。大多数时候，你可以完全信任它的判断——但总有一种情况，你比总厨更懂自己的菜。

Cost(P)=∑i=1NTcompute(i)+∑(u,v)∈E[P(u)≠P(v)]⋅size(u→v)BP(u),P(v)\text{Cost}(P) = \sum_{i=1}^{N} T_{\text{compute}}^{(i)} + \sum_{(u,v) \in E} \big[ P(u) \ne P(v) \big] \cdot \frac{\text{size}(u \to v)}{B_{P(u), P(v)}}
公式说明：某个放置方案的代价，等于每台设备的计算耗时，加上所有跨设备数据边按实际带宽折算后的传输时间总和。
  1. 计算耗时分母$T_{\text{compute}}^{(i)}$ 是设备 $i$ 上所有操作耗费的总时间，同步地贡献给总耗时。
    ∑i=1NTcompute(i)\sum_{i=1}^{N} T_{\text{compute}}^{(i)}
  2. 跨设备判定$P(u) \ne P(v)$ 如果成立，说明边 $(u, v)$ 连接了不同设备，数据必须走通信链路，此时取1；否则不产生通信成本。
    [P(u)≠P(v)]\big[ P(u) \ne P(v) \big]
  3. 传输时间折算$\text{size}(u \to v)$ 是要传的数据量，$B_{P(u),P(v)}$ 是这两台设备间的实际带宽。
    size(u→v)BP(u),P(v)\frac{\text{size}(u \to v)}{B_{P(u), P(v)}}

jit手动控制

有时候，你很清楚自己餐厅里哪个灶台是最快、最稳的，而总厨却不知道——因为它只看到抽象的图，看不到你某个设备挂了其他任务。这时你要发声：这道菜，我指定给二号厨师做。

在JAX中，你通过jit装饰器来“按住”总厨的手。jit的作用是把一个Python函数编译成一个XLA可执行程序，你可以在编译时给XLA下达若干约束，包括设备指定。比如jax.jit(fn, device=devices[1])，意思是：这个函数编译出来的所有操作，都放到1号设备上执行。这就是设备参数指定的威力。

jit还有另一层意义：编译边界的指定。如果你不加jit，JAX会一行一行地解释执行你的Python函数，每次操作都可能触发一次独立的设备调度——相当于厨师做一道菜要重新看一遍菜谱，哪怕这道菜已经做过一百遍。用了jit，整个函数被编译成一个可执行块，XLA可以在块内做融合优化（比如把多个小操作合并成一个kernel），大幅减少设备调度的次数和中间张量的内存占用。这也带来了一个重要的约束：jit编译要求静态形状——函数的输入shape和dtype在编译时必须确定，否则无法生成固定代码。

当一个函数被jit编译后，所有在函数内部产生的中间张量都只存在于设备端，根本不会传回Python主机。这就是计算与通信融合：你从Python侧看到的，只是“喂入一批数据，取出一个结果”，而内部的千军万马全部在设备端自行调度完毕。

jit让你能对“单个函数内的所有计算”做整体编排，但它的默认心智模型仍是“一个函数在一个地方跑”。当数据大到一张台子放不下，你就需要更高阶的分工工具——pmap。

fjit(x)=Exec(Compile(Gf))(x),Gf=Trace(f,shape(x))f_{\text{jit}}(x) = \text{Exec}\big( \text{Compile}(G_f) \big) (x), \quad G_f = \text{Trace}(f, \text{shape}(x))
公式解读：对函数 $f$ 应用 jit，等于先用输入的形状信息把函数“描”成一张静态计算图，再把它编译成设备可执行程序，最后执行。
  1. 形状追踪$\text{Trace}(f, \text{shape}(x))$ 表示让JAX用假张量走过一遍函数，记录下每个操作的形状，得到计算图。
    Gf=Trace(f,shape(x))G_f = \text{Trace}(f, \text{shape}(x))
  2. 编译可执行$\text{Compile}(G_f)$ 把计算图交给XLA，生成针对特定设备优化的机器码。
    Compile(Gf)\text{Compile}(G_f)
  3. 设备端执行$\text{Exec}(\cdots)(x)$ 在设备上执行编译产物，输入输出通过设备传输，中间过程对外不可见。
    fjit(x)=Exec(Compile(Gf))(x)f_{\text{jit}}(x) = \text{Exec}\big(\text{Compile}(G_f)\big)(x)

pmap分片

餐厅迎来了史上最大的一桌客人：一批512张订单同时涌进来。一位厨师就算有三头六臂也做不完。这时候总厨下令：把这512张订单按厨师人数均分，每人拿一叠，同时开做。这就是数据并行，而JAX中的pmap就是这道指令。

pmap（parallel map）做的事情可以用一句话概括：沿着张量的第一个维度，把数据均匀切分成N片，每片分给一个设备，然后在所有设备上执行同一个函数。比如你有4台设备，输入张量形状是(8, 512)，pmap会在设备0拿到(2, 512)、设备1拿到(2, 512)……每个设备只处理属于自己的那两行数据。这就是分片维度控制的核心——你不需要手动切片，pmap会替你完成。

这里非常重要的细节是局部计算映射：你用pmap包裹的函数，其内部逻辑是用“单片数据”的视角来写的——仿佛你只有一张台子、只处理那一小叠订单。pmap负责把这份逻辑复制到N台设备上，并让每台设备针对自己的分片运行。设备之间原本独立的计算互不干扰，直到需要汇总结果（比如求全局梯度）时，才触发跨设备归约——AllReduce操作会把各设备的局部结果合并成全局结果，再广播回去。

pmap让“一份大订单切成小叠分给众人”变得优雅而简洁。但它也把分布式训练的一个隐藏成本推到了聚光灯下——设备之间的传菜，可不是免费的。

x∈RB×d  →pmap  [x[0:B/N],  x[B/N:2B/N],  …,  x[(N−1)B/N:B]]x \in \mathbb{R}^{B \times d} \;\xrightarrow{\text{pmap}}\; \big[ x_{[0: B/N]}, \; x_{[B/N: 2B/N]}, \; \ldots, \; x_{[(N-1)B/N : B]} \big]
公式说明：一个批次为 $B$ 的张量，在 $N$ 台设备上被均匀切成 $N$ 片，每台设备只拿 $B/N$ 行数据。
  1. 分片切片$x_{[0 : B/N]}$ 表示第0号设备拿到的数据切片——把第一维从0切到 $B/N$。
    x[0:B/N]x_{[0 : B/N]}
  2. 依次分配每台设备 $k$ 分配到区间 $[kB/N, (k+1)B/N]$ 内的数据。
    x[kB/N:(k+1)B/N]x_{[kB/N : (k+1)B/N]}
  3. 统一函数执行每个设备 $d_k$ 对同一函数 $f$ 执行，得到局部结果并汇总。
    yk=f(xk),Y=concat(y0,y1,…,yN−1)y_k = f(x_k), \quad Y = \text{concat}(y_0, y_1, \ldots, y_{N-1})

通信开销优化

厨师们埋头各做各的订单，但总有些菜需要拼盘——比如把四个人的半成品合成一道完整的菜。这个过程必须有人离开灶台，走到别人的工作台前传递食材。每一次设备间通信都有代价，而分布式训练中最沉重的一次传菜发生在梯度同步。

设想你用pmap在4张GPU上并行训练：每个GPU算出了自己那一份数据的梯度。为了让模型保持一致，需要把4份梯度加起来取平均，然后分发到所有GPU上，更新各自的模型参数。这个操作叫AllReduce（全归约）。在通信拓扑上，最简单的做法是选一台主设备收集所有梯度、算完再广播回去——但主设备会成为瓶颈。更聪明的做法是环状AllReduce：让4台设备围成一个环，数据像接力棒一样沿着环逐跳传递，每个设备收到前一个设备的部分求和后再传给下一个。这样所有设备几乎同时完成汇总，没有冷落谁。

通信成本可以量化。对于环状AllReduce，总耗时中包含两项：一项是固定的启动延迟（α），它与数据量无关——就好比厨师每一次开口叫同事，都要先花一点“组织语言”的时间；另一项是传输时间（β·M），它与消息大小成正比——传的菜越多，端菜走路的时间越长。

理解了通信成本，就理解了扩GPU规模时为什么“更多厨师不一定更快”——如果每轮通信的时间已经超过了计算节省的时间，加设备反而变慢。而把通信隐藏到计算背后（通信计算重叠），是优化分布式训练吞吐量的关键技术。

Tallreduce=2(N−1)α+2(N−1)NβMT_{\text{allreduce}} = 2(N-1)\alpha + \frac{2(N-1)}{N} \beta M
公式解读：$N$ 台设备做一次环状全归约的耗时，由 $2(N-1)$ 次延迟启动和 $\frac{2(N-1)}{N}$ 倍消息体量的带宽传输共同构成。
  1. 延迟项$2(N-1)\alpha$ —— 每次通信有固定延迟α，环状归约要经历 $2(N-1)$ 次这样的握手，即使消息为零也要付出。
    2(N−1)α2(N-1)\alpha
  2. 带宽项$\frac{2(N-1)}{N} \beta M$ —— 消息总量 $M$ 乘以带宽系数 β，再乘以一个随设备数缓慢增长的系数。
    2(N−1)NβM\frac{2(N-1)}{N} \beta M
  3. 设备数影响当 $N$ 增大，带宽项趋近于 $2\beta M$，说明环状AllReduce的带宽成本有上限——这正是它可扩展的原因。
    lim⁡N→∞2(N−1)N=2\lim_{N \to \infty} \frac{2(N-1)}{N} = 2

显存利用率

厨房再大，台面也是有限的。当模型大到一张GPU放不下——比如一个大模型光参数就要20GB，而单卡只有16GB显存——你必须考虑：如何把模型这副“大案板”切成几块，分放到不同台子上。

设备放置直接影响显存利用率。在数据并行（pmap）的模式下，每个设备都保存一份完整的模型参数，相当于每位厨师都配了一套完整的刀具——这当然耗台面。而如果采用模型并行，把模型的不同层（或同一层的不同权重块）放到不同设备上，每台设备只需要保存自己负责的那部分参数，台面压力就大大减轻。这就是分片减少冗余的逻辑：不做无谓的复制，只保留必需的部分。

但是，显存管理的挑战还有一层：训练过程中的显存峰值并不是“模型参数大小”那么简单。运行时，显存里同时住着四类房客：模型参数、梯度、优化器状态（比如Adam的动量）、以及前向传播时为了反向传播而缓存的中间激活值。这四者叠加，才是真正的峰值占用。有时候模型不大，但batch开得太大，激活值直接把显存挤爆——就像台面上备了太多半成品食材，结果连放成品盘子的地方都没有了。

理解了这四项，你就明白为什么大型模型训练几乎必然要走“分片+设备放置”这条路——它不只是加速，更是让放不下的模型能“挤”进显存的关键。

Mpeak=Mmodel+Mgrad+Mopt+MactM_{\text{peak}} = M_{\text{model}} + M_{\text{grad}} + M_{\text{opt}} + M_{\text{act}}
公式含义：训练时的显存峰值，等于模型参数、梯度、优化器状态和激活值四者占用之和；任何一项超标都会导致“厨房台面溢出”。
  1. 模型参数$M_{\text{model}}$ 是模型本身的权重占用，由参数量乘以每个参数字节数决定。
    Mmodel=num_params×bytes_per_paramM_{\text{model}} = \text{num\_params} \times \text{bytes\_per\_param}
  2. 梯度存储$M_{\text{grad}}$ —— 反向传播算出的梯度，同样大小，因为每个参数都有对应梯度。
    Mgrad≈MmodelM_{\text{grad}} \approx M_{\text{model}}
  3. 优化器状态$M_{\text{opt}}$ —— Adam之类的优化器会保存每个参数的额外状态，往往能到参数量的两倍。
    Mopt=2×Mmodel(Adam)M_{\text{opt}} = 2 \times M_{\text{model}} \quad (\text{Adam})
  4. 激活值累积$M_{\text{act}}$ —— 前向过程的中间张量，与batch大小和网络深度相关，最容易被忽视。
    Mact∝batch_size×depthM_{\text{act}} \propto \text{batch\_size} \times \text{depth}

数据与模型分片

做饭有两种分工流派。第一种：每个厨师做整道菜，但服务不同的客人——这叫数据并行，每个设备持有完整的模型，处理不同的数据样本。第二种：每个厨师只做这道菜的某一个工序——有人切菜、有人掌勺、有人装盘，设备各持模型的一部分，依次协作完成一个样本的处理，这叫模型并行。

在实际的大型训练任务中，两种流派常常混搭。以GPT-3这类千亿参数模型为例：模型太大，任何一张卡都无法单独放下完整权重，所以必须做模型并行——把Transformer的每一层切到不同设备上，前一层算完传给下一层。这还不够，因为每层内部的矩阵也大得惊人，于是还会把单个矩阵切成块分布到多张卡上做并行矩阵乘法（这就是2D分片）。与此同时，为了利用大batch训练，还会对数据维度做数据并行——多组设备各自处理一批不同的样本。数据并行与模型并行叠加，就是混合并行。一句话总结：数据并行解决“样本太多”的问题，模型并行解决“模型太大”的问题。

用矩阵分块的视角来看模型并行会非常清晰。假设一个大矩阵 $W$（比如全连接层的权重）横跨4台设备存储，可以按行切成4块，每台设备保存一行块；前向传播时，输入 $x$ 也要做相应切片分发。

分片策略的选择不是非此即彼，而是根据模型大小、batch大小、设备拓扑动态组合的艺术。而你真正要掌握的，是怎样用JAX的两个工具——pmap做数据分片，手动jit配合设备放置做模型分片——来搭出属于你自己的混合并行方案。

W=[W1W2⋮WN],y=xW=∑k=1NxkWkW = \begin{bmatrix} W_1 \\ W_2 \\ \vdots \\ W_N \end{bmatrix}, \quad y = x W = \sum_{k=1}^{N} x_k W_k
公式解读：把权重矩阵按行分到 $N$ 台设备，每台设备计算 $x_k \cdot W_k$，最后把结果相加得到完整的 $y$。
  1. 行分块$W_k$ 是 $W$ 的第 $k$ 个行块，驻留在第 $k$ 号设备上。
    Wk=W  [k⋅HN:(k+1)⋅HN,  :]W_k = W\;[k \cdot \tfrac{H}{N} : (k{+}1) \cdot \tfrac{H}{N}, \; :]
  2. 局部MatMul第 $k$ 台设备拿输入 $x$ 的对应切片，与自己的 $W_k$ 做矩阵乘法。
    yk=xWky_k = x W_k
  3. 全归约聚合把 $N$ 台设备的局部结果相加，得到最终输出——这本质上就是一个AllReduce。
    y=y1+y2+⋯+yNy = y_1 + y_2 + \cdots + y_N

扩展实践

当你在4张GPU上跑通了训练，下一个自然的冲动是：把它扩展到8张、16张、甚至一个跨越8台机器的64卡集群。这时，设备放置的复杂度从“一张厨房里的布局”升级为“一个连锁餐厅网络”的管理。

第一个必须考虑的问题是多级分布策略：一个64卡集群，通常由8台机器组成，每台机器内有8张GPU。机器内部的GPU通过NVLink等高速总线互联，带宽可达数百GB/s；而机器之间走的是万兆网络或InfiniBand，带宽差一个数量级。一个聪明的放置方案，必须拓扑感知——尽量让高频通信的参与者落在同一台机器内部，而不是跨机器传输。想象一下，如果总厨把需要频繁交换数据的两位厨师安排在不同的楼栋，传菜全靠外卖骑手，那效率必然惨不忍睹。

在工程实践上，JAX生态提供了一些探查工具来辅助设备放置决策：jax.devices()可以列出所有可见设备，jax.local_device_count()告诉你每台机器上有几张卡，jax.debug.visualize风格的剖析API能输出HLO的计算图布局。更进一步的性能分析往往依赖perfetto或NVIDIA的Nsight等profiling工具，它们帮助你看清楚每张卡上计算占多少时间、通信占多少时间，从而精准定位瓶颈。

下面是一个实操思路：当建立八机64卡集群时，先测出每张GPU到其他63张GPU的通信带宽，生成一张“距离矩阵”，然后让数据并行（通信量大的部分）沿着高带宽链路放置，模型并行（通信量相对小的部分）可以跨机放置。这套逻辑已经体现在很多大模型训练框架的“设备网格”抽象中。

设备放置的最终命题，不是“把所有计算分到所有设备上”，而是“在正确的位置做正确的事，并用最省的通信把各部分连起来”。这个过程，既是科学，也是手艺。

S(N)=1(1−p)+pNS(N) = \frac{1}{(1-p) + \dfrac{p}{N}}
这是阿姆达尔定律：$p$ 是可并行部分的比例，$N$ 是设备数。它揭示了扩展的上限——即使设备无限多，串行部分 $(1-p)$ 仍决定了加速比天花板。
  1. 并行比例$p$ 是总任务中能被多设备分摊的比例，例如矩阵乘法和数据加载。
    0≤p≤10 \le p \le 1
  2. 串行开销$(1-p)$ 是无论如何都无法并行化的部分，例如单点聚合、调度、Python解释器开销。
    1−p1-p
  3. 加速比极限$S(N)$ 是 $N$ 台设备相对单设备的加速倍数；当 $N \to \infty$，加速比趋近于 $1/(1-p)$。
    lim⁡N→∞S(N)=11−p\lim_{N \to \infty} S(N) = \frac{1}{1-p}

AI 应用场景

场景一：大语言模型的多机多卡训练

阶段：模型训练（分布式并行）

解决的问题：传统单卡训练无法容纳百亿级参数的LLM，显存直接溢出；即便勉强塞下，训练一个epoch也要数月。利用设备放置策略，把模型权重按矩阵分块分布到多张GPU上（模型并行），同时把数据分片到多组设备（数据并行），实现训练速度的线性扩展。

具体案例：某公司训练一个130亿参数的对话模型，硬件为16台A100（共128卡）。工程师用pmap对数据维度做128路并行，同时把Transformer每层的权重矩阵按列切分到8卡（模型分片）。通过拓扑感知放置，将同一台机器内使用的8卡组成一个“分片组”，利用NVLink高带宽完成矩阵乘法中的AllReduce；跨机器的梯度同步走InfiniBand。最终吞吐量从单卡训练的0.3TFLOPs提升到280TFLOPs，显存占用从“溢出”降到每卡约38GB。

场景二：金融风控模型的快速实验迭代

阶段：模型推理与批量预测

解决的问题：风控团队每天需要处理数千万条用户的特征张量做风险评分，单GPU推理链路延迟高、吞吐不足。利用jit把推理函数编译并按设备指定放置到多张GPU上，同时用pmap对用户批次做分片并行推理。

具体案例：某支付平台用JAX重写风控模型推理服务。原来用Python逐batch推理，P99延迟200ms。工程师用jax.jit编译整个打分函数，将模型参数预先加载到4张GPU的常量内存中，再用pmap把10000条用户的特征张量分成4片并行打分。优化后P99延迟降至25ms，吞吐量提升8倍，且因为jit消除了大量Python到设备间的张量搬运，显存占用下降了40%。

Python 演示

python45 行
import jaximport jax.numpy as jnp # 1. 查看这台机器上有哪些"厨师"（设备）devices = jax.devices()print("可用设备数量:", len(devices))for i, d in enumerate(devices):    print(f"  设备[{i}]: {d}") n_dev = len(devices) # 2. 定义一个简单的神经网络层：y = x @ W + b#    它等价于厨房里的一道固定工序def linear_layer(x):    W = jnp.ones((16, 16)) * 2.0    b = jnp.zeros(16)    return jnp.dot(x, W) + b # 3. 用 pmap 做数据并行分片#    输入的第一个维度必须能被设备数整除（视为batch维度）batch_size = max(n_dev, 2)data = jnp.ones((batch_size, 16))try:    p_layer = jax.pmap(linear_layer)    out = p_layer(data)    print("pmap 输出形状:", out.shape)except Exception as e:    print("当前环境不支持 pmap 多设备:", e) # 4. 用 jit 编译并指定放到第一个设备上执行#    相当于总厨手动指定"这道菜只由1号厨师做"fast_fn = jax.jit(linear_layer)x = jnp.ones((4, 16))y = fast_fn(x)print("jit 输出形状:", y.shape)print("结果设备:", y.devices) # 5. 手动将数组复制到特定设备（显式设备放置）if len(devices) >= 2:    target = devices[1]    x2 = jax.device_put(jnp.ones((4, 16)), target)    print("已手动放置到:", x2.devices)else:    x2 = jax.device_put(jnp.ones((4, 16)), devices[0])    print("已手动放置到唯一设备:", x2.devices)
目的：演示JAX中设备放置的三种手段——pmap数据分片、jit编译控制、device_put手动放置，帮助理解多设备调度机制。
  1. devices = jax.devices()获取全部可见设备列表，这是做任何放置决策的起点。
  2. p_layer = jax.pmap(linear_layer)将函数编译为“并行版本”，首维自动分片到全部设备。
  3. out = p_layer(data)调用并行函数时，数据沿首维被切分到各设备执行。
  4. fast_fn = jax.jit(linear_layer)编译整个函数，XLA在内部完成计算与通信的融合优化。
  5. y.devices查看输出张量实际存放在哪个设备上，帮助你确认放置是否生效。
  6. jax.device_put(x, target)显式把一个张量放到指定设备，适合手动控制模型分片的场景。

常见误区

误区我以为多加几张GPU，训练速度就会线性变快。反正每张卡分到更少的数据，每个人都轻松一些，总体应该更快才对。

正解设备增加的同时，通信开销也在增加。每轮梯度同步的AllReduce成本是 $2(N-1)\alpha + 2\beta M$，当设备规模变大而单卡计算量变小时，通信逐渐占据主导。在厨房里，如果传菜喊话的时间已经超过了做菜的时间，加再多的厨师也无济于事，甚至会互相挡路。判断是否该加设备的准则是：确保每台设备上的计算时间远大于其参与通信的时间。

为何常见：大家天然觉得“人多力量大”，忽略了协作本身有成本。加上各类框架宣传的“线性扩展”理想曲线，很容易误以为加速比不受通信影响。

误区我以为 `pmap` 就像Python里写个for循环，把数据切成N份，分别丢给N个设备算一遍，再把结果拼起来。感觉它就是个进阶版map函数。

正解pmap 确实在逻辑上等价于“分片→单程序多数据执行→拼接结果”，但本质区别在于所有分片是并行同时执行的，且在语义上被编译成一个统一的XLA计算，设备间的数据依赖会被整体分析。for循环是串行的——第一份算完才能算第二份——而pmap让所有设备同时开火。另一个关键点：pmap中的通信操作（如跨设备归约）是显式的原语，会被编译进同一个执行计划。

为何常见：从名字上看，pmap自带"map"字样，容易让人认为只是并行版的Python map；而实际编码时它的调用方式确实和map相似，掩盖了底层的编译与通信语义差异。

误区XLA不是会搜索最优放置吗？那我什么都不用管，它一定能找到最好的方案。如果慢，那肯定是硬件不行。

正解XLA的代价模型基于静态估算，它并不知道运行时的真实情况——比如某个GPU同时还在跑另一个任务、某条NVLink链路因为别的进程占用而拥堵、或者你的数据加载（DataLoader）是CPU瓶颈。这些动态因素会让XLA“纸上谈兵”式的估算失真。实践中的做法是：先用XLA自动放置跑通，再用性能剖析工具观察各设备的利用率，发现问题后用jit(device=...)手动把高通信量的计算挪到物理距离近的设备上。

为何常见：JAX的官方宣传强调“自动并行”，很多人就此理解为完全免干预；加上对XLA内部搜索机制的不了解，容易形成过度自信。

误区我理解模型并行就是把一个大模型切成两块，一张卡存前半部分，另一张卡存后半部分，推理的时候第一张算完传给第二张就行。训练时应该也是一样的逻辑吧？

正解推理确实可以这样简单流水线化，但训练阶段每个样本都需要前向+反向两次传播，而且反向传播的顺序与前向相反。如果只是简单地把层切到不同设备，前向时第二张卡要等第一张卡算完，反向时又要等第二张卡算完梯度——通信量和等待时间都会成倍增加。真实的模型并行还需要对张量做更细致的矩阵分块、在分块间做AllReduce，并配合梯度累积来减少同步频率。这也是为什么实践中往往用pmap处理数据分片，再配合手写切分逻辑做模型分片，两者结合才高效。

为何常见：推理的流水线直觉简单直观，容易让人认为训练只是“多跑一个反向”；但实际上反向传播的依赖顺序让朴素切分层的方式在训练中严重低效。

面试题

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

单选题 · easy

  • A.将函数编译成单个设备的加速代码
  • B.沿张量首维分片，多设备并行执行相同函数← 正确答案
  • C.自动将模型参数复制到所有设备
  • D.优化Python代码的执行速度

答案：B。pmap的核心是数据并行：沿第一个维度切分数据，分发给多台设备同步执行同一个函数。A描述的是jit，C和D则并非pmap的职责。就像总厨把大订单切成小叠分给每位厨师同时做。

Q2. jit在设备放置中最直接的用途是什么？

单选题 · easy

  • A.把所有计算放到CPU上
  • B.让函数在每次调用时重新编译
  • C.编译整个函数并在指定设备上执行← 正确答案
  • D.自动切分数据到所有设备

答案：C。jit把整个函数编译成XLA可执行程序，可以通过device参数指定执行设备；它不会自动切分数据（那是pmap），也不会放到CPU上，更不会每次重新编译。相当于总厨把整本菜谱读完后指定一位厨师做整套菜。

Q3. 环状AllReduce的时间模型中，$2(N-1)\alpha$ 这一项代表什么？

单选题 · mid

  • A.数据传输的总带宽消耗
  • B.所有设备计算耗时的总和
  • C.与消息大小无关的固定通信延迟← 正确答案
  • D.设备间传输数据的总字节数

答案：C。$\alpha$是单次通信的启动延迟，$2(N-1)$是环状结构需要的通信轮数，两者相乘得到与消息大小无关的固定延迟开销。就像厨师每次开口传话都要先花一秒钟组织语言，这跟菜量无关。

Q4. 用pmap处理形状为(32, 256)的输入，设备数为8台时，每台设备实际拿到的数据形状是？

单选题 · mid

  • A.(32, 256)
  • B.(4, 256)← 正确答案
  • C.(32, 32)
  • D.(8, 256)

答案：B。pmap沿第一个维度均匀切分，32 ÷ 8 = 4，因此每台设备获得(4, 256)的数据。这是数据并行最基本的计算：切分的维度是第一维，其他维度保持不变。

Q5. 训练时模型的梯度同步需要跨设备汇总梯度，这通常通过什么通信原语实现？

单选题 · easy

  • A.AllReduce← 正确答案
  • B.Broadcast
  • C.Scatter
  • D.Gather

答案：A。梯度同步需要把各设备的梯度求和并广播回去，这个过程正是AllReduce（全归约）。Broadcast是单源广播，Scatter是分发，Gather是收集，它们单独都无法完成“求和+广播”的双重任务。

Q6. 在数据并行训练中，通信与计算重叠的目的是？

单选题 · hard

  • A.减少计算量
  • B.让通信时间被计算时间覆盖← 正确答案
  • C.增加设备利用率
  • D.减少显存占用

答案：B。通信计算重叠的核心思想是：在反向传播计算梯队的同时，预取/发送已经算好的梯度，使通信和计算并行进行，从而让通信时间“藏”在计算内部而非额外暴露。B最准确地描述了这一目标；C是结果的副产品而非直接目的。

Q7. 关于JAX的设备放置，以下说法正确的是？

多选题 · mid

  • A.XLA会自动推断张量的最优设备放置← 正确答案
  • B.用户可以通过jit(device=...)手动指定设备← 正确答案
  • C.pmap可以自动进行模型并行分片
  • D.手动放置一定比自动放置快

答案：AB。XLA确实会自动分析计算图并选择放置方案，用户也可以用jit的device参数手动指定。C错误：pmap做的是数据并行，不是模型并行。D错误：手动放置受限于人的认知，未必优于XLA的全局搜索，所以不能说“一定”更快。

Q8. 以下哪些策略可以有效减少多GPU训练中的通信开销？

多选题 · hard

  • A.降低梯度同步频率← 正确答案
  • B.增加设备间数据传输的数据量
  • C.将高频通信的计算放到物理拓扑相近的设备← 正确答案
  • D.减少数据并行度，全部用模型并行

答案：AC。降低同步频率（如梯度累积）减少了AllReduce次数，确实是常见优化；拓扑感知放置能让通信走更快的链路，也能减低成本。B显然错误——增加数据量会放大通信开销。D也不对：模型并行本身也有通信成本，不是消除通信的银弹。

Q9. 设备放置策略直接影响以下哪些训练指标？

多选题 · mid

  • A.模型收敛精度
  • B.训练吞吐量← 正确答案
  • C.显存峰值占用← 正确答案
  • D.学习率大小

答案：BC。设备放置决定了通信开销和显存分布，直接作用于训练吞吐量（每秒处理的样本数）和显存峰值占用。收敛精度和学习率由优化算法决定，与设备放置无直接关系——只要数值计算正确（同精度），放置方式不会改变数学结果。

Q10. 一个模型的权重矩阵W为(1024, 1024)，被按行切成4块放到4台设备，前向传播计算$y = xW$时，最后的输出y如何得到？

单选题 · hard

  • A.取任意一台设备的计算结果
  • B.将4台设备的局部结果做拼接
  • C.将4台设备的局部结果相加（AllReduce）← 正确答案
  • D.将输入x直接广播到所有设备

答案：C。按行分块后，$xW = x_1W_1 + x_2W_2 + x_3W_3 + x_4W_4$，需要把4台设备计算的部分结果求和才能得到完整的输出y。拼接是错误的——因为每块只是部分贡献，不是互斥的输出片段；这个求和操作正是AllReduce在模型并行中的典型形态。