Change
6d4aa4a
6d4aa4abec2b7e1743f0803a7bbbba492edad8a2 · commit on GitHub
pytorch-blog-feed: changed (294253 bytes, HTTP 200)
raw/pytorch-blog-feed/response.xml modified
- Source
- pytorch-blog-feed
- Lines added
- +1,368
- Lines removed
- -703
- Stored bytes at this commit
- 294,253
- Timestamp
- origin
- Raw artifact at this commit
- raw/pytorch-blog-feed/response.xml
Recorded headers
| observed_at | 2026-10-21T04:47:21.001Z |
|---|---|
| origin_date | 2026-10-21T03:26:27.000Z |
| status | 200 |
| final URL | https://pytorch.org/blog/feed/ |
| etag | "dcb5d572d1ff95849defaf75275efbdb" |
| last-modified | Tue, 20 Oct 2026 19:32:47 GMT |
| date | Wed, 21 Oct 2026 04:47:20 GMT |
| age | 4853 |
| cache-control | public, max-age=60, s-maxage=43200, stale-while-revalidate=86400, stale-if-error=604800 |
| cf-cache-status | null |
| content-encoding | null |
| content-length | 294253 |
@
@@ -12,7 +12,7 @@ <atom:link href="https://pytorch.org/blog/feed/" rel="self" type="application/rss+xml" /> <link>https://pytorch.org/blog/</link> <description></description>-
<lastBuildDate>Mon, 19 Oct 2026 22:53:51 +0000</lastBuildDate>+
<lastBuildDate>Tue, 20 Oct 2026 19:32:47 +0000</lastBuildDate> <language>en-US</language> <sy:updatePeriod> hourly </sy:updatePeriod>@
@@ -28,6 +28,1368 @@ <height>32</height></image> <item>+
<title>Introducing TorchTPU</title>+
<link>https://pytorch.org/blog/introducing-torchtpu/</link>+
+
<dc:creator><![CDATA[The TorchTPU Team†]]></dc:creator>+
<pubDate>Tue, 20 Oct 2026 16:02:30 +0000</pubDate>+
<category><![CDATA[Blog]]></category>+
<guid isPermaLink="false">https://pytorch.org/?p=188414</guid>+
+
<description><![CDATA[<p>TL;DR TorchTPU is an open source collaboration between Google and Meta that introduces a native and principled PyTorch backend for Google TPUs. It enables seamless out-of-the-box model execution and TPU...</p>+
<p>The post <a href="https://pytorch.org/blog/introducing-torchtpu/">Introducing TorchTPU</a> appeared first on <a href="https://pytorch.org">PyTorch</a>.</p>+
]]></description>+
<content:encoded><![CDATA[<h2><img fetchpriority="high" decoding="async" class="alignnone wp-image-188419 size-full" src="https://pytorch.org/wp-content/uploads/2026/10/image-with-text-scaled.jpg" alt="" width="2560" height="846" srcset="https://pytorch.org/wp-content/uploads/2026/10/image…+
<h2>TL;DR</h2>+
<p>TorchTPU is an open source collaboration between Google and Meta that introduces a native and principled PyTorch backend for Google TPUs. It enables seamless out-of-the-box model execution and TPU migration while delivering industry-competitive inference and training performance through eager exe…+
<h2>Introducing TorchTPU</h2>+
<p>Google TPUs have been used very successfully, both within and outside of Google, to train and deploy large ML models at massive scale. They are specifically designed for ML and, as such, they provide significant benefits in terms of performance, energy, and cost. While TPUs have been historically…+
<p>To program Google’s TPU at a level below the Python framework, developers utilize the XLA compiler stack. In 2020, Google introduced the first compiled version of PyTorch for TPU, PyTorch/XLA, based on Lazy Tensors. While foundational, this initial integration struggled to keep pace with th…+
<p>More specifically, as we started TorchTPU we deliberately formulated a set of principles to follow:</p>+
<ul>+
<li><strong>PyTorch Citizenship</strong>:+
<ul>+
<li>No JAX dependencies in the runtime;</li>+
<li>Solid eager mode;</li>+
<li>Torch.compile() based on Torch Dynamo (not Lazy Tensor) and with similar performance to JAX compile;</li>+
<li>Wide support for PyTorch distributed idioms (DDP, FSDP, DTensor, …) by providing TPU collectives at the lowest level in the PyTorch stack (e.g. Process Group API);</li>+
<li>Seamless integration into the PyTorch ecosystem;</li>+
<li>Community engagement.</li>+
</ul>+
</li>+
<li><strong>Focus on the User:</strong>+
<ul>+
<li>Porting a workload to TPU should be as simple as changing <code>model.to("cuda")</code> to <code>model.to("tpu")</code>. (One cannot achieve peak TPU performance with that simple change, since a model’s shapes, sharding and ops may be skewed toward a GPU architecture; however, the model should r…+
<li>Provide user-friendly and actionable messages for runtime/compiler errors;</li>+
<li>Integrate with the tools and frameworks that the PyTorch community already uses (e.g., torch.profiler, DTensor, Ray).</li>+
</ul>+
</li>+
<li><strong>Improve the State of the Art</strong>:+
<ul>+
<li>Build a solution that can compete with PyTorch on GPUs for large scale workloads;</li>+
<li>Introduce a JIT compilation approach for eager execution, with runtime optimizations that can bridge the performance gap between op-by-op execution and compiled execution;</li>+
<li>Improve the TPU programming stack by introducing bounded dynamism, by avoiding XLA recompilation due to shape changes, and by exploring dynamic TPU kernels that can be precompiled ahead of time.</li>+
</ul>+
</li>+
</ul>+
<p>TorchTPU is in active development, it is used in production for both inference and training workloads for some of the largest Google Cloud customers, and can be found as an Open Source project in this <a href="https://github.com/google-pytorch/torch_tpu">GitHub repo</a>. At the time of this writi…+
<ul>+
<li>Added support for both eager and compiled modes, including a new DeferAndFuse eager mode;</li>+
<li>Added support for ~450 Aten ops (skipping ops with structural or composite delegates), including forward and backward pass for all Core ATen ops required for the major deep learning architectures;</li>+
<li>Added ability to author custom kernels via JAX/Pallas and via Helion DSL;</li>+
<li>Validated out-of-box execution on the 5,000 most downloaded Hugging Face models, including:+
<ul>+
<li>LLMs: Qwen, DeepSeek, Gemma, LLama, Mistral, GPT, OPT, BLOOM, Falcon, StableML</li>+
<li>Audio/Speech: Whisper, Wav2Vec2, HuBERT, MMS</li>+
<li>Vision: DinoV2, SigLIP, CLIP, Swin Transformer, ConvNeXt, ViT, ResNet, DETR</li>+
<li>Text Encoders & Search: BERT, ALBERT & Electra, RoBERTa, DeBERTa, T5, MT5, Flan-T5</li>+
</ul>+
</li>+
<li>Validated DTensor-based sharding;</li>+
<li>Validated multi-host inference (vLLM and SGLang) and fine-tuning / training workloads (TorchTitan), with Ray-based orchestration;</li>+
<li>Validated integration with vLLM, SGLang, TorchTitan, Ray, and TorchRec.</li>+
</ul>+
<p>In the remainder of this document we discuss key aspects of the TorchTPU architecture and some initial results.</p>+
<h2>Architecture Deep Dive</h2>+
<h3>A TPU Primer</h3>+
<p><a href="https://cloud.google.com/tpu">Google TPUs</a> are specialized <a href="https://en.wikipedia.org/wiki/Application-specific_integrated_circuit">ASICs</a> designed specifically to perform the massive tensor operations that are the foundation of deep learning. From a <a href="https://cloud.g…+
<ul>+
<li>One or more <strong>TensorCores</strong>, each comprising one or more matrix-multiply units (MXUs), a vector unit (VPU), and a scalar unit (see <a href="https://dl.acm.org/doi/pdf/10.1145/3360307">this link</a>). Based on the TPU generation, an MXU can include 256×256 multiply-accumulators,…+
<li>One or more <strong>SparseCores</strong>, which are <a href="https://cloud.google.com/tpu/docs/system-architecture-tpu-vm#sparsecore">specialized units</a> for sparse operations (e.g., gather, scatter, embeddings) and for collectives. They can run in parallel with a TensorCore, enabling compute-…+
<li>A large <strong>High Bandwidth Memory</strong> (HBM) – 288 GB at 8.6 TB/s in <a href="https://cloud.google.com/blog/products/compute/tpu-8t-and-tpu-8i-technical-deep-dive">TPU 8i</a>, which can be used to store ML data and share that with the CPU.</li>+
<li>Bidirectional <strong>Inter-Chip Interconnect</strong> (ICI) – 19.2 TBps in <a href="https://blog.google/innovation-and-ai/infrastructure-and-cloud/google-cloud/eighth-generation-tpu-agentic-era/">TPU 8t and 8i</a>, allowing fast communication between TPU chips within user-configured <a href="ht…+
<li>Bidirectional <strong>Data-Center-Network</strong> (DCN) – 400 GBps in <a href="https://blog.google/innovation-and-ai/infrastructure-and-cloud/google-cloud/eighth-generation-tpu-agentic-era/">TPU 8t</a>, allowing communication between TPU chips across Slices.</li>+
</ul>+
<p>TPUs are programmed with the <a href="https://github.com/openxla/xla">XLA compiler</a> stack through two different entry points: <a href="https://openxla.org/stablehlo/spec">StableHLO</a> (higher-level IR for ML ops) and <a href="https://docs.jax.dev/en/latest/pallas/design/design.html#lowering-p…+
<p><!-- IMAGE PLACEHOLDER: IMAGE 1 of 4: TPU programming stack sketch (StableHLO / Mosaic / XLA, "See picture below") - download from the Google Doc, upload to the WP media library, then replace this comment with the <img> --></p>+
<p><img decoding="async" class="wp-image-188536 size-large aligncenter" src="https://pytorch.org/wp-content/uploads/2026/10/Fig-1-618x1024.png" alt="" width="618" height="1024" srcset="https://pytorch.org/wp-content/uploads/2026/10/Fig-1-618x1024.png 618w, https://pytorch.org/wp-content/uploads/2026…+
<p>A C++ program can access a local TPU by instantiating a <a href="https://openxla.org/xla/pjrt/cpp_api_overview">PjRT Client</a>, whose API provides methods to compile MLIR (StableHLO or Mosaic) into a TPU executable, to launch that executable, and to copy data to-from the TPU. Note that TPU execu…+
<p>Notable features of the TPU programming model:</p>+
<ul>+
<li>A TensorCore device can be thought of as a single-threaded processor with one or more massive matrix-multiply units (MXU). Consequently, to reach peak performance it is key that a model’s shapes be sufficiently wide to exploit the full potential of the MXU. This differs from GPUs where lau…+
<li>The XLA compiler can instruct a TensorCore to offload sparse and collective operations to SparseCore, for more efficient and/or parallel execution. This happens either automatically or through user-supplied MLIR annotation in a compiled StableHLO module.</li>+
<li>A TPU device can be accessed by only one user-space process. Following our goal of seamless integration in the PyTorch ecosystem, for distributed TorchTPU workloads we adopt a one-process-per-device model in contrast with JAX’s one-process-per-host model. However, because a TPU can be accessed b…+
<li>StableHLO follows the <a href="https://en.wikipedia.org/wiki/Static_single-assignment_form">SSA paradigm</a> and, hence, doesn’t allow for in-place mutations, i.e., a write operation always produces a new value (at least conceptually). This is particularly significant for TorchTPU, which does pr…+
<li>The XLA compiler and PjRT runtime assume that tensor shapes are known at compile time. Recognizing this as a limitation for PyTorch workloads, we have been adding support for bounded dynamism, i.e., when the maximum bounds of a shape’s dimensions are known at compile time, but the actual dimensi…+
<li>Collective ops are lowered to XLA as StableHLO ops, versus using a separate library such as NCCL.</li>+
</ul>+
<h3>Usability and Autoload Mechanism</h3>+
<p>We designed TorchTPU with usability in mind. When the user imports PyTorch on a TPU machine where TorchTPU wheel is installed, running the PyTorch code on TPU involves only trivial changes and can be done in a couple of lines of code thanks to PyTorch’s autoload mechanism:</p>+
<p>See the table below for a side-by-side comparison.</p>+
<table>+
<thead>+
<tr>+
<th><strong>Before (GPU)</strong></th>+
<th>After (TPU)</th>+
<th><strong>Alternative (TPU or GPU)</strong></th>+
</tr>+
</thead>+
<tbody>+
<tr>+
<td>+
<pre><code class="language-python">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)
+
...</code></pre>+
</td>+
<td>+
<pre><code class="language-python">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)
+
...</code></pre>+
</td>+
<td>+
<pre><code class="language-python">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)
+
...
+
</code></pre>+
</td>+
</tr>+
</tbody>+
</table>+
<h3>Numerical Precision</h3>+
<p>TorchTPU adapts standard PyTorch precision APIs for TPU hardware. Like all high-performance machine learning architectures, TPUs are optimized to deliver peak performance and energy efficiency using specific native data types preferred by their particular generation. For TPUs the native matrix mu…+
<p>TorchTPU maps PyTorch precision modes to underlying XLA/StableHLO configurations to expose this hardware capability:</p>+
<ul>+
<li><em>medium</em>: Uses lower precision for multiplication (converting inputs to native BF16 for the MXU), while accumulation is performed in FP32 to maintain stability. TorchTPU defaults to this mode to fully leverage TPU hardware acceleration, as single-pass BF16 multiplication with FP32 accumul…+
<li><em>high</em>: Uses a 3-pass algorithm where inputs are split into high and low components to compute three cross-terms, recovering most FP32 precision at roughly 3× the compute cost.</li>+
<li><em>highest</em>: Uses a 6-pass algorithm that computes all cross-terms to approximate full IEEE FP32 accuracy, incurring roughly 6× compute overhead on the MXU.</li>+
</ul>+
<p>For the vast majority of deep learning workloads, we recommend using the <em>medium</em> (default) precision. Explicitly selecting <code>high</code> or <code>highest</code> in TorchTPU triggers a performance warning (<code>TORCH_WARN_ONCE</code>). Only switch to <em>high</em> or <em>highest</em> …+
<p>See the table below for a side-by-side comparison.</p>+
<table>+
<thead>+
<tr>+
<th>Dimension</th>+
<th>PyTorch on CUDA</th>+
<th>PyTorch on TPU (TorchTPU)</th>+
</tr>+
</thead>+
<tbody>+
<tr>+
<td><strong>Default Precision Mode</strong></td>+
<td>highest (pure FP32)</td>+
<td>medium (Precision.DEFAULT)</td>+
</tr>+
<tr>+
<td><strong>Fast Math Accelerator Primitive</strong></td>+
<td>Tensor Cores (TF32)</td>+
<td>Systolic MXUs (bfloat16)</td>+
</tr>+
<tr>+
<td><strong>Fast Math Bits</strong></td>+
<td>1 sign bit + 10 mantissa bits + 8 exponent bits (19-bit TF32 format)</td>+
<td>1 sign bit + 7 mantissa bits + 8 exponent bits (16-bit BF16 format)</td>+
</tr>+
<tr>+
<td><strong>Accumulation Precision</strong></td>+
<td>FP32 (in highest and high), TF32 (in medium)</td>+
<td>FP32 across all modes</td>+
</tr>+
<tr>+
<td><strong>Higher/Highest Precision Mechanism</strong></td>+
<td>Dedicated hardware FP32/TF32 units</td>+
<td>Multi-pass algorithmic decomposition (3× or 6× MXU passes)</td>+
</tr>+
<tr>+
<td><strong>Performance Impact of high/highest</strong></td>+
<td>Baseline speed or modest slowdown (~1–2×)</td>+
<td>~3× MXU compute cost for high; ~6× compute overhead for highest</td>+
</tr>+
<tr>+
<td><strong>Scoping API</strong></td>+
<td>Global torch.set_float32_matmul_precision</td>+
<td>Global torch.set_float32_matmul_precision + scoped torch.tpu.precision context manager</td>+
</tr>+
</tbody>+
</table>+
<h3>Eager execution, StableHLO lowering, and JIT compilation</h3>+
<p>Following <a href="https://docs.pytorch.org/tutorials/advanced/extend_dispatcher.html">this PyTorch guidance</a>, we have extended the PyTorch dispatcher with a new C++ backend, registered as <em>PrivateUse1</em>. Additionally, we have used the standard <code>at::Tensor</code> mechanism to define…+
<p>One unique characteristic of TorchTPU’s eager execution is that it can dispatch Aten ops to the TPU with multiple modalities:</p>+
<ul>+
<li><em>Immediately</em>, effectively implementing op-by-op dispatch;</li>+
<li><em>Deferred</em> all the way to an automatically-determined <strong>Materialization Point</strong> (e.g., a <code>.to("cpu")</code> statement) and aggregated/compiled together with all the other ops that were deferred up to that point;</li>+
<li>Anything in between.</li>+
</ul>+
<p>While this flexibility certainly brings additional complexity to the design, we have been deliberate about adding it from the start since we wanted to experiment with different policies so as to improve eager performance over the basic op-by-op dispatch modality. (Additionally, this flexibility t…+
<p>In summary, TorchTPU provides 3 eager modalities, which can be set with either environment variables or through a context manager:</p>+
<ul>+
<li><strong>Strict</strong> (Default): We dispatch one Aten op at a time, and they execute on the TPU asynchronously with the CPU. This is the equivalent of what you’d expect when running PyTorch eagerly on GPU.</li>+
<li><strong>Debug</strong> (<code>export TORCH_TPU_LAUNCH_BLOCKING=1</code>): We dispatch one Aten op at the time, but TPU execution is synchronous with the CPU. This is slow and aimed at debugging, and is equivalent to using <code>CUDA_LAUNCH_BLOCKING</code> on GPU with the additional benefit of en…+
<li><strong>DeferAndFuse</strong> (<code>export TORCH_TPU_DEFER_AND_FUSE=1</code>): We defer and fuse multiple Aten ops into larger chunks. This is a new eager/compiled <em>hybrid</em> modality introduced by TorchTPU and can be 50% or more faster than Strict Eager while reducing peak HBM memory (op …+
</ul>+
<p>Now that we have described the basics of TorchTPU eager mode, let’s sketch what happens when we receive a callback invocation from the PyTorch dispatcher. Let’s assume we are using the DeferAndFuse mode and let’s consider <code>at::mul.out</code> as an example:</p>+
<ul>+
<li>Create a Deferred Op instance that contains:+
<ul>+
<li>A StableHLO builder object, which when invoked returns a StableHLO function implementing the semantics of <code>aten::mul</code>, as shown below (note that any implicit broadcasting and dtype promotion are handled elsewhere)+
<pre><code class="language-mlir">func.func @main(%arg0: tensor<4x16xbfloat16>, %arg1: tensor<4x16xbfloat16>) -> tensor<4x16xbfloat16> {
+
%0 = stablehlo.multiply %arg0, %arg1 : tensor<4x16xbfloat16>
+
return %0 : tensor<4x16xbfloat16>
+
}</code></pre>+
</li>+
<li>An optional map of static op parameters (e.g., <code>training</code> and <code>momentum</code> for <code>aten::native_batch_norm</code>) to be used for the computation of a cache key (discussed below);</li>+
</ul>+
</li>+
<li>Return a new <code>at::Tensor</code> with <code>c10::StorageImpl</code> pointing (through indirection) to that Deferred Op instance.</li>+
</ul>+
<p>Upon receiving a callback from the PyTorch dispatcher corresponding to <code>mytensor.to("cpu")</code> or <code>mytensor.item()</code> or similar, TorchTPU will:</p>+
<ol>+
<li>Analyze the region R of deferred ops accumulated so far that lead to the computation of <code>mytensor</code>;</li>+
<li>Apply a liveness analysis to determine which ops to materialize within R, in addition to that producing <code>mytensor</code>; importantly ops with side effects (e.g., collectives) are always preserved (as so is their invocation order) while ops that are not live outside of R can be detected det…+
<li>Create a compilation cache key for the graph G of ops to materialize, so as to quickly and uniquely identify that graph in the future;</li>+
<li>Use the computed cache keys to check in a <strong>Compilation Cache</strong> if graph G has been already compiled.+
<ul>+
<li>If so, dispatch the executable for G on the TPU;</li>+
<li>If not, use a thread pool to (asynchronously) compile G and then dispatch its execution on the TPU.</li>+
</ul>+
</li>+
<li>Wait for the TPU value of <code>mytensor</code> to be ready and then copy it from the TPU HBM to the CPU storage associated with the destination <code>at::Tensor</code>.</li>+
</ol>+
<p>On the surface, it may seem that the approach described above would be inefficient because the CPU starts TPU execution only upon reaching a <code>.to("cpu")</code> statement. For real workloads we can expect that after few initial iterations/steps, which will populate the compilation cache, CPU …+
<h3>Compilation Cache</h3>+
<p>TorchTPU employs a Compilation Cache to avoid unnecessary XLA compilations. Similar to modern computer storage systems, this cache encompasses multiple tiers working in tandem:</p>+
<ul>+
<li><strong>Tier-1</strong>: Process-local in-memory cache for sharing compilation within the same process. This is the fastest of the three tiers.</li>+
<li><strong>Tier-2</strong>: Machine-local shared memory that enables cross-process sharing on the same host. This is useful for multi-device training/inference.</li>+
<li><strong>Tier-3</strong>: Persistent cache backed by a remote file system (e.g., Google Cloud Storage), allowing reuse across hosts, programs, and job restarts. This is useful when running TPU programs in execution environments like Colab and data centers.</li>+
</ul>+
<p>The tier-1 cache is always enabled. By default, tier-2 is enabled for distributed programs and disabled for single-device programs. The user can override this behavior as needed. Tier-3 is disabled by default, but can be explicitly enabled by specifying a shared storage path.</p>+
<p>Each entry in the cache is a compiled TPU executable for a StableHLO computation graph. In the tier-1 cache, each executable is in the form of a C++ object that can be readily scheduled for execution on a TPU device. Since objects cannot be shared across processes directly, the tier-2 and tier-3 …+
<p>For performance reasons, we do not use the StableHLO graphs themselves as cache keys. Instead, we fingerprint the computation, compiler options, and the versions of major components into three 64-bit integers and use the tuple as the cache key:</p>+
<ul>+
<li>A shapeless key that captures the structure of the computation graph, but not the sizes of the tensor inputs/outputs,</li>+
<li>A dimensions key that captures the tensor sizes, and</li>+
<li>A compile options key that captures the options passed to the XLA compiler and the versions of TorchTPU, libtpu, and the PjRt runtime.</li>+
</ul>+
<p>The shapeless key and the dimensions key together are also known as the graph key. While in theory we could get cache key collisions between different graphs, compiler options, and component versions, we consider the 192-bit key space large enough that such collisions will not arise in practice.<…+
<h3>View semantics</h3>+
<p>One particular area of contrast between PyTorch’s operational semantics and StableHLO semantics is the handling of mutability and shared-memory views:</p>+
<ul>+
<li>PyTorch assumes that each Tensor is a pointer into a linear byte array, with stride metadata at the Tensor level to describe the translation between a logical index tuple into a linear offset into that array. Multiple Tensors can share the same underlying byte array so that in-place writes can b…+
<li>In contrast, StableHLO is a purely functional language, where all tensor values are immutable after construction. The physical layout is abstracted from the user and it reflects tiling decisions made by the compiler, consequently direct stride manipulation is not supported. Additionally, in-plac…+
</ul>+
<p>Achieving correct view semantics on TPU in the presence of writes is not straightforward; however, TorchTPU accomplished it with a little machinery involving <strong>view decomposition</strong>, <strong>view inversion</strong> and deferred ops. This semantic bridge allows TorchTPU to natively sup…+
<h3>Streams and SparseCore</h3>+
<p>When optimizing model performance, there are 3 classes of operations that a PyTorch user may want to execute simultaneously:</p>+
<ul>+
<li>Compute ops;</li>+
<li>Collective ops;</li>+
<li>Data transfers between host and device.</li>+
</ul>+
<p>On GPUs this is typically accomplished by using separate command queues so that the device can dispatch the work to different HW units (SMs, DMAs, etc). In PyTorch device command queues are exposed as <a href="https://docs.pytorch.org/docs/2.13/generated/torch.Stream.html">Streams</a>, and they a…+
<p>TorchTPU also implements the PyTorch Stream API, but achieving the same functionality available on GPU is not possible in all cases. Considering that a TPU TensorCore is fundamentally a single-threaded execution processor, we can only dispatch one computation at a time on a given TensorCore. On t…+
<h3>Compiled Mode</h3>+
<p>Our primary goal is to maintain a user experience consistent with other <code>torch.compile()</code> backends and, consequently, to share the same compiler frontend, namely Dynamo. Per <a href="https://docs.pytorch.org/docs/stable/torch.compiler_custom_backends.html">this PyTorch guidance</a> we …+
<pre><code class="language-python">compiled_model = torch.compile(model, backend="tpu")</code></pre>+
<p>Dynamo invokes our backend by passing <code>torch.fx.GraphModule</code> plus a list of example input tensors. At that point we trace the provided graph as if we were in eager mode, but with all ops being deferred and no actual execution on the TPU. That allows us to generate and compile a single …+
<h3>Bounded Dynamism</h3>+
<p>A main difference between TPUs and GPUs is that XLA traditionally operates on static shapes to generate optimized hardware layouts and fused kernels. In practice, this is not a severe limitation and when a tensor dimension (such as batch size or sequence length) varies during workload execution, …+
<p>To avoid both the manual code changes and recompilation overhead, active development is ongoing to natively support dynamic shapes via XLA Bounded Dynamism in TorchTPU. With bounded dynamism, dynamic buffers are allocated to a maximum upper bound within the main model module, and the graph is com…+
<h3>Kernels</h3>+
<p>We distinguish two types of kernels:</p>+
<ol>+
<li>Kernels provided by the PyTorch framework such as <a href="https://docs.pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html"><code>torch.nn.functional.scaled_dot_product_attention</code></a>.</li>+
<li>Custom kernels provided by the user to accelerate a specific part of a model; those are typically model specific;</li>+
</ol>+
<p>For 1) we had initially implemented an approach where a TPU kernel is authored in Pallas and then exported to Mosaic with symbolic shapes, plus a TorchTPU step to specialize the kernel’s shapes at its invocation time. The mechanism turned out to be less flexible than anticipated and therefore we …+
<p>For 2) we have developed three solutions:</p>+
<ul>+
<li>(Production ready) <em>The user writes their TPU kernel in Pallas</em> and integrates it with PyTorch via a TorchTPU API, as follows+
<pre><code class="language-python"># 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)</code></pre>+
</li>+
<li>(Production ready) The <em>user leverages kernels from <a href="http://openxla.org/tokamax">Tokamax</a></em>, a library of high-performance TPU kernels, now available for seamless interoperability with PyTorch. Tokamax provides highly optimized implementations of key operators, including mixture…+
<li>(Experimental) The user rewrites their GPU kernel in <a href="https://pytorch.org/blog/helion/">Helion</a>, which enables targeting multiple accelerators with the same PyTorch-based DSL, as shown below+
<pre><code class="language-python"># 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)</code></pre>+
</li>+
</ul>+
<h3>Distributed Support</h3>+
<p>To ensure compatibility with existing PyTorch distributed APIs and frameworks, we decided to implement TPU collectives at the Process Group API level, by inheriting from <code>c10d::Backend</code>; that allows us to intercept collective ops in both eager and compile mode and seamlessly support Py…+
<p>Intercepting collectives at the Process Group level takes us almost to the finish line, with the exception of a few notable TPU idiosyncrasies w.r.t. GPUs:</p>+
<ul>+
<li><strong>Different Initialization Sequence</strong>. Consider a host with multiple devices attached. In the case of GPU, a process launched by <code>torchrun</code> can see all the attached devices, even though it will use only one of them (e.g, with <code>device=f"cuda:{local_rank}"</code>), and…+
<li><strong>Default Process Group Knowledge</strong>. Consider two subgroups (4, 0, 1, 3) and (5, 2, 6, 7) running an <code>all_reduce</code> among themselves. Note that <a href="https://docs.pytorch.org/docs/stable/distributed.html#torch.distributed.all_reduce"><code>torch.distributed.all_reduce</c…+
<pre><code class="language-mlir">%0 = "stablehlo.all_reduce"(%arg0) <{replica_groups=dense<[[4,0,1,3], [5,2,6,7]]>: tensor<2x4xi64>}> ({...</code></pre>+
</li>+
<li><strong>MPMD Support</strong>. Historically, the TPU stack was built with SPMD in mind. Thanks to tools such as <a href="https://openxla.org/shardy/getting_started_jax">Shardy</a>, distributed JAX programs are written for a single host, with no communication code included, and then are sharded a…+
<ul>+
<li>User-inserted statements conditional on the process’ rank;</li>+
<li>DTensor’s uneven splitting of a shared tensor when a tensor’s dimension is not divisible by the number of devices.</li>+
</ul>+
<p>Unpredictable deadlocks or memory corruption can happen when TPUs StableHLO modules that include collective ops and the modules are not identical (this is due to potentially mismatching sync flag IDs or buffer offsets across ranks in the generated TPU code). TorchTPU deals with that as follows:</…+
<ul>+
<li>(Default) <strong>Skinny Collectives</strong>. By default, TorchTPU compiles each collective op in its own “skinny” StableHLO module, which guarantees consistency across ranks for buffer and sync flag IDs. That works, but can incur a noticeable performance penalty w.r.t. what’s discussed in the …+
<li>(Optional) <strong><code>@spmd_safe</code> Decorator</strong>. If the user knows that a specific section of code is SPMD (e.g., there is no if statement conditional on <code>torch.distributed.get_rank()</code>), they can mark that section of code with a TorchTPU’s <code>@spmd_safe</code> decorat…+
<li>(Optional) <strong>Compilation Handshake</strong> (<a href="https://github.com/google-pytorch/torch_tpu/issues/2886">RFC</a>). This mechanism runs when a given <code>torch.compile</code> statement triggers compilation of a module containing communication operations. It initiates a consensus roun…+
<ul>+
<li><code>OFF</code> – never perform the handshake (current default);</li>+
<li><code>COMPILE_STAGE</code> – perform the handshake before starting an FX graph compilation;</li>+
<li><code>DISPATCH_STAGE</code> – perform the handshake before starting a compiled FX graph execution. This mode should be used only if a user expects <strong>rank specific recompilations</strong> since it involves larger overhead than the previous option.</li>+
</ul>+
</li>+
</ul>+
</li>+
</ul>+
<h3>Distributed Initialization</h3>+
<p>Initializing a distributed workload in TorchTPU follows standard PyTorch distributed semantics, with the caveat that TorchTPU requires two TPU-specific variables a standard PyTorch torchrun would not know how to set:</p>+
<ul>+
<li><code>TORCH_TPU_TOPOLOGY</code>: The physical mesh geometry of the TPU slice (e.g., “2,2,1” for a 2x2x1 slice configuration).</li>+
<li><code>TORCH_TPU_SLICEBUILDER_ADDRESSES</code>: A list of open ports/addresses (one per device) used to construct the TPU Inter-Chip Interconnect (ICI) network mesh (e.g., “localhost:8471,localhost:8472,…”).</li>+
</ul>+
<p>An alternative is to use <code>torch.tpu.distributed.set_tpu_launch_env()</code> from Python, which can set those variables in a Google Cloud environment. Below we show sample code to launch a distributed TorchTPU workload.</p>+
<p><strong>A. Worker implementation (<code>train.py</code>)</strong></p>+
<pre><code class="language-python">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()</code></pre>+
<p><strong>B. Launching the workload using torchrun (CLI)</strong></p>+
<pre><code class="language-bash"># 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</code></pre>Diff display stops at 400 lines. The line counts above are from the whole diff. 47 lines shown here cut at 300 characters. The raw artifact at this commit is linked above.