单张 TPU v5e 跑 12B 模型：重打包 Gemma 4 QAT 权重实测

核心洞察

这篇文章最有意思的地方在于，它证明了单张 TPU v5eCloud TPU v5e，Google Cloud 第五代 TPU 芯片，单芯片 16GB HBM、800 GiB/s 带宽、峰值 BF16 197 TFLOPs，以较低成本提供推理/训练能力。 芯片上跑 12B 模型是可行的，而且质量几乎不掉。作者没有用花哨的技巧，就是把 Google 已经训练好的 QAT量化感知训练（Quantization-Aware Training），在训练过程中模拟量化误差，使模型在量化后保持精度。 权重原样打包，不重新取整。实测下来比 Google 自己导出的 4-bit 版本还高 2.4 分，速度一样。如果你手头只有一张 v5e，想跑尽可能大的模型，这篇值得细看。

---

这篇文章是一份完整的操作指南，讲的是如何把 Google 量化感知训练（QAT）过的 Gemma 4Gemma 4，Google的轻量级开源模型系列，边缘变体（E2B/E4B）可在1.5GB RAM以下运行。 权重重新打包，让 vLLM一个高性能的LLM推理和服务库，为DeepSeek-OCR提供优化的推理能力，支持流式输出和批量处理。 能在单张 Google Cloud TPU v5e 芯片上跑起来。每个构建版本都从分类、数学、工具调用、吞吐量和长提示词几个维度打了分。所有逐条输出、日志和脚本都已提交。

在单张 v5e 芯片上，重新打包后的 QAT 构建能跑 Gemma 4 从 E2B 到 26B 的所有尺寸。到 12B 为止，在 3880 条分类测试集上跟 bf16 打平，比 Google 自己导出的 4-bit 版本在同等速度下高出最多 2.4 分。12B 也因此成为这张芯片上能跑的最大模型：权重 11.31 GiB，16 并发下每秒输出 675 个 token，GSM8KGrade School Math 8K benchmark for evaluating mathematical reasoning capabilities of language models. 得分 0.964，BFCLBerkeley 函数调用基准（Berkeley Function Calling Leaderboard），用于评估模型调用工具/函数的能力。 工具调用得分 0.955。

代码仓库地址：https://github.com/xbill9/gemma4-dev/tree/main/jev-tpu-v5e1

---

核心结论

1. 重新打包后的 QAT 权重可在单张 TPU v5e 芯片（15.75 GiB HBM）上服务 Gemma 4 从 E2B 到 26B 的所有尺寸，到 12B 为止在 3880 条分类测试集上与 bf16 打平。

2. 12B w8a8emb4 是单张 v5e 芯片上能跑的最大模型：权重 11.31 GiB，16 并发下每秒输出 675 个 token，GSM8K 得分 0.964，BFCL 工具调用得分 0.955。

3. 不做二次取整的重打包版本比 Google 自己导出的 4-bit 版本测试集得分更高（E2B +2.4 分、E4B +1.3 分、12B +0.6 分），且两者速度相同。

4. 从 QAT 权重转 int8 优于从 bf16 转 int8：E2B 上测试集高 1.5 分，12B 上 BFCL 高 7.2 分（0.955 对 0.882）。

5. 启用 --kv-cache-dtype fp8 可将 12B 的 KV 缓存从 9728 token 翻倍至 18944 token，约 3600 token 长提示词下 16 并发吞吐量从 95 tok/s 提升至 157 tok/s（1.65 倍），首 token 延迟从 22.0s 降至 13.3s。

为什么要重新打包

一张 TPU v5e 芯片（v5litepod-1）有 15.75 GiB 的 HBM。Gemma 4 E4B 在 bf16 下是 14.9 GiB，12B 是 22.4 GiB。所以 E2B 以上的模型都必须用 4-bit 或 8-bit 权重才能塞进这张芯片。

Google 为每个 Gemma 4 尺寸都训练了 4-bit 版本，并以“未量化”的 bf16 检查点形式发布（-qat-q4_0-unquantized）。这些权重里每 32 个一组，已经落在 16 级的网格上。重新打包就是把这些值存成 vLLM 能直接服务的格式，不做二次取整：

v5e 原生支持 int8 乘 int8，所以在这张芯片上 int8 构建是最快的。

---

开始之前你需要准备

• 一个 Google Cloud 项目，在 us-west4-a 有 TPU v5e flex-start 配额，gcloud CLI 已登录
• 一个 Cloud Storage 桶，用来存检查点和结果
• Hugging Face token 存在 Secret Manager 里，名字叫 hf-token
• 克隆仓库：git clone https://github.com/xbill9/gemma4-dev

---

第一步：重新打包 QAT 权重

4-bit 重打包会恢复每组训练时的步长，把组存成 int4，每个值都保留它在网格上的训练位置：

huggingface-cli download google/gemma-4-12B-it-qat-q4_0-unquantized \
  --local-dir ~/models/gemma-4-12B-it-qat-q4_0-unquantized
python3 ../jev-tpu-31b/repack_q4_0.py repack ~/models/gemma-4-12B-it-qat-q4_0-unquantized \
  ~/models/gemma-4-12B-it-qat-q4_0-w4a16-ct

int8 构建用同样的 QAT 值，通过 ../jev-tpu-31b/w8a8_from_qat.py 逐通道转成 int8。已发布的构建在 Hugging Face 上，路径是 xbill9/gemma-4--it-qat-。

在 TPU 后端上服务这些模型，需要对 vLLM 的 tpu_inference 打三个补丁，放在 jev-tpu-v5e1/patches/ 里：JAX 路径上的 int8 W8A8 方法、在芯片上保持打包状态的 int4 嵌入表（只解包当前步读取的行）、以及 int4 的 lm_head。

---

第二步：在单张 v5e 芯片上服务

每次运行是一个 flex-start 排队资源，启动后给固定版本的 vLLM 镜像打补丁，依次服务一系列构建，然后上传所有结果：

gcloud alpha compute tpus queued-resources create jev-tpu-v5e1-$RUN --zone us-west4-a \
  --accelerator-type v5litepod-1 --runtime-version v2-alpha-tpuv5-lite \
  --node-id jev-tpu-v5e1-$RUN-node \
  --provisioning-model flex-start --max-run-duration 4h --valid-until-duration 2h \
  --metadata jev-code=<bundle>.tgz,jev-run=$RUN,jev-qr=jev-tpu-v5e1-$RUN \
  --metadata-from-file startup-script=tpu/startup_quant.sh,jev-arms=arms.txt,jev-patches=patches.txt

每个构建一行配置，指定检查点、要测什么、以及服务参数：

/work/models/gemma-4-12B-it-qat-w8a8-int8-emb4=12b-w8a8-emb4=read+load+gen=tools,gmu=0.92

READY /work/models/gemma-4-12B-it-qat-w8a8-int8-emb4 after 405s

---

单张 v5e 芯片能跑什么

测试集是 Bespoke Labs 基准集里的 3880 条公开记录，逐条与 bf16 配对比较（分数，95% 区间）。吞吐量是 1、4、16 并发下的每秒输出 token 数：

到 12B 为止，每个构建都跟 bf16 打平；26B 重打包版低了 1.1 分。E4B、12B 和 26B 的 bf16 版本塞不进这张芯片，所以它们的 bf16 参照跑在 v6e 上。

12B w8a8emb4 占 11.31 GiB 权重，剩下的全给 KV 缓存：

TPU KV cache size: 9,728 tokens, Maximum concurrency for 4,096 tokens per request: 2.38x

26B 混合专家重打包版以 13.58 GiB 服务，缓存 2176 token。

---

重打包版对比 Google 的 4-bit 导出

Google 的 -qat-w4a16-ct 导出会对 QAT 权重的每一组重新取整。重打包保留了原始值，测试集得分更高：

在同一台 VM 上依次服务，两者速度一样：

e2b-google load at concurrency 16: 1910.8 tok/s, range 1884.7 to 1910.8
e2b-repack load at concurrency 16: 1910.2 tok/s, range 1878.3 to 1910.9

16 并发下：E4B 是 1020 对 1009，12B 是 392 对 388。

---

第三步：评测数学和工具调用

gen_eval.py 对已服务的模型跑两个任务：

JEV_GEN_MAX_TOKENS=2048 python3 gen_eval.py run http://localhost:8000 "$MODEL" gsm8k out/gsm8k.jsonl
python3 gen_eval.py run http://localhost:8000 "$MODEL" bfcl_simple out/bfcl_simple.jsonl

12b-w8a8-emb4 gsm8k: {"task": "gsm8k", "max_tokens": 2048, "n": 1319, "right": 1272, "accuracy": 0.9644, "errors": 0, "truncated": 4, "wall_s": 584.8}

• GSM8K：全部 1319 道测试题，零样本思维链，贪心解码，最终数字匹配就算对。
• BFCL v3 simple：400 条记录，只提供一个工具，服务时加 --tool-call-parser gemma4，函数名和参数都正确才算对。

---

数学和工具调用表现如何

工具调用在每个尺寸上都稳住了：E2B 的每个重打包版都在 bf16 的 1.3 分以内。GSM8K 上 E2B 重打包版比 bf16 低 0.9 到 2.0 分，而 QAT 权重本身以 bf16 存储时也低了 1.3 分（−2.7 到 0.0），所以重打包最多只增加了 0.8 分的损失。E2B 上 4-bit 重打包的数学保持得最好。E4B 和 12B 在这张芯片上没有这两个任务的 bf16 参照。

---

int8 从 QAT 权重来，还是从 bf16 来

第三方 int8 构建直接从 bf16 版本取整，格式一样，速度一样。但从 QAT 权重构建的 int8 在 E2B 上测试集高了 1.5 分（+0.5 到 +2.6）。12B 上 BFCL 高了 7.2 分（0.955 对 0.882）：从 bf16 取整的版本在 400 条记录里有 16 条把字符串参数多包了一层引号，QAT 版本一条都没有。

---

第四步：长提示词

w4a16_client.py loadlong 会流式发送约 3600 token 的提示词，每条都有唯一前缀，所以没有一条能从缓存命中：

12b-w8a8-emb4 long 3584 at concurrency 16: 3611 prompt tokens, 94.7 out tok/s, 1430.0 total tok/s, ttft 21.96 s

E2B w8a8emb4 在长提示词下保持 bf16 的 1.38 倍速度，首 token 也更快。12B 的缓存一次大约只能装两个这种请求。

---

提示：fp8 KV 缓存让 12B 的空间翻倍

--kv-cache-dtype fp8 把缓存存成每值一字节：

TPU KV cache size: 18,944 tokens, Maximum concurrency for 4,096 tokens per request: 4.62x

单请求在两种缓存下速度一样。fp8 缓存代价是 0.7 个测试集分，GSM8K 和 BFCL 基本保住。

---

横向对比

---

那到底选哪个

想要单张 v5e 芯片上最强的模型，选 12B w8a8emb4，大量长提示词同时进来时加 --kv-cache-dtype fp8。它有自己的服务目录 tpu-vllm-v5e1-12b-w8a8emb4，带一个 MCP 服务器，在那里以 57 / 216 / 713 输出 token 每秒服务。E4B 选 w8a8emb4。E2B 要速度选 w8a8，要数学最接近 bf16 选 q4w4a16 重打包版。

---

清理

每次运行在最后一个构建完成后会自己删除排队资源：

gcloud alpha compute tpus queued-resources list --zone us-west4-a

---

总结

这篇文章的目标是搞清楚 Google 的 QAT Gemma 4 权重在单张 TPU v5e 芯片上能做什么。方案的关键在于把 QAT 值重新打包成 int4 和 int8 格式而不重新取整，以及对 vLLM 的 TPU 后端做三处补充。结果如下：

• 重打包版能在单张 v5e 芯片上服务 Gemma 4 从 E2B 到 26B 的所有尺寸，到 12B 为止测试集与 bf16 打平
• 12B w8a8emb4 以 11.31 GiB 塞进去，每秒输出 675 token，GSM8K 0.964，BFCL 0.955
• 重打包版测试集比 Google 的 4-bit 导出高最多 2.4 分，速度一样
• 从 QAT 权重转 int8 比从 bf16 转 int8 更好：E2B 高 1.5 个测试集分，12B 高 7.2 个 BFCL 分
• fp8 KV 缓存把 12B 的缓存翻倍到 18944 token，长提示词吞吐量提升 1.65 倍
• E2B 构建在 GSM8K 上比 bf16 低 0.9 到 2.0 分，大部分来自 QAT 训练本身
• 26B 以 2176 token 上下文服务，比 bf16 低 1.1 分

范围说明：每个构建都跑在单张 TPU v5e 芯片（v5litepod-1，flex-start）上，区域 us-west4-a，vLLM 版本 0.29.1rc1.dev468+g0b7f11a1e，镜像 vllm/vllm-tpu@sha256:19a1a052…，加上 jev-tpu-v5e1/patches/ 里的三个补丁，每个构建跑一次。E4B、12B 和 26B 的 bf16 测试集参照跑在 v6e 上，E4B 和 12B 在这张芯片上没有 GSM8K 或 BFCL 的 bf16 参照。吞吐量是三轮 256 输出 token 的中位数。配对是逐条记录的，差值只有在 95% 区间排除零时才算数。重打包的检查点是非官方的，基于 Google 在 Apache 2.0 下的发布衍生。部分分析和写作使用了 AI 辅助（Claude）；每个数字都来自已提交的输出文件。

在单张 TPU v5e 芯片上服务 Google QAT Gemma 4 的方案，通过增量式的逐步方法得到了验证。

---

参考链接

• 代码、补丁、逐条结果和日志：https://github.com/xbill9/gemma4-dev/tree/main/jev-tpu-v5e1
• 12B 服务目录：https://github.com/xbill9/gemma4-dev/tree/main/tpu-vllm-v5e1-12b-w8a8emb4
• 12B int8 加 int4 词表：https://huggingface.co/xbill9/gemma-4-12B-it-qat-w8a8-int8-emb4
• 12B 4-bit 重打包：https://huggingface.co/xbill9/gemma-4-12B-it-qat-q4_0-w4a16-ct
• Google QAT 源检查点：https://huggingface.co/google/gemma-4-12B-it-qat-q4_0-unquantized
• Google QAT 公告：https://blog.google/innovation-and-ai/technology/developers-tools/quantization-aware-training-gemma-4/
• GSM8K：https://github.com/openai/grade-school-math
• Berkeley Function Calling Leaderboard：https://huggingface.co/datasets/gorilla-llm/Berkeley-Function-Calling-Leaderboard
• Bespoke Labs 公开测试集：https://github.com/bespokelabsai/nimble/blob/0e67403/docs/PUBLIC_BENCHMARKS.md
• vLLM TPU 文档：https://docs.vllm.ai/projects/tpu/en/latest/
• Cloud TPU v5e：https://cloud.google.com/tpu/docs/v5e

常见问题（FAQ）

为什么重新打包的 Gemma 4 QAT 权重比 Google 官方 4-bit 导出效果更好？

因为重打包保留了 QAT 训练时的原始网格值，不做二次取整；而 Google 导出会对每组重新取整，引入误差。实测重打包版在分类任务上比官方导出高最多 2.4 分，速度相同。

单张 TPU v5e 芯片能运行 Gemma 4 的哪些尺寸？最大能跑多大？

单张 v5e 芯片可运行 Gemma 4 从 E2B 到 26B 的所有尺寸。其中 12B 是能跑的最大模型，权重 11.31 GiB，16 并发下每秒输出 675 个 token，GSM8K 得分 0.964。

重新打包 QAT 权重需要哪些准备工作？

需要：Google Cloud 项目并在 us-west4-a 有 TPU v5e flex-start 配额；gcloud CLI 已登录；一个 Cloud Storage 桶；Hugging Face token 存在 Secret Manager 中名为 hf-token；并克隆代码仓库。