Hugging Face 联合 Google 与 Meta PyTorch 团队推出开源后端 TorchTPU,使 PyTorch 模型能原生识别并在 Google Cloud TPU 上运行,简化从 GPU 到 TPU 的迁移流程。该项目兼顾即时执行与 torch.compile 编译优化,已验证数千个 Hugging Face 模型但性能与稳定性仍需用户在目标 TPU 上检验。
Hugging Face于2026年10月20日发布技术文章,介绍与Google、Meta PyTorch团队合作推进的TorchTPU。这个开源后端把Google Cloud TPU注册为PyTorch可以直接识别的设备,使开发者能够使用较为熟悉的PyTorch接口运行Hugging Face模型,而不必像过去那样围绕PyTorch/XLA重写大量代码。
这项变化的重点不只是“支持另一种芯片”,而是试图把TPU纳入PyTorch开发者已经习惯的工作流:模型可以在即时执行模式下运行,也可以通过torch.compile进行编译和优化;分布式训练则继续使用PyTorch的相关接口。对长期以GPU为默认目标的开源模型生态来说,硬件迁移的门槛因此可能进一步下降。
从PyTorch/XLA到原生TPU后端
TPU是Google为机器学习矩阵运算设计的专用加速器,既被Google用于训练和服务Gemini等模型,也通过Google Cloud向外部开发者提供。过去,PyTorch用户在TPU上通常要依赖PyTorch/XLA。这个路径能够工作,但建立在惰性张量和图记录机制之上,代码执行、调试以及图重新编译的行为,与普通PyTorch的即时执行模式存在明显差异。
Hugging Face在《Running Native PyTorch on TPUs with Zero Code Changes》中称,TorchTPU的目标是让用户尽量只改变设备配置,就能把模型迁移到TPU。文章给出的示例使用device_map="tpu"加载模型,并把输入张量放到"tpu"设备;Transformers、Diffusers和Kernels等Hugging Face生态组件也被纳入这一支持范围。
不过,“无需重写”更准确地说是迁移入口变得简单,而不是所有模型都能在任何条件下获得与GPU相同的性能。模型形状、算子覆盖率、批处理方式、分片策略和编译设置,仍会影响最终结果。
即时执行与编译优化并存
根据Google Cloud的PyTorch on Cloud TPU文档,TorchTPU提供标准的torch.device("tpu")接口。用户可以先以即时执行模式运行代码,用于交互式调试和逐步检查;在需要稳定吞吐量时,则可以使用torch.compile,由XLA编译器进一步优化计算图。
TorchTPU的技术路线并不是把PyTorch强行改造成JAX。其设计目标包括不在运行时依赖JAX、保留PyTorch的张量语义和即时执行方式,并在底层将PyTorch的ATen操作降低到StableHLO。项目还尝试在即时执行过程中进行算子融合,以减少逐算子执行在加速器上产生的开销,同时通过动态形状边界减少不必要的重新编译。
这套设计也延伸到分布式训练。TorchTPU在PyTorch的ProcessGroup等较低层处理集合通信,使DDP、FSDP和DTensor等上层机制有机会直接复用,而不是要求开发者重新学习一套TPU专用的分片模型。
Hugging Face模型覆盖与性能数据仍需区分看待
TorchTPU的开源代码库已经公开,采用Apache 2.0许可证。Google和合作团队表示,项目已经对约5,000个下载量较高的Hugging Face模型进行开箱运行验证,覆盖Llama、DeepSeek、Mistral、Gemma、Qwen等语言模型,以及Whisper、CLIP、ViT和其他视觉、语音模型。
这些数字说明项目正在扩大兼容性范围,但并不等于每个模型都已经适合生产环境。验证一个模型能够运行,与在真实业务中达到稳定的吞吐量、延迟和成本目标,是两个不同的问题。Hugging Face在另一条硬件后端路径上推出的TorchNeuron公开测试也体现了类似方向:通过统一的PyTorch入口,减少开发者针对不同加速器维护专用代码的需要,但硬件差异并不会因此消失。
关于性能,TorchTPU团队在技术论文中称,项目支持即时和编译模式、约450个ATen操作,以及多主机推理和TorchTitan训练工作负载。论文还提到,TorchTPU已经用于部分Google Cloud客户的训练和推理场景。由于这些结果主要来自项目方的测试和部署经验,读者应把它们视为工程进展和厂商披露,而不是适用于所有模型的独立基准结论。
TorchTitan、vLLM和SGLang加入同一软件栈
Google与Meta还在TorchTitan集成文章中介绍了面向大模型训练的TPU适配版本。相关工作包括使用Pallas编写的注意力、损失和专家层内核,以及通过SimpleFSDP和Graph Trainer把参数分片、前向、反向和优化步骤更完整地交给编译器处理。
在推理侧,Google开发者博客宣布了基于TorchTPU的vLLM-TorchTPU和SGLang-TorchTPU。这些引擎面向多轮和智能体工作负载,支持前缀缓存、KV缓存卸载、长上下文、分离式prefill与decode,以及视觉语言模型等场景。Google称,相关项目已经在部分预览客户中运行,并通过公开的InferenceX项目披露性能结果;其中的成本和吞吐量数字仍属于特定硬件、模型和配置下的项目方或合作方数据。
TPU能否真正成为PyTorch的另一种标准设备
TorchTPU的战略意义,在于它同时面对两个长期存在的障碍:一方面,TPU拥有面向大规模机器学习的硬件和云基础设施;另一方面,大量研究人员和企业团队已经围绕PyTorch建立了模型、训练脚本和部署工具。若每一次更换加速器都要求重写这些软件,硬件选择就会受到生态锁定影响。
TorchTPU试图把竞争焦点从“开发者是否愿意学习另一套框架”转向“不同硬件后端能否在同一套PyTorch接口下竞争”。这并不意味着GPU、TPU和其他加速器会拥有相同的性能、价格或算子覆盖,但它可能降低开发者尝试替代硬件的初始成本。
下一阶段的关键,将是TorchTPU能否继续扩大模型和算子覆盖,能否让编译缓存、分布式通信和调试体验在生产环境中保持稳定,以及独立用户能否复现项目方披露的效率结果。对开发者而言,最现实的做法仍是先在目标TPU实例上验证自己的模型和工作负载,再判断“只改设备配置”是否足以满足实际需求。