第八代TPU深度拆解：训练与推理实战指南

最近芯片圈子里聊得最多的，除了GPU就是Google的第八代TPU。很多训练团队不是没看过那几张泄露截图，真正的问题是：这些参数到了自己手上到底意味着什么，能跑多大的模型，能不能接进现在的PyTorch/JAX管线。这篇文章就是把第八代TPU的训练与推理芯片参数掰开来看——哪些是能直接拿来当参考的资料，哪些还只是推测，以及根据我过去在TPU v4、v5p和v6上跑训练和推理踩出来的经验，它落到真实工作负载时应该怎么配置、怎么避坑。

适合两类人看：一类是正在调训练基础设施的工程师，另一类是准备做算力采购的决策者。如果你是第一次接触TPU，建议先看第3章和第4章，能帮你少走一半弯路。先说清楚：目前Google对外公开的最新型号是TPU v7 Ironwood，业内常说的“第八代TPU”，是依据路线图和供应链信息推算的下一代云端训练与推理芯片。所以凡是参数表里的精确数字，我都标注为“推测”或“流传版本”，不吹成官方规格，免得后面发布时被打脸。

1. 第八代TPU是什么来头：代际演进里的关键变化

1.1 从TPU v1到v7，Google到底在迭代什么

很多人以为TPU是突然冒出来的“黑马”，其实这条线已经走了快十年。Google在2015年发布TPU v1，当时只做推理，专门跑矩阵乘法和卷积，单靠这一招就在语音和搜索场景里省下了大量GPU成本。到v2、v3才把训练能力补齐，引入了更灵活的可编程矩阵单元。之后v4、v5e、v5p到v6 Trillium，每一代都在做同一件事：把算力、内存带宽、芯片间互联三项指标同步抬高。

这一点非常关键。TPU的设计思路一直不是“堆通用计算单元”，而是围绕Transformer这类密集矩阵运算做定制。模型结构越固定，它在矩阵乘单元上的利用率越高。第八代TPU如果按传闻推演，大概率还是在做这种“同步换代”：单芯片浮点峰值提高，HBM容量和带宽加大，ICI互联带宽继续拉高，然后通过软件栈把这三者一起吃满。它不是某个单点参数突然暴涨，而是一次系统工程层面的整体升级。

1.2 第八代TPU要解决什么问题

当前大模型训练和推理的痛点，其实很集中。

第一个是长上下文带来的KV Cache问题。模型上下文从32K涨到128K甚至1M之后，KV Cache所占的内存空间和带宽远远超过参数本身。如果芯片内存容量不够大，推理时就要频繁把KV Cache换进换出，整体吞吐直接崩掉。第八代TPU如果配到192GB级别的HBM，一台节点能同时服务的用户数和上下文长度都会明显上升。

第二个是MoE模型的全局通信。现在800B、1T参数级别的模型几乎都是Mixture-of-Experts架构，光有算力不行，expert之间要频繁做All-to-All通信。芯片间互联带宽不够，训练时一大半时间都在等token传输。每代TPU把ICI带宽往上拉，本质上就是为了不在这里卡脖子。

第三个是单位token成本。大模型从预训练到持续微调再到推理，每一个阶段都是钱。Google做TPU的目的不只是让自家Gemini跑得更快，更关键的是把单token成本压到别人追不上的水平。评价第八代TPU，不能只看峰值算力，要看“每瓦能跑多少token”“每美元能训多少参数”。

1.3 一个需要先辟谣的点

现在自媒体聊AI芯片，最喜欢说的是“PFLOPS翻倍”。但这个指标特别容易被带偏。同一块芯片，BF16、FP8、INT8三种精度下的峰值完全不同，稀疏计算还能再翻一倍。很多人拿INT8稀疏的峰值去对比别人BF16的稠密值，数字是好看了，实际工作负载完全不是一回事。

看第八代TPU的参数，首先要分清口径：训练大模型主要看BF16和FP8，推理看INT8或FP8，稠密和稀疏要分开算。下文参数表里我尽量把口径写清楚，标注“非官方推测”的地方，仅供参考。

2. 详版参数拆解：算力、内存、互联与功耗

2.1 算力矩阵：BF16、FP8、INT8与稀疏

下面这个表格是我综合多个渠道整理的推测值，不是Google官方公布的数字。但即便如此，也能看出下一代TPU的设计方向。

先解释精度。BF16是今天大模型训练的默认精度，动态范围和FP32差不多，但只占两个字节，显存压力和通信压力都小。FP8是这两年才普及的，训练时可以和BF16混合使用，把主流矩阵运算降到8位，速度更快、省显存，但对数值稳定性要求很高，不是无脑全部改成FP8就行。INT8则主要用于推理。

峰值算力还有一个残酷的现实：它只是纸面极限。实际运行时要受矩阵维度、数据布局、互联带宽共同影响。举个例子，你在服务端做小batch解码时，矩阵乘法的尺寸可能只有几百乘几百，矩阵单元利用率可能连30%都不到。所以真正有参考意义的不是峰值PFLOPS，而是“在真实模型、真实并发下能打出的有效算力”。这个有效算力，往往要靠跑Profile才能拿到。

2.2 内存容量与带宽：决定batch与KV Cache

推理场景里，HBM带宽比算力更能决定体验。模型参数要过内存，KV Cache也要过内存。decode阶段每个token的生成都依赖KV Cache的读取，带宽越大，每秒能吐出的token就越多。

KV Cache到底有多吃内存？可以用一个公式快速估算：

每个token的KV Cache大小 = 2（K和V两组） × 层数 × KV头数 × head_dim × 精度字节数 示例：70B模型，层数80，KV头数8，head_dim 128，BF16精度（2字节）： 2 × 80 × 8 × 128 × 2 = 327,680 Bytes ≈ 320 KiB / token 上下文长度32K时：327,680 × 32,768 = 10.7 GiB

一个70B模型光参数就占140GB，再加上32K上下文的KV Cache，192GB HBM正好能扛住。如果是128K上下文，一个用户就要占40多GB，这时内存容量不够，服务端只能降并发或者做KV Cache换出。第八代TPU把HBM提到192GB级别，并同步提高带宽，对长上下文推理的提升甚至比算力翻倍还明显。

简单算一下decode上限：如果每token占用320KB，16TB/s的内存带宽，理论峰值是每秒五万个token左右。当然这是极端理想值，实际要考虑多路并发和内存读写开销，但数量级摆在那里：内存带宽越高，并发推理的吞吐天花板越高。

2.3 互联带宽与POD规模

单芯片再强，放到集群里也要看芯片之间怎么通信。TPU的组网方式和GPU集群不一样：GPU通常靠NVLink加InfiniBand，TPU走的是ICI加OCS光电路交换。第八代TPU大概率会把单芯片ICI带宽拉到4Tbps级别，这样带来的好处非常直接。

一是张量并行通信更快。比如把70B模型切成多卡推理，每一层Transformer里的矩阵乘都要跨卡同步结果。通信带宽低，训练和推理都跑不快。二是MoE的All-to-All没那么痛了。专家分布在多张卡上，每个token都要找相关专家，通信模式是“All-to-All”，比纯数据并行通信量高一个数量级。ICI带宽不够，MoE集群就是一群芯片互相等待。

TPU的“Slice”概念也要提一下。每次申请资源，你拿到的不是一个任意组合的集群，而是一个固定拓扑的slice，比如v8-8、v8-16、v8-32。这种设计让Google可以预先规划好光交换连接，也简化了软件层的并行配置。缺点是你不能像拼接木桶那样随意加卡，申请多大、拓扑多规则，都得按平台约束来。

2.4 功耗、液冷与TCO

算力涨上去，功耗不可能不涨。上一代旗舰级芯片已经在液冷方案里跑，第八代如果单芯片功耗上到800W甚至更高，整个数据中心的散热架构都得围着它重新设计。

这里说的TCO，不能只看单卡价格，还要看整柜功耗上限带来的连锁后果。一个机柜如果只能提供40kW电力，那单卡500W和单卡900W能塞进去的芯片数量完全不同。反过来，如果900W芯片的吞吐是500W芯片的两倍，单卡功耗高反而可能更划算。做采购决策时，最好自己建一个简单公式：总吞吐 = 单卡吞吐 × 卡数，总成本 = 液冷改造 + 电价 + 芯片租赁费用。用这个公式去测，才不会看到PFLOPS就冲动下单。

3. 训练场景配置指南：从申请资源到跑起LoRA

3.1 TPU资源申请的正确姿势

在Cloud上申请TPU，现在主流是用Queued Resource，也就是异步排队资源。提交请求后不用一直等着，资源准备好后再去连接。命令大致长这样：

gcloud alpha compute tpus queued-resources create my-v8-16 \ --accelerator-type=v8-16 \ --runtime-version=tpu-vm-jax-0.4.39 \ --node-prefix=my-v8-vm \ --zone=us-central2-b \ --project=my-project

几个关键点说一下。accelerator-type里的数字代表slice规模，v8-16就是16颗第八代TPU组成的slice。runtime-version必须是TPU平台提供的镜像版本，而不是你自己拍脑袋装的环境。TPU VM的运行镜像把内核驱动、XLA编译器、JAX或TensorFlow版本打包在一起，升级一次就是整体升级，不能像普通GPU服务器那样随便装CUDA、换驱动。

资源起来后，可以用SSH连进去确认设备可见性：

import jax print(jax.devices())

看到8个或16个TPU设备条目，说明环境正常。这里有个容易忽略的细节：TPU VM虽然是虚拟机，但它不是普通计算实例，重启后本地盘数据很可能丢，模型权重、数据集和编译缓存都要放到持久盘或对象存储里。

3.2 训练环境里最容易踩的三个坑

我在TPU上踩过的坑，基本集中在三个地方。

第一个是动态shape。TPU走的XLA编译器会把完整计算图编译成二进制，然后整图执行。图里面只要有一个维度的值不确定，比如KV Cache序列长度随batch变化，编译器就可能反复触发重新编译。最夸张的一次，我在TPU上写自回归生成逻辑，因为没有固定输入长度，XLA每几步就重新编译一次，整个训练速度慢到没法看。解决办法是提前把所有输入padding到固定长度，让shape在整轮执行中保持稳定。

第二个是首次编译时间。第一次启动训练时，XLA要把整张图编译出来，一个7B模型的JIT编译时间可以到7到15分钟。很多新手以为机器卡死了，直接重启，结果下次还是重新编译一遍。正确做法是提前把编译好的缓存保存起来，重启后直接复用。

第三个是版本对齐。PyTorch/XLA、JAX、TPU runtime之间是强绑定的。你随便升级一个轮子，很可能出现“算子回退”，也就是某个操作没有被XLA识别，默默跑在CPU或GPU上，性能掉几个数量级还不报错。排查方法也简单：训练开始时看一眼日志里的算子统计，确认没有非TPU的可疑回退。

3.3 LoRA微调与大模型训练的实际建议

先说结论：低于1B的模型，我个人不建议专门上TPU，XLA的编译开销摊下来不划算，用GPU或者CPU更灵活。7B到70B的LoRA微调可以做，但建议优先考虑JAX或者PyTorch/XLA，不要直接拿GPU上的PEFT库硬跑。

在TPU上做LoRA，核心思路是把低秩矩阵插进attention层，然后用XLA把整个模型编译成一张图。如果你在GPU上已经用惯了PEFT，迁移过来会有点别扭：不是所有自定义OP都有XLA实现，很多适配模块需要自己写。实际操作中，我会把需要训练的LoRA参数集中在attention两端，训练时冻结其他层权重，用混合精度BF16，优化器用AdamW的XLA实现。

全参数训练又是另一套玩法。70B以上的模型，建议直接用GSPMD或FSDP类分布式策略把参数切分到多张卡上。TPU的优势是ICI带宽高，张量并行通信成本相对低，但前提是你把sharding配置写对。我的经验是：先用数据并行把一个batch跑通，再切到张量并行，最后才考虑混合并行，一步步逼近目标性能。

训练时序列长度最好固定成max_len，不足的padding。这样XLA的图和内核cache都能稳定复用。模型checkpoint也要注意：TPU上保存的XLA参数和GPU上的torch.load格式不对应，要么统一用框架自带的保存方式，要么显式转换。别等训练跑了一半才发现权重无法恢复。

4. 推理部署实战：引擎选型与流式管线

4.1 推理引擎怎么选

TPU推理不像GPU那样有vLLM、TensorRT-LLM一堆现成优化栈。想直接用开源推理框架跑TPU，目前基本不现实，这也算生态上最短板的地方。实际可选的路线大概四条：

选型时要重点问一个问题：你现有的模型代码能不能被XLA完整编译。能，性能通常不错；不能，出现算子回退后，性能可能比GPU差很多。所以在TPU上搞推理，第一步不是调参，而是先跑一次完整的XLA编译，看看日志里有没有失败或回退的算子。

4.2 流式推理管线怎么搭

大模型推理的特点是“预填充重、解码轻”。用户输入一段长prompt，模型要快速处理完，这是计算密集阶段；之后每个token的生成是顺序的，存的是内存带宽密集操作。流式部署要解决的核心问题是：不能等整段生成完再返回，首token时延要压下去，后续token要持续流出来。

我的做法是分三层处理。第一层做连续批处理，把多个请求攒到一起，尽量让矩阵运算Batch变大。第二层做prefill和decode分离，也就是预填充用大batch高算力跑，decode用小batch低算力跑，避免互相拖累。第三层做前缀缓存，把system prompt这类公共部分提前算好放进KV Cache，用户请求进来时直接复用。

部署形态上，用GKE加TPU节点池比较顺手。一个简化版的Deployment声明长这样：

apiVersion: apps/v1 kind: Deployment metadata: name: tpu-llm-server spec: replicas: 2 template: spec: nodeSelector: cloud.google.com/gke-tpu-accelerator: "v8" containers: - name: server image: gcr.io/my-project/tpu-serve:latest ports: - containerPort: 8080 resources: limits: google.com/tpu: 8

这份配置只是示意，实际生产还要挂PV、下载模型权重、做编译预热。这里特别提醒：TPU服务上线前一定先做“模型预热”，把XLA编译产物和模型权重都加载好。如果用户请求来了才开始编译，光等编译时间就够把线上SLA打穿好几轮。

流式推理的性能指标，主要看TTFT（首token时延）和TPOT（每个输出token的时延）。TPU的批量调度天然偏向吞吐，如果业务类型是请求量稀疏、每请求间隔很长，TPU的算力利用率会非常难看。把连续批处理窗口设成100到200毫秒，通常能平衡时延和吞吐。

4.3 哪些场景别用TPU推理

虽然第八代TPU听着很强，但并不是所有推理场景都适合。

• 请求稀疏、峰值流量不确定的小服务，用CPU或小GPU成本更低。
• 单模型特别小，比如1B以下，XLA的固定编译开销会把优势吃干净。
• 强依赖GPU自定义算子或CUDA生态库的项目，迁移成本远大于收益。
• 对延迟要求极低、需要个位数毫秒响应级别，这类实时RPC服务不要指望TPU，批处理机制会拖累你。

一个很现实的经验：TPU推理适合“稳定、大量、长上下文”的负载。如果业务是这个形态，第八代TPU的大内存和高带宽会非常能打。如果不是，先别盲目迁移。

5. 和GPU阵营对比：参数之外的工程账

5.1 参数对比表

拿第八代TPU和当前主流GPU旗舰做一个粗略比较。注意，两边参数口径不同，这里只看数量级：

参数上两边其实拉不开绝对差距。真正拉开的是工程账。

5.2 从GPU迁移到TPU的成本怎么算

很多人只比较硬件参数，忽略了迁移成本。这个成本要细算。

第一是框架迁移成本。如果你的模型本来就是JAX或TensorFlow生态，迁移到TPU相对顺滑。如果是PyTorch模型，能用PyTorch/XLA跑通就已经算幸运了。遇到自定义CUDA kernel，基本没有捷径，得重新用XLA算子实现一遍。这个工作量不是几周能解决的。

第二是DevOps成本。GPU服务器你可以当普通机器用，TPU VM更像一个封闭的定制计算节点，环境由平台镜像控制，很多常规运维操作要改思路。日志采集、监控指标、自动扩缩容，都要适配TPU的调度模型。

第三是成本曲线。GPU市场已经有大量现成的计费模式，TPU则按Google云的TPU小时计费。如果业务波动大，TPU固定slice会让你在低峰期也在付钱。反过来，负载稳定时TPU的能效和单位吞吐成本可能更优。

一个比较稳妥的判断方法：把你的生产模型在现有GPU集群上的“每美元产出token数”测出来，再把同等负载放到TPU测试集群上跑一周，对比两个数字。别猜，直接逐步实测。

5.3 什么时候选TPU

综合来看，下面几种场景我愿意主动选TPU：

• 训练或者推理规模已经到POD级别，有几十甚至几百张卡同步工作，TPU的ICI互联和OCS组网优势能体现出来。
• 模型架构相对固定，比如公司自研的MoE大模型，短期不会有大规模结构改动，值得投入XLA编译优化。
• 长上下文推理场景多，比如百万token级别的文档处理、代码仓分析，需要大HBM和超高内存带宽。
• 团队有JAX基础，或者愿意为XLA生态投入工程人力。

反过来，如果项目还处在快速迭代期，模型结构三天两头变，那就继续留在CUDA生态，灵活比峰值更重要。

6. 最后分享几条实操经验

最后把这几条年在TPU上摸出来的经验一次性说清楚，每一条都是真金白银换来的。

第一，申请资源前先查目标区域可用性。看起来配额很足，但某些热门zone里v8-16这样的slice可能一直pending。多准备一个备选区域，比等配额靠谱。

第二，第一次XLA编译非常慢，别急着重启。把日志存下来，看到图编译阶段卡住其实是正常现象。等几分钟再回来看，很多情况是已经编译完了。生产服务必须预热，否则用户的第一批请求全都会撞上编译时延。

第三，训练中断恢复时，本地环境会重置。权重文件、数据缓存、编译缓存都要放到持久盘或对象存储。我以前吃过亏，训练到一半VM被调度回收，本地权重全没了，只能从头继续。

第四，用FP8训练前，先在小batch上对比BF16的loss收敛曲线。有些模型对精度很敏感，直接切FP8会损失效果。不要看训练速度上去了就以为万事大吉。

第五，做INT8推理时，不要只盯着整体准确率。关键层的PPL指标、生成样本的连续性都要看。量化掉的某些attention layer可能只影响部分切片，但会让特定场景的生成质量明显劣化。

第六，也是最重要的一条：别只看PFLOPS和HBM大小，先拿你真实模型的Profile在它上面跑一周，能效、时延、成本全部记录下来，比任何参数表都管用。参数是纸面能力，跑通才是真本事。