Cloud TPU release notes

Generally available: Compute Engine supports Google's custom-developed accelerator Tensor Processing Unit (TPU), providing a converged experience across AI accelerators on Google Cloud. You can use the Compute Engine instance API and managed instance group (MIG) API to create and manage TPU VMs. You can perform standard VM configurations such as using a custom OS or configure boot disk size. Compute Engine APIs support the creation and management of TPU slices across all consumption options, enabling small-scale experimentation and large-scale training and inference workloads.

For more information, see TPU resources in Compute Engine.

Generally available: Cloud TPU now offers TPU availability in AI zones. To learn more, see About AI zones.

Generally available: TPU7x is generally available (GA). TPU7x is the first release within the Ironwood family, Google Cloud's seventh generation TPU. TPU7x supports large-scale AI training and inference, providing performance and cost-effectiveness for demanding workloads such as large language (LLMs), mixture of experts (MoEs), and diffusion models. For more information, see the TPU7x (Ironwood) documentation.

Preview: TPU7x is available in Preview. TPU7x is the first release within the Ironwood family, Google Cloud's seventh generation TPU. TPU7x supports large-scale AI training and inference, providing performance and cost-effectiveness for demanding workloads such as large language models (LLMs), mixture of experts (MoEs), and diffusion models. For more information, see the TPU7x (Ironwood) documentation.

Public preview: You can request Cloud TPUs using future reservations in calendar mode. This mode, powered by the Dynamic Workload Scheduler, lets you check TPU availability up to 120 days in advance and request capacity based on your schedule. You can use calendar mode to reserve TPUs for 1 to 90 days. Requesting a short-term reservation with calendar mode is a good fit for training and experimentation workloads that require precise start times and have a defined duration. For more information, see Request a short-term reservation using calendar mode.

Public preview: You can enable reservation sharing for Cloud TPU. This feature lets you share a reservation across multiple projects. You can also share a reservation with Vertex AI for training or serving workloads. For more information, see Share a Cloud TPU reservation.

Flex-start for Cloud TPU, powered by Dynamic Workload Scheduler, is available in Preview. Flex-start is a flexible and cost-effective consumption option for AI workloads. Flex-start enables you to dynamically provision TPUs for up to 7 days using the queued resources API, without long-term reservations. This option is ideal for quick experimentation, small-scale testing, dynamic inference provisioning, and model fine-tuning. For more information about Flex-start for Cloud TPU, see Request Cloud TPUs using Flex-start.

This Release Note announces General Availability of Trillium AKA v6e. Trillium is the 6th generation and latest Cloud TPU. It is fully integrated with our AI Hypercomputer architecture to deliver compelling value to our Google Cloud Platform AI customers.

We used Trillium TPUs to train the new Gemini 2.0, Google's most capable AI model yet, and now enterprises and startups alike can take advantage of the same powerful, efficient, and sustainable infrastructure. Today, Trillium is generally available for Google Cloud customers and this week we will be delivering our first large tranches of Trillium capacity to some of our biggest Google Cloud Platform customers.

Here are some of the key improvements that Trillium delivers over the prior generations, v5e and v5p:

• Over 4x improvement in training performance.

• Up to 3x increase in inference throughput.

• A 67% increase in energy efficiency.

• An impressive 4.7x increase in peak compute performance per chip.

• Double the High Bandwidth Memory (HBM) capacity.

• Double the Interchip Interconnect (ICI) bandwidth.

• 100,000 Trillium chips per Jupiter network fabric with 13 Petabits/sec of bisection bandwidth, capable of scaling a single distributed training job to hundreds of thousands of accelerators.

• Trillium provides up to 2.1x increase in performance per dollar over Cloud TPU v5e and up to 2.5x increase in performance per dollar over Cloud TPU v5p in training dense LLMs like Llama2-70b and Llama3.1-405b.

• GKE integration enables seamless AI workload orchestration using Google Compute Engine MIGs including XPK for faster iterative development.

• Multislice training with Trillium scales from one to hundreds of thousands of chips across pods using DCN.

• Training and serving fungibility enables use of same Cloud TPU quota for both training and inference.

• Support for collection scheduling with collection SLOs being defended.

• Full-host VM support to enable inference support for larger models (70B+ parameters).

• Official Libtpu releases that guarantees stability across all three frameworks (Jax/Pytorch-XLA/Tensorflow).

These enhancements enable Trillium to excel across a wide range of AI workloads, including:

• Scaling AI training workloads like LLMs including dense and Mixture of Experts (MoE) models

• Inference performance and collection scheduling

• Embedding-intensive models acceleration

• Delivering training and inference price-performance

You can now request Cloud TPUs as queued resources in the Google Cloud Console. Queuing your request for TPU resources can help alleviate stockout issues. If the resources you request are not immediately available, your request is added to a queue until the request succeeds or you delete it. You can also specify a time range in which you want to fulfill the resource request. For more information, see Manage queued resources.

Creating a Multislice TPU environment is now available in the Google Cloud Console. You can use Multislice to run training jobs using multiple TPU slices within a single Pod or on slices in multiple Pods. You must use a queued resource request to create a Multislice environment. For more information, see Cloud TPU Multislice overview.

Cloud TPU now supports TensorFlow 2.16.1. For more information see the TensorFlow 2.16.1 release notes.

Cloud TPU now supports TensorFlow 2.14.1. For more information see the TensorFlow 2.14.1 release notes.

Cloud TPU now supports TensorFlow 2.15.0, which adds support for PJRT. For more information see the TensorFlow 2.15.0 release notes.

Cloud TPU now supports TensorFlow 2.13.1. For more information see the TensorFlow 2.13.1 release notes.

Cloud TPU now supports TensorFlow 2.14.0. For more information see the TensorFlow 2.14.0 release notes.

You can now create Cloud Tensor Processing Unit (TPU) nodes in Google Kubernetes Engine (GKE) to run AI workloads, from training to inference models. GKE manages your cluster by automating TPU resource provisioning, scaling, scheduling, repairing, and upgrading. GKE provides TPU infrastructure metrics in Cloud Monitoring, TPU logs, and error reports for better visibility and monitoring of TPU node pools in GKE clusters. TPUs are available with GKE Standard clusters. GKE supports TPU v4 in version 1.26.1.gke-1500 and later, and supports TPU v5e in version 1.27.2-gke.1500 and later. To learn more, see TPUs in GKE introduction.

Cloud TPU now supports TensorFlow 2.12.1. For more information see the TensorFlow 2.12.1 release notes.

Cloud TPU now supports TensorFlow 2.13.0. For more information see the TensorFlow 2.13.0 Release Notes.

You can now view historical logs of maintenance events on your TPU in system event audit logs. For additional information see the maintenance events documentation.

Cloud TPU now supports Tensorflow 2.11.1. For more information see the TensorFlow 2.11.1 release notes.

Cloud TPU now supports Tensorflow 2.12.0. For more information see the TensorFlow 2.12 release notes.

Cloud TPUs now support the PyTorch 2.0 release, via PyTorch/XLA integration. On top of the underlying improvements and bug fixes in PyTorch's 2.0 release, this release introduces several features, and PyTorch/XLA specific bug fixes.

Beta Features

PJRT runtime

• Checkout our newest document; PjRt is the default runtime in 2.0.
• New Implementation of xm.rendezvous with XLA collective communication which scales better (#4181)
• New PJRT TPU backend through the C-API (#4077)
• Use PJRT to default if no runtime is configured (#4599)
• Experimental support for torch.distributed and DDP on TPU v2 and v3 (#4520)

FSDP

• Add auto_wrap_policy into XLA FSDP for automatic wrapping (#4318)

Stable Features

Lazy Tensor Core Migration

• Migration is completed, checkout this dev discussion for more detail.
• Naively inherits LazyTensor (#4271)
• Adopt even more LazyTensor interfaces (#4317)
• Introduce XLAGraphExecutor (#4270)
• Inherits LazyGraphExecutor (#4296)
• Adopt more LazyGraphExecutor virtual interfaces (#4314)
• Rollback to use xla::Shape instead of torch::lazy::Shape (#4111)
• Use TORCH_LAZY_COUNTER/METRIC (#4208)

Improvements & Additions

• Add an option to increase the worker thread efficiency for data loading (#4727)
• Improve numerical stability of torch.sigmoid (#4311)
• Add an api to clear counter and metrics (#4109)
• Add met.short_metrics_report to display more concise metrics report (#4148)
• Document environment variables (#4273)
• Op Lowering
  • _linalg_svd (#4537)
  • Upsample_bilinear2d with scale (#4464)

Experimental Features

TorchDynamo (torch.compile) support

• Checkout our newest doc.
• Dynamo bridge python binding (#4119)
• Dynamo bridge backend implementation (#4523)
• Training optimization: make execution async (#4425)
• Training optimization: reduce graph execution per step (#4523)

PyTorch/XLA GSPMD on single host

• Preserve parameter sharding with sharded data placeholder (#4721)
• Transfer shards from server to host (#4508)
• Store the sharding annotation within XLATensor(#4390)
• Use d2d replication for more efficient input sharding (#4336)
• Mesh to support custom device order. (#4162)
• Introduce virtual SPMD device to avoid unpartitioned data transfer (#4091)

Ongoing development

• Ongoing Dynamic Shape implementation
  • Implement missing XLASymNodeImpl::Sub (#4551)
  • Make empty_symint support dynamism. (#4550)
  • Add dynamic shape support to SigmoidBackward (#4322)
  • Add a forward pass NN model with dynamism test (#4256)
• Ongoing SPMD multi host execution (#4573)

Bug fixes & improvements

• Support int as index type (#4602)
• Only alias inputs and outputs when force_ltc_sync == True (#4575)
• Fix race condition between execution and buffer tear down on GPU when using bfc_allocator (#4542)
• Release the GIL during TransferFromServer (#4504)
• Fix type annotations in FSDP (#4371)

Cloud TPU now supports TensorFlow patches: 2.8.4, 2.9.3, and 2.10.1. See the TensorFlow release notes for details:

• 2.8.4 release notes
• 2.9.3 release notes
• 2.10.1 release notes

Cloud TPU now supports Tensorflow 2.11.0. For more information see TensorFlow 2.11 release notes.

Cloud TPU now supports Tensorflow 2.10.0. For more information see TensorFlow 2.10 release notes.

Cloud TPU now supports TensorFlow 2.6.5 and TensorFlow 2.7.3.

For more information see TensorFlow 2.6.5 and TensorFlow 2.7.3 release notes.

Cloud TPU now supports Tensorflow 2.8.2 and 2.9.1. For more information see TensorFlow 2.8.2 release notes and TensorFlow 2.9.1 release notes.

Cloud TPU now supports Tensorflow 2.6.3. For more information see TensorFlow 2.6.3 release notes.

Cloud TPU now supports Tensorflow 2.5.3 and 2.7.1. For more information see TensorFlow 2.5.3 release notes and TensorFlow 2.7.1 release notes.

Cloud TPU now supports Tensorflow 2.8.0. For more information, see TensorFlow 2.8.0 Release Notes.

Cloud TPU team just released TF-2.4.4, TF-2.5.2 and TF-2.6.2 on Cloud TPUs. The TensorFlow release notes for these releases are shown below.

• TF 2.4.4
• TF 2.5.2
• TF 2.6.2

Cloud TPU now supports Tensorflow 2.7.0. For more information, see Tensorflow 2.7.0 Release Notes

Cloud TPU team just released TF-2.3.4, TF-2.4.3 and TF-2.5.1 on Cloud TPUs. The TensorFlow release notes for these releases are shown below.

• Tensoflow-2.3.4 Release notes

• Tensoflow-2.4.3 Release notes

• Tensoflow-2.5.1 Release notes

Cloud TPU now supports Tensorflow 2.6.0. For more information, see Tensorflow 2.6.0 Release Notes.

In TF 2.6.0, TensorFlow has introduced a new version of the TF/XLA bridge using the MLIR compiler infrastructure. The MLIR bridge is enabled by default. To explicitly disable it at runtime, add the following code snippet to your model's code:

tf.config.experimental.disable_mlir_bridge()

The Cloud TPU team has released support for TensorFlow 2.4.2. The corresponding Tensorflow release notes are:

Tensoflow-2.4.2 Release notes

Cloud TPU team just released TF-2.1.4, TF-2.2.3 and TF-2.3.3 on Cloud TPUs. The TensorFlow release notes for these releases are shown below.

• Tensoflow-2.1.4 Release notes

• Tensoflow-2.2.3 Release notes

• Tensoflow-2.3.3 Release notes

Cloud TPU now supports Tensorflow 2.5.0. For more information, see Tensorflow 2.5.0 Release Notes

New Cloud TPU VMs make training your ML models on TPUs easier than ever

The new Cloud TPU VM architecture makes it easier than ever before to use our industry-leading TPU hardware. The Cloud TPU VMs provide direct access to TPU host machines, offering a new and improved user experience for developing and deploying TensorFlow, PyTorch, and JAX on Cloud TPUs. Instead of accessing Cloud TPUs remotely over the network, Cloud TPU VMs let you set up your own interactive development environment on each TPU host machine. Now you can write and debug an ML model line-by-line using a single TPU VM, and then scale it up on a Cloud TPU Pod slice to take advantage of the super-fast TPU interconnects. You have root access to every TPU VM you create, so you can install and run any code you wish in a tight loop with your TPU accelerators. You can use local storage, execute custom code in your input pipelines, and more easily integrate Cloud TPUs into your research and production workflows. Google supports Cloud TPU integrations with TensorFlow, PyTorch, and JAX, and you can even write your own integrations via a new libtpu shared library on the VM. For more information, see https://cloud.google.com/blog/products/compute/introducing-cloud-tpu-vms.

Cloud TPU now supports Shared VPC

Shared VPC allows an organization to connect resources from multiple projects to a common VPC network to communicate with each other securely and efficiently using internal IPs from that network. This release enables connecting to Cloud TPU Nodes from Shared VPC networks.

Compute Engine TPU Metrics and Logs In-Context

New Monitoring tab for TPUs provides key TPU Metrics and access to logs at a glance. You can see a variety of key TPU metrics including MXU utilization, CPU, memory, sent/received traffic, and more. In addition, it provides instant access to TPU logs which give insight into important events.

PyTorch/XLA 1.6 Release (GA)

Highlights

Cloud TPUs now support the PyTorch 1.6 release, via PyTorch/XLA integration. With this release we mark our general availability (GA) with the models such as ResNet, FairSeq Transformer and RoBERTa, and HuggingFace GLUE task models that have been rigorously tested and optimized.

In addition, with our PyTorch/XLA 1.6 release, you no longer need to run the