AI 算力圈最现实的焦虑,常常是模型能不能顺利跑起来。

Google 开源了 TorchTPU。

它给 PyTorch 提供了一个面向 Google TPU 的原生后端,让 PyTorch 模型和操作可以跑在 TPU 设备上。

一句话,Google 想把 PyTorch 开发者请进 TPU 世界。

一个 pip 包,背后是一整套 TPU 入口

一个 pip 包,把 PyTorch 和 TPU 接到了一起。

TorchTPU 的项目说明很直白。

它包含面向 TPU 的 PyTorch ATen kernel 实现,也包含用于优化 XLA 计算执行的编译缓存。

它支持分布式训练,覆盖数据并行和张量并行。

它还给出模型推理和训练示例,并提供调试、基准测试工具。

当前公开版本是 PyPI 上的 0.2.0,支持 Python 3.11+。

安装命令里有一个细节很关键:需要用 --extra-index-url 安装 CPU 版 PyTorch,避免默认装上 CUDA 版本。

说实话,这种细节非常工程化。

它说明项目目标很明确:让开发者少踩坑,尽快把环境跑起来。

仓库还准备了用户指南和交互式 notebooks。

Notebooks 使用 Marimo,类似 Jupyter,官方建议在 TPU VM 上运行,v6e GCE 实例最容易上手。

对开发者来说,仓库连 Bazel 构建、Bazelisk、clangd 环境都写得很细。

这类仓库真正锋利的地方,是把复杂硬件适配压缩成一个可安装入口。

模型清单很直白:Llama、Qwen、UNet 都被点名

Google 把模型清单摆出来了。

素材列出的示例包括 Llama、Qwen、ResNet、DLRM、minGPT、UNet。

这份清单很有针对性。

Llama 和 Qwen 指向大语言模型。

ResNet、UNet 指向视觉模型。

DLRM 指向推荐系统。

minGPT 则让人一眼联想到小规模 GPT 训练。

Google 等于把常见任务摆到桌面上:这些模型,可以拿来看,也可以拿来跑。

老实讲,这种清单比一句性能口号更有说服力。

开发者通常不信 PPT。

他们更信能运行的代码、能复现的示例、能看懂的文档。

XLA 才是真正的暗线

真正的暗线,藏在 XLA 里。

TorchTPU 多次提到 XLA。

它提供编译缓存,用于优化 XLA 计算执行。

它也集成 torch.compile,支持图模式执行。

这里的想象空间很大。

如果 PyTorch 模型通过 torch.compile 进入图模式,再由 XLA 编译到 TPU,性能优化就不再只靠人工逐个适配。

它更像一条自动化流水线。

热门评论里有人写:

XLA is pretty good + potentially a better compiler infrastructure than inductor which will lead to improved performance on all hardware.

这句话有点激进。

但它点中了重点。

TorchTPU 表面是给 TPU 做适配,实际也把 XLA 推到 PyTorch 开发者面前。

最有意思的反转:TensorFlow 被拖出来鞭尸

评论区最扎眼的一句,关于 TensorFlow。

热门评论里写道:

Also I guess this is the final nail in the coffin of Tensorflow.

这话当然夸张。

但它能出现,本身就说明一个变化。

Google 正在认真把 PyTorch 生态接进自家 TPU。

评论区仍有人在问 TPU 在哪、驱动在哪。

现在 Google 给出的答案,是更贴近主流开发者的 PyTorch 后端。

这点我挺认同。

硬件再强,如果开发者不愿意迁移,也只是机房里的昂贵设备。

有意思的是,Google 这次没有只讲自家技术栈。

它把门打开,让 PyTorch 模型直接进来。

不止 Google,大家都在抢 PyTorch 后端

TorchTPU 并不是评论区里唯一的主角。

有人说,自己正在为 Tenstorrent 加速器做开源 libtpu 实现,并称已经能运行大部分 torchtpu。

也有人提到,Amazon 为 AWS Trainium 做了 torch-neuronx。

这些信息放在一起,味道就变了。

从 Tenstorrent 到 Amazon Trainium,评论区已经给出同一个信号:想让硬件被用起来,先做好 PyTorch 后端。

评论区里,大家讨论的重点已经不只是硬件参数。

后端在哪,编译器在哪,示例模型在哪,开发者能不能今天安装、明天跑通。

框架层的竞争,已经开始抢开发者入口。

真正卡脖子的,可能不止代码,还有 TPU 在哪

评论区也有人问得很直接:

Pardon my ignorance, but who else besides Google has TPUs?

还有人问:

Is the Linux driver available?

这些问题很关键。

开源后端降低了软件门槛,但没有自动解决硬件可得性。

热门评论里有人说,Google 在通过云服务,也通过精选合作伙伴出售 TPU 铲子。

也有人提到,Nebius 已经从炒作峰值跌了近一半。

这提醒市场:算力叙事很热,供给、生态和商业化节奏也会反过来影响估值。

换句话说,TorchTPU 让 TPU 看起来更亲近 PyTorch。

但 TPU 能不能从 Google 体系走向更大市场,还要看硬件、云服务和生态伙伴怎么接。

Google 开源 TorchTPU,表面是给 PyTorch 加一个 TPU 后端。

更深一层,它是在把 TPU、XLA 和主流模型生态拧到一起。

当芯片竞争从参数表打到框架层,真正的问题也变了:开发者会不会为了一个后端,换掉整套算力习惯?

【锐评】:Google 开源 TorchTPU,表面是技术适配,实际是抢 PyTorch 开发者入口。

参考链接:
https://github.com/google-pytorch/torch_tpu