在 Colab 上把 PyTorch 跑上 TPU：五个 notebook 的实测
来源：dev.to — 2026-10-24
📋 概述
五个 Colab notebook 在单个 v5e 芯片（16GB 高带宽显存）上跑 Gemma 4 E2B：Hugging Face 的 generate 只有每秒 5.0 到 5.3 个 token，改成编译过的 decode step 提到 31.3 到 33.4，再加 4-bit 每层表能到 50.2 到 50.4；换成环形缓存速度 38.6，但能吃下 131,328 token 的完整上下文。同芯片的 JAX 引擎是 130.8 到 132.9，也就是说编译加 4-bit 表之后 PyTorch 仍落后 2.6 倍，而优化之前差了 14.8 倍。所有方案在同一个 19 token 提示、64 个贪心 token 下写出的 56 token 答案一致。
🔑 核心要点
- 编译一次 decode step 就把吞吐从每秒 5 token 提到 50 token
- 环形缓存让上下文撑到 131,328 token，代价是降到 38.7 tok/s
- 同芯片 JAX 仍有 2.6 倍优势，优化之前是 14.8 倍
- 所有方案在同样的提示下输出同一个 56 token 答案
- 全部在 Colab v5e-1 单芯片上复现，每一步一个 notebook
💡 金句
“One compiled decode step takes PyTorch from 5 to 50 tokens per second, and a ring cache takes it to the full context.”
译：一次编译过的解码步骤把 PyTorch 从每秒 5 个 token 提到 50 个，而环形缓存把它带到完整上下文。
← 返回 dev.to 首页