TPU：推理 vs 训练

可视化演示

展示脉动阵列矩阵乘法数据流、训练vs推理负载特征对比及训练推理混部资源调度策略，支持模式切换与参数调节。

场景差异

一家餐厅的后厨，藏着两种截然不同的活法。研发新菜的试菜间里，师傅不慌不忙，电子秤精确到 0.1 克，一次试做几十份样品，只为把味道拿捏到极致；而前厅的出餐窗口，客人催得紧，师傅只有三分钟，拿起标准量勺“一勺盐”就下锅。同样是做饭，一个要“精”，一个要“快”。TPU 面对的训练和推理任务，正是这样两种活法。

TPU 在处理训练任务时，就像试菜间里的师傅：模型权重是“菜谱”，每次做菜（前向传播）都要按当下的菜谱来，尝完发现问题（计算损失）再倒推调整（反向传播）。训练阶段要求高精度，因为误差会在反复迭代中像滚雪球一样累积；一次喂入的数据要足够多（大批量），因为后面还有昂贵的“改菜谱”环节（梯度更新）等着摊薄成本。

推理则完全是另一套逻辑：菜谱已经定稿，只需按章执行一次前向传播，不用改菜谱、不用存档。模型被量化为 INT8，数值从 16 位浮点变成 8 位整数，计算量和内存都大幅缩减。在批处理策略上，训练偏爱大 batch，用吞吐换单位成本；推理偏爱小 batch 甚至 batch=1，因为每个请求背后都有一位等不及的“客人”——延迟才是推理的第一指标。单个请求的平均等待时间与批大小的关系可以写成：

$$L = \displaystyle\frac{T_{\text{batch}}}{B}$$

1. 延迟 L：单个推理请求从进入队列到拿到结果的总等待时间，单位是毫秒。

$$L_{\text{infer}}$$

2. 批处理时长 T_batch：TPU 处理完一整批样本的总耗时，包含排队、计算和传输。

$$T_{\text{batch}}$$

3. 批大小 B：一次送入处理器的样本数量，训练时可能是上百，推理时常为 1。

$$B$$

4. 平均分摊关系：批内每个样本的等待时间约等于总时长除以批大小，所以批越大，单个样本等得越久。

$$\frac{T_{\text{batch}}}{B}$$

所以训练师傅尽可以一次端上 100 份样品慢慢研究，出餐师傅却只能一份一份地炒。同样是 TPU，负载特征不同，活法自然不同——这正是我们理解后面所有架构设计与调度策略的起点。

硬件架构

这家 AI 餐厅的镇店之宝，是一口超级大炒锅——MXU（矩阵乘法单元）。普通 CPU 像单口小灶，一次只能炒一两道菜；GPU 像一排家用灶台，能同时开几十个火；而 MXU 是一口能“一次翻炒”整座食材山的巨型炒锅：把几百个数一次性倒进去，翻炒（乘加）结果瞬间出炉。它专精一个动作，却把速度做到了极致。

光有炒锅还不够，还得有趁手的量具。TPU 为训练和推理配备了两套“秤”：训练用的是 bfloat16 高精度浮点——别被“16 位”骗了，它的指数范围和 FP32 一样大，能表示的“数量级”没缩水，只是把尾数精度压到 8 位，对梯度更新来说完全够用；推理用的是 INT8 整数——把每个数粗犷地记成 8 位，存储减半、速度翻倍。

MXU 的威力来自“脉动阵列”设计：成百上千个乘加器排列成阵列，数据像波浪一样在阵列里流动，每个时钟周期完成一次大规模矩阵乘法。此外，TPU 集群内部还有高带宽互联（如 TPU v3 Pod），让成百上千颗芯片能快速交换梯度。矩阵乘法的规模可以直接用维度表示：

$$C_{M \times N} = A_{M \times K} \times B_{K \times N}$$

1. 矩阵 A：左侧输入矩阵，M 行代表批里的 M 个样本。

$$A_{M \times K}$$

2. 矩阵 B：右侧权重矩阵，K 行 N 列，N 常为输出神经元数或类别数。

$$B_{K \times N}$$

3. 归约维度 K：两矩阵做内积时逐元素相乘再累加的维度，MXU 把这段累加一次性流水完成。

$$K$$

4. 输出 C：M 行 N 列的结果矩阵，正是 MXU 一次“翻炒”的产出。

$$C_{M \times N}$$

超级炒锅配电子秤，能做出精雕细琢的宴席菜；换上一把量勺，就能以两倍速度端出快餐。训练和推理各自用顺手的“量具”，这口锅的潜力才被完全释放。

训练工作流

试菜间的工作从来不是一次性的：按菜谱做一道菜（前向），请品鉴师打分（损失），倒推是哪一步放多了盐、哪一步火候不够（反向），然后修改菜谱再来一遍（更新）。如此循环往复几百甚至几百万次，菜谱才终于接近完美。训练神经网络的过程，跟试菜如出一辙。

一次典型的训练迭代分三步：前向传播——数据从输入层流过每一层网络，经矩阵乘法（MXU 的工作）和激活函数得到预测；计算损失——拿预测和正确答案对比，得到一个“差多远”的分数；反向传播——用链式法则从损失倒推每一层参数应有的调整方向（梯度）。方向知道了，就沿着它迈一小步。

多个厨师协作时还要“对表”：数据并行训练中，每张 TPU 芯片各自算出一份梯度，然后通过高速互联做 AllReduce 汇总——大家取平均值，再同步更新参数。这一步如果太频繁，通信开销会吃掉计算收益，所以训练才偏爱大 batch：一次多算点，同步次数就少了。参数更新遵循最经典的梯度下降法则：

$$\theta_{t+1} = \theta_t - \eta \; \nabla_{\theta} \mathcal{L}(\theta_t)$$

1. 当前参数 θ_t：第 t 步迭代时模型的全部权重，就是当前的“菜谱”。

$$\theta_t$$

2. 梯度 ∇L：损失函数对参数的偏导，指明“口味偏咸还是偏淡、该往哪个方向调”。

$$\nabla_{\theta} \mathcal{L}(\theta_t)$$

3. 学习率 η：每次调整的步幅大小，步子太大会跑偏，太小则收敛太慢。

$$\eta$$

4. 更新规则：新菜谱 = 旧菜谱 - 学习率 × 调整方向，一步步逼近最优。

$$\theta_{t+1} = \theta_t - \eta \nabla_{\theta} \mathcal{L}$$

训练的关键词是“迭代与对齐”：前向、反向、同步、更新，周而复始。每一次试菜都让菜谱离完美更近一步，而大批量和高速互联，就是为了让这场漫长的修炼尽量快一些。

推理工作流

出餐窗口的师傅从不研发新菜，她们手里只有一本定稿菜谱，动作只有一个：按菜谱做，快速端出去。为了再快一点，她们把电子秤换成了量勺——盐“一勺”而不是“3.27 克”。这个“换秤”的动作，在 AI 世界里叫量化（Quantization）：把模型权重从高精度浮点映射成 INT8 整数。

推理是训练的反面：它只做一次前向传播，没有反向传播，不需要保存梯度。它的考核指标是延迟——用户点下按钮后多久看到结果。在 TPU 上，推理任务通常以 INT8 计算：权重和激活值都被压缩成 8 位整数，内存占用减半，同一颗芯片的算力翻倍。Google 第一代 TPU 就是专为 INT8 推理设计的，峰值算力高达 92 TOPS。

精度损失怎么办？聪明的方法叫“校准”：用少量真实数据确定缩放因子，把浮点范围最合理地映射到整数范围；更高级的做法是“量化感知训练”，让模型在训练阶段就提前适应“量勺”的粗糙，于是推理时口味几乎不变。量化的数学本质就是把浮点数 r 通过缩放和取整变成整数 q：

$$q = \operatorname{round}\left(\frac{r}{s}\right) + z$$

1. 原始浮点值 r：量化前的权重或激活值，对应“精确克数”。

$$r$$

2. 缩放因子 s：浮点范围与整数范围的比例，决定“一勺相当于几克”。

$$s$$

3. 四舍五入 round：把连续值取整到离散整数，量勺没法舀出小数克。

$$\operatorname{round}\left(\frac{r}{s}\right)$$

4. 零点偏移 z：让浮点零点在整数域也有精确对应，校正偏移。

$$z$$

量勺虽然粗犷，但配合训练阶段“适应量勺”的菜谱，端出来的菜顾客根本尝不出差别。推理的哲学就是：牺牲一点点无关紧要的精度，换来翻倍的速度和减半的成本。

存储管理

厨房里最怕什么？不是锅不够大，而是半成品没地方放、传菜通道太窄。试菜间里，师傅会把每一道工序的半成品都留样——因为如果最终菜咸了，得回头查是哪一步的盐加多了。训练也一样：反向传播需要回看前向传播每一层的中间结果，这些中间结果叫“激活值”，必须全部留在内存里。

出餐窗口就轻松多了：菜谱（权重）定稿后放在手边，做完一道菜直接上桌，半成品（激活值）算完即弃。所以训练芯片的内存压力是推理的好几倍——一个大 batch 训练，激活值占用随批大小线性膨胀，甚至超过权重的占用；而推理只需要让权重常驻内存即可。

为了缓解数据搬运压力，TPU 采用“权重驻留”（weight-stationary）的数据流设计：权重尽量固定在 MXU 的片上寄存器里，输入数据从内存流过，反复使用同一份权重，大幅减少重复搬移。这一设计直接针对一个朴素事实——搬运数据比计算更贵，内存带宽往往是瓶颈。训练时的激活值总占用可以这样估算：

$$M_{\text{act}} = B \times \sum_{l=1}^{L} H_l \, W_l \, C_l$$

1. 批大小 B：一次送入的样本数，激活值显存随它线性增长。

$$B$$

2. 第 l 层特征图尺寸 H_l W_l C_l：第 l 层中间结果的高、宽、通道数，即“半成品的体积”。

$$H_l W_l C_l$$

3. 逐层求和 Σ：把所有层的半成品体积加总，得到单样本的激活值总量。

$$\sum_{l=1}^{L} H_l W_l C_l$$

4. 总占用 M_act：训练为反向传播留存的总显存，批越大、网络越深，占得越多。

$$M_{\text{act}}$$

试菜间的大冰箱和出餐台的小架子，决定了两种场景对内存的配置完全不同。理解了这一点，你就明白为什么推理专用芯片可以做得更小、更省电，而训练芯片必须带着“大冰箱”上阵。

性能对比

饭店老板要算效率账。试菜间忙了一下午，做了 30 份样品，产出是 3 道新菜；出餐窗口忙了一下午，出了 300 份菜，收入是试菜间的几十倍。两种“忙”都占用了同一个厨房，但账本上该记的指标完全不同——训练看“每秒钟处理了多少样本”，推理看“每个请求多少毫秒返回”。

从硬件利用率的视角看，训练是“大而稳”的洪流：大批量数据把 MXU 填得满满当当，利用率能拉到 80% 以上。推理则是“短而急”的脉冲：单个请求算完就停，如果一次只来一个，芯片大部分时间在空转。为了救利用率，推理服务引入了动态 batching——把同时到达的多个请求拼成一锅一起算，既保吞吐又不让单个请求等太久。

精度档位的选择直接影响性能账。同一颗 TPU 上，INT8 的运算速度比 bfloat16 快一倍（8 位 vs 16 位，同一 MXU 每周期处理的数据量翻倍），功耗也更低。这就是为什么服务方宁可花心思做量化校准，也要在推理链路上用 INT8。量化带来的理论加速比是：

$$S = \frac{T_{\text{bf16}}}{T_{\text{int8}}} \approx 2$$

1. 分子 T_bf16：用 bfloat16 计算某个矩阵乘所需的时间。

$$T_{\text{bf16}}$$

2. 分母 T_int8：用 INT8 计算同一个矩阵乘所需的时间。

$$T_{\text{int8}}$$

3. 加速比 S：同一硬件两种精度的时间比，约等于 2，即推理每周期吞吐翻倍。

$$S \approx 2$$

两本账摆在一起，精明的老板会问：试菜间闲下来的时候，能不能让出餐窗口的活分给它？训练和推理的错峰互补，正是混部调度要解决的问题。

混部调度

现实中的餐厅，出餐高峰只在饭点，试菜可以安排在下午闲时；但厨房租金是全天付的，锅闲着就是亏钱。聪明的老板开始排“共享时刻表”：客人催得急的时间段，试菜师傅让出炒锅；后半夜没人来，试菜师傅加班研发。这就是资源混部调度（Co-location）的思想——让训练和推理共享同一批 TPU。

混部的关键，是认清两种任务的“地位差异”：推理面向用户，有服务等级协议（SLO）——比如“99% 的请求必须在 10 毫秒内返回”，这是硬指标，必须高优先级、不可抢占；训练是内部任务，晚几个小时完成没人投诉，可以被暂停、被压缩、被抢占，属于弹性负载。调度器的职责就是：推理优先拿资源，训练用剩余容量“填缝”。

具体怎么做？资源池化：把一批 TPU 放进统一资源池，按需切分成虚拟切片分配给不同任务。动态伸缩：推理高峰自动收缩训练占用量，低峰再放开。监控闭环：实时统计推理的 p99 延迟，一旦逼近 SLO 阈值，立刻抢占训练任务。最终效果可以用芯片利用率来衡量：

$$U = 1 - \frac{T_{\text{idle}}}{T_{\text{total}}}$$

1. 闲置时间 T_idle：芯片没有任务可跑的空转时间，是纯浪费。

$$T_{\text{idle}}$$

2. 总时间 T_total：统计周期内的总时间。

$$T_{\text{total}}$$

3. 利用率 U：芯片实际干活的时间占比，混部的目标就是让 U 趋近 1。

$$U = 1 - \frac{T_{\text{idle}}}{T_{\text{total}}}$$

同一个厨房、同一口炒锅，靠一张排班表服务两种生意。推理赚口碑，训练练内功，各取所需，皆大欢喜——这就是混部调度的魅力所在。

工程实践

有一家互联网“餐厅”，白天是线上推荐服务的高峰，用户请求像散客一样络绎不绝；晚上是模型的例行更新，大批训练任务像包场宴席。过去它准备了两套厨房：推理集群白天忙死、夜里闲死，训练集群夜里忙死、白天闲死——整体利用率只有 45%，老板看着账单直叹气。

后来它把两套厨房合并成一个共享 TPU 池，写下了三条规则：第一，推理任务永远最高优先级，INT8 计算，p99 延迟必须达标；第二，训练任务使用剩余容量，随时准备被抢占，好在支持断点续训，被打断也不慌；第三，调度器每 100 毫秒看一次推理延迟水位，逼近阈值就压缩训练，回落后再放开。每日还有一次“错峰”：白天训练只占 30% 容量，夜间自动提升到 70%。

运行一个月后：芯片平均利用率从 45% 提升到 82%；同样的日均请求量，需要的芯片数量减少了约 30%；推理 p99 延迟稳定在 20ms 以内，完全符合服务等级协议。混部的硬约束可以写成一条概率公式：

$$P(L_{\text{infer}} \le L_{\text{SLO}}) \ge 0.99$$

1. 实测延迟 L_infer：线上真实采集的推理请求响应时间。

$$L_{\text{infer}}$$

2. 承诺阈值 L_SLO：服务等级协议中规定的最大允许延迟。

$$L_{\text{SLO}}$$

3. 达标概率 P：满足延迟要求的请求所占比例，混部要求不低于 99%。

$$P(L_{\text{infer}} \le L_{\text{SLO}})$$

账本上多出来的 37 个百分点，就是混部调度交出的一份漂亮答卷。试菜与出餐，终于在同一屋檐下相安无事——训练打磨模型的锋芒，推理守护用户的体验，一片 TPU 两种精彩。

AI 应用场景

Python 演示

import random  def simulate(shared=True, steps=2000, seed=42):    random.seed(seed)    total = 8    busy = 0    infer_req = 0    infer_miss = 0    train_done = 0    train_queue = []     for t in range(steps):        n_req = sum(1 for _ in range(4) if random.random() < 0.25)        infer_req += n_req        if t % 2 == 0:            train_queue.append(15)         if shared:            n_serve = min(n_req, total)            remain = total - n_serve        else:            n_serve = min(n_req, 2)            remain = 6         can_train = min(len(train_queue), remain)        infer_miss += n_req - n_serve        busy += n_serve + can_train         active = train_queue[:can_train]        waiting = train_queue[can_train:]        active = [r - 1 for r in active if r > 1]        train_done += can_train - len(active)        train_queue = active + waiting     infer_success = 100.0 * (infer_req - infer_miss) / infer_req    utilization = 100.0 * busy / (total * steps)    return infer_success, utilization, train_done  def main():    print('%-8s %-12s %-12s %-10s' % ('模式', '推理成功率', '芯片利用率', '训练完成'))    for mode, shared in (('分部署', False), ('共享混部', True)):        rate, util, done = simulate(shared=shared)        print('%-8s %-11.2f%% %-11.2f%% %-9d' % (mode, rate, util, done))  main()

1. import random引入随机数库，模拟推理请求的到达波动。
2. n_req = sum(1 for _ in range(4) if random.random() < 0.25)每个时间片最多来 4 个推理请求，每个以 25% 概率到达，模拟流量峰谷。
3. if shared:区分“共享混部”与“固定分部署”两条路线。
4. n_serve = min(n_req, total)共享池给推理最高优先级，8 个单元全部可服务。
5. remain = total - n_serve服务完推理后，剩下的容量才轮到训练。
6. can_train = min(len(train_queue), remain)训练用“填缝”形式占满剩余容量。
7. active = [r - 1 for r in active if r > 1]正在执行的训练任务消耗一个时间片，完成的移出队列。
8. utilization = 100.0 * busy / (total * steps)统计整个周期的芯片利用率，用来对比两种方案。
9. simulate(shared=True)跑共享混部方案。
10. simulate(shared=False)跑固定分部署方案，作为对照。

常见误区

面试题