TorchTPU 现已开源:在 TPU 上运行原生 PyTo... 笔记

TorchTPU 现已开源:在 TPU 上运行原生 PyTorch

PyTorch Conference 2026 推出了用于在 Google Cloud 及其他基础设施上运行 PyTorch 的新产品与集成方案。TorchTPU 作为 Google TPU 的原生 PyTorch 栈,现已在 GitHub 上开源。这使得用户能够借助 vLLM、SGLang 和 Hugging Face 等熟悉工具,在 TPU 上实现模型的无缝训练、推理与服务部署。该集成为用户提供包括 TPU、GPU 和 CPU 在内的多种计算选项,并保持一致的操作模型。TorchTPU 支持以最小代码改动和熟悉的 API 在 TPU 上运行现有的 PyTorch 工作负载;它支持 eager 模式调试,并通过 torch.compile 优化 TPU 利用率与内存使用。Hugging Face 模型现可原生运行于 TorchTPU,通过设置 device_map="tpu" 简化部署。推理优化通过 TorchTPU 与服务引擎的集成以及跨 TPU 和 GPU 的舰队路由实现。vLLM 在 TPU 上提供高吞吐服务与内存效率,支持连续批处理与量化等特性。SGLang 通过在请求间复用 KV 缓存,加速 TPU 上的代理式与结构化工作负载。llm-d 1.0 作为 CNCF Sandbox 项目,将 PyTorch 服务引擎与 Kubernetes 连接,以支持分布式 LLM 推理。TorchTitan 提供原生的 PyTorch 分布式高性能大规模训练,已在数千个 TPU 上得到验证。VeRL 这一流行的强化学习框架现已包含 TPU 后端,支持可扩展的后训练与强化学习工作负载。Ray 已添加原生 TorchTPU 支持,可在单一工作流中编排跨多种加速器的 AI 流水线。LiteRT 作为 Google 的端侧 AI 推理解决方案,现已加入 PyTorch 生态系统,使 PyTorch 模型能够在多样化设备上运行。TPU 代理套件也已在 GitHub 上进入早期预览阶段,用于自动化 PyTorch 模型在 TPU 上的接入、诊断与性能调优。
CdXz5zHNQW_YCinlopqHk.jpeg