deniz.in

Markets

Weather

Loading weather

· via dev.to (home feed)

Google open-sources TorchTPU, PyTorch/XLA archive set for November 15

Google released TorchTPU, a native PyTorch backend for TPUs, and will archive the deprecated PyTorch/XLA repository on November 15, 2026, leaving torch_xla users weeks to plan a migration.

Google open-sources TorchTPU, PyTorch/XLA archive set for November 15

Google ships a native PyTorch path to TPUs

Google published TorchTPU on GitHub on October 20, 2026, a backend that makes TPUs a device PyTorch can target natively, in the same way it targets GPUs and CPUs. According to a dev.to post by Gabriel Anhaia, the release also ends the previous route: PyTorch/XLA is now deprecated, and its repository is scheduled to be archived on November 15, 2026. That gives teams running torch_xla under four weeks before the project becomes read-only.

The dev.to post says Google introduced TorchTPU at PyTorch Conference 2026 and, per the Google Cloud blog it cites, developed it with Meta's PyTorch team. The backend offers eager execution for debugging and torch.compile for production performance.

What version 0.2.0 includes

The repository is Apache 2.0 licensed, and the current release on PyPI is 0.2.0, requiring Python 3.11 or newer. Google's feature pages list TPU v5e, v5p, v6e and v7 as validated hardware; v8i and v8t appear on the near-term roadmap but are not supported yet.

Serving support ships alongside training. Google links a vLLM backend built with Inferact and an SGLang backend built with RadixArk, and says Hugging Face Transformers v5 can deploy to TPUs using device_map="tpu". On the training side, Google reports peak model-FLOPs utilization of up to 52% when training Qwen 3 8B with TorchTitan, with near-linear scaling out to 8,192 Ironwood TPUs. These are Google's own figures rather than an independent benchmark.

The port starts with one string

The FAQ characterizes migrating an existing workload as changing .to("cuda") to .to("tpu"). One installation detail matters: the README warns that a default torch install pulls the CUDA build, so Google recommends installing torch_tpu 0.2.0 against the CPU wheel from download.pytorch.org.

The execution model docs describe the default strict eager mode, where each operation becomes a single-op graph compiled at -O1 and cached, letting the cache warm quickly. torch.compile instead hands TorchTPU the full FX graph and compiles it at -O2. Two environment variables sit in between: TORCH_TPU_LAUNCH_BLOCKING makes errors surface at the operation that caused them, and TORCH_TPU_DEFER_AND_FUSE fuses deferred operations into larger graphs, which Google claims can deliver up to 50% speedup over default eager with no model changes.

Known limits to check first

The documentation is candid about what stalls the device. Python branches that depend on tensor values, such as branching on my_tensor.all(), force a copy back to the host and drain the execution queue, and they also break the graph in compiled or fused mode. Operations with data-dependent output shapes, like torch.masked_select, behave similarly.

Dynamic shapes are handled through compiled buckets rounded to multiples of 128, one by default and up to eight. With the default dynamic=None setting, a growing dimension is treated as symbolic; in Google's decode-loop example with an expanding KV cache, only the first three steps pay compile time. The stated constraints are at most eight buckets per trace, no dynamic=True, no output shapes in the style of torch.nonzero, and no custom Pallas kernels in that path. These bounded-dynamic executables are not written to the AOTAutograd cache, so they recompile on every process start.

Feature maturity varies. Mixed eager and compiled runs, forward and backward passes, Pallas kernels and SimpleFSDP are validated; FSDPv2 and SymInts are experimental, and FSDPv2 is flagged as low performance.

If you run PyTorch/XLA today

The FAQ draws a technical line between the stacks: PyTorch/XLA relied on JAX runtime components plus its own APIs, while TorchTPU is a native backend with no JAX runtime dependency — and Google notes TorchTPU is not meant to displace JAX. After November 15 the old repository stays readable and clonable but accepts no new issues or pull requests. The dev.to post suggests auditing a codebase with git grep for torch_xla, xla_device and xm. references, then checking each hit against TorchTPU's feature matrix.

A preview agent suite rides along

Google also shipped a preview TPU agent suite: a package of skills and subagents for coding agents, fronted by a single tpu-agent skill that routes to specialists for onboarding, profiling, hardware checks and kernels. Claude Code and Antigravity are supported, with Codex and Cursor listed as coming soon; outside contributions are not accepted for now. The installer writes skills, per-agent prompts and a managed CLAUDE.md block (AGENTS.md under Antigravity), so reviewing the diff before committing is advisable.

Why it matters

PyTorch users on TPUs now face a hard, near-term deadline: an archived, unmaintained stack versus a 0.2.0 release with documented gaps. The timing is competitive as well. TechCrunch, cited by the dev.to post, frames the launch as Google joining Amazon in challenging Nvidia's grip on AI developers — arriving one day after AWS announced its own PyTorch backend for Trainium. If TPUs, GPUs and CPUs all become native PyTorch targets, choosing hardware starts to look like a line-of-code decision rather than a framework commitment.

  • #pytorch
  • #google-tpu
  • #machine-learning
  • #open-source
  • #deep-learning

Related posts