TorchTPU 作为 Google 推出的 PyTorch 原生 TPU 运行时，在编译优化层面与传统的 GPU 加速存在显著差异。要在 TPU 上充分发挥 torch.compile 的性能优势，关键在于理解其内存分配策略与算子融合机制。本文将从工程实践角度出发，剖析 Fused Eager 与静态编译模式的底层差异，并给出可落地的参数配置与监控建议。
TorchTPU 的三种执行模式与内存特征
TorchTPU 采用了独特的「Eager First」设计哲学，通过 PrivateUse1 接口将 TPU 深度集成到 PyTorch 的执行栈中。这种架构使得开发者无需修改核心逻辑，仅需将设备初始化切换为 tpu 即可获得原生 PyTorch 体验。在这一基础上，TorchTPU 提供了三种逐层递进的执行模式，每种模式对应不同的内存分配特征与适用场景。
第一种是 Debug Eager 模式，该模式下每一次操作都会触发与 CPU 的同步，主要用于定位 shape 不匹配、NaN 值传播或 OOM 崩溃等问题。由于每步都需要等待 TPU 执行完成再返回 CPU，其执行效率较低，但内存占用相对可控，适合调试阶段使用。第二种是 Strict Eager 模式，它保持了单操作分发的特性，但改为异步执行，CPU 与 TPU 可以并行工作，仅在用户代码显式调用同步点时才阻塞。这种模式最贴近默认的 PyTorch 行为，同时保持了较高的吞吐量。第三种是 Fused Eager 模式，这是 TorchTPU 的核心创新点：系统通过自动反射分析操作流，将连续的算子动态融合成更大的计算块后再提交给 TPU。根据官方数据，Fused Eager 相比 Strict Eager 可以实现 50% 到 100% 以上的性能提升，而且无需用户介入任何配置。融合后的计算块减少了 TensorCore 与内存之间的数据传输次数，从而显著降低了内存带宽开销，这对于受限于 HBM 访问延迟的 TPU 架构尤为关键。
torch.compile 的静态编译路径与 XLA 集成
当开发者需要进一步挖掘性能上限时，torch.compile 提供了完整的静态编译能力。TorchTPU 在这一路径上选择了独特的架构：使用 Torch Dynamo 捕获 FX 图后，并不经过 Torch Inductor，而是直接将算子映射到 StableHLO（XLA 的中间表示），再交给 XLA 编译器生成优化的 TPU 二进制文件。这一设计是经过深思熟虑的：XLA 在 TPU 拓扑上的通信与计算重叠优化经过了大量生产验证，而 StableHLO 作为桥接层既保留了 PyTorch 算子的语义，又能让 XLA 进行硬件感知的低级别优化。
在内存分配层面，静态编译引入了几个需要关注的特性。首次编译时，系统需要为融合后的算子分配工作空间（workspace）用于中间结果的缓存，这会导致第一次执行的内存峰值明显高于后续运行。XLA 自身拥有一套复杂的内存规划器，会根据计算图的数据依赖关系和融合策略来分配缓冲区。理解这一点对于预估集群显存需求至关重要：当 batch size 较大或模型层数较深时，首次编译阶段的峰值内存可能成为 OOM 的触发点。完成编译后，这些工作空间会被缓存到磁盘或共享内存中，后续运行可以直接加载已编译的二进制，从而避免重复的内存分配开销。
内存配置实践参数与监控要点
针对不同的部署场景，TorchTPU 提供了一系列可调节的参数来平衡内存占用与编译开销。首先是编译模式的选择，torch.compile 支持 default、reduce-overhead 和 max-autotune 三种模式。在 TPU 环境下，default 模式通常能够在编译时间与运行时性能之间取得较好平衡；reduce-overhead 模式通过降低 Python 调用开销来减少延迟，适合推理场景；max-autotune 模式则会尝试更多的融合策略组合，但编译时间显著增加，且可能因为搜索空间过大而影响内存稳定性，建议仅在离线调优阶段使用。
其次是编译缓存的配置。TorchTPU 的 Compilation Cache 可以配置为单主机模式或多主机持久化模式。在多节点训练场景中，建议将缓存路径指向共享存储（如 GCS 或 NFS），这样不同节点可以共享已编译的算子图，避免重复编译带来的内存峰值。但需要注意的是，缓存的共享依赖模型结构和输入 shape 的一致性，如果不同节点的 batch size 或序列长度存在差异，缓存命中率的下降会导致额外的编译开销。
对于内存敏感的部署场景，可以通过设置环境变量 TORCHTPU_FLAGS 来控制内存分配策略。例如，torchtpu_allow_tf32 控制是否启用 TF32 计算精度，启用后可减少一半的显存占用但可能略微影响数值精度；torchtpu_lazy_rematerialization 则可以启用重计算技术，通过在前向传播时丢弃部分中间结果、在反向传播时重新计算的方式来换取更低的峰值内存。此外，XLA 的 autoFDO（反馈导向优化）功能可以记录运行时的算子执行热度，在后续编译中优先融合高频路径，这也能间接改善内存分配效率。
在实际监控方面，建议关注以下几个关键指标：编译缓存命中率（可通过 XLA 的调试接口查询）、首次迭代的峰值内存使用量、以及编译过程的耗时占比。如果发现缓存命中率持续低于 60%，需要检查是否存在动态 shape 变化（如可变序列长度或动态 batch size），这类场景会触发频繁的重新编译，导致内存使用波动。对于动态序列长度的处理，TorchTPU 团队正在推进 bounded dynamism 特性，计划通过在 XLA 中引入更先进的动态 shape 处理机制来减少 recompilation 开销，该特性预计在 2026 年内可用。
算子融合的硬件感知调优
TPU 的硬件特性对算子融合策略有着直接影响。当前一代 TPU 的 TensorCore 在矩阵乘法效率上对输入维度有偏好：当注意力头的维度设置为 128 或 256 时，能更好地利用 TensorCore 的并行计算能力，而传统的 64 维设置则无法发挥峰值性能。这一硬件特性意味着，即使使用相同的 torch.compile 配置，仅仅是修改模型架构中的 head_dim 参数，就可能带来显著的性能差异。类似地，SparseCore 用于处理不规则的内存访问模式如 embedding 查找和 gather/scatter 操作，在设计包含大规模嵌入层的模型时，需要考虑将这类操作显式地路由到 SparseCore 以获得最优的内存利用率。
对于需要编写自定义算子的开发者，TorchTPU 支持通过 @torch_tpu.pallas.custom_jax_kernel 装饰器将 Pallas 内核直接集成到编译流程中，从而实现细粒度的算子融合控制。Helion 内核的支持也正在开发中，未来将进一步扩展自定义算子的表达能力。
总结
TorchTPU 在内存分配与算子融合方面提供了分层的优化策略：从开箱即用的 Fused Eager 自动融合，到需要显式配置的 torch.compile 静态编译，开发者可以根据调试、迭代速度或极致性能的不同需求灵活选择。理解首次编译的内存峰值来源、合理配置编译缓存、关注 TPU 硬件特性（如 head_dim 与 SparseCore），是在生产环境中稳定部署 TorchTPU 工作负载的关键。随着动态 shape 支持和预编译内核库的完善，TorchTPU 的内存效率与编译速度仍有进一步提升空间。
资料来源：Google Developers Blog, "TorchTPU: Running PyTorch Natively on TPUs at Google Scale" (2026-04-07)