Google已开源TorchTPU。这是一层软件，可让开发者在Google张量处理单元（TPU）上运行PyTorch——这是一款广泛用于构建和运行AI模型的工具包。TorchTPU由Google与Meta的PyTorch团队共同开发，旨在通过相对少量的代码改动，将现有工作负载迁移到TPU硬件上，而不必重写代码。
Google表示，TorchTPU支持两种执行方式：一是即时执行，按调用顺序运行操作，便于调试；二是面向生产工作负载的编译执行。开发者可以将设备设置从“cuda”改为“tpu”，并使用TPU分布式后端；Google也推荐将PyTorch的加速器检测API作为另一种选择。TorchTPU已在GitHub上发布，但其路线图仍包括支持动态输入形状、提升编译速度、扩大算子覆盖范围，以及支持强化学习和量化等工作。
对开发者来说有哪些变化？
此次发布让TPU能够接入许多PyTorch开发者已经在使用的软件。Google表示，TorchTPU可与vLLM和SGLang等推理服务工具，以及TorchTitan、Ray和TorchRec等训练与编排项目配合使用。Hugging Face表示，在Transformers v5中，开发者可以通过设置device_map="tpu"来选择TPU。其Diffusers库也支持在TPU上运行，包括编译模式和将任务分配到多个加速器上执行。
这与Google此前通过PyTorch/XLA在TPU上运行PyTorch的做法不同。Google表示，PyTorch/XLA需要对代码做更大幅度的修改、编译时间较长，也不支持PyTorch的即时执行模式。据InfoQ报道，Google搭建了一套测试流程，评估约5,000个Hugging Face模型，覆盖500多种架构。这表明测试范围较广，但并不意味着所有模型或工作负载都能不经修改地运行。
与其他硬件选项相比如何？
Google表示，TorchTPU的DeferAndFuse模式——一种结合即时执行与编译执行的混合模式——相比Strict Eager模式，速度至少快50%，同时还能降低加速器的峰值内存用量。在另一项测试中，Google使用标准PyTorch运行了一个拥有300亿参数的Muse Glimmer模型。InfoQ报道称，在序列长度为5,120时，该模型在一颗TPUv7芯片上的推理吞吐量为每秒25,700个token；序列长度为10,240时，则为每秒22,400个token。使用四颗TPUv7芯片训练时，这两个序列长度下的速度分别达到每秒5,600和4,600个token。这些数据来自Google，并非独立对比结果。
TechCrunch将TorchTPU与Amazon的TorchNeuron相提并论。TorchNeuron于同一周发布，旨在让PyTorch工作负载能够运行在Amazon的Trainium芯片上。两者都希望让用户更容易使用Nvidia GPU以外的硬件。TorchTPU的发布为PyTorch用户增加了一个硬件选项，但这本身并不能证明迁移工作负载会降低成本，或胜过开发者当前使用的配置。
为什么重要
对于已经在使用PyTorch的开发者来说，TorchTPU有望让他们更轻松地在Google TPU上测试或部署工作负载，无需从头重构。迁移对特定模型能否提升性价比或速度，目前仍不明确；文中引用的性能数据均来自Google自身测试。
