Gemma 4 E2B in Pure JAX on a Colab TPU: Google's 4-Bit Export Against an Exact Repack

本文为原文前 6,000 字符的节选翻译，完整内容请查看原文。

This article provides a step by step guide to a Colab notebook that serves Gemma 4 E2B on a single TPU v5e chip with a pure-JAX engine and compares two 4-bit builds of the same model against the weights Google trained. Every number below was measured in the notebook on a Colab v5e-1 runtime, and the executed notebook is committed. Google ships E2B in 4 bits as gemma-4-E2B-it-qat-w4a16-ct. Its export rounds every weight a second time, onto a grid the model never trained on. A repack that stores the trained grid instead lands 342.6 times closer to the original model in next-token predictions, matches its top token 99.29% of the time against 85.96%, writes the same 128-token story word for word, runs at the same speed and downloads 0.8 GB less.

本文提供了一份 Colab 笔记本的分步指南，该笔记本在单个 TPU v5e 芯片上使用纯 JAX 引擎运行 Gemma 4 E2B 模型，并将该模型的两个 4 位版本与谷歌训练的权重进行了对比。以下所有数据均在 Colab v5e-1 运行时环境中测量得出，且已提交了执行后的笔记本。谷歌发布的 4 位版本为 gemma-4-E2B-it-qat-w4a16-ct。其导出过程对每个权重进行了二次舍入，使其落入模型从未训练过的网格上。相比之下，一种存储了训练网格的重打包版本在预测下一个 token 时，与原始模型的距离缩短了 342.6 倍，其首选 token 的匹配率从 85.96% 提升至 99.29%，能够逐字写出相同的 128 个 token 的故事，运行速度相同，且下载体积减少了 0.8 GB。

The Marathon: This notebook is an entry in the AI GDE Marathon: JAX on TPU Tutorial, a series of Apache 2.0 Colab notebooks covering JAX on the TPU backend. Colab’s TPU runtimes are the same v5e and v6e silicon the source measurements came from, so every cell measures on the reader’s chip and prints the source result beside it as context.

马拉松：本笔记本是“AI GDE 马拉松：JAX on TPU 教程”系列的一部分，该系列包含一系列基于 Apache 2.0 协议的 Colab 笔记本，涵盖了 TPU 后端上的 JAX 应用。Colab 的 TPU 运行时环境与原始测量所使用的 v5e 和 v6e 芯片相同，因此每个单元格都会在读者的芯片上进行测量，并将原始结果打印在旁边作为参考。

Why Two 4-Bit Builds of One Model? Gemma 4 E2B was trained with quantization-aware training (QAT): during training every weight was held on a 4-bit grid, one scale per group of 32 values, so the model learned to work with exactly those values. Google publishes the result twice: Build Hugging Face What it stores QAT reference google/gemma-4-E2B-it-qat-q4_0-unquantized the trained grid values, in bf16 Google’s export google/gemma-4-E2B-it-qat-w4a16-ct int4, each group re-rounded with step = largest weight ÷ 7.5 The repack xbill9/gemma-4-E2B-it-qat-q4_0-w4a16-ct int4, the trained levels and the trained step. Both 4-bit builds use the same compressed-tensors W4A16 format and the same shapes, so any loader that reads one reads the other. The difference is which numbers are inside.

为什么同一个模型会有两个 4 位版本？Gemma 4 E2B 采用了量化感知训练 (QAT)：在训练过程中，每个权重都保持在 4 位网格上，每 32 个值使用一个缩放比例，因此模型学会了精确地使用这些值。谷歌发布了两次结果：QAT 参考版本（google/gemma-4-E2B-it-qat-q4_0-unquantized）存储了 bf16 格式的训练网格值；谷歌导出版本（google/gemma-4-E2B-it-qat-w4a16-ct）为 int4 格式，每组权重通过“步长 = 最大权重 ÷ 7.5”重新舍入；重打包版本（xbill9/gemma-4-E2B-it-qat-q4_0-w4a16-ct）为 int4 格式，存储了训练后的层级和步长。这两个 4 位版本使用相同的压缩张量 W4A16 格式和相同的形状，因此任何能读取其中一个的加载器都能读取另一个。区别在于内部存储的数值不同。

At This Point You Should Have… A Google account with Colab access to a TPU runtime. About 40 minutes, most of it downloading 26 GB and compiling. No Hugging Face token is needed: all three checkpoints are public and ungated, and the engine is a public GitHub repo.

此时你应该准备好：一个拥有 TPU 运行时访问权限的 Google Colab 账户。大约 40 分钟的时间，其中大部分用于下载 26 GB 的数据和编译。无需 Hugging Face 令牌：所有三个检查点都是公开且无需授权的，引擎也是一个公开的 GitHub 仓库。

Step 1 — Open the Notebook on a TPU: Open the notebook from the link above, then Runtime → Change runtime type → v5e-1 TPU, then Runtime → Run all. Cell 1.1 records the chip: JAX 0.7.2 1 device(s): TPU v5 lite (tpu) HBM 16.91 GB host RAM 50.5 GB free disk 195 GB. Cell 1.2 clones the engine, tpu-jax, at a pinned commit. It is a Gemma 4 E2B decoder in pure JAX with no PyTorch in the path, and its 4-bit path unpacks each weight with plain XLA operations before multiplying.

第一步 — 在 TPU 上打开笔记本：通过上述链接打开笔记本，然后选择“运行时”→“更改运行时类型”→“v5e-1 TPU”，接着选择“运行时”→“全部运行”。单元格 1.1 记录了芯片信息：JAX 0.7.2，1 个设备：TPU v5 lite (tpu)，HBM 16.91 GB，主机内存 50.5 GB，可用磁盘空间 195 GB。单元格 1.2 克隆了引擎 tpu-jax 的指定提交版本。这是一个纯 JAX 实现的 Gemma 4 E2B 解码器，路径中不包含 PyTorch，其 4 位路径在乘法运算前使用普通的 XLA 操作解压每个权重。

Step 2 — Download the Three Checkpoints: Each checkpoint is pinned to a Hub revision, so the bytes you download are the bytes measured here. Cell 2.2 then indexes every tensor: qat 1951 tensors 0 packed int4 lm_head stored: False; stock 2504 tensors 276 packed int4 lm_head stored: True; repack 2503 tensors 276 packed int4 lm_head stored: False. stock lm_head identical to embed_tokens: True (0.805 GB); stock minus repack on disk: 0.805 GB. Both 4-bit builds pack the same 276 linear layers. Google’s export also stores lm_head.weight, a byte-for-byte copy of the token embedding, although the config ties the two. That one tensor is the whole 0.8 GB difference in download size.

第二步 — 下载三个检查点：每个检查点都固定在特定的 Hub 版本上，因此你下载的字节数与此处测量的一致。单元格 2.2 随后索引了每个张量：QAT 版本有 1951 个张量，0 个打包的 int4，lm_head 未存储；Stock 版本有 2504 个张量，276 个打包的 int4，lm_head 已存储；重打包版本有 2503 个张量，276 个打包的 int4，lm_head 未存储。Stock 版本的 lm_head 与 embed_tokens 完全相同（0.805 GB）；Stock 版本与重打包版本在磁盘上的大小差为 0.805 GB。两个 4 位版本都打包了相同的 276 个线性层。谷歌的导出版本还存储了 lm_head.weight，这是 token 嵌入的逐字节副本，尽管配置中已将两者绑定。正是这一个张量造成了 0.8 GB 的下载大小差异。

Step 3 — Are the Weights the Trained Ones? Cell 3.1 reads 24 tensors from the first, middle and last layers, every projection kind and 5.7 million groups of 32, from all three files. It unpacks each 4-bit build with the engine’s own function and compares the result with the QAT values. Google’s export uses largest weight ÷ 7.5 as the step in every group, and its values sit 6.7% from the trained ones. The repack sits 0.19% away, three values in four are bit-for-bit the trained value, and the rest differ only by the bf16 rounding of the stored step. The last line explains why the trained step has to be stored. If a group’s largest weight sits on level m of the trained step d, the ÷ 7.5 rule gives a step of m·d / 7.5, which never equals d because m is a whole number from 1 to 8. The largest weight sits on level 7 in 39.5% of groups and level 8 in the other 60.5%, so a fixed ÷ 8 rule would be right for only some of them.

第三步 — 权重是训练时的权重吗？单元格 3.1 从三个文件中读取了第一层、中间层和最后一层的 24 个张量，涵盖了所有投影类型以及 570 万组（每组 32 个值）。它使用引擎自带的函数解压每个 4 位版本，并将结果与 QAT 值进行比较。谷歌的导出版本在每组中使用“最大权重 ÷ 7.5”作为步长，其数值与训练值相差 6.7%。重打包版本仅相差 0.19%，四分之三的数值与训练值完全一致，其余部分仅因存储步长的 bf16 舍入而有所不同。最后一行解释了为什么要存储训练步长。如果某组的最大权重位于训练步长 d 的第 m 级，那么 ÷ 7.5 的规则给出的步长为 m·d / 7.5，由于 m 是 1 到 8 之间的整数，该值永远不会等于 d。最大权重在 39.5% 的组中位于第 7 级，在其余 60.5% 的组中位于第 8 级，因此固定的 ÷ 8 规则仅对其中一部分有效。

Step 4 — Does the Difference Reach the Output? Cells 4.2 to 4.4 load each model in turn, feed all three the same 8 WikiText-2 sequences of 512 tokens, and keep each model’s predicted next-token distribution at every position. Cell 4.5 compares each 4-bit build with the QAT model on the same 4,096 positions. The repack is closer to the QAT model at every one of the 4,096 positions, and the gap holds on every sequence: the stock build’s best sequence is still 200 times further off than the repack’s worst. Google’s export also predicts the text 6.4% worse than the model it was made from.

第四步 — 差异会影响输出吗？单元格 4.2 到 4.4 依次加载每个模型，向三个模型输入相同的 8 个 WikiText-2 序列（每个 512 个 token），并保留每个模型在每个位置预测的下一个 token 分布。单元格 4.5 在相同的 4,096 个位置上比较了每个 4 位版本与 QAT 模型。在所有 4,096 个位置上，重打包版本都比 Stock 版本更接近 QAT 模型，且这一差距在每个序列中都成立：Stock 版本表现最好的序列与 QAT 模型的距离，仍然是重打包版本表现最差序列距离的 200 倍。此外，谷歌的导出版本在文本预测上的表现比其原始模型差 6.4%。

Step 5 — Speed, Memory and What Each One Writes: Cell 5.1 decodes a fixed prompt, “The history of the Roman Empire”, for 128 greedy tokens, once to compile and three times timed.

第五步 — 速度、内存及输出内容：单元格 5.1 对固定提示词“The history of the Roman Empire”进行 128 个贪婪 token 的解码，运行一次用于编译，并进行三次计时测试。