Google 与 Meta 发布 TorchTPU：PyTorch 原生 TPU 后端

Google 与 Meta 发布 TorchTPU：PyTorch 原生 TPU 后端 · AIHOT

PyTorch：Blog

·

The TorchTPU Teamâ€

2026-10-21 00:00

2 小时前

精选

AI 评分

70

Google 与 Meta 发布 TorchTPU：PyTorch 原生 TPU 后端

Introducing TorchTPU

AI 导读

Google 与 Meta 发布开源项目 TorchTPU，为 Google TPU 提供原生 PyTorch 后端，模型迁移只需把 .to("cuda") 改为 .to("tpu")。

推荐理由

Google 与 Meta 官方详解 TorchTPU çš„æž¶æž„è®¾è®¡ã€ç²¾åº¦ç­–ç•¥å’Œåˆ†å¸ƒå¼æ”¯æŒï¼Œè¯»è€ å¯æ®æ­¤è¯„ä¼°æŠŠçŽ°æœ‰ PyTorch 工作流迁到 TPU 的改动成本。

另外 1 家也报道了

正文 · AI 翻译

è¯‘æ–‡å°šä¸å®Œæ•´ï¼Œå®Œæ•´å† å®¹è¯·åˆ‡æ¢åˆ°åŽŸæ–‡ã€‚

TL;DR

TorchTPU 是 Google 与 Meta 合作开发的开源项目，为 Google TPU å¼•å ¥äº†åŽŸç”Ÿä¸”ç¬¦åˆè®¾è®¡åŽŸåˆ™çš„ PyTorch åŽç«¯ã€‚å®ƒæ”¯æŒå¼€ç®±å³ç”¨çš„æ— ç¼æ¨¡åž‹æ‰§è¡Œä¸Ž TPU 迁移，并通过 eager 执行优化、torch.compile é›†æˆå’ŒåŽŸç”Ÿç¡¬ä»¶åŠ é€Ÿï¼Œæä¾›å ·æœ‰è¡Œä¸šç«žäº‰åŠ›çš„æŽ¨ç†ä¸Žè®­ç»ƒæ€§èƒ½ã€‚

推出 TorchTPU

Google TPU 在 Google å† å¤–éƒ½å·²è¢«éžå¸¸æˆåŠŸåœ°ç”¨äºŽå¤§è§„æ¨¡è®­ç»ƒå’Œéƒ¨ç½²å¤§åž‹æœºå™¨å­¦ä¹ æ¨¡åž‹ã€‚å®ƒä»¬ä¸“ä¸ºæœºå™¨å­¦ä¹ è€Œè®¾è®¡ï¼Œå› æ­¤åœ¨æ€§èƒ½ã€èƒ½è€—å’Œæˆæœ¬æ–¹é¢å ·æœ‰æ˜¾è‘—ä¼˜åŠ¿ã€‚å°½ç®¡ TPU 历史上一直与 TensorFlow 和 JAX ç”Ÿæ€ç³»ç»Ÿç›¸å ³è”ï¼Œæˆ‘ä»¬å¾ˆé«˜å ´æŽ¨å‡º TorchTPU ——这是 Google 于 2025 年中启动、与 Meta 的 PyTorch 团队合作开展的新项目，旨在让 PyTorch 原生运行在 TPU 上。

为了在 Python 框架之下对 Google TPU è¿›è¡Œç¼–ç¨‹ï¼Œå¼€å‘è€ ä½¿ç”¨ XLA ç¼–è¯‘å™¨æ ˆã€‚2020 年，Google 推出了首个面向 TPU 的 PyTorch 编译版本 PyTorch/XLA，基于 Lazy Tensorsã€‚å°½ç®¡å ·æœ‰å¥ åŸºæ„ä¹‰ï¼Œè¿™ä¸€æ—©æœŸé›†æˆéš¾ä»¥è·Ÿä¸Š PyTorch ç”±ç¤¾åŒºé©±åŠ¨çš„å¿«é€Ÿæ¼”è¿›ã€‚å ³é”®åœ¨äºŽï¼ŒPyTorch/XLA éœ€è¦ä¾µå ¥å¼ä¿®æ”¹æºä»£ç ä»¥æ˜¾å¼å®šä¹‰å¯ç¼–è¯‘åŒºåŸŸï¼Œç¼–è¯‘æ—¶å»¶è¿Ÿä¸¥é‡ï¼Œä¸æ”¯æŒ PyTorch 原生的命令式（“eager”）执行模式，并且在支持 PyTorch ç”Ÿæ€å† æ–°å ´çš„åˆ†å¸ƒå¼è®¡ç®—åŽŸè¯­æ–¹é¢æ»žåŽã€‚ä¸Ž PyTorch/XLA 和 TorchAX 等以往将 PyTorch å¼ºè¡Œå¡žå ¥ JAX 运作方式的尝试不同，TorchTPU 是作为 PyTorch 的一个符合设计原则的 TPU 后端而构建的：eager æ‰§è¡Œæ˜¯é¦–è¦ä¼˜å ˆçº§ï¼ŒçŽ°æœ‰ PyTorch å·¥ä½œè´Ÿè½½åªéœ€æžå°‘çš„ä»£ç æ”¹åŠ¨å³å¯å¼€ç®±å³ç”¨ï¼ˆä¾‹å¦‚ï¼Œ .to("cuda") → .to("tpu") ）。

æ›´å ·ä½“åœ°è¯´ï¼Œåœ¨å¯åŠ¨ TorchTPU 时，我们刻意制定了一套遵循的原则：

• PyTorch Citizenship : 运行时中不依赖 JAX； 稳健的 eager 模式； 基于 Torch Dynamo（而非 Lazy Tensor）的 Torch.compile()，性能与 JAX compile 相近； 通过在 PyTorch æ ˆçš„æœ€åº•å±‚ï¼ˆä¾‹å¦‚ Process Group API）提供 TPU 集合通信，广泛支持 PyTorch 分布式惯用法（DDP、FSDP、DTensor 等）； æ— ç¼é›†æˆåˆ° PyTorch 生态系统中； 社区参与。

• Focus on the User: 将工作负载移植到 TPU 应当像把 model.to("cuda") 改为 model.to("tpu") ä¸€æ ·ç®€å•ã€‚ï¼ˆä» é è¿™ä¸€ç®€å•æ”¹åŠ¨æ— æ³•è¾¾åˆ° TPU çš„å³°å€¼æ€§èƒ½ï¼Œå› ä¸ºæ¨¡åž‹çš„å½¢çŠ¶ã€åˆ†ç‰‡å’Œç®—å­å¯èƒ½åå‘ GPU 架构；但模型应当能够开箱即用。） 为运行时/编译器错误提供用户友好且可操作的提示信息； 与 PyTorch ç¤¾åŒºå·²åœ¨ä½¿ç”¨çš„å·¥å ·å’Œæ¡†æž¶é›†æˆï¼ˆä¾‹å¦‚ torch.profiler、DTensor、Ray）。

• Improve the State of the Art : 构建一个能够在大规模工作负载上与 GPU 上的 PyTorch 竞争的解决方案； 为 eager æ‰§è¡Œå¼•å ¥ JIT 编译方法，通过运行时优化弥合逐算子执行与编译执行之间的性能差距； é€šè¿‡å¼•å ¥æœ‰ç•ŒåŠ¨æ€æ€§ã€é¿å å› å½¢çŠ¶å˜åŒ–å¯¼è‡´ XLA 重新编译，以及探索可提前预编译的动态 TPU å† æ ¸ï¼Œæ”¹è¿› TPU ç¼–ç¨‹æ ˆã€‚

TorchTPU 正在积极开发中，它已在生产环境中用于部分最大 Google Cloud 客户的推理和训练工作负载，并且可以在该 GitHub 仓库 ä¸­æ‰¾åˆ°å ¶å¼€æºé¡¹ç›®ã€‚æ’°å†™æœ¬æ–‡æ—¶ï¼Œæˆ‘ä»¬å·²ï¼š

• 新增对 eager å’Œç¼–è¯‘æ¨¡å¼çš„æ”¯æŒï¼ŒåŒ æ‹¬æ–°çš„ DeferAndFuse eager 模式；

• 新增对约 450 个 Aten ç®—å­çš„æ”¯æŒï¼ˆè·³è¿‡å ·æœ‰ç»“æž„åŒ–æˆ–å¤åˆå§”æ‰˜çš„ç®—å­ï¼‰ï¼ŒåŒ æ‹¬ä¸»è¦æ·±åº¦å­¦ä¹ æž¶æž„æ‰€éœ€çš„å ¨éƒ¨ Core ATen ç®—å­çš„å‰å‘å’Œåå‘ä¼ æ’­ï¼›

• 新增通过 JAX/Pallas 和 Helion DSL ç¼–å†™è‡ªå®šä¹‰å† æ ¸çš„èƒ½åŠ›ï¼›

• Validated out-of-box execution on the 5,000 most downloaded Hugging Face models, including: LLM：Qwen、DeepSeek、Gemma、LLama、Mistral、GPT、OPT、BLOOM、Falcon、StableML 音频/语音：Whisper、Wav2Vec2、HuBERT、MMS 视觉：DinoV2、SigLIP、CLIP、Swin Transformer、ConvNeXt、ViT、ResNet、DETR æ–‡æœ¬ç¼–ç å™¨ä¸Žæœç´¢ï¼šBERT、ALBERT 与 Electra、RoBERTa、DeBERTa、T5、MT5、Flan-T5

• 已验证基于 DTensor 的分片；

• 已验证多主机推理（vLLM 和 SGLang）以及微调/训练工作负载（TorchTitan），并采用基于 Ray 的编排；

• 已验证与 vLLM、SGLang、TorchTitan、Ray 和 TorchRec 的集成。

在本文档的剩余部分，我们将讨论 TorchTPU æž¶æž„çš„å ³é”®æ–¹é¢ä»¥åŠä¸€äº›åˆæ­¥ç»“æžœã€‚

æž¶æž„æ·±å ¥æŽ¢è®¨

TPU å ¥é—¨

Google TPU 是专门设计的 ASIC ï¼Œç”¨äºŽæ‰§è¡Œä½œä¸ºæ·±åº¦å­¦ä¹ åŸºç¡€çš„å·¨é‡å¼ é‡è¿ç®—ã€‚ä»Ž 硬件架构角度 来看，一个 TPU 芯片 åŒ å«ï¼ˆå–å†³äºŽä»£é™ ï¼‰ï¼š

• 一个或多个 TensorCore ï¼Œæ¯ä¸ªåŒ å«ä¸€ä¸ªæˆ–å¤šä¸ªçŸ©é˜µä¹˜æ³•å•å ƒï¼ˆMXUï¼‰ã€ä¸€ä¸ªå‘é‡å•å ƒï¼ˆVPUï¼‰å’Œä¸€ä¸ªæ ‡é‡å•å ƒï¼ˆå‚è§ 此链接 ï¼‰ã€‚æ ¹æ® TPU ä»£é™ ï¼Œä¸€ä¸ª MXU å¯åŒ å« 256×256 ä¸ªä¹˜åŠ å™¨ï¼Œå®ƒä»¬ä»¥ 脉动阵列 的形式排列。

• 一个或多个 SparseCore ，它们是用于稀疏操作（例如 gather、scatter、embeddings）和集合通信的 ä¸“ç”¨å•å ƒ 。它们可以与 TensorCore 并行运行，从而实现计算-计算和计算-é€šä¿¡çš„é‡å ã€‚

• 大容量 é«˜å¸¦å®½å† å­˜ （HBM）——在 TPU 8i 中为 288 GB、8.6 TB/s，可用于存储 ML 数据并与 CPU å ±äº«ã€‚

• 双向 芯片间互连 （ICI）——在 TPU 8t 和 8i 中为 19.2 TBpsï¼Œå è®¸ç”¨æˆ·é ç½®çš„ TPU Slice å† çš„ TPU 芯片之间快速通信，TPU Slice 是 TPU Pod çš„é€»è¾‘å­ç»„ï¼Œç”±é¢„é ç½®æ•°é‡çš„ TPU 芯片组成，这些芯片以 3D torus（适用于 TPU 7、TPUv 8t）或 Boardfly（TPU 8iï¼‰é ç½®è¿žæŽ¥ã€‚

• 双向 数据中心网络 （DCN）——在 TPU 8t 中为 400 GBpsï¼Œå è®¸è·¨ Slice 的 TPU 芯片之间通信。

TPU é€šè¿‡ä¸¤ä¸ªä¸åŒçš„å ¥å£ç‚¹ä½¿ç”¨ XLA 编译器 æ ˆè¿›è¡Œç¼–ç¨‹ï¼š StableHLO （用于 ML 算子的较高级 IR）和 Mosaic ï¼ˆç”¨äºŽè‡ªå®šä¹‰å† æ ¸çš„è¾ƒä½Žçº§ IR ï¼‰ã€‚å®žé™ çš„ TPU ISA æ˜¯ä¸“æœ‰çš„ã€å¯èƒ½å‘ç”Ÿå˜åŒ–ä¸”ä¸å ¬å¼€ã€‚å®žé™ ä¸Šï¼Œå¯ä»¥å°† StableHLO 以及（更是如此）Mosaic 视为等同于 PTX。参见下图的高层示意图。

C++ 程序可以通过实例化 PjRT Client 来访问本地 TPUï¼Œå ¶ API 提供了将 MLIR（StableHLO 或 Mosaic）编译为 TPU 可执行文件、启动该可执行文件以及在 TPU 之间复制数据的方法。请注意，TPU 执行是异步的，CPU å¯ä»¥é€šè¿‡ç­‰å¾ å ³è”çš„ PjRTBuffer æ¥ç­‰å¾ TPU 结果就绪。

TPU 编程模型的显著特点：

• TensorCore è®¾å¤‡å¯ä»¥è¢«è§†ä¸ºä¸€ä¸ªå•çº¿ç¨‹å¤„ç†å™¨ï¼Œå¸¦æœ‰ä¸€ä¸ªæˆ–å¤šä¸ªå¤§è§„æ¨¡çŸ©é˜µä¹˜æ³•å•å ƒï¼ˆMXUï¼‰ã€‚å› æ­¤ï¼Œè¦è¾¾åˆ°å³°å€¼æ€§èƒ½ï¼Œå ³é”®åœ¨äºŽæ¨¡åž‹çš„å½¢çŠ¶è¦è¶³å¤Ÿå®½ï¼Œä»¥å 分发挥 MXU 的潜力。这与 GPU 不同，在 GPU ä¸Šå¹¶è¡Œå¯åŠ¨å¤šä¸ªè¾ƒå°çš„å† æ ¸æœ‰åŠ©äºŽæœ€å¤§åŒ–ç¡¬ä»¶åˆ©ç”¨çŽ‡ã€‚

• XLA 编译器可以指示 TensorCore 将稀疏运算和集合运算卸载到 SparseCore，以实现更高效和/或并行的执行。这可以自动发生，也可以通过编译后的 StableHLO 模块中用户提供的 MLIR 注解来实现。

• 一个 TPU è®¾å¤‡åªèƒ½è¢«ä¸€ä¸ªç”¨æˆ·ç©ºé—´è¿›ç¨‹è®¿é—®ã€‚éµå¾ªæˆ‘ä»¬æ— ç¼é›†æˆåˆ° PyTorch ç”Ÿæ€ç³»ç»Ÿçš„ç›®æ ‡ï¼Œå¯¹äºŽåˆ†å¸ƒå¼ TorchTPU 工作负载，我们采用每设备一进程的模型，这与 JAX 的每主机一进程模型形成对比。然而，由于一个 TPU 只能被一个用户空间进程访问，RL å ±ç½®æ— æ³•å¼€ç®±å³ç”¨åœ°æ”¯æŒï¼Œæˆ‘ä»¬æ­£åœ¨ç ”ç©¶æ”¾å®½æˆ–ç»•è¿‡è¿™ä¸€ TPU 限制的方法。

• StableHLO 遵循 SSA 范式 ï¼Œå› æ­¤ä¸å è®¸åŽŸåœ°ä¿®æ”¹ï¼Œå³å†™æ“ä½œæ€»æ˜¯äº§ç”Ÿä¸€ä¸ªæ–°å€¼ï¼ˆè‡³å°‘åœ¨æ¦‚å¿µä¸Šå¦‚æ­¤ï¼‰ã€‚è¿™å¯¹ TorchTPU 尤为重要，它确实通过视图分解和反转提供了熟悉的 PyTorchâ€œè§†å›¾â€è¯­ä¹‰ï¼ˆåŒ æ‹¬ä¿®æ”¹ï¼‰ï¼Œå¦‚æœ¬æ–‡æ¡£åŽé¢æ‰€è¿°ã€‚

• XLA 编译器和 PjRT è¿è¡Œæ—¶å‡å®šå¼ é‡å½¢çŠ¶åœ¨ç¼–è¯‘æ—¶å·²çŸ¥ã€‚è®¤è¯†åˆ°è¿™æ˜¯ PyTorch å·¥ä½œè´Ÿè½½çš„ä¸€ä¸ªé™åˆ¶ï¼Œæˆ‘ä»¬ä¸€ç›´åœ¨æ·»åŠ å¯¹æœ‰ç•ŒåŠ¨æ€æ€§çš„æ”¯æŒï¼Œå³å½“å½¢çŠ¶ç»´åº¦çš„æœ€å¤§è¾¹ç•Œåœ¨ç¼–è¯‘æ—¶å·²çŸ¥ï¼Œä½†å®žé™ ç»´åº¦ä» åœ¨è¿è¡Œæ—¶å·²çŸ¥æ—¶ã€‚

• 集合操作被降级为 StableHLO æ“ä½œä¼ é€’ç»™ XLA，而不是使用 NCCL 等单独的库。

å¯ç”¨æ€§ä¸Žè‡ªåŠ¨åŠ è½½æœºåˆ¶

我们在设计 TorchTPU æ—¶è€ƒè™‘äº†å¯ç”¨æ€§ã€‚å½“ç”¨æˆ·åœ¨å®‰è£ äº† TorchTPU wheel 的 TPU æœºå™¨ä¸Šå¯¼å ¥ PyTorch 时，在 TPU 上运行 PyTorch ä»£ç åªéœ€è¿›è¡Œå¾®ä¸è¶³é“çš„æ›´æ”¹ï¼Œå¹¶ä¸”å¾—ç›ŠäºŽ PyTorch çš„è‡ªåŠ¨åŠ è½½æœºåˆ¶ï¼Œåªéœ€å‡ è¡Œä»£ç å³å¯å®Œæˆï¼š

è¯·å‚é˜ ä¸‹è¡¨è¿›è¡Œå¹¶æŽ’æ¯”è¾ƒã€‚

之前（GPU）

之后（TPU）

替代方案（TPU 或 GPU）

import torch
import torch.distributed as dist

def main():
    device = torch.device("cuda")
    dist.init_process_group(backend="nccl")
    rank = dist.get_rank()
    word_size = dist.get_world_size()
    ...
    model = model.to(device)
    ...

import torch
import torch.distributed as dist

def main():
    device = torch.device("tpu")
    dist.init_process_group(backend="tpu_dist")
    rank = dist.get_rank()
    word_size = dist.get_world_size()
    ...
    model = model.to(device)
    ...

import torch
import torch.distributed as dist

def main():
    device = torch.accelerator.current_accelerator()
    dist.init_process_group()
    rank = dist.get_rank()
    word_size = dist.get_world_size()
    ...
    model = model.to(device)
    ...

数值精度

TorchTPU 为 TPU ç¡¬ä»¶é€‚é æ ‡å‡† PyTorch 精度 APIã€‚ä¸Žæ‰€æœ‰é«˜æ€§èƒ½æœºå™¨å­¦ä¹ æž¶æž„ä¸€æ ·ï¼ŒTPU é’ˆå¯¹å ¶ç‰¹å®šä»£é™ æ‰€åå¥½çš„åŽŸç”Ÿæ•°æ®ç±»åž‹è¿›è¡Œäº†ä¼˜åŒ–ï¼Œä»¥å®žçŽ°å³°å€¼æ€§èƒ½å’Œèƒ½æ•ˆã€‚å¯¹äºŽ TPUï¼ŒåŽŸç”ŸçŸ©é˜µä¹˜æ³•æ ¼å¼ä¸º bfloat16（BF16），它在高动态范围与高效吞吐之间取得平衡。由于 TPU çŸ©é˜µä¹˜æ³•å•å ƒï¼ˆMXU）被设计为使用 BF16 è¾“å ¥æ‰§è¡Œç®—æœ¯è¿ç®—ï¼Œå¹¶ä»¥ 32 ä½æµ®ç‚¹è¿›è¡Œä¸­é—´ç´¯åŠ ï¼Œè¿™ä¸€é ç½®æ˜¯ TPU é«˜æ€§èƒ½è®¡ç®—çš„åŸºçŸ³ã€‚å¯¹äºŽéžæ”¶ç¼©è¿ç®—ï¼ˆä¾‹å¦‚é€å ƒç´ åŠ æ³•ã€æ¿€æ´»ã€layernorm 和 softmax），TorchTPU 使用 TPU å‘é‡æ‰§è¡Œå•å ƒï¼ˆVPU）执行 FP32 å¼ é‡è¿ç®—ï¼Œè¯¥å•å ƒåŽŸç”Ÿæ‰§è¡Œå®Œæ•´çš„ IEEE-754 FP32 算术。

TorchTPU 将 PyTorch ç²¾åº¦æ¨¡å¼æ˜ å°„åˆ°åº•å±‚ XLA/StableHLO é ç½®ï¼Œä»¥æš´éœ²è¿™ä¸€ç¡¬ä»¶èƒ½åŠ›ï¼š

• medium ï¼šä¹˜æ³•ä½¿ç”¨è¾ƒä½Žç²¾åº¦ï¼ˆå°†è¾“å ¥è½¬æ¢ä¸º MXU 的原生 BF16ï¼‰ï¼Œè€Œç´¯åŠ ä»¥ FP32 执行以保持稳定性。TorchTPU 默认使用此模式以å 分利用 TPU ç¡¬ä»¶åŠ é€Ÿï¼Œå› ä¸ºå•é BF16 ä¹˜æ³•é åˆ FP32 ç´¯åŠ å¯¹å‡ ä¹Žæ‰€æœ‰æ·±åº¦å­¦ä¹ å·¥ä½œè´Ÿè½½éƒ½èƒ½æä¾›è¶³å¤Ÿç²¾åº¦ã€‚

• high ：使用 3 éç®—æ³•ï¼Œå°†è¾“å ¥æ‹†åˆ†ä¸ºé«˜ã€ä½Žåˆ†é‡ä»¥è®¡ç®—ä¸‰ä¸ªäº¤å‰é¡¹ï¼Œä»¥çº¦ 3 倍计算成本恢复大部分 FP32 精度。

• highest ：使用 6 遍算法计算所有交叉项，以近似完整的 IEEE FP32 精度，在 MXU 上产生约 6 倍计算开销。

å¯¹äºŽç»å¤§å¤šæ•°æ·±åº¦å­¦ä¹ å·¥ä½œè´Ÿè½½ï¼Œæˆ‘ä»¬å»ºè®®ä½¿ç”¨ medium （默认）精度。在 TorchTPU 中显式选择 high 或 highest 会触发性能警告（ TORCH_WARN_ONCE ï¼‰ã€‚åªæœ‰åœ¨æ˜Žç¡®è¯†åˆ«åˆ°ä¸Žç²¾åº¦ç›¸å ³çš„æ”¶æ•›é—®é¢˜æ—¶æ‰åˆ‡æ¢åˆ° high 或 highest ï¼Œå› ä¸ºè¿™äº›æ¨¡å¼ä¼šå¸¦æ¥æ˜¾è‘—çš„æ€§èƒ½å¼€é”€ã€‚é™¤å ¨å±€è®¾ç½®å¤–ï¼ŒTorchTPU 还提供线程本地上下文管理器（ torch.tpu.precision ）用于细粒度作用域（例如针对数值敏感的层）。

并排对比见下表。

维度

CUDA 上的 PyTorch

TPU 上的 PyTorch（TorchTPU）

默认精度模式

highest（纯 FP32）

medium（Precision.DEFAULT）

å¿«é€Ÿæ•°å­¦åŠ é€Ÿå™¨åŽŸè¯­

Tensor Cores（TF32）

脉动 MXU（bfloat16）

快速数学位数

1 符号位 + 10 尾数位 + 8 指数位（19 位 TF32 æ ¼å¼ï¼‰

1 符号位 + 7 尾数位 + 8 指数位（16 位 BF16 æ ¼å¼ï¼‰

ç´¯åŠ ç²¾åº¦

FP32（highest 和 high 中），TF32（medium 中）

所有模式均为 FP32

更高/最高精度机制

专用硬件 FP32/TF32 å•å ƒ

多遍算法分解（3 倍或 6 倍 MXU 遍）

high/highest 的性能影响

åŸºçº¿é€Ÿåº¦æˆ–é€‚åº¦å‡æ ¢ï¼ˆçº¦ 1–2 倍）

high 约 3 倍 MXU 计算成本；highest 约 6 倍计算开销

作用域 API

å ¨å±€ torch.set_float32_matmul_precision

å ¨å±€ torch.set_float32_matmul_precision + 作用域 torch.tpu.precision 上下文管理器

Eager 执行、StableHLO 降级和 JIT 编译

遵循 这份 PyTorch 指南 ，我们通过一个新的 C++ 后端扩展了 PyTorch dispatcherï¼Œå¹¶å°†å ¶æ³¨å†Œä¸º PrivateUse1 ã€‚æ­¤å¤–ï¼Œæˆ‘ä»¬ä½¿ç”¨äº†æ ‡å‡†çš„ at::Tensor æœºåˆ¶æ¥å®šä¹‰å¼ é‡å­˜å‚¨çš„è®¾å¤‡ç‰¹å®šå®žçŽ°ï¼ˆé€šè¿‡åœ¨ c10::StorageImpl å† ä½¿ç”¨ c10::UniqueVoidPtr ）。

TorchTPU 的 eager 执行的一个独特特性是，它可以以多种模式将 Aten 算子分发到 TPU：

• 立即 ，即有效地实现逐算子分发；

• 延迟 ，一直延迟到自动确定的 物化点 （例如，一条 .to("cpu") è¯­å¥ï¼‰ï¼Œå¹¶ä¸Žæˆªè‡³è¯¥ç‚¹è¢«å»¶è¿Ÿçš„æ‰€æœ‰å ¶ä»–ç®—å­ä¸€èµ·èšåˆ/编译；

• ä»‹äºŽä¸¤è€ ä¹‹é—´çš„ä»»ä½•æƒ å†µã€‚

è™½ç„¶è¿™ç§çµæ´»æ€§ç¡®å®žç»™è®¾è®¡å¸¦æ¥äº†é¢å¤–çš„å¤æ‚æ€§ï¼Œä½†æˆ‘ä»¬ä»Žä¸€å¼€å§‹å°±åˆ»æ„åŠ å ¥å®ƒï¼Œå› ä¸ºæˆ‘ä»¬å¸Œæœ›è¯•éªŒä¸åŒçš„ç­–ç•¥ï¼Œä»¥åœ¨åŸºæœ¬çš„é€ç®—å­åˆ†å‘æ¨¡å¼ä¹‹ä¸Šæå‡ eager 性能。（此外，正如本文档后面所讨论的，这种灵活性在处理视图语义和 torch.compile 时非常方便。）

总之，TorchTPU 提供了 3 种 eager 模式，可以通过环境变量或上下文管理器来设置：

• Strict （默认）：我们一次分发一个 Aten 算子，它们在 TPU 上与 CPU 异步执行。这相当于在 GPU 上以 eager 模式运行 PyTorch 时的预期行为。

• Debug （ export TORCH_TPU_LAUNCH_BLOCKING=1 ）：我们一次分发一个 Aten 算子，但 TPU 执行与 CPU åŒæ­¥ã€‚è¿™å¾ˆæ ¢ï¼Œæ—¨åœ¨ç”¨äºŽè°ƒè¯•ï¼Œç›¸å½“äºŽåœ¨ GPU 上使用 CUDA_LAUNCH_BLOCKING ，并额外提供了 TorchTPU 的边界检查。

• DeferAndFuse （ export TORCH_TPU_DEFER_AND_FUSE=1 ）：我们将多个 Aten 算子延迟并融合成更大的块。这是 TorchTPU å¼•å ¥çš„ä¸€ç§æ–°çš„ eager/编译 混合 模式，比 Strict Eager å¿« 50% 或更多，同时降低峰值 HBM å† å­˜ï¼ˆç®—å­èžåˆå‡å°‘äº† HBM 中的临时存储）。

现在我们已经描述了 TorchTPU eager 模式的基础知识，让我们勾勒一下当收到来自 PyTorch dispatcher 的回调调用时会发生什么。假设我们使用的是 DeferAndFuse 模式，并以 at::mul.out 为例：

• Create a Deferred Op instance that contains: A StableHLO builder object, which when invoked returns a StableHLO function implementing the semantics of aten::mul , as shown below (note that any implicit broadcasting and dtype promotion are handled elsewhere) func.func @main(%arg0: tensor<4x16xbfloat16>, %arg1: tensor<4x16xbfloat16>) -> tensor<4x16xbfloat16> { %0 = stablehlo.multiply %arg0, %arg1 : tensor<4x16xbfloat16> return %0 : tensor<4x16xbfloat16> } ä¸€ä¸ªå¯é€‰çš„é™æ€ç®—å­å‚æ•°æ˜ å°„ï¼ˆä¾‹å¦‚ï¼Œå¯¹äºŽ aten::native_batch_norm 使用 training 和 momentum ），用于计算缓存键（下文讨论）；

• 返回一个新的 at::Tensor ï¼Œå ¶ c10::StorageImpl （通过间接方式）指向该 Deferred Op 实例。

当收到来自 PyTorch dispatcher 的对应于 mytensor.to("cpu") 或 mytensor.item() æˆ–ç±»ä¼¼å† å®¹çš„å›žè°ƒæ—¶ï¼ŒTorchTPU 将：

• 分析到目前为止累积的、导致 mytensor 计算的延迟算子区域 R；

• 应用活跃性分析来确定 R 中哪些算子需要物化，以及产生 mytensor çš„é‚£ä¸ªç®—å­ï¼›é‡è¦çš„æ˜¯ï¼Œå ·æœ‰å‰¯ä½œç”¨çš„ç®—å­ï¼ˆä¾‹å¦‚ï¼Œé›†åˆé€šä¿¡ï¼‰æ€»æ˜¯è¢«ä¿ç•™ï¼ˆå ¶è°ƒç”¨é¡ºåºä¹Ÿæ˜¯å¦‚æ­¤ï¼‰ï¼Œè€Œ R ä¹‹å¤–ä¸æ´»è·ƒçš„ç®—å­å¯ä»¥åœ¨ä¸åŒçš„å·¥ä½œè´Ÿè½½å‰¯æœ¬ä¹‹é—´ç¡®å®šæ€§åœ°æ£€æµ‹åˆ°ï¼Œå› ä¸º Python çš„å† å­˜ç®¡ç†ä½¿ç”¨å¼•ç”¨è®¡æ•°ã€‚

• 为要物化的算子图 G 创建一个编译缓存键，以便将来快速且唯一地识别该图；

• Use the computed cache keys to check in a Compilation Cache if graph G has been already compiled. 如果是，则在 TPU 上分发 G 的可执行文件； å¦‚æžœä¸æ˜¯ï¼Œåˆ™ä½¿ç”¨çº¿ç¨‹æ± ï¼ˆå¼‚æ­¥åœ°ï¼‰ç¼–è¯‘ G，然后在 TPU ä¸Šåˆ†å‘å ¶æ‰§è¡Œã€‚

• ç­‰å¾ mytensor 的 TPU å€¼å°±ç»ªï¼Œç„¶åŽå°†å ¶ä»Ž TPU HBM å¤åˆ¶åˆ°ä¸Žç›®æ ‡ at::Tensor å ³è”çš„ CPU 存储。

è¡¨é¢ä¸Šçœ‹ï¼Œä¸Šè¿°æ–¹æ³•ä¼¼ä¹Žæ•ˆçŽ‡ä¸é«˜ï¼Œå› ä¸º CPU 只有在遇到 .to("cpu") 语句时才会启动 TPU 执行。对于真实工作负载，我们可以预期，在少数初始迭代/步骤填å 编译缓存之后，CPU æ‰§è¡Œåº”ä» æ¶‰åŠ PyTorch è¿½è¸ªå’Œç®—å­åˆ†å‘ï¼Œå› æ­¤ä¼šå§‹ç»ˆé¢†å ˆäºŽ TPU æ‰§è¡Œï¼Œè€ŒåŽè€ å°†å¿™äºŽå†—é•¿çš„æ•°å€¼è®¡ç®—ã€‚å°½ç®¡å¦‚æ­¤ï¼Œä¸ºäº†åŠ é€Ÿå¯¹å¤šæ¬¡é‡å¤çš„ç®—å­åºåˆ—ï¼ˆæƒ³æƒ³ä¸€å † transformer å—ï¼‰çš„ç‰©åŒ–ï¼Œæˆ‘ä»¬æ·»åŠ äº†é‡å¤ç®—å­å¯å‘å¼ï¼ˆRepeated Ops Heuristic），它会检测此类重复，并在到达 TPU 到 CPU ä¼ è¾“ä¹‹å‰å°±è§¦å‘ç‰©åŒ–ç‚¹ã€‚

编译缓存

TorchTPU é‡‡ç”¨ç¼–è¯‘ç¼“å­˜æ¥é¿å ä¸å¿ è¦çš„ XLA ç¼–è¯‘ã€‚ä¸ŽçŽ°ä»£è®¡ç®—æœºå­˜å‚¨ç³»ç»Ÿç±»ä¼¼ï¼Œè¯¥ç¼“å­˜åŒ å«å¤šä¸ªååŒå·¥ä½œçš„å±‚çº§ï¼š

• 第 1 层 ï¼šè¿›ç¨‹æœ¬åœ°å† å­˜ç¼“å­˜ï¼Œç”¨äºŽåœ¨åŒä¸€è¿›ç¨‹å† å ±äº«ç¼–è¯‘ç»“æžœã€‚è¿™æ˜¯ä¸‰å±‚ä¸­æœ€å¿«çš„ã€‚

• 第 2 层 ï¼šæœºå™¨æœ¬åœ°å ±äº«å† å­˜ï¼Œæ”¯æŒåŒä¸€ä¸»æœºä¸Šçš„è·¨è¿›ç¨‹å ±äº«ã€‚è¿™å¯¹äºŽå¤šè®¾å¤‡è®­ç»ƒ/推理很有用。

• 第 3 层 ：由远程文件系统（例如 Google Cloud Storageï¼‰æ”¯æŒçš„æŒä¹ ç¼“å­˜ï¼Œå è®¸è·¨ä¸»æœºã€ç¨‹åºå’Œä½œä¸šé‡å¯å¤ç”¨ã€‚è¿™åœ¨ Colab 和数据中心等执行环境中运行 TPU 程序时很有用。

第 1 å±‚ç¼“å­˜å§‹ç»ˆå¯ç”¨ã€‚é»˜è®¤æƒ å†µä¸‹ï¼Œç¬¬ 2 å±‚å¯¹åˆ†å¸ƒå¼ç¨‹åºå¯ç”¨ï¼Œå¯¹å•è®¾å¤‡ç¨‹åºç¦ç”¨ã€‚ç”¨æˆ·å¯ä»¥æ ¹æ®éœ€è¦è¦†ç›–æ­¤è¡Œä¸ºã€‚ç¬¬ 3 å±‚é»˜è®¤ç¦ç”¨ï¼Œä½†å¯ä»¥é€šè¿‡æŒ‡å®šå ±äº«å­˜å‚¨è·¯å¾„æ˜¾å¼å¯ç”¨ã€‚

缓存中的每个条目都是针对 StableHLO 计算图编译的 TPU 可执行文件。在第 1 层缓存中，每个可执行文件都是一个 C++ 对象，可以直接调度到 TPU è®¾å¤‡ä¸Šæ‰§è¡Œã€‚ç”±äºŽå¯¹è±¡æ— æ³•ç›´æŽ¥è·¨è¿›ç¨‹å ±äº«ï¼Œç¬¬ 2 层和第 3 层缓存将可执行文件保存为序列化的二进制字符串。

å‡ºäºŽæ€§èƒ½åŽŸå› ï¼Œæˆ‘ä»¬ä¸ä½¿ç”¨ StableHLO 图本身作为缓存键。相反，我们将计算、编译器选项和主要组件的版本指纹化为三个 64 ä½æ•´æ•°ï¼Œå¹¶ä½¿ç”¨è¯¥å ƒç»„ä½œä¸ºç¼“å­˜é”®ï¼š

• ä¸€ä¸ªæ— å½¢çŠ¶é”®ï¼Œæ•èŽ·è®¡ç®—å›¾çš„ç»“æž„ï¼Œä½†ä¸æ•èŽ·å¼ é‡è¾“å ¥/输出的大小，

• ä¸€ä¸ªç»´åº¦é”®ï¼Œæ•èŽ·å¼ é‡å¤§å°ï¼Œä»¥åŠ

• ä¸€ä¸ªç¼–è¯‘é€‰é¡¹é”®ï¼Œæ•èŽ·ä¼ é€’ç»™ XLA 编译器的选项以及 TorchTPU、libtpu 和 PjRt 运行时的版本。

æ— å½¢çŠ¶é”®å’Œç»´åº¦é”®åˆåœ¨ä¸€èµ·ä¹Ÿç§°ä¸ºå›¾é”®ã€‚è™½ç„¶ç†è®ºä¸Šä¸åŒå›¾ã€ç¼–è¯‘å™¨é€‰é¡¹å’Œç»„ä»¶ç‰ˆæœ¬ä¹‹é—´å¯èƒ½å‘ç”Ÿç¼“å­˜é”®å†²çªï¼Œä½†æˆ‘ä»¬è®¤ä¸º 192 位键空间足够大，实践中不会出现此类冲突。

视图语义

PyTorch 的操作语义与 StableHLO è¯­ä¹‰ä¹‹é—´ä¸€ä¸ªç‰¹åˆ«ä¸åŒçš„é¢†åŸŸæ˜¯å¯å˜æ€§å’Œå ±äº«å† å­˜è§†å›¾çš„å¤„ç†ï¼š

• PyTorch 假设每个 Tensor 都是指向线性字节数组的指针，并在 Tensor 层级带有 stride å ƒæ•°æ®ï¼Œç”¨äºŽæè¿°é€»è¾‘ç´¢å¼•å ƒç»„åˆ°è¯¥æ•°ç»„ä¸­çº¿æ€§åç§»é‡ä¹‹é—´çš„è½¬æ¢å ³ç³»ã€‚å¤šä¸ª Tensor å¯ä»¥å ±äº«åŒä¸€åº•å±‚å­—èŠ‚æ•°ç»„ï¼Œå› æ­¤å¯ä»¥ä»Žä»»æ„ Tensor æ‰§è¡ŒåŽŸåœ°å†™å ¥ï¼Œå¹¶ä¸”è¯¥å†™å ¥ä¼šä»Žä»»ä½•å ¶ä»– Tensor å¯è§ã€‚å¾—ç›ŠäºŽé™„åŠ çš„ stride å ƒæ•°æ®ï¼Œè§†å›¾æ“ä½œå¯ä»¥åœ¨ dispatch 阶段在 CPU ä¸Šä»¥è¾ƒä½Žå¼€é”€å®Œæˆï¼Œåªéœ€æ›´æ–°è¯¥å ƒæ•°æ®å³å¯ã€‚è™½ç„¶è¿™äº›è§†å›¾å¤§å¤šå±žäºŽå°‘æ•°å‡ ç§å¸¸è§æ“ä½œï¼Œä½†ä½¿ç”¨ as_strided 运算符可以定义任意复杂的模式。

• 相比之下，StableHLO æ˜¯ä¸€ç§çº¯å‡½æ•°å¼è¯­è¨€ï¼Œå ¶ä¸­æ‰€æœ‰å¼ é‡å€¼åœ¨æž„é€ åŽå‡ä¸å¯å˜ã€‚ç‰©ç†å¸ƒå±€å¯¹ç”¨æˆ·æ˜¯æŠ½è±¡çš„ï¼Œå®ƒåæ˜ ç¼–è¯‘å™¨åšå‡ºçš„åˆ†å—å†³ç­–ï¼Œå› æ­¤ä¸æ”¯æŒç›´æŽ¥æ“ä½œ stride。此外，StableHLO ä¸æ”¯æŒåŽŸåœ°å†™å ¥ï¼›ç›¸åï¼Œå¯ä»¥é€šè¿‡ dynamic_update_slice æ“ä½œåº”ç”¨é€‰æ‹©æ€§æ›´æ–°ï¼Œå°†ä¸€ç»„æ–°å€¼åˆå¹¶åˆ°çŽ°æœ‰å¼ é‡çš„å­åˆ‡ç‰‡ä¸­ã€‚

åœ¨å­˜åœ¨å†™å ¥çš„æƒ å†µä¸‹ï¼Œåœ¨ TPU 上实现正确的视图语义并不简单；然而，TorchTPU å€ŸåŠ©ä¸€äº›æœºåˆ¶å®žçŽ°äº†è¿™ä¸€ç‚¹ï¼ŒåŒ æ‹¬ 视图分解 、 视图反转 和延迟操作。这一语义桥接使 TorchTPU 能够原生支持原本不可能支持的 PyTorch 模式，例如用于分布式训练的拼接参数数组。

流与 SparseCore

在优化模型性能时，PyTorch 用户可能希望同时执行 3 类操作：

• 计算操作；

• 集合通信操作；

• ä¸»æœºä¸Žè®¾å¤‡ä¹‹é—´çš„æ•°æ®ä¼ è¾“ã€‚

在 GPU ä¸Šï¼Œè¿™é€šå¸¸é€šè¿‡ä½¿ç”¨å•ç‹¬çš„å‘½ä»¤é˜Ÿåˆ—æ¥å®žçŽ°ï¼Œä»¥ä¾¿è®¾å¤‡å¯ä»¥å°†å·¥ä½œåˆ†æ´¾åˆ°ä¸åŒçš„ç¡¬ä»¶å•å ƒï¼ˆSM、DMA 等）。在 PyTorch 中，设备命令队列以 Streams çš„å½¢å¼æš´éœ²ï¼Œå®ƒä»¬é€šå¸¸æ—¢ç”¨äºŽç”¨æˆ·ä»£ç ï¼Œä¹Ÿç”¨äºŽæ¡†æž¶ä»£ç ï¼ˆä¾‹å¦‚åœ¨ FSDPv2 中）。

TorchTPU 也实现了 PyTorch Stream APIï¼Œä½†å¹¶éžåœ¨æ‰€æœ‰æƒ å†µä¸‹éƒ½èƒ½å®žçŽ° GPU 上可用的相同功能。考虑到 TPU TensorCore 本质上是一个单线程执行处理器，我们在给定 TensorCore 上一次只能分派一个计算。相反，GPU åŒ å«è®¸å¤šç›¸åŒçš„è®¡ç®—æ ¸å¿ƒï¼ˆSM），可以并行运行。所幸，一个 TPU TensorCore å ³è”ä¸€ä¸ªæˆ–å¤šä¸ª DMA 引擎以及一个或多个 SparseCore 引擎。不出所料，DMA å¼•æ“Žå¯ç”¨äºŽåœ¨ä¸»æœºä¸Žè®¾å¤‡ä¹‹é—´è¿›è¡Œæ•°æ®ä¼ è¾“ï¼Œå¹¶ä¸Ž TensorCore 操作并行执行。更有趣的是，SparseCore å¼•æ“Žæ˜¯ä¸€ä¸ªç‹¬ç«‹å•å ƒï¼Œé’ˆå¯¹ç¨€ç–æ“ä½œï¼ˆä¾‹å¦‚ gather/scatter）进行了优化，并且还可以支持集合通信操作。原则上它可以支持完整的 StableHLO 操作集，但需要注意的是，它在稀疏数学运算上会优于 TensorCoreï¼Œè€Œåœ¨ç¨ å¯†æ•°å­¦è¿ç®—ä¸Šä¼šé€ŠäºŽ TensorCoreï¼Œè€ŒåŽè€ æ­£æ˜¯ TensorCore å† MXU çš„å¼ºé¡¹ã€‚ä¸ºäº†è®©ç”¨æˆ·å®Œå ¨æŽ§åˆ¶å·¥ä½œè¿è¡Œçš„ä½ç½®ï¼Œå¯ä»¥å°† TorchTPU 流手动固定到 SparseCore，并且从性能和操作支持完整性的角度来看，该流与运行在 TensorCore ä¸Šçš„é»˜è®¤æµå¹¶ä¸ç›¸åŒã€‚æµè¿˜å¯ä»¥ç”¨äºŽåè°ƒå¼‚æ­¥æ•°æ®ä¼ è¾“ï¼ˆä½¿ç”¨ DMA 引擎），方式与 GPU 上类似。

编译模式

æˆ‘ä»¬çš„é¦–è¦ç›®æ ‡æ˜¯ä¿æŒä¸Žå ¶ä»– torch.compile() åŽç«¯ä¸€è‡´çš„ç”¨æˆ·ä½“éªŒï¼Œå¹¶å› æ­¤å ±äº«ç›¸åŒçš„ç¼–è¯‘å™¨å‰ç«¯ï¼Œå³ Dynamoã€‚æ ¹æ® 这份 PyTorch 指南 ，我们注册了一个新的 TorchDynamo åŽç«¯ï¼Œå ¶ç”¨æ³•å¦‚ä¸‹ï¼š

compiled_model = torch.compile(model, backend="tpu")

Dynamo é€šè¿‡ä¼ å ¥ torch.fx.GraphModule ä»¥åŠä¸€ç»„ç¤ºä¾‹è¾“å ¥å¼ é‡æ¥è°ƒç”¨æˆ‘ä»¬çš„åŽç«¯ã€‚æ­¤æ—¶ï¼Œæˆ‘ä»¬åƒåœ¨ eager æ¨¡å¼ä¸‹ä¸€æ ·å¯¹æ‰€æä¾›çš„å›¾è¿›è¡Œè¿½è¸ªï¼Œä½†æ‰€æœ‰ç®—å­éƒ½è¢«å»¶è¿Ÿæ‰§è¡Œï¼Œä¸”ä¸ä¼šåœ¨ TPU ä¸Šå®žé™ è¿è¡Œã€‚è¿™ä½¿æˆ‘ä»¬èƒ½å¤Ÿä¸ºæ•´ä¸ª FX 图生成并编译单个 StableHLO 程序。最后，我们的 Dynamo 后端返回一个 Python å¯è°ƒç”¨å¯¹è±¡ï¼Œè¯¥å¯¹è±¡å°è£ äº†å¯¹æ‰€ç”Ÿæˆ TPU 可执行文件的调用。请注意，作为一项临时限制，在 TPU 机器上调用 torch.compile 时，我们会自动设置 dynamic=False ï¼Œå› ä¸ºæ”¯æŒ SymInt 需要有界动态性（下文将讨论），但该特性目前尚未默认启用。

有界动态性

TPU 与 GPU 的一个主要区别在于，XLA ä¼ ç»Ÿä¸ŠåŸºäºŽé™æ€å½¢çŠ¶è¿è¡Œï¼Œä»¥ç”Ÿæˆä¼˜åŒ–çš„ç¡¬ä»¶å¸ƒå±€å’Œèžåˆå† æ ¸ã€‚åœ¨å®žè·µä¸­ï¼Œè¿™å¹¶éžä¸¥é‡é™åˆ¶ï¼›å½“å¼ é‡ç»´åº¦ï¼ˆå¦‚æ‰¹å¤§å°æˆ–åºåˆ—é•¿åº¦ï¼‰åœ¨å·¥ä½œè´Ÿè½½æ‰§è¡ŒæœŸé—´å‘ç”Ÿå˜åŒ–æ—¶ï¼Œå¼€å‘è€ é€šå¸¸é€šè¿‡é¢„åˆ†é æ›´å¤§çš„å¼ é‡å¹¶åœ¨è®¡ç®—å‰åŽæ‰‹åŠ¨å¡«å /切片来解决。

ä¸ºé¿å æ‰‹åŠ¨ä¿®æ”¹ä»£ç å’Œé‡æ–°ç¼–è¯‘çš„å¼€é”€ï¼Œæˆ‘ä»¬æ­£åœ¨ç§¯æžå¼€å‘ï¼Œä»¥é€šè¿‡ TorchTPU 中的 XLA æœ‰ç•ŒåŠ¨æ€æ€§åŽŸç”Ÿæ”¯æŒåŠ¨æ€å½¢çŠ¶ã€‚åœ¨æœ‰ç•ŒåŠ¨æ€æ€§ä¸‹ï¼ŒåŠ¨æ€ç¼“å†²åŒºåœ¨ä¸»æ¨¡åž‹æ¨¡å—å† æŒ‰æœ€å¤§ä¸Šé™åˆ†é ï¼Œå›¾åªç¼–è¯‘ä¸€æ¬¡ã€‚è®¡ç®—ä½¿ç”¨æœ‰ç•Œç¼“å†²åŒºåœ¨ TPU 上执行，而 XLA 会跟踪逻辑维度大小并自动处理填å 和切片，不会触发重新编译。

å† æ ¸

æˆ‘ä»¬åŒºåˆ†ä¸¤ç±»å† æ ¸ï¼š

• 由 PyTorch æ¡†æž¶æä¾›çš„å† æ ¸ï¼Œä¾‹å¦‚ torch.nn.functional.scaled_dot_product_attention 。

• ç”¨æˆ·æä¾›çš„è‡ªå®šä¹‰å† æ ¸ï¼Œç”¨äºŽåŠ é€Ÿæ¨¡åž‹çš„ç‰¹å®šéƒ¨åˆ†ï¼›è¿™äº›é€šå¸¸ä¸Žå ·ä½“æ¨¡åž‹ç›¸å ³ï¼›

对于第 1 类，我们最初实现了一种方法：用 Pallas 编写 TPU å† æ ¸ï¼Œç„¶åŽä»¥ç¬¦å·å½¢çŠ¶å¯¼å‡ºåˆ° Mosaic，再通过一个 TorchTPU æ­¥éª¤åœ¨è°ƒç”¨æ—¶ç‰¹åŒ–å† æ ¸çš„å½¢çŠ¶ã€‚è¯¥æœºåˆ¶è¢«è¯æ˜Žä¸å¦‚é¢„æœŸçµæ´»ï¼Œå› æ­¤æˆ‘ä»¬æ­£åœ¨è¯•éªŒä¸€ç§æ–°æ–¹æ³•ï¼Œä»¥ç¼–å†™æ”¯æŒåŠ¨æ€å½¢çŠ¶çš„ Mosaic å† æ ¸ã€‚

对于第 2 类，我们开发了三种解决方案：

• (Production ready) The user writes their TPU kernel in Pallas and integrates it with PyTorch via a TorchTPU API, as follows # Define a simple vector addition Pallas kernel def add_vectors_kernel(x_ref, y_ref, o_ref): o_ref[...] = x_ref[...] + y_ref[...] # Wrap the Pallas kernel into a JAX function @torch.tpu.pallas.jax_op("pallas::add_vectors") def add_vectors(x: jax.Array, y: jax.Array) -> jax.Array: out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype) return pl.pallas_call(add_vectors_pallas_kernel, out_shape)(x, y) # Use it like a standard PyTorch function x = torch.tensor([1.0, 2.0, 3.0], device="tpu") y = torch.tensor([4.0, 5.0, 6.0], device="tpu") z = add_vectors(x, y)

• （生产就绪） 用户可利用来自 Tokamax çš„å† æ ¸ ，这是一个高性能 TPU å† æ ¸åº“ï¼ŒçŽ°å·²å¯ç”¨äºŽä¸Ž PyTorch æ— ç¼äº’æ“ä½œã€‚Tokamax æä¾›å ³é”®ç®—å­çš„é«˜åº¦ä¼˜åŒ–å®žçŽ°ï¼ŒåŒ æ‹¬æ··åˆä¸“å®¶æž„å»ºæ¨¡å—ï¼ˆå¦‚åˆ†ç»„çŸ©é˜µä¹˜æ³•å’Œ token gather/scatter）、注意力变体（如分页注意力、门控 deltanet、Kimi delta 注意力和多头潜在注意力）以及 TopK è®¡ç®—ã€‚ç®—å­ç»„åˆå°†éšæ—¶é—´æ¼”è¿›ï¼Œä»¥æ”¯æŒåœ¨æœ€æ–°ç¡¬ä»¶åŠ é€Ÿå™¨ä¸Šè¿è¡Œçš„æœ€æ–° OSS 模型（例如，发布时支持 TPU v8iï¼‰ã€‚è¿™äº›å† æ ¸ç”±ç¨³å¥çš„é›†æˆåŸºç¡€åŸºç¡€è®¾æ–½æä¾›æ”¯æŒï¼Œä¾‹å¦‚è‡ªåŠ¨è°ƒä¼˜ï¼ˆå¦‚ tile 大小搜索）、低开销的基于 XProf çš„åŸºå‡†æµ‹è¯•å’Œå›žå½’æµ‹è¯•ï¼Œä»¥é™ä½Žè¿ç»´å¤æ‚æ€§ã€‚è™½ç„¶å† æ ¸æœ¬èº«ç”¨ JAX/Pallas 编写，但该库的基础设施采用 torch.tpu.pallas.jax_op API 来透明地管理互操作，并暴露原生 PyTorch API ä»¥æ–¹ä¾¿å¼€å‘è€ ã€‚

• (Experimental) The user rewrites their GPU kernel in Helion , which enables targeting multiple accelerators with the same PyTorch-based DSL, as shown below # Define the kernel directly in Helion @helion.kernel(backend="pallas") def add_vectors(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: out = torch.empty_like(x) # Device execution loop over automatically partitioned tiles for tile in hl.tile(out.size()): out[tile] = x[tile] + y[tile] return out # Use it like a standard PyTorch function x = torch.tensor([1.0, 2.0, 3.0], device="tpu") y = torch.tensor([4.0, 5.0, 6.0], device="tpu") z = add_vectors(x, y)

分布式支持

为确保与现有 PyTorch 分布式 API å’Œæ¡†æž¶å ¼å®¹ï¼Œæˆ‘ä»¬å†³å®šåœ¨ Process Group API 层面实现 TPU 集合通信，通过继承 c10d::Backend ；这使我们能够在 eager å’Œç¼–è¯‘æ¨¡å¼ä¸‹æ‹¦æˆªé›†åˆæ“ä½œï¼Œå¹¶æ— ç¼æ”¯æŒ PyTorch 分布式框架，如 DTensor、FSDPv2、SimpleFSDP、FairScale 等。与 GPU 上的 PyTorch ç±»ä¼¼ï¼Œæˆ‘ä»¬éµå¾ªï¼ˆå¹¶å®žé™ å¼ºåˆ¶æ‰§è¡Œï¼‰æ¯è®¾å¤‡ä¸€è¿›ç¨‹çš„æ–¹æ³•ã€‚ä»Žå®žçŽ°è§’åº¦çœ‹ï¼ŒTPU æ ˆä¸æä¾›å•ç‹¬çš„ TPU 间通信库，而是将 TPU 集合操作表示为 StableHLO 操作 （ stablehlo.all_gather 、 stablehlo.all_scatter 、 stablehlo.all_to_all 等），这与上述机制自然契合。

在 Process Group å±‚é¢æ‹¦æˆªé›†åˆæ“ä½œå‡ ä¹Žè®©æˆ‘ä»¬åˆ°è¾¾ç»ˆç‚¹ï¼Œé™¤äº†ä¸Ž GPU 相比 TPU 的一些显著特性：

• 不同的初始化序列 。考虑一个连接了多个设备的主机。在 GPU çš„æƒ å†µä¸‹ï¼Œç”± torchrun å¯åŠ¨çš„è¿›ç¨‹å¯ä»¥çœ‹åˆ°æ‰€æœ‰è¿žæŽ¥çš„è®¾å¤‡ï¼Œå³ä½¿å®ƒåªä¼šä½¿ç”¨å ¶ä¸­ä¸€ä¸ªï¼ˆä¾‹å¦‚ï¼Œä½¿ç”¨ device=f"cuda:{local_rank}" ），并且可以在调用 torch.distributed.init_process_group() 之前将模型移动到该设备。然而，在 TPU çš„æƒ å†µä¸‹ï¼Œä¸€ä¸ªè¿›ç¨‹åªèƒ½çœ‹åˆ°ä¸€ä¸ªè®¾å¤‡ï¼Œå¹¶ä¸”ä¸çŸ¥é“ä¸»æœºè¿žæŽ¥äº†å¤šå°‘è®¾å¤‡ã€‚è¦æ­£ç¡®åˆå§‹åŒ– TPU 运行时，进程需要在调用 torch.distributed.init_process_group 之前知道 (a) 运行分布式的意图，(b) TPU 数量（world size），以及 (c) TPU ç±»åž‹ï¼ˆå› ä¸ºå®ƒå½±å“æ‹“æ‰‘ï¼‰ã€‚

• Default Process Group Knowledge . Consider two subgroups (4, 0, 1, 3) and (5, 2, 6, 7) running an all_reduce among themselves. Note that torch.distributed.all_reduce takes as input only the subgroup in which the collective op must execute. In contrast, stablehlo.all_reduce requires information about all subgroups (i.e., PyTorch’s default process group), even those the current process is not part of, as shown in the following MLIR snippet: %0 = "stablehlo.all_reduce"(%arg0) <{replica_groups=dense<[[4,0,1,3], [5,2,6,7]]>: tensor<2x4xi64>}> ({...

• MPMD Support . Historically, the TPU stack was built with SPMD in mind. Thanks to tools such as Shardy , distributed JAX programs are written for a single host, with no communication code included, and then are sharded at massive scale with user-provided directives, while the XLA compiler is fully responsible for inserting collective ops at the right points while maximizing the overlap between computation and communication. Unfortunately, the PyTorch ecosystem doesn’t adopt a clean SPMD approach (a notable exception is SimpleFSDP ) and different processes may execute different ML ops leading to a collective op for reasons such as: ç”¨æˆ·æ’å ¥çš„ä»¥è¿›ç¨‹ rank 为条件的语句； å½“å¼ é‡ç»´åº¦ä¸èƒ½è¢«è®¾å¤‡æ•°é‡æ•´é™¤æ—¶ï¼ŒDTensor å¯¹å ±äº«å¼ é‡çš„ä¸å‡åŒ€æ‹†åˆ†ã€‚ å½“åŒ å«é›†åˆæ“ä½œçš„ TPU StableHLO æ¨¡å—ä¸ä¸€è‡´æ—¶ï¼Œå¯èƒ½ä¼šå‘ç”Ÿä¸å¯é¢„æµ‹çš„æ­»é”æˆ–å† å­˜æŸåï¼ˆè¿™æ˜¯ç”±äºŽç”Ÿæˆçš„ TPU ä»£ç ä¸­è·¨ rank çš„åŒæ­¥æ ‡å¿— ID æˆ–ç¼“å†²åŒºåç§»å¯èƒ½ä¸åŒ¹é ï¼‰ã€‚TorchTPU 对此的处理如下： (默认) Skinny Collectives ã€‚é»˜è®¤æƒ å†µä¸‹ï¼ŒTorchTPU ä¼šå°†æ¯ä¸ªé›†åˆé€šä¿¡æ“ä½œç¼–è¯‘åˆ°å ¶ç‹¬ç«‹çš„â€œskinny” StableHLO 模块中，这保证了各 rank ä¹‹é—´ç¼“å†²åŒºå’ŒåŒæ­¥æ ‡å¿— ID çš„ä¸€è‡´æ€§ã€‚è¿™æ ·åšè™½ç„¶å¯è¡Œï¼Œä½†ç›¸æ¯”ä¸‹ä¸€æ¡ä¸­è®¨è®ºçš„æ–¹æ¡ˆä¼šå¸¦æ¥æ˜Žæ˜¾çš„æ€§èƒ½æŸå¤±ï¼ˆä¾‹å¦‚ç«¯åˆ°ç«¯ 30-40%ï¼‰ï¼Œå› ä¸ºå®ƒä¸å è®¸ XLA ç¼–è¯‘å™¨å°†è®¡ç®—ä¸Žé€šä¿¡é‡å ã€‚ç›®å‰æ­£åœ¨è¿›è¡Œçš„ä¸€é¡¹å·¥ä½œæ˜¯å°† Skinny Collectives 运行在绑定到 SparseCore çš„ç‹¬ç«‹æµä¸Šï¼Œä»Žè€Œä½¿å ¶èƒ½å¤Ÿä¸Žè¿è¡Œåœ¨ TensorCore ä¸Šçš„æ•°å­¦è®¡ç®—é‡å ã€‚æˆ‘ä»¬é¢„è®¡è¯¥ä¼˜åŒ–å°†æŒ½å›ž Skinny Collectives 的大部分性能损失。 (可选) @spmd_safe è£ é¥°å™¨ ã€‚å¦‚æžœç”¨æˆ·çŸ¥é“æŸæ®µä»£ç æ˜¯ SPMD 的（例如，不存在以 torch.distributed.get_rank() 为条件的 if 语句），他们可以用 TorchTPU 的 @spmd_safe è£ é¥°å™¨æ ‡è®°è¯¥æ®µä»£ç ã€‚åœ¨ torch.compile 下，这会消除 TorchTPU åŽŸæœ¬ä¼šåœ¨é›†åˆé€šä¿¡å‘¨å›´æ’å ¥çš„å›¾æ–­ç‚¹ï¼ˆå‚è§ä¸Šæ–‡ Skinny Collectivesï¼‰ï¼Œå¹¶å› æ­¤ç”Ÿæˆä¸€ä¸ªå† åµŒé›†åˆé€šä¿¡çš„å¤§åž‹ StableHLO 模块，XLA ç¼–è¯‘å™¨å¯ä»¥é€šè¿‡ç§»åŠ¨é›†åˆé€šä¿¡ä½ç½®æ¥é«˜æ•ˆä¼˜åŒ–è¯¥æ¨¡å—ï¼Œä»Žè€Œæœ€å¤§åŒ–è®¡ç®—ä¸Žé€šä¿¡çš„é‡å ï¼›è¿™å·²è¢«è¯æ˜Žåœ¨ SimpleFSDP 或手动放置集合通信操作时能带来显著的性能收益。这种方法的一个挑战是 MPMD ä»£ç å¯èƒ½ä¸æ˜“å¯Ÿè§‰ï¼Œç”¨æˆ·å¯èƒ½åœ¨ä»£ç å¹¶éžä¸¥æ ¼ SPMD çš„æƒ å†µä¸‹ä½¿ç”¨ @spmd_safe è£ é¥°å™¨ï¼Œä»Žè€Œå¯¼è‡´è„†å¼±çš„æƒ å†µã€‚ (Optional) Compilation Handshake ( RFC ). This mechanism runs when a given torch.compile statement triggers compilation of a module containing communication operations. It initiates a consensus round across all replicas to verify that the compiled FX graphs are identical, which in turn allows replicas to disable graph breaks around collectives for the given FX graphs, similarly to the @spmd_safe mechanism described above.The handshake stage is controlled by an environment variable TORCH_TPU_SPMD_HANDSHAKE_STAGE which can be set to: OFF – 从不执行握手（当前默认）； COMPILE_STAGE – 在开始 FX 图编译之前执行握手； DISPATCH_STAGE – 在开始已编译的 FX å›¾æ‰§è¡Œä¹‹å‰æ‰§è¡Œæ¡æ‰‹ã€‚ä» å½“ç”¨æˆ·é¢„æœŸä¼šå‡ºçŽ° 特定于 rank 的重新编译 æ—¶æ‰åº”ä½¿ç”¨æ­¤æ¨¡å¼ï¼Œå› ä¸ºå®ƒæ¯”ä¸Šä¸€ä¸ªé€‰é¡¹æ¶‰åŠæ›´å¤§çš„å¼€é”€ã€‚

分布式初始化

在 TorchTPU ä¸­åˆå§‹åŒ–åˆ†å¸ƒå¼å·¥ä½œè´Ÿè½½éµå¾ªæ ‡å‡† PyTorch 分布式语义，但需注意 TorchTPU 需要两个 TPU ç‰¹æœ‰çš„å˜é‡ï¼Œè€Œæ ‡å‡† PyTorch torchrun 不知道如何设置它们：

• TORCH_TPU_TOPOLOGY ：TPU 切片的物理 mesh å‡ ä½•ç»“æž„ï¼ˆä¾‹å¦‚ï¼Œå¯¹äºŽ 2x2x1 åˆ‡ç‰‡é ç½®ä¸ºâ€œ2,2,1”）。

• TORCH_TPU_SLICEBUILDER_ADDRESSES ：用于构建 TPU 芯片间互连（ICI）网络 mesh 的开放端口/地址列表（每个设备一个）（例如，“localhost:8471,localhost:8472,…”）。

另一种方法是在 Python 中使用 torch.tpu.distributed.set_tpu_launch_env() ，它可以在 Google Cloud 环境中设置这些变量。下面我们展示启动分布式 TorchTPU å·¥ä½œè´Ÿè½½çš„ç¤ºä¾‹ä»£ç ã€‚

A. Worker 实现（ train.py ）

import os
import torch
import torch.distributed as dist

def worker_main():
  # 1. Initialize distributed process group on TPU using the tpu_dist backend
  dist.init_process_group(backend="tpu_dist")

  # 2. Query rank, world size, and construct target TPU device
  rank = dist.get_rank()
  world_size = dist.get_world_size()
  device = torch.device("tpu")

  print(f"[Rank {rank}/{world_size}] Initialized on {device}")

  # 3. Perform a distributed collective operation across TPU workers
  x = torch.ones(4, 4, device=device) * (rank + 1)
  dist.all_reduce(x, op=dist.ReduceOp.SUM)

  print(f"[Rank {rank}] Reduced sum: {x[0, 0].item()}")

  # 4. Clean up process group resources
  dist.destroy_process_group()

if __name__ == "__main__":
  worker_main()

B. 使用 torchrun 启动工作负载（CLI）

# Single host launch (spawns 1 process per TPU chip on an 8-chip node)
eval $(python -m torch.tpu.distributed | sed 's/^/export /')
torchrun --nproc_per_node=8 train.py

最后，请注意给定节点的默认进程组是该节点的 切片 。

最优设备 Mesh æ˜ å°„

由于 TPU 独特的连接性，将模型的并行策略（TC、CP、FSDP、DP）与物理 TPU 硬件互连拓扑对齐非常重要，以最大化带宽并最小化网络延迟。回顾一下，TPU 使用分层网络架构：

• ICI ï¼šç›´æŽ¥çš„å ‰çº¤/铜缆互连，以 Torus mesh 连接芯片。提供数十到数百 GB/秒的带宽。

• DCN ：用于连接不同 TPU åˆ‡ç‰‡çš„æ ‡å‡†ä¸»æœºçº§ä»¥å¤ªç½‘ã€‚å ¶å¸¦å®½æ˜¾è‘—ä½ŽäºŽ ICI，延迟显著高于 ICI。

一个好的规则是使用 TorchTPU 的 topology_aware_mesh() APIï¼Œå¹¶å°†é€šä¿¡æœ€å¯†é›†çš„å¹¶è¡Œç»´åº¦æ˜ å°„åˆ°æœ€å¿«ã€æœ€è¿žç»­çš„ç‰©ç† ICI 维度，也就是将 mesh 维度从网络强度最低到网络强度最高进行排序。示例见下表：

并行类型

通信模式

频率 / 数据量

最优物理放置

å¼ é‡å¹¶è¡Œï¼ˆTP）

all_reduce、reduce_scatter

每个 transformer 层（前向 + 反向）

最快的 ICI 环（同一 tray/host，连续的 1-hop 芯片）

上下文 / 序列并行（CP/SP）

all_to_all、Ring Attention

高（每个 attention 层）

本地 ICI 轴（slice å† ï¼‰

专家并行（EP）

all_to_all

每个 MoE router 步骤

本地 ICI 轴

FSDP / å®Œå ¨åˆ†ç‰‡ DP

all_gather、reduce_scatter

每个模块/层 1 次

复制数据并行（DP）

all_reduceï¼ˆä» æ¢¯åº¦ï¼‰

每个训练步骤一次

ICI 或 DCN（多 Slice）

ä»¥ä¸‹æ˜¯ä¸€ä¸ªä»£ç ç¤ºä¾‹ï¼š

from torch.distributed.device_mesh import DeviceMesh

# Single-Slice 2D Mesh Example (DP + TP)

# Logical mesh shape: (dp=4, tp=2) on an 8-chip slice ordered from least
# communication-intensive to most communication-intensive:
mesh_shape = (4, 2)
dim_names = ("dp", "tp")

# Computes optimal rank layout aligned with physical ICI rings
optimal_rank_tensor = torch.tpu.topology_aware_mesh(
    mesh_shape=mesh_shape,
)

# Instantiate the DeviceMesh
mesh_2d = DeviceMesh("tpu", mesh=optimal_rank_tensor,
                     mesh_dim_names=dim_names, preserve_rank_order=True)

dp_mesh = mesh_2d["dp"]
tp_mesh = mesh_2d["tp"]

作为未来工作，我们将探索一种 APIï¼Œå è®¸ç”¨æˆ·æ›´ç›´æŽ¥åœ°è‡ªå®šä¹‰å ¶ mesh，从而在特殊场景下实现最优的集合通信性能。

分布式 Pallas å† æ ¸

TorchTPU 支持自定义 Pallas å† æ ¸ï¼Œå¯é€šè¿‡ ICI 上的硬件 RDMA ç›´æŽ¥è¿›è¡ŒèŠ¯ç‰‡é—´é€šä¿¡ï¼Œæ— éœ€ä¸»æœº CPU å‚ä¸Žã€‚è¿™ä½¿ä½ å¯ä»¥ç¼–å†™é«˜æ•ˆçš„ç‚¹å¯¹ç‚¹é€šä¿¡å† æ ¸ï¼ˆä¾‹å¦‚ Ring Attention、自定义 All-to-All / Shuffles、流水线并行移位或 Halo 交换），并将它们作为原生 PyTorch 函数调用。示例如下：

# 1. Pallas Kernel Body: Asynchronously sends local buffer to neighbor rank (rank + 1)
def pallas_send_to_next_kernel(x_ref, y_ref, send_sem, recv_sem):
  rank = jax.lax.axis_index("x")
  world_size = jax.lax.axis_size("x")
  dest_id = (rank + 1) % world_size
  # Trigger hardware Remote DMA over TPU interconnect
  async_copy = pltpu.make_async_remote_copy(
      src_ref=x_ref,
      dst_ref=y_ref,
      send_sem=send_sem,
      recv_sem=recv_sem,
      device_id=(dest_id,),
  )
  async_copy.start()
  async_copy.wait()

# 2. JAX Wrapper with ShardMap and DMA Semaphores
def build_jax_fn(mesh):
  def local_pallas_call(x):
    return pl.pallas_call(
        pallas_send_to_next_kernel,
        out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype),
        grid_spec=pltpu.PrefetchScalarGridSpec(
            num_scalar_prefetch=0,
            in_specs=[pl.BlockSpec(x.shape)],
            out_specs=pl.BlockSpec(x.shape),
            scratch_shapes=(
                pltpu.SemaphoreType.DMA,
                pltpu.SemaphoreType.DMA,
            ),  # send & recv semaphores
        ),
    )(x)
  return jax.shard_map(
      local_pallas_call,
      mesh=mesh,
      in_specs=P("x"),
      out_specs=P("x"),
      check_vma=False,
  )

# 3. Create the TorchTPU callable wrapper
def get_send_to_next_op(world_size: int):
  mesh = torch.tpu.pallas.get_pallas_mesh(axis_names=("x",), mesh_shape=(world_size,))
  jax_fn = build_jax_fn(mesh)
  return torch.tpu.pallas.jax_op(
      "pallas::send_to_next",
      build_jax_fn(mesh),
      mesh=mesh,
      input_partition_specs=[P("x")],
      output_partition_specs=[P("x")],
  )

下表总结了如何编写通信型 Pallas å† æ ¸ã€‚

阶段

å ³é”®æ­¥éª¤ä¸Žæ“ä½œ

1. Pallas å† æ ¸ä¸»ä½“

查询逻辑 rank ä¸Žç›®æ ‡ ID

调用

pltpu.make_async_remote_copy(...)

控制异步调用生命周期：

desc.start()

/

desc.wait()

1. åˆ†ç‰‡æ˜ å°„ä¸Žä¿¡å·é‡

在

scratch_shapes

ä¸­åˆ†é  DMA 信号量

在 JAX Mesh 与 PartitionSpec 上使用

jax.shard_map

åŒ è£

1. TorchTPU 注册

使用

torch.tpu.pallas.jax_op

提供 mesh 和

input_partition_specs

直接在 PyTorch å¼ é‡ä¸Šè°ƒç”¨ï¼ˆEager 与

torch.compile

）

è§£å†³è®¾å¤‡æ˜ å°„ä¸ä¸€è‡´é—®é¢˜ 。此前在 TorchTPU 中，使用底层芯片间操作（如 lax.axis_index 和 async_remote_dma / make_async_remote_copy ）的 Pallas å† æ ¸åŸºäºŽåŽŸå§‹ç¡¬ä»¶è®¾å¤‡ ID 而非 PyTorch rank 来路由通信，与 JAX 保持一致。当硬件设备 ID 与 rank 顺序不一致时，环通信和跨设备 DMA ä¼šè¢«è·¯ç”±åˆ°éžé¢„æœŸçš„ç‰©ç†è®¾å¤‡ï¼Œè€Œåœ¨éš”ç¦»å­ç»„ä¸Šçš„æ‰§è¡Œä¼šå› åˆ†åŒºæ•°é‡ä¸åŒ¹é è€Œå¤±è´¥ã€‚

TorchTPU 通过强制逻辑寻址解决了这些不一致问题：

• Rank å¯¹é½æ˜ å°„ ：TorchTPU 现在将 Pallas å† æ ¸ç›´æŽ¥æ˜ å°„åˆ° PyTorch rank，而非原始物理设备 ID，通过 get_pallas_mesh API 确保路由一致，该 API 将 pallas ä¸­çš„åˆ†åŒºå’Œè½´ç›´æŽ¥æ˜ å°„åˆ° PyTorch Process Group 中的 rank 顺序。

• 逻辑设备寻址： TorchTPU 现在强制使用 DeviceIdType.LOGICAL å’Œè§„èŒƒçš„è®¾å¤‡åˆ†é ï¼ˆ da(0, p) = p ），保证每个逻辑 rank å§‹ç»ˆæ˜ å°„åˆ°å ¶æ­£ç¡®çš„ç‰©ç†ç›®æ ‡ã€‚

统一 ProcessGroup 集成。 TorchTPU 现在提供获取 rank 对齐 JAX Mesh çš„æœºåˆ¶ï¼Œä½¿å† æ ¸èƒ½å¤Ÿç»§æ‰¿ PyTorch ProcessGroup æ‹“æ‰‘ã€‚è¿™å¯ä»¥åœ¨å ¨é‡ä¸–ç•Œå’Œéš”ç¦»å­ç»„é ç½®ä¸‹é«˜æ•ˆæ‰§è¡Œï¼Œæ— éœ€æ‰‹åŠ¨ä¿®è¡¥çŠ¶æ€æˆ–ä¸­æ–­è®¡ç®—å›¾ã€‚

开发模型

TorchTPU 以开放方式开发，是 Google ä¸Žç¤¾åŒºåä½œçš„æˆæžœï¼Œä¼˜å ˆè€ƒè™‘é€æ˜Žåº¦å’Œå‚ä¸Žåº¦ã€‚è¯¥å¼€å‘æ¨¡å¼é¼“åŠ±æ›´å¹¿æ³›çš„ PyTorch 生态和行业合作伙伴贡献力量。

我们预计 TorchTPU å°†æœåŠ¡äºŽå¤šæ ·åŒ–çš„ç”¨æˆ·ç¾¤ä½“ï¼ŒåŒ æ‹¬ç ”ç©¶æœºæž„å’Œè¡Œä¸šé¢†å¯¼è€ ã€‚æˆ‘ä»¬è¿˜è®¡åˆ’å°†å ¶ä¸Ž Hugging Face ç­‰å¹³å°é›†æˆï¼Œä¸ºå¹¿æ³›çš„å¼€æ”¾æºä»£ç æ¨¡åž‹æä¾›å¼€ç®±å³ç”¨çš„ TPU 支持。

ä¸ºæ”¯æŒå¼€å‘è€ ï¼ŒTorchTPU æä¾›å ¨é¢çš„æ–‡æ¡£ï¼ŒåŒ æ‹¬ （参见 http://torchtpu.readthedocs.io/ ） ：

• 教程 ï¼šå¼•å¯¼å¼å­¦ä¹ è·¯å¾„ï¼ŒåŒ æ‹¬è™šæ‹Ÿæœºé ç½®å’Œæ¨¡åž‹ç§»æ¤ã€‚

• é æ–¹ ï¼šé¢å‘é—®é¢˜çš„è§£å†³æ–¹æ¡ˆï¼Œé€‚ç”¨äºŽæ€§èƒ½ä¼˜åŒ–å’Œå† æ ¸å®žçŽ°ç­‰ä»»åŠ¡ã€‚

• 参考 ：技术规范和 API æ˜ å°„ã€‚

• 说明 ï¼šæ·±å ¥æŽ¢è®¨æž¶æž„å’Œç¡¬ä»¶åŸºç¡€ã€‚

文档使用可直接运行的 notebook 构建。

该项目利用 Bazel 进行稳健的构建和测试管理，测试框架仍在演进，并计划采用 Pytest。CI/CD 基础设施正在集成到官方 PyTorch CI HUD ä¸­ï¼Œä»¥ä¿æŒé«˜å¯é æ€§å¹¶ä¿ƒè¿›ç¤¾åŒºè´¡çŒ®ã€‚

TorchTPU è®¾è®¡ä¸ºå¯æ‰©å±•ï¼Œèƒ½å¤Ÿä¸Žå ³é”® PyTorch ç”Ÿæ€ç»„ä»¶æ— ç¼é›†æˆã€‚å®ƒæ”¯æŒé€šè¿‡ Pallas 和 Helion è¿›è¡Œè‡ªå®šä¹‰å† æ ¸å¼€å‘ï¼Œå¹¶é€šè¿‡ vLLM、SGLang 和 TorchTitan ç­‰é¡¹ç›®ä¸ºåˆ†å¸ƒå¼è®­ç»ƒå’ŒæŽ¨ç†å·¥ä½œè´Ÿè½½æä¾›ä¸€ç­‰æ”¯æŒã€‚è¿™ç§æ¨¡å—åŒ–æ–¹æ³•å è®¸ç”¨æˆ·åˆ©ç”¨çŽ°æœ‰çš„ PyTorch 惯用法，同时针对 TPU 特定的硬件架构进行优化。

æ¡ˆä¾‹ç ”ç©¶

æ¡ˆä¾‹ç ”ç©¶ 1：Hugging Face 模型覆盖

TorchTPU 提供与 PyTorch ç”Ÿæ€çš„å¼€ç®±å³ç”¨å ¼å®¹æ€§ï¼Œä½¿å¼€å‘è€ æ— éœ€ä¿®æ”¹ä»£ç æˆ–é‡æ–°ç¼–å†™å³å¯åœ¨ Google Cloud TPU ä¸Šè¿è¡Œå¼€æ”¾æºä»£ç æ¨¡åž‹ã€‚ä¸ºéªŒè¯è¿™ä¸€æ‰¿è¯ºï¼Œæˆ‘ä»¬æž„å»ºäº†ä¸€ä¸ªè‡ªåŠ¨åŒ–æµ‹è¯•æ¡†æž¶ï¼Œä½¿ç”¨ Hugging Face 模型覆盖 NLP、视觉、音频、多模态和生成式扩散领域。该持续集成流水线以 CPU 和 GPU 基线为基准，对 TorchTPU 的覆盖范围、稳定性和数值等价性进行基准测试。我们在从 v5e 开始的 TPU ç¡¬ä»¶ä»£é™ ä¸Šï¼Œä»¥ eager 和编译模式持续评估约 5,000 个 Hugging Face 模型，涵盖数百种架构。

æˆ‘ä»¬é‡‡ç”¨åŸºäºŽè¯­æ–™åˆ†å±‚å’Œå‡è¡¡æž¶æž„é€‰æ‹©çš„ç¨³å¥æµ‹è¯•æ–¹æ³•ã€‚æˆ‘ä»¬çš„æ–¹æ³•ä¾èµ–äºŽä¸¤çº§é€‰æ‹©æµç¨‹ï¼Œæ—¨åœ¨å¹³è¡¡å®žé™ é‡‡ç”¨ä¸Žå¹¿æ³›çš„æž¶æž„å¤šæ ·æ€§ï¼Œé¿å å¯¹ç»†å¾®å˜ä½“è¿›è¡Œå†—ä½™æµ‹è¯•ï¼ŒåŒæ—¶ç¡®ä¿çœŸå®žç”Ÿäº§æƒé‡å¾—åˆ°å åˆ†è¯„ä¼°ã€‚åœ¨ç¬¬ä¸€çº§ä¸­ï¼Œæˆ‘ä»¬æ ¹æ®ç¤¾åŒºé‡‡ç”¨åº¦ã€ä¸‹è½½é€Ÿåº¦å’Œè¶‹åŠ¿è¯„åˆ†ä¼˜å ˆé€‰æ‹©æŽ’åé å‰çš„æ¨¡åž‹â€”â€”ä¾‹å¦‚ Qwen、Gemma、Mistral、DeepSeek、GLM、Kimi、Llama、Whisper、Stable Diffusion、FLUX 和 Vision Transformerâ€”â€”ä»¥ç¡®ä¿å ³é”®ç¤¾åŒºçƒ­é—¨æ¨¡åž‹å®Œç¾Žè¿è¡Œã€‚ç¬¬äºŒçº§æ˜¯ä»Žç¬¬ä¸€çº§ä¸­é€‰å‡ºçš„èšç„¦å­é›†ï¼Œç¡®ä¿æˆ‘ä»¬ä»Žè¶ è¿‡ 500 个不同的模型家族中至少选择 1 ä¸ªä»£è¡¨æ€§æ¨¡åž‹ã€‚æˆ‘ä»¬ä½¿ç”¨ç¬¬äºŒçº§åœ¨å ¨é¢çš„ PyTorch 算子和编译器模式集合上验证相对于 GPU 基线的数值准确性。

验证过程从第 1 阶段开始，该阶段侧重于冒烟测试和降低验证，以保证运行时稳定性。在此初始阶段，流水线使用合成权重实例化模型，验证 eager 模式的前向执行，并确认图追踪、FX 到 StableHLO 的图生成以及 XLA ç¼–è¯‘å‡èƒ½æˆåŠŸå®Œæˆä¸”æ— é”™è¯¯ï¼ˆä¾‹å¦‚ï¼Œæ— å´©æºƒæˆ–å¼‚å¸¸ï¼‰ã€‚ä¸€æ—¦ç¡®è®¤è¿è¡Œæ—¶ç¨³å®šæ€§ï¼Œç¬¬ 2 é˜¶æ®µä¾¿ä½¿ç”¨é¢†åŸŸç‰¹å®šçš„è¯„ä¼°æŒ‡æ ‡ï¼Œå¯¹ç §æ—¢å®šåŸºçº¿è¯„ä¼°æ•°å€¼æ­£ç¡®æ€§ã€‚å¯¹äºŽå› æžœ LLM，数学保真度通过困惑度（ Perplexity = e L = exp((1/ N ) Σ i =1 N L i )）来衡量，若 TPU 与基线困惑度之间的相对差异保持在 0.10 ä»¥å† ï¼Œåˆ™æµ‹è¯•é€šè¿‡ã€‚å¯¹äºŽè§†è§‰ã€éŸ³é¢‘å’Œç¼–ç å™¨æž¶æž„ä¸­çš„é«˜ç»´ç‰¹å¾ç©ºé—´ï¼Œæˆ‘ä»¬é‡‡ç”¨åŒé˜ˆå€¼ç›¸å¯¹ MSE æŒ‡æ ‡ï¼ˆ Relative MSE = MSE( Y TPU , Y Baseline ) / (Var( Y Baseline ) + 10 −7 )），要求相对 MSE ≤ 0.10 或绝对 MSE < 0.05，以防止基线输出接近常数。

下表总结了截至本文撰写时的模型覆盖结果：

æŒ‡æ ‡ç±»åˆ«

范围 / 定义

结果

已验证模型总数

端到端测试的不同开源模型

4,917 个模型

独特架构家族

跨 Transformer、Diffuser 等的不同架构类别

500+ 个架构家族

第 1 层（流行度）– é€šè¿‡çŽ‡ï¼ˆå³æ— å´©æºƒï¼‰

下载量最高的社区 LLM、视觉、音频和多模态模型

98%

– 4,917 个中的 4,821 个

第 2 层（架构）– 数值一致性

æ ‘å† ä¸åŒæž¶æž„çš„ä»£è¡¨æ€§æ¨¡åž‹

97%

– 657 个中的 638 个（相对于 GPU）

对于有界动态性支持（针对 SymInts），我们在 839 个模型（来自 5K è¯­æ–™åº“çš„å› æžœ LLM 列表）上评估了正确性，并确认了 834 个（99%）模型的数值正确。

æ¡ˆä¾‹ç ”ç©¶ 2：开箱即用的性能

TorchTPU 致力于成为 PyTorch 的真正原生 TPU åŽç«¯ã€‚åœ¨æˆ‘ä»¬çœ‹æ¥ï¼Œè¿™ä¸€æœŸæœ›åŒ æ‹¬åœ¨è¿è¡Œçº¯ PyTorchã€è®¾å¤‡æ— å ³ä»£ç æ—¶å ·æœ‰è‰¯å¥½çš„æ€§èƒ½ã€‚æˆ‘ä»¬é€šè¿‡ä¸¤ä¸ªç¤ºä¾‹æ¥å±•ç¤ºè¿™ä¸€ç‚¹ã€‚

ç¬¬ä¸€ä¸ªç¤ºä¾‹åŒ å«åœ¨ TorchTPU ä»“åº“ä¸­ï¼Œæ˜¯ä¸€ä¸ªè‡ªåŒ å«çš„ã€çº¯ PyTorch 实现的 Muse Glimmer 30B æ¨¡åž‹ï¼Œæˆ‘ä»¬ä» ä½¿ç”¨æ ‡å‡† PyTorch 模块和算子编写，例如 nn.Linear 、 nn.RMSNorm 、 F.scaled_dot_product_attention 和 torch.distributed 集合通信。我们 有意 æ²¡æœ‰ä½¿ç”¨ä»»ä½•æ¨¡åž‹ç‰¹å®šçš„è‡ªå®šä¹‰å† æ ¸ï¼ˆä½†ä»ä½¿ç”¨äº† PyTorch å† éƒ¨å† æ ¸ï¼‰ã€‚åœ¨æ­¤ç¤ºä¾‹ä¸­ï¼Œæˆ‘ä»¬å±•ç¤ºäº†ï¼š

• Inference (prefill) performance on a single TPUv7 chip at 序列长度为 5,120 时， ~25.7k TPS ， 59% MFU ，以及 序列长度为 10,240 时， ~22.4k TPS ， 51% MFU 。 请注意，此示例使用 FSDP（通过原生 PyTorch 集合通信实现），world size 为 2ï¼Œå› ä¸ºå•ä¸ª TPUv7 芯片向 PyTorch æš´éœ²ä¸¤ä¸ªè®¾å¤‡ï¼ˆå¯¹åº”ä¸¤ä¸ªåŒ å«çš„ TensorCore）。

• Training on 4 TPUv7 chips (i.e. a full single host), using FSDP, at 每个 TPU 在序列长度为 5,120 时达到 ~5.6k TPS ， 38% MFU ，以及 每个 TPU 在序列长度为 10,240 时达到 ~4.6k TPS ， 32% MFU 。 在第二个示例中，我们启用了 激活检查点 ï¼Œå®ƒä» åœ¨ HBM ä¸­å­˜å‚¨å‰å‘ä¼ æ’­æ¿€æ´»çš„ä¸€ä¸ªå­é›†ï¼Œå¹¶åœ¨åå‘ä¼ æ’­è®¡ç®—éœ€è¦æ—¶é‡æ–°è®¡ç®—/重物化未存储的中间激活。

æ‰€è€ƒè™‘ç¤ºä¾‹çš„ä¸€ä¸ªå†³å®šæ€§ç‰¹å¾æ˜¯ï¼ŒåŒä¸€ä»½ä»£ç åœ¨é™¤æŒ‡å®šåŠ é€Ÿå™¨ç±»åž‹ï¼ˆä¾‹å¦‚ç¼–å†™ device="tpu" 而非 device="cuda" ）之外 没有任何 å®žè´¨æ€§æ”¹åŠ¨çš„æƒ å†µä¸‹ï¼Œå³å¯åœ¨ TPU 和 GPU ä¸Šè¿è¡Œï¼Œå¹¶ä¸”åœ¨æ ‡å‡† PyTorch 优化与分片技术（如 torch.compile 和 FSDP）下在 TPU 上表现非常出色。

为进一步验证 TorchTPU 在纯 PyTorch ä»£ç ä¸Šçš„æ€§èƒ½ï¼Œæˆ‘ä»¬è¿˜ä»Ž TorchTitan 项目中选取了一个 外部产出 的 Muse Glimmer 30B 实现，该实现最初是为 GPU/CUDA ç¼–å†™çš„ã€‚åœ¨å¯¹å ¶åšæœ€å°é™åº¦é€‚é åˆ° TorchTPU åŽï¼ˆæˆ‘ä»¬ä» å°† TorchTPU 尚不支持的 FlexAttention 替换为原生 F.scaled_dot_product_attention ，并应用了 torch.compile 和 SimpleFSDP），TorchTPU 移植版实现了：

• Training on 4 TPUv7 chips: 在序列长度为 5,120 时，每 TPU 约 5.3k TPS ， MFU 为 36% ，以及 在序列长度为 10,240 时，每 TPU 约 4.2k TPS ， MFU 为 29% 。

总之，所有这些示例表明，纯 PyTorch 模型实现，甚至最初为 GPU 编写的实现，只需简单的 PyTorch æ”¹åŠ¨å¹¶é åˆæ ‡å‡† PyTorch 优化，即可在 TPU ä¸Šè¡¨çŽ°è¶³å¤Ÿå‡ºè‰²ï¼Œè€Œæ— éœ€å€ŸåŠ©è‡ªå®šä¹‰å† æ ¸ç­‰é«˜çº§çš„ã€è®¾å¤‡ç‰¹å®šçš„çˆ¬å±±æŠ€æœ¯ã€‚æˆ‘ä»¬è®¤ä¸ºè¿™æ˜¯ TorchTPU 采用的一项重要特性。当然，对于追求峰值性能的生产工作负载，我们绝对建议采用基于 Pallas å† æ ¸çš„æœ€ä¼˜æ–¹æ¡ˆï¼Œä»¥åŠ vLLM-TorchTPU、SGLang-TorchTPU 或 TorchTitan-TorchTPU ç­‰ç”Ÿäº§çº§æŠ€æœ¯æ ˆä½œä¸ºç¤ºä¾‹ã€‚

æ¡ˆä¾‹ç ”ç©¶ 3：将训练扩展到 Pod 规模

我们使用 TorchTitan 验证了 TorchTPU 可扩展至 1 个 TPU v7 POD。对于此测试，我们选择了 Qwen3 MoE 235B。作为初步冒烟测试，我们确认在 64 个 TPU v7 é ç½®ä¸Šï¼ŒTorchTitan é åˆ TorchTPU 的模型性能与 MaxText 上的 JAX 相当（PyTorch/TorchTitan 为 1828 TPS，JAX/MaxText 为 1937 TPS）。

下文讨论的数据基于以下设置计算得出：

• 模型： qwen3_moe_235b_scan

• 本地批大小（即每个 TensorCore）：16，序列长度：4096

• 并行策略：单切片运行使用 FSDP（通过 SimpleFSDP），多切片运行使用 HSDP

• 精度：BF16 é åˆ FP32 混合精度归约（梯度归约：f32，RoPE：f32，优化器状态：f32）

• å† æ ¸ï¼ˆPallas）：Tokamax Splash Attentionï¼ˆå¸¦æ¥å¤§å¹ æ€§èƒ½æå‡ï¼‰ã€èžåˆäº¤å‰ç†µæŸå¤±å† æ ¸ï¼Œä»¥åŠ MoE 层的分组矩阵乘法

• 扫描优化：94 ä¸ªè§£ç å™¨å±‚ä½œä¸ºå•ä¸ªæ‰«æå¾ªçŽ¯åœ¨å †å æƒé‡ä¸Šè¿è¡Œ

• è§£ç å™¨å±‚è¾“å ¥çš„æ¿€æ´»å¸è½½ï¼ˆä½œä¸ºæ‰«æä¼˜åŒ–çš„ä¸€éƒ¨åˆ†ï¼‰

• é›†åˆé€šä¿¡æµæ°´çº¿åŒ–ï¼Œä»¥å°†é›†åˆé€šä¿¡ä¸Žç›¸é‚»å±‚çš„è®¡ç®—é‡å ï¼ˆæ‰«æçš„ä¸€éƒ¨åˆ†ï¼‰

• 完整激活检查点

• 使用 TorchTitan 的 Graph Trainer è¿›è¡Œè”åˆå‰å‘å’Œåå‘ä¼ æ’­

• 专家不均衡：使用 TorchTitan çš„è¾ åŠ©æŸå¤±é˜²æ­¢ MoE 不均衡

• 训练数据集：Hugging Face C4 en æ•°æ®é›†ï¼Œé åˆè‡ªå®šä¹‰æ•°æ®é›†åˆ†ç‰‡

拓扑

TPU v7 芯片数量

TensorCore 数量

总批大小

每 TensorCore TPS

MFU

步时间

1 åˆ‡ç‰‡é åˆ FSDP

1024

2048

32768

2,236

28.8%

29.3s

2 åˆ‡ç‰‡é åˆ HSDP

4096

65536

2,214

28.5%

29.6s

4 åˆ‡ç‰‡é åˆ HSDP

8192

131072

2,095

26.9%

31.3s

8 åˆ‡ç‰‡é åˆ HSDP

16384

262144

2,104

27.1%

31.2s

在所有拓扑下，我们的数据证实了 弱扩展 ：整体计算量呈线性扩展，如下所示，总每秒 token 数线性增长，而 MFU 基本保持平稳。

我们的数据还显示，在所考虑的各种并行级别下， 损失曲线 å‡ ä¹Žå®Œå ¨ç›¸åŒã€‚

在多切片 HSDP é ç½®ä¸­ï¼Œå‰å‘ä¼ æ’­ä¸æ¶‰åŠä»»ä½•è·¨åˆ‡ç‰‡é€šä¿¡ï¼Œæ¯ä¸ªåˆ‡ç‰‡å¤„ç†è‡ªå·±é‚£ä»½æ‰¹æ¬¡æ•°æ®ã€‚ä½†åœ¨åå‘ä¼ æ’­ä¸­ï¼Œä¸€ä¸ªåˆ‡ç‰‡å ˆè®¡ç®—ç¬¬ t 层的梯度，然后通过 all-reduce（在 DCN 上，使用 Megascale-XLA）对各切片的所有这些梯度求和。计算随后继续处理第 t-1 层的梯度。得益于 XLA，TorchTPU 能够将第 t+1 层的 all-reduce 与第 t å±‚çš„è®¡ç®—å®Œå ¨é‡å ï¼Œå¦‚ä¸‹å›¾æ‰€ç¤ºã€‚å®žçŽ°è¿™ç§é‡å æ˜¯æ¯ä¸ª TensorCore 的 TPS ä»Žä¸€ç‰‡åˆ°å «ç‰‡å‡ ä¹Žä¿æŒå¹³ç¨³çš„ä¸»è¦åŽŸå› ï¼Œå¦‚ä¸Šè¡¨æ‰€ç¤ºã€‚

æ¡ˆä¾‹ç ”ç©¶ 4：在 vLLM 上使用 Qwen3.5 进行 LLM 推理（SemiAnalysis InferenceX）

大语言模型（LLMï¼‰æŽ¨ç†â€”â€”å°¤å ¶æ˜¯é’ˆå¯¹å¤§è§„æ¨¡æ··åˆä¸“å®¶ï¼ˆMoE）架构——是 AI äº‘åŸºç¡€è®¾æ–½ä¸­è¦æ±‚æœ€é«˜ã€æˆæœ¬æœ€å ³é”®çš„å·¥ä½œè´Ÿè½½ä¹‹ä¸€ã€‚åœ¨ SemiAnalysis 发布的一项独立基准评估中（ TPU Inference Externalization Full Steam Ahead — InferenceX ），Google 的 TPU v7 æ­é  vLLM-TorchTPU è½¯ä»¶æ ˆï¼Œäº¤ä»˜äº†å ·æœ‰è¡Œä¸šç«žäº‰åŠ›çš„æ¯ç¾Žå ƒæ€§èƒ½ã€‚è¯¥åŸºå‡†æä¾›äº†å¯¹ TorchTPU 软件在大规模混合专家（MoE）工作负载上就绪程度的第三方推理审计，表明面向 TPU 的新 vLLM è½¯ä»¶æ ˆåœ¨ç¤¾åŒºå ¬å¼€å¯ç”¨çš„ TPU 系统上提供了强大的成本效率和高吞吐量。

在此基础上，Google 最近提交了 Qwen 3.5 在 Ironwood TPU 上的首批 AgentX 性能结果。AgentX æ•æ‰äº†æ™ºèƒ½ä½“å·¥ä½œè´Ÿè½½çš„å››ä¸ªè¦ç´ â€”â€”å¤šè½®ã€é•¿ä¸Šä¸‹æ–‡ã€é«˜å‰ç¼€å¤ç”¨å’Œå­æ™ºèƒ½ä½“çªå‘ã€‚TPU å·²èƒ½å¤Ÿè¯æ˜Žå ¶å³å°†æŽ¨å‡ºçš„ vLLM-TorchTPU è½¯ä»¶æ ˆèƒ½å¤ŸæˆåŠŸå¤„ç†æ™ºèƒ½ä½“å·¥ä½œè´Ÿè½½ã€‚ä» ç”¨ 4 块芯片，TPU v7 å°±å®žçŽ°äº†å¸•ç´¯æ‰˜å‰æ²¿çš„å ¨é¢é¢†å ˆã€‚

• 在 InferenceX 8K/1K 上，每百万 token 总体成本最低，为 2.01 美分（低于我们 9 月首次亮相时的 3.6 美分）；

• 在 AgentX 上，每百万 token 总体成本最低，为 0.83 ç¾Žåˆ†ï¼Œå¹¶å ·å¤‡ SOTA äº¤äº’æ€§ï¼ˆç”¨æˆ·å“åº”èƒ½åŠ›ï¼‰ã€‚å¸•ç´¯æ‰˜å‰æ²¿å ¨é¢é¢†å ˆã€‚

è¿™äº›åŸºå‡†ç»“æžœåæ˜ äº†åœ¨ TPU 系统上实现的强大、真实的在线服务效率，突显了在 Google 定制设计的 ML 芯片上部署大规模模型的巨大经济可行性。

要实现这些有竞争力的性能和 TCO 结果，需要通过自定义 Pallas å† æ ¸å’Œæœ‰é’ˆå¯¹æ€§çš„è¿è¡Œæ—¶è°ƒä¼˜æ¥è§£å†³å ³é”®çš„è½¯ç¡¬ä»¶å¯¹é½ç“¶é¢ˆï¼Œä¸‹è¡¨åˆ—å‡ºäº†å‡ ä¸ªç¤ºä¾‹ï¼š

优化

瓶颈

实现

将专家并行 MoE 层作为单个 SPMD 程序

在 DP8 + EP 下，vLLM 的 MoE 层将每个 rank 的 token 收集到每个 rank，对它们进行路由、置换，运行两个分组矩阵乘法和一次合并，然后将结果 reduce-scatter 回去。在并发度为 256 的 prefill 密集型步骤中，该 reduce-scatter 每层耗时 1,190 µsï¼Œå æ‰€æœ‰è®¾å¤‡æ—¶é—´çš„ 22%ï¼Œä¸”æœŸé—´æ²¡æœ‰å ¶ä»–æ“ä½œå¹¶è¡Œè¿è¡Œã€‚

一个 SPMD ç¨‹åºåœ¨å ¨éƒ¨å «ä¸ª TensorCore 上运行该层。每个 TensorCore å°†å ¶è¡Œé‡åŒ–ä¸º FP8，并在 SparseCore 上异步地对它们进行 all-gather。同时，它用 Pallas top-k 为自己的 token é€‰æ‹©ä¸“å®¶ï¼Œå¹¶åœ¨ä¸“å®¶ç´¢å¼•å’Œè¡Œç¼©æ”¾å› å­çš„ç¬¬äºŒæ¬¡å°åž‹ all-gather ä¸­å ±äº«å®ƒä»¬ã€‚éšåŽï¼Œä¸€ä¸ª Pallas å† æ ¸åœ¨è·¯ç”±åˆ°æ¯ä¸ªæœ¬åœ°ä¸“å®¶çš„è¡Œä¸Šè¿è¡Œè¯¥ä¸“å®¶çš„ä¸¤ä¸ªçŸ©é˜µä¹˜æ³•ï¼ŒæŒ‰ç´¢å¼•èŽ·å–è¿™äº›è¡Œã€‚å®ƒé€šè¿‡è¿œç¨‹ DMA 将每个结果行（以 FP8 å½¢å¼ï¼Œå¸¦æ¯è¡Œç¼©æ”¾å› å­ï¼‰æŽ¨é€åˆ°æ‹¥æœ‰è¯¥ token 的 TensorCore（

pltpu.make_async_remote_copy

ï¼‰ã€‚è¿™äº›æŽ¨é€ä¸ŽåŽç»­ä¸“å®¶çš„çŸ©é˜µä¹˜æ³•é‡å ã€‚ç„¶åŽï¼Œæ¯ä¸ª TensorCore å¯¹ä¸ºå ¶ token åˆ°è¾¾çš„è¡Œæ±‚å’Œï¼Œå¹¶æŒ‰è·¯ç”±å™¨åˆ†æ•°åŠ æƒã€‚æ¯å±‚ä¸­ï¼Œè¿™å–ä»£äº†è·¯ç”±è¡Œçš„ argsort、SparseCore 的 permute 和 combine、两个分组矩阵乘法以及 combine reduce-scatter。

Sequence-on-Lane æ³¨æ„åŠ›å¸ƒå±€ï¼ˆæ— å¡«å FP8 KV 缓存）

TPU çš„å‘é‡å•å ƒä»¥ 32 位字工作，KV ç¼“å­˜ä¹ŸæŒ‰æ­¤å¸ƒå±€ï¼Œå› æ­¤å››ä¸ª FP8 å€¼å ±äº«æ¯ä¸ªå­—ã€‚æ ‡å‡† KV 布局用来自多个头的 K 和 V å¡«å 一个字。在 8 è·¯å¼ é‡å¹¶è¡Œä¸‹ï¼Œæ¯ä¸ªæ ¸å¿ƒåªæŒæœ‰æ¨¡åž‹ä¸¤ä¸ª KV å¤´ä¸­çš„ä¸€ä¸ªï¼Œå› æ­¤æ¯ä¸ªå­—çš„ä¸€åŠæ˜¯å¡«å ï¼šæ¯ä¸ªæ ¸å¿ƒä¸Šæ¯ä¸ª token 的 1,024 å­—èŠ‚ç¼“å­˜ä» å¯¹åº” 512 字节数据。

Sequence-on-lane 将一页的 token 沿 128 个 lane 排列，将头维度沿子 lane å‘ä¸‹æŽ’åˆ—ï¼Œå› æ­¤ä¸€ä¸ªå­—å®¹çº³åŒä¸€ token 的四个值，没有任何浪费。这将 KV å† å­˜å’Œæ¯ä¸ªè§£ç æ­¥éª¤è¯»å–çš„ KV 字节减半。Sequence-on-lane 是 vllm-torchtpu 现在支持的通用布局。

Gated DeltaNet 优化

60 层中有 45 层使用 Gated DeltaNet，它每个头保持固定的 128 × 128 状态，而不是不断增长的 KV ç¼“å­˜ã€‚è§£ç å¿ é¡»åœ¨æ¯ä¸€æ­¥è¯»å–å’Œå†™å ¥æ¯ä¸ªåºåˆ—çš„çŠ¶æ€ï¼›åœ¨å¹¶å‘åº¦ä¸º 64 æ—¶ï¼Œç§»é™¤å† æ ¸çš„æ‰€æœ‰ç®—æœ¯åªè®©å®ƒå¿«äº† 6%。逐 token 计算时，预填å æ˜¯ä¸€ä¸ªé¡ºåºé€’å½’ï¼Œä½¿çŸ©é˜µå•å ƒç©ºé—²ã€‚

一个 Pallas å† æ ¸èžåˆäº†å·ç§¯å’Œé€’å½’ï¼Œæ‰¹å¤„ç†è§£ç åºåˆ—ï¼Œå¹¶ä»¥ 64 token 的块处理预填å ，大部分重写为矩阵乘法。状态与 KV ç¼“å­˜ä½äºŽåŒä¸€ä¸ªåˆ†é¡µæ± ä¸­ã€‚å† æ ¸ç›´æŽ¥ä»Žæ± é¡µæµå¼ä¼ è¾“å®ƒï¼Œå¹¶åŽŸåœ°è¯»å–é€’å½’çŠ¶æ€ï¼Œé€šè¿‡æ”¹å˜å¯»å€è¿™äº›å­—èŠ‚çš„æ–¹å¼è€Œä¸æ˜¯è½¬æ¢å®ƒä»¬ã€‚å› æ­¤ï¼Œå ±äº«æ± ä¸ä¼šå¸¦æ¥åžåé‡æŸå¤±ï¼šåœ¨æˆ‘ä»¬çš„æµ‹è¯•ä¸­ï¼Œå®ƒåœ¨æ¯ç§é ç½®ä¸‹éƒ½æä¾›äº†ç­‰äºŽæˆ–é«˜äºŽç‹¬ç«‹ç¼“å­˜çš„åžåé‡ã€‚

这些结果突显了现代服务基础设施中的 TPU 优化，表明 TCO 和软硬件对齐决定了真实世界的服务效率。最终，带有 TorchTPU 后端的新 vLLM-TPU ä½¿å·¥ç¨‹å›¢é˜Ÿèƒ½å¤Ÿåœ¨å ¬å¼€å¯ç”¨çš„ TPU åŸºç¡€è®¾æ–½ä¸Šä½¿ç”¨æ ‡å‡† PyTorch 惯用法部署现成的开放权重模型。

æ¡ˆä¾‹ç ”ç©¶ 5：使用 TorchTitan 进行 LLM 微调

TorchTPU 与 TorchTitan 原生集成——TorchTitan 是 PyTorch é¢å‘å¤§è§„æ¨¡åˆ†å¸ƒå¼è®­ç»ƒçš„å‚è€ƒå¹³å°â€”â€”åœ¨ç¨ å¯†å¤§è¯­è¨€æ¨¡åž‹ï¼ˆMuseGlimmer 30B、Llama 3 8B–405B）、稀疏混合专家模型（Qwen 3 235B MoE）以及生成式扩散 Transformer（FLUX.1-dev 12B）上提供统一的高性能训练环境。在 TPU v6e（Trillium）和 TPU v7x（Ironwood）集群上以 BF16 精度进行评估，多切片执行通过 torchrun 在 Google Kubernetes Engine（GKEï¼‰ä¸Šæ— ç¼ç¼–æŽ’ã€‚TorchTitan 作为开源项目可在此 GitHub 仓库 中找到。

分布式策略与分片架构 ã€‚é«˜æ•ˆçš„å¤šåŠ é€Ÿå™¨è®­ç»ƒè¦æ±‚æ¨¡åž‹å¹¶è¡Œç­–ç•¥ä¸Ž TPU çš„ç‰©ç†äº’è¿žæ‹“æ‰‘ä¹‹é—´å®žçŽ°æ— ç¼å¯¹é½ã€‚ä¸ºæ­¤ï¼Œæˆ‘ä»¬é‡‡ç”¨ 2D+ DeviceMesh æ˜ å°„ï¼Œå°†é€šä¿¡å¯†é›†çš„ç»´åº¦â€”â€”ä¾‹å¦‚å¼ é‡å¹¶è¡Œå’Œä¸“å®¶å¹¶è¡Œâ€”â€”ä¸Žè¿žç»­çš„é«˜å¸¦å®½èŠ¯ç‰‡é—´äº’è¿žï¼ˆICIï¼‰å¯¹é½ï¼ŒåŒæ—¶å°†å¤åˆ¶çš„æ•°æ®å¹¶è¡Œæµé‡è·¯ç”±åˆ°è¾ƒæ ¢çš„è·¨åˆ‡ç‰‡æ•°æ®ä¸­å¿ƒç½‘ç»œï¼ˆDCN）。这一策略通过我们与 SimpleFSDP çš„é›†æˆæ¥å®žçŽ°ï¼ŒåŽè€ æä¾›ç¬¦åˆ SPMD 规范的操作。为进一步降低延迟，我们将 XLA è¿è¡Œæ—¶é ç½®ä¸ºæŠŠé›†åˆé€šä¿¡æ“ä½œï¼ˆå¦‚ All-Gather 和 Reduce-Scatter）直接卸载到片上 SparseCore 引擎，从而有效地将通信与计算密集的 TensorCore MXU æ•°å­¦è¿ç®—é‡å ã€‚æœ€åŽï¼Œæˆ‘ä»¬é€šè¿‡ä½¿ç”¨ HLO diffing å·¥å ·å’Œ @spmd_safe æ³¨è§£æ¥ç¡®ä¿ç¡®å®šæ€§æ‰§è¡Œå¹¶é˜²æ­¢åŒæ­¥æ­»é”ï¼Œè¿™äº›å·¥å ·å’Œæ³¨è§£æ¶ˆé™¤äº† torch.compile ä¸‹ä¸å¿ è¦çš„å›¾ä¸­æ–­ã€‚

性能与多节点扩展 。通过系统性的逐层优化，TorchTPU 上的 TorchTitan 为 FLUX 模型实现了卓越的性能和多节点扩展效率。从开箱即用的 eager 基线步进时间 11.49s（0.26 samples/s）开始，应用 DeferAndFuse eager 模式将步进延迟改善至 9.95s（1.15× åŠ é€Ÿï¼‰ã€‚å¼•å ¥ SimpleFSDP åˆ†ç‰‡è¿›ä¸€æ­¥å°†æ‰§è¡ŒåŠ é€Ÿè‡³ 6.52s（1.76× åŠ é€Ÿï¼‰ï¼Œè€Œé›†æˆè‡ªå®šä¹‰ Pallas SplashAttention å† æ ¸åˆ™å°†æ­¥è¿›æ—¶é—´é™è‡³ 5.55s（2.07× åŠ é€Ÿï¼‰ã€‚å¯ç”¨å®Œæ•´çš„ torch.compile(backend=”tpu”) 带来显著的性能飞跃，达到 3.02s（3.80× åŠ é€Ÿï¼‰ï¼Œè€Œåº”ç”¨ç›®æ ‡æ¨¡åž‹çª¥å­”é‡å†™åˆ™äº§ç”Ÿæœ€ç»ˆä¼˜åŒ–åŽçš„æ­¥è¿›æ—¶é—´ 2.78s（1.08 samples/s）。在多主机集群中，TorchTitan 展现出接近理想的弱扩展和强扩展：在 FLUX.1-dev 上，使用 SimpleFSDP 从单个 v6e-8 节点（8 芯片）扩展到 v6e-256 切片（64 ä¸ªèŠ‚ç‚¹å ± 256 èŠ¯ç‰‡ï¼‰ï¼Œæ¯æ ·æœ¬æ­¥è¿›æ—¶é—´ä» ä»Ž 0.92s ç•¥å¾®å¢žåŠ åˆ° 1.00s（总步进时间 3.01sï¼‰ï¼Œåœ¨å ‰äº’è¿ž ICI 网络上展现出接近线性的效率。在更大规模下，TorchTitan 在 TPU v7x（Ironwood）上以 BF16 精度、8K 上下文窗口在 1,024 个芯片上运行 Llama 3 405B，实现 500 tokens/sec/chip 和 50+% MFUã€‚åŒæ ·ï¼Œåœ¨ 8,192 个 TPU v7x 芯片上评估 Qwen3 235B MoE 也展现出高效扩展，在 BF16 精度、4K 序列长度下达到 28% MFU。

å ³é”®æž¶æž„ä¼˜åŒ– ã€‚ä¸ºåœ¨å¤šæ ·åŒ–çš„æ¨¡åž‹æž¶æž„ä¸­ç»´æŒå³°å€¼è®¡ç®—å¯†åº¦å¹¶æ¶ˆé™¤ç¡¬ä»¶ç“¶é¢ˆï¼Œæˆ‘ä»¬é‡‡ç”¨äº†å¤šé¡¹ç³»ç»Ÿå’Œå† æ ¸ä¼˜åŒ–ï¼š

• SimpleFSDP： 一种对编译器友好的 FSDP 版本。FSDP2 é€šè¿‡æ¨¡å—é’©å­æ³¨å ¥å ¶ all-gather 和 reduce-scatter，这会把 torch.compile 打断成逐算子分发；SimpleFSDP åˆ™å°†å®ƒä»¬è¡¨è¾¾ä¸ºè¢«è¿½è¸ªå›¾ä¸­çš„æ™®é€šæ“ä½œã€‚è¿™æ ·é›†åˆé€šä¿¡å¯¹ XLA 保持可见，TPU 编译器便能围绕它们进行融合、与 MXU è®¡ç®—é‡å ï¼Œå¹¶å¯å°†å ¶å¸è½½åˆ° SparseCore。

• Scan 高阶算子 （torch._higher_order_ops.scan）：将重复的 transformer å—ä½œä¸ºå¯¹å †å æƒé‡çš„ä¸€æ¬¡å¾ªçŽ¯æ¥è¿è¡Œã€‚éšç€æ·±åº¦å¢žé•¿ï¼Œè¿™ä½¿ç¼–è¯‘æ—¶é—´å’Œèšåˆå‚æ•°å† å­˜ä¿æŒæ’å®šã€‚é€å±‚çš„é›†åˆé€šä¿¡ä½äºŽå¾ªçŽ¯å† éƒ¨ï¼Œå› æ­¤ XLA å¯ä»¥å°†å®ƒä»¬ä¸Žè®¡ç®—é‡å ï¼šå»¶è¿Ÿéšè—è°ƒåº¦å™¨å°†æ¯ä¸€å±‚çš„ all-gather 隐藏在该层的计算之后，集合通信流水线器可选择性地预取下一层的 all-gatherï¼Œè€Œåœ¨åå‘ä¼ æ’­ä¸­ï¼Œå®ƒå¯ä»¥è®©æŸä¸€å±‚çš„ reduce-scatter 在下一层计算时完成。

• Pallas SplashAttention å† æ ¸ ：将缺乏向量瓦片对齐的默认注意力实现替换为自定义 Pallas å† æ ¸ï¼Œæ”¯æŒå®Œæ•´çš„ Grouped-Query Attention（GQA）、偏置和 dropout，在扩散注意力块上带来 +17.5% çš„åŠ é€Ÿã€‚

• Fused Pallas Cross-Entropy Loss : Avoids out-of-memory (OOM) faults when materializing [B×S×V] logits for large vocabularies by porting Tokamax’s online reduction kernel to compute loss in-place.

• Graph trainer: Captures the full training step, forward, loss, backward and optimizer update, as a single graph that compiles into one XLA program. This gives the TPU compiler a global view for fusion, collective scheduling and memory planning, and lets TPU-specific passes decide per activation whether to save, recompute or offload to host, and donate optimizer buffers so state is updated in place rather than double-buffered.

• Activation Offloading : The activations that the backward pass needs can be moved to host memory after the forward and copied back just before backward uses them, which frees HBM for larger models and longer sequences. The graph trainer inserts the offload and reload ops into the traced graph. Under scan, they sit inside the loop body, so saved activations move one layer at a time. XLA turns these copies into async transfers, and in backward copies the next layer’s activations back while the current layer computes. This is essential for the largest models, such as Llama3-405B and Qwen3-235B, to fit at their target sequence lengths.

Ongoing development for TorchTitan on TorchTPU targets four strategic priorities: expanding list of models optimized for TPU; integrating FP8 and FP4 low-bit quantization via torchao ; and supporting end-to-end Reinforcement Learning (RL) by coupling training loops with vLLM/SGLang sampling backends in VeRL.

Case Study 6: Hardware-Accelerated Recommendation Models with TorchRec and SparseCore

Recommendation and ranking architectures—such as Meta’s Deep Learning Recommendation Model (DLRM)—face critical memory and compute bottlenecks due to high-cardinality embedding tables. On standard GPU clusters, managing these tables may require complex software-based sharding via DistributedModelParallel (DMP) and separate optimizer updates, which often introduce latency. To bridge this gap, Google collaborated with the Meta PyTorch team to upstream a native TPU backend directly into the TorchRec repository, allowing recommendation models to bypass host-side compute bottlenecks and tap directly into TPU SparseCore (SC) hardware.

• Architectural Integration. The resulting implementation creates a seamless data pipeline where high-cardinality lookups are accelerated by dedicated hardware. A multi-threaded C++ PyBind backend runs on the host CPU to parallelize partition mapping and index padding, outputting structured batches that map directly to the TPU SparseCore units. Within the TPU, we implemented In-Place Optimizer Fusion: rather than materializing gradient tensors (.grad) in High Bandwidth Memory (HBM)—a process that traditionally complicates memory management—we map backpropagation directly to SparseDenseMatmulGrad. This ensures parameter updates (SGD, AdaGrad, or Adam) occur in-place, eliminating gradient memory overhead entirely. Finally, low-level ATen operators, such as SparseDenseMatmulOp (for forward lookups) and RaggedAllToAll (for optical ICI cross-chip collectives), provide hardware-accelerated kernels that execute directly on the silicon.

• Performance & Results. When evaluated on an industry-standard DLRM v2 model across a distributed 32-chip TPU v7x (Ghostfish) cluster, this integration achieved 100% performance parity with Google’s highly optimized JAX baseline, reaching an identical 18.3ms step latency. Notably, this performance matching was established using baseline SparseCore kernels before applying advanced optimizations. Convergence and mathematical fidelity were equally robust, yielding identical loss trajectories compared to the JAX implementation.

Active development focuses on further saturating bandwidth through TensorCore/SparseCore pipelining—interleaving MLP execution with embedding lookups—and advanced table-stacking techniques to unify memory buffers and reduce kernel overhead. Additionally, ongoing work aims to expose lower-level SparseCore embedding lookup operations (e.g. SparseMapRow, SparseGather etc) via ATen kernels, enabling users to implement custom lookup algorithms and advanced sharding strategies. This integration solidifies the TPU as a high-efficiency engine for the next generation of recommendation architectures.

Future Work

With TorchTPU now publicly available on GitHub , this initial open-source release represents only the beginning of our development journey. As we engage directly with the broader PyTorch community, our roadmap for the next 6 to 12 months focuses on production hardening, ecosystem expansion, and continually integrating community feedback to bridge the gap between TPU hardware capabilities and PyTorch-native developer workflows. We are prioritizing four strategic pillars to drive this evolution:

• Foundational Hardening & User Velocity : We are refining the “True Eager” experience by building a robust library of precompiled, dynamic kernels for ATen operators. We are also making XLA Bounded Dynamism fully production-ready, allowing the torch.compile stack to handle SymInts without extra recompilation, and streamlining our GKE packaging to lower the barrier for deployment across PyTorch versions. Furthermore, we are continuously expanding operator coverage and advancing our quantization support—specifically targeting FP8 and FP4 formats for next-generation TPUv8 architectures, and implementing nested compiler regions to reduce compilation time.

• Pod-Scale Distributed Training & Advanced Training : Building on our successful POD-scale training validation of Qwen3, we are expanding to additional model architecture and sharding techniques at and beyond that scale. We are enabling end-to-end Reinforcement Learning (RL) training loops by tightly coupling scalable training with high-throughput vLLM/SGLang sampling backends orchestrated via Ray Core. To enhance distributed stability, we are also validating a consensus-based compilation handshake to safely optimize collective operations with compute without fragile @spmd_safe annotations.

• Emerging Architectures, Inference & Serving : We are committing to proactively support cutting-edge model architectures as they emerge in the open-source ecosystem—including reasoning models, multimodal and generative diffusion models (e.g., Wan 2.2, Flux Klein), and novel Mixture-of-Experts (MoE) variants. Our efforts focus on optimizing the serving stack in vLLM and SGLang for low-latency production, incorporating native FP8/FP4 quantization across future TPU generations, and ensuring day-one TPU support for new architectural paradigms.

• High-Efficiency Communication & Diagnostics : We are optimizing collective operations within and across TPU slices and adding zero-copy DMA to enable hardware-assisted memory offloading. To support debugging at scale, we are transitioning to a WebGPU-based Trace Viewer 2.0 and deploying “flight recorder” continuous profiling (XProfz) to identify intermittent bottlenecks in production environments.

TorchTPU is built on the belief that a successful ecosystem thrives through transparent, collaborative engineering. As we open-source our four foundational repositories (TorchTPU, vLLM-TorchTPU, SGLang-TorchTPU, and TorchTitan) at the October 2026 PyTorch Conference, we actively seek feedback on model architecture prioritization and workload sequencing. We invite developers, researchers, and partners to collaborate with us via GitHub or our active TPU Builders program.

†The TorchTPU Team

Claudio Basile, Michael Voznesensky, Gunhyun Park, Ben Wilson, Parker Holloway, Zhanyong Wan, Vlad Belous, Manish Khanna, Paul Ganssle, Kevin Gleason, Andrew Dame, Brendan Folie, Bhavya Bahl, Haifeng Jin, Pedro Goncalves Mokarzel, Arian Arfain, Luke Boyer, Chi Zhang, Yiling Yuan, Adam Mainz, Aditya Patel, Huangze Tang, Jay Gan, Lizzie Liang, Will Froom, Tushar Agarwal, Karlo Basioli, Theotime Combes, Emilio Cota, Minos Papadopoulos, Tudor Pipernea, Yarong Mu, Qi Zhou, Yaoshiang Ho, Aleksey Vlasenko, Alexander Brauckmann, Ziyin Huang, Jack Oh, Jerry Hsu, Tanvi Bajpai, Jialei Chen, Wenjun Gu, Chintan Shah, Chris Chan, Sayce Falk, Chris Jones, Rob Torres, Kat Ko, Lee Howes, Freya Azad, Andi Gavrilescu, Joe Pamer, Bill Jia

Plus many, many others…

The post Introducing TorchTPU appeared first on PyTorch .

来源： PyTorch：Blog · pytorch.org

Google / Gemini

Meta / Llama

部署工程

产品更新

#

部署/工程

Google

Meta