Cloud TPU 上的 PyTorch

借助 PyTorch on Cloud TPU，您可以使用标准 PyTorch 语义、分布式 API 以及您已在使用的生态系统工具（包括 Hugging Face、vLLM、SGLang、llm-d、TorchTitan、TorchRec 和 Ray），在 Cloud TPU 硬件上训练、微调和运行开源模型。

该堆栈的基础是 TorchTPU，这是一个与 Meta 的 PyTorch 团队共同开发的开源 PyTorch 原生后端。TorchTPU 为机器学习生命周期的每个阶段提供库、框架和工具，涵盖核心运行时执行、自定义内核、生产推理、大规模训练和集群编排。

开始使用

• 在 TPU 上运行 PyTorch 模型：从 TorchTPU 和 TorchTPU 文档开始。
• 提供 LLM 和多模态模型：使用 vLLM-TorchTPU 或 SGLang-TorchTPU。
• 在 Kubernetes 舰队中路由请求：使用 llm-d。
• 预训练或微调基础模型：使用 TPU 上的 TorchTitan（或用于推荐模型的 TorchRec）。
• 扩缩数据处理、训练、RL 和服务流水线：使用 Ray。
• 使用编码辅助工具来加载、分析和调优模型：使用 TPU Agent Suite（预览版）。

TorchTPU：适用于 Cloud TPU 的 PyTorch 原生后端

TorchTPU 为 Cloud TPU 上的 PyTorch 提供设备后端、图编译器集成和自定义内核 API。TorchTPU 是与 PyTorch 团队合作开发的，可让您在 Eager 模式下以交互方式调试模型，使用 torch.compile 编译模型代码，并使用 torch.distributed (tpu_dist) 跨切片扩缩。

如需安装 TorchTPU 并在 Cloud TPU 上运行张量运算，请执行以下操作：

pip install torch_tpu

import torch

device = torch.device("tpu")

x = torch.randn((4, 4), device=device)
y = torch.ones((4, 4), device=device)
result = torch.matmul(x, y).cpu()

print(f"Executed on device: {x.device}")
print(f"Result shape on CPU: {result.shape}, sum: {result.sum():.4f}")

动机

虽然 Cloud TPU 历来与 JAX 和 TensorFlow 搭配使用，但以前在 TPU 上运行 PyTorch 需要进行延迟张量图记录 (PyTorch/XLA)，这通常意味着需要添加代码来标记可编译区域。TorchTPU 默认使用 PyTorch 的标准“每个设备一个进程”执行模型，以 Eager 模式运行操作，因此现有工作负载和库无需重写结构即可运行。

注意： TorchTPU 取代了 PyTorch/XLA。我们建议您使用 TorchTPU 启动新的工作负载，并将现有的 PyTorch/XLA 工作负载迁移到 TorchTPU。

设计

TorchTPU 为 PyTorch 带来了 PyTorch 原生 TPU 支持，其设计目标主要有三个：

• 标准 PyTorch 语义：每个设备运行一个进程，并保留 PyTorch 张量视图和就地突变，而无需 JAX 运行时依赖项。
• 默认即刻执行：标准 PyTorch 模型会立即在即刻模式下运行，以便进行交互式调试和逐步检查。
• 编译器和内核优化：除了逐个算子的即刻执行之外，TorchTPU 还提供即刻模式下的算子融合、torch.compile 与动态形状边界的集成（以减少重新编译）以及自定义内核 API。

对于图编译，torch.compile 会使用 TorchDynamo 捕获 PyTorch 计算图，并通过 OpenXLA 和 StableHLO 将其编译为 TPU 机器代码。集合操作在 torch.distributed 进程组级别 (tpu_dist) 运行。

主要优势

• 算子、模型和 Hugging Face 支持：TorchTPU 支持前向和后向传递中的标准 PyTorch (ATen) 算子，并持续针对热门 Hugging Face 模型进行测试。您可以通过传递 device_map="tpu" 在 Transformers v5 中部署模型，使用 Eager 模式、torch.compile 和张量并行运行 Diffusers 流水线，以及在 TensorBoard 和 XProf 中使用 torch.profiler 分析工作负载。
• 单芯片到多切片伸缩：tpu_dist 后端支持 DistributedDataParallel (DDP)、FullyShardedDataParallel (FSDP2) 和 DTensor。此外，torch.tpu.distributed.topology_aware_mesh 会将并行性维度映射到物理 TPU 互连。
• 可配置的即刻执行模式：提供三种即刻执行模式：
  • Strict（默认）：使用标准命令式 PyTorch 语义，单独且异步地分派运算符。
  • Debug：通过张量边界检查同步运行。
  • DeferAndFuse：融合 Eager 算子序列，以减少高带宽内存 (HBM) 流量。
• 全图编译和缓存：torch.compile(backend="tpu") 会将图编译为 TPU 可执行文件，并在进程、主机和作业重启之间缓存这些文件。
• 自定义 TPU 内核：您可以调用自定义 Pallas 内核（包括直接在芯片之间通信的内核）作为标准 PyTorch 函数，使用 Tokamax 库中的预构建内核，或使用 Helion 编写可移植的基于图块的内核。

如需了解详情，请参阅 TorchTPU 文档，在 PyTorch on TPU 讨论论坛中提问，或在 GitHub 上报告问题。

生产环境推理和服务

对于生产部署，Cloud TPU 上的 PyTorch 与开源推理引擎和 Kubernetes 原生分布式服务基础设施集成。

vLLM-TorchTPU 和 SGLang-TorchTPU：推理引擎

vLLM-TorchTPU 和 SGLang-TorchTPU 是适用于 vLLM 和 SGLang 推理引擎的 TorchTPU 后端。这两个后端都运行 Qwen、DeepSeek、GLM、Kimi、GPT-OSS 和 Gemma 4 等开放模型。如需了解详情，请参阅博文宣布在 TPU 上推出新的 vLLM 和 SGLang 推理引擎。

注意： vLLM-JAX 和 SGLang-JAX 仍会继续维护和支持基于 JAX 的服务工作流。

动机

提供推理、多模态和工具使用工作负载通常需要自定义模型定义、推测性解码方法和量化方案。在 TorchTPU 上构建 vLLM 和 SGLang 后端，可让您直接在 PyTorch 中自定义和调试服务堆栈。

设计

两个后端都保留了标准 vLLM 和 SGLang 调度器、KV 缓存管理器和 API 服务器，同时在底层使用 TorchTPU 执行和自定义 Pallas 算子。大多数模型都会运行其标准 PyTorch 模型定义，而计算密集型块会在编译后的计算图中调用 Pallas 内核，以便独立更新特定于硬件的内核。

主要优势

• 并行处理、量化和推测性解码：支持跨主机的张量、数据、专家、流水线和上下文并行处理，以及多模态（图片和视频）和嵌入输入。支持的推测解码方法包括多词元预测 (MTP)、EAGLE-3、DFlash、DSpark 和 n-gram 匹配。支持的量化格式包括 8 位 (FP8) KV 缓存以及 8 位 (FP8) 和 4 位（MXFP4 和 NVFP4）浮点模型检查点。
• 两层前缀缓存、工具调用和结构化输出：自动前缀缓存（包括混合注意力模型和循环模型）会将活跃前缀保留在 HBM 中，并将较旧的前缀卸载到主机内存，以避免在多轮请求中重新计算。内置了工具调用和结构化输出功能
• 使用 TPU Sync 重叠 KV 缓存传输：支持跨主机解耦预填充和解码服务（将提示处理和 token 生成分离到不同的工作器池），并使用 TPU Sync 进行主机 KV 卸载。TPU Sync 是一个开源 C++ 库，可直接在后台线程上读取 PyTorch 设备缓冲区，而无需通过主机 CPU 内存复制数据。

如需详细了解 vLLM-TorchTPU，请参阅 vLLM-TorchTPU 代码库，在 vLLM Slack 频道中提问，或在 GitHub 上报告 vLLM-TorchTPU 问题。

如需详细了解 SGLang-TorchTPU，请参阅 SGLang-TorchTPU 代码库，或在 GitHub 上报告 SGLang-TorchTPU 问题。

llm-d：Kubernetes 原生分布式推理

llm-d 是一个开源云原生计算基金会 (CNCF) 沙盒框架，用于在 Kubernetes 上进行分布式 LLM 推理，由 Red Hat、IBM、 Google Cloud以及由 500 多个贡献者和 35 个组织组成的社区共同开发。

动机

单实例推理引擎负责管理每个副本的执行，而生产集群必须处理流量高峰、多租户优先级、集群范围的上下文共享和零停机时间发布。llm-d 提供集群级路由和操作层，用于协调推理服务器集群。

设计

llm-d 在 Kubernetes 上的 vLLM 和 SGLang 等模型服务器前端运行。llm-d 提供以下功能：

• 根据前缀缓存状态和副本负载路由请求。
• 协调解耦的预填充和解码工作器池。
• 支持 MoE 分片和异步批量推理。
• 管理跨设备内存、主机内存和存储空间的分层 KV 缓存分流。

主要优势

• 多轮会话路由：将多轮对话路由到包含其缓存上下文的副本。轮次优先调度和有状态 Responses API 支持让智能体能够跨轮次重复使用上下文，从而减少重新计算。
• 多租户服务质量：优先级和公平性流控制可让对延迟时间敏感的请求和突发代理工作负载共享一个舰队，而不会出现资源匮乏。
• 集群操作和滚动更新：提供 v1 配置 API、路由器高可用性以及协调的滚动更新，将预填充和解码副本作为单个逻辑单元进行管理。
• GKE 上的原生 TorchTPU 应用：您可以在一个 llm-d 端点后面与其他加速器一起应用 Cloud TPU，并使用标准 Kubernetes 工具管理 vLLM-TorchTPU 和 SGLang-TorchTPU 副本。

如需了解详情，请在 llm-d Slack 频道中提问，或在 GitHub 上报告问题。

训练、微调和领域库

为了预训练、微调和调整大型模型，TorchTPU 提供了针对 TPU 硬件进行调优的参考训练配方和特定领域库。

TPU 上的 TorchTitan：大规模训练参考

TPU 上的 TorchTitan 扩展了 TorchTitan（用于生成式 AI 训练的 PyTorch 模块化参考平台），具有 TPU 分片传递、自定义内核和训练方案。如需了解详情，请参阅博文将 TorchTPU 集成到 TorchTitan 中并针对 TPU 优化模型。

动机

如需在多主机 TPU 切片上扩缩训练，需要跨芯片互连对参数进行分片，并使通信与计算重叠。Google 和 Meta 在 TorchTitan 中共同开发了 TPU 分片和编译传递，因此您无需设计自定义 TPU 分片策略或手动安排通信。

设计

模型是用单设备 PyTorch 编写的，不含加速器专用代码。 TorchTitan 会根据配置文件配置模型层，让您无需修改模型定义或训练循环，即可将各个层替换为自定义 Pallas 内核（例如 Splash Attention）。

TorchTitan 的 SimpleFSDP 实现将参数分片操作直接放置在编译后的计算图中，以便 XLA 编译器可以将设备通信与计算重叠，而全步计算图编译则将整个训练步骤捕获为单个 XLA 程序。对于多切片运行，混合分片 (HSDP) 会通过芯片间互连 (ICI) 在每个切片内进行分片，并通过数据中心网络 (DCN) 跨切片进行复制。此布局将跨切片流量限制为每步一次梯度缩减。

主要优势

• 预训练和 RL 后训练 recipe：包括密集 LLM（Llama 3、Qwen 3 和 MuseGlimmer 30B）、混合专家 (MoE) 模型（Qwen 3 235B，在多达 8,192 个 TPU7x 芯片上进行了验证）和扩散转换器 (FLUX.1) 的参考 recipe。对于强化学习训练后（例如群组相对策略优化 [GRPO]），VeRL（预览版）会将参数更新委托给 TorchTitan，并将更新后的权重推送到服务引擎，以生成训练轨迹（展开）。
• 预配置的编译器和并行设置：配方配置包括预配置的编译器和并行设置，因此您可以通过配置文件而不是代码更改来调整模型 FLOP 利用率 (MFU)。
• 新架构的初始配置：Hugging Face Transformers 建模后端通过 TorchTitan 训练器运行 Hugging Face Hub 模型检查点（包括没有专家并行性的 MoE 模型）。对于完整模型移植，TorchTitan 为工程师和编码智能体提供了分步入门技能。
• 独立组件重用：您可以端到端采用 TorchTitan，也可以将各个模块、自定义运算符和编译器标志导入现有训练代码库。

如需了解详情，请参阅 TorchTitan on TPU 入门指南，在 PyTorch 讨论论坛中提问，或在 GitHub 上报告问题。

TorchRec：推荐系统

TorchRec 是一个 PyTorch 领域库，为推荐模型提供分布式嵌入表和分片基元。在 Cloud TPU 上，TorchRec 会将稀疏嵌入表查找分流到专用 TPU SparseCore 处理器，同时在 TPU TensorCore 上运行密集矩阵运算。

如需详细了解 TorchRec，请参阅 TorchRec 文档，在 PyTorch on TPU 讨论论坛中提问，或在 GitHub 上报告问题。

集群编排和开发者工具

为了扩缩多节点流水线并自动执行模型初始配置，TorchTPU 生态系统与 Ray 集成，并提供可访问 TPU 文档的编码智能体扩展程序。

Ray：分布式训练、服务和强化学习

Ray 使用 Ray Train、Ray Serve、Ray Data、Ray Core 以及 verl 等强化学习框架，在集群中扩缩 Python 和 AI 工作负载。

动机

多切片训练和多阶段流水线（例如结合了数据预处理、训练和推出生成的 RL 工作流）需要自动切片预配、拓扑发现和容错工作器协调。torchrun 等独立启动实用程序不会预配 TPU 切片或协调多阶段 actor 流水线，因此 Ray 会管理 TorchTPU 工作负载的集群级调度和放置。

设计

当您提交指定 TPU 代系、拓扑和工作器数量的作业时，Ray 会预留完整的 TPU 切片，以便切片中的所有工作器一起调度，发现芯片互连地址，并为每个 TPU 设备启动一个 PyTorch 工作器。Ray Train 的 PyTorch 后端会在作业启动前自动初始化 tpu 设备和 tpu_dist 进程组。

主要优势

• 拓扑匹配的布置：将跨主机的 worker 布置与物理 TPU 世代和切片布局相匹配。
• 自动故障恢复：在 GKE 上的多个 TPU 切片中运行分布式训练作业，并在工作器或切片发生故障时自动恢复。
• 直接到 TPU 的数据预提取：Ray Data 的 iter_torch_batches 方法会将批次作为 PyTorch 张量直接预提取到 TPU 设备内存中。
• 多主机推理编排：管理 vLLM-TorchTPU 和 SGLang-TorchTPU 的多主机服务部署。

TPU Agent Suite：（预览版）适用于 TPU 的编码智能体

TPU Agent Suite 是一款开源扩展程序，适用于 Claude Code 和 Google Antigravity 等编码辅助（未来还将支持更多编码智能体），专注于在 TPU v6e 和 TPU7x 上进行 Transformer 推理（密集大语言模型和 MoE 模型）。

动机

通用编码辅助通常缺乏有关 TPU 硬件拓扑、XLA 编译器标志、XProf 遥测和 Pallas 内核限制的最新上下文。TPU Agent Suite 可让您的助理访问官方 TPU 文档和结构化工作流。这有助于防止出现无效的标志建议，并自动执行迭代性能分析和基准测试。

设计

该套件会向您现有的编码辅助添加一个 TPU 知识库和六个特定于任务的子智能体。路由智能体不会自行运行命令，而是将每个提示委托给适当的子智能体。在 TPU 虚拟机或 GKE 集群上执行任何命令都需要获得您的批准，并且调优循环会在返回结果之前运行有限次数的记录尝试。

主要优势

• 模型加载和准确率验证：TPU 模型加载智能体可解决不受支持的算子问题，选择服务路径（vLLM-TorchTPU 或标准 PyTorch），并将逐层输出与 CPU 参考模型进行比较，以验证数值准确率。
• 崩溃和瓶颈诊断：TPU 诊断智能体可分析崩溃日志、检查 XProf 性能分析文件和 XLA 编译器轨迹，并找出运行时瓶颈。
• 内核、编译器和分片优化：TPU 内核开发智能体可协助处理自定义 Pallas 内核，TPU 性能优化智能体可对 XLA 编译器标志组合进行基准测试，TPU 分片智能体可规划多切片并行布局。
• 文档查找和引用：TPU 知识库智能体可回答技术问题，并明确引用官方 TPU 文档和教程，同时会在主题未涵盖时进行说明。

如需了解详情，请参阅 TPU Agent Suite 文档，在 GitHub 讨论中分享反馈和提出问题，或在 GitHub 上报告问题。

后续步骤

• 开始阅读 TorchTPU 文档。
• 试用 TPU 上的 TorchTitan 配方。
• 使用 vLLM-TorchTPU 和 SGLang-TorchTPU 部署推理工作负载。
• 在 PyTorch on TPU 讨论论坛中提问。
• 如需与 Google 产品和工程团队讨论特定用例，请发送电子邮件至 [e-mail]。