zak

For future reference, the team looked into this, and it appears that the interruptions you experienced were specific to your project and a small number of other projects. The vast majority of TRC projects should see much longer Cloud TPU uptimes when they are able to create on-demand TPUs.

I'm sorry that you had such a frustrating time and that we weren't able to sort it out via email while it was happening. If you decide to try TRC again and run into issues like this, please be sure to engage with TRC support!

Could you share a few technical details about the issues you've encountered with TF / JAX / PyTorch on Cloud TPUs? The overall Cloud TPU user experience improved a whole lot when we enabled direct access to TPU VMs, and I believe the newer JAX and PyTorch integrations are improving very rapidly. I'd love to know which issues are currently causing the most friction.

As mentioned in another comment, it sounds like you're using preemptible TRC TPU quota. If you use on-demand TRC TPU quota instead, that should improve your uptime substantially.

It sounds like you're primarily using preemptible TPU quota, which doesn't come with any availability or uptime expectations at all.

By default, the TRC program grants both on-demand quota and preemptible quota. If you are able to create a TPU VM with your on-demand quota, it should last quite a bit longer than a few hours. (There are situations in which on-demand TRC TPU VMs can be interrupted, but these ought to be rare.) If your on-demand TPU VMs are being interrupted frequently, please email TRC support and provide the names of the TPU hosts that were interrupted so folks can try to help.

When there is very high demand for Cloud TPUs, it's certainly possible for preemptible TPU VMs to be interrupted frequently. It would be an interesting engineering project to make a very robust training system that could make progress even with low TPU VM uptime, and I hope someone does it! Until then, though, you should have a better experience with on-demand resources when you're able to create them. Reserved capacity is even better since it provides an expectation of both availability and uptime.

A few quick comments:

> But it’s important for hobbyists and tinkerers to be able to participate in the AI ecosystem

Totally agree! This was a big part of my original motivation for creating the TPU Research Cloud program. People sometimes assume that e.g. an academic affiliation is required to participate, but that isn't true; we want the program to be as open as possible. We should find a better way to highlight the work of TRC tinkerers - for now, the GitHub and Hugging Face search buttons near the top of https://sites.research.google/trc/publications/ provide some raw pointers.

I'm sorry to hear that you've personally had a hard time getting TPU v3 capacity in europe-west4-a. In general, TPU availability varies by region and by hardware generation, and we've experimented with different ways of prioritizing projects. It's possible that something was misconfigured on our end if your TPU lifetimes were so short. Could you email Jonathan the name of the project(s) you were using and any other data you still have handy so we can figure out what was going wrong?

Also, thanks for the kind words for Jonathan and the rest of the TRC team. They haven't lost any power or control, and they are allocating a lot more Cloud TPU capacity than ever. However, now that everyone wants to train LLMs, diffusion models, and other exciting new things, demand for TPU compute is way up, so juggling all of the inbound TRC requests is definitely more challenging than it used to be.

Actually, the TPU Research Cloud program is still going strong! We've expanded the compute pool significantly to include Cloud TPU v4 Pod slices, and larger projects still use hundreds of chips at a time. (TRC capacity has not been reclaimed for internal use.)

Check out this list of recent TRC-supported publications: https://sites.research.google/trc/publications/

Demand for Cloud TPUs is definitely intense, so if you're using preemptible capacity, you're probably seeing more frequent interruptions, but reserved capacity is also available. Hope you email the TRC support team to say hello!

No, Cloud TPUs support JAX, PyTorch, and TensorFlow, and the new TPU VM architecture provides enough low-level access that users could add support for additional frameworks themselves if they are willing to put in substantial effort.

If you haven't used Cloud TPUs in a while, I'd encourage you to try them now with TPU VMs and the latest versions of JAX, PyTorch / XLA, or TensorFlow. We've gotten a lot of positive feedback from customers and TRC users, so we think the overall experience has improved a lot, though there's always more we want to do.

People especially seem to find Cloud TPUs easy to use in comparison to alternatives when they are scaling up ML training workloads. Once you have a model running on a single TPU core, it is relatively straightforward from a systems perspective to scale it out to thousands of cores. You still need to work through the ML challenges of scaling, but that is more tractable when you aren't simultaneously struggling with systems-level issues.

In particular, you don't need to master a sequence of different networking technologies as you scale up, and the TPU interconnect is so much faster at scale than other technologies (10X last time I checked) that you don't have to work as hard to avoid network bottlenecks. Support for model parallelism on Cloud TPUs is improving across the ML frameworks, too.

To be clear, training ML models at scales that we currently consider large is still very challenging on any platform - for example, the logbooks that Meta recently published are fascinating: https://github.com/facebookresearch/metaseq/blob/main/projec...

We aim for Cloud TPUs to simplify the process of training models at these scales and far beyond: https://ai.googleblog.com/2022/04/pathways-language-model-pa...

We try to make it easy to switch back and forth between Cloud TPUs and other hardware platforms using JAX, PyTorch, and TensorFlow. This is a difficult technical challenge, but the XLA compiler helps a lot, and switching is easier now than it has ever been.

Thanks very much! We've come a long way, but there is always more interesting work required to keep up with the deep learning frontier and enable Cloud TPU customers and TRC users to expand it further.

I started the TRC program alongside the Cloud TPU program to make interesting amounts of ML compute available to a broad group of creative people, not only to academic researchers.

The TRC program welcomes hobbyists, artists, students, independent learners, technical writers, and a variety of others. We love it when the TPU Research Cloud enables people to do something that wouldn't have been possible otherwise.

I definitely recommend applying to the TRC program - please feel free to say directly that you are a hobbyist. The sign-up form is short, and it's likely that you can access a lot of compute if you are flexible and persistent.

Yes, there are a couple of ways to use Cloud TPUs at lower priority and lower cost. If you are a hobbyist, I highly recommend trying out Cloud TPUs for free via the TPU Research Cloud: https://sites.research.google/trc/

If you are flexible, you (or your scripts) can access a lot of compute power at odd hours!

No, Vectorflow is not supported out of the box, and I'm not sure the workloads it targets are the right fit for Cloud TPU hardware. However, be sure to check out the "Ranking and recommendation" section of the linked blog post above - Cloud TPUs are able to accelerate the ML models with very large embeddings that are increasingly common in state-of-the-art ranking and recommendation systems.

The MLPerf 1.0 results provided an apples-to-apples comparison of large-scale TPU and GPU systems across several ML workloads: https://cloud.google.com/blog/products/ai-machine-learning/g...

In MLPerf 1.1, we showcased model training at larger scale: https://cloud.google.com/blog/topics/tpus/google-showcases-c...

The deep learning workloads that people find most interesting and the underlying hardware and software systems are all changing very rapidly. In addition to following MLPerf, we generally recommend that people run rigorous performance and cost comparisons on the actual workloads that they care about accelerating.

Thanks, Frank! You personally helped more Cloud TPU and TRC users than I can count, and you always came through something needed to get done and fast. I really appreciated it!

I started pitching the Cloud TPU program in 2016. Many, many people have contributed since then to build the products that are available today.

Google is a large and complicated place, but we're getting closer to providing the magical interactive supercomputing experience we've wanted for a long time.

The deep learning landscape is evolving very rapidly, and there is a lot of interest in scaling up further, so the next few years will be exciting.

Thanks, and congratulations to many others across many teams who have supported the Cloud TPU program over the years!

We love Colab and would love to upgrade the Colab TPU integration to support TPU VMs! No timeframe yet, but the right folks across JAX / Colab / Cloud TPU are very aware of this issue.

In the previous Cloud TPU architecture, PyTorch and JAX users had to create a separate CPU VM for every remote TPU host and arrange for these CPU hosts to communicate indirectly with the TPU hosts via gRPC. This was cumbersome and made debugging difficult.

With TPU VMs, none of this is necessary. You can SSH directly into each TPU host machine and install arbitrary software on a VM there to handle data loading and other tasks with much greater flexibility.

The blog post provides an example of training cost improvement using PyTorch / XLA on TPU VMs in the "Local execution of input pipeline" section. Hopefully we will be able to provide more tutorials on using PyTorch / XLA with TPU VMs soon.

With TPU VMs, workloads that require lots of CPU-TPU communication can now do that communication locally instead of going over the network, which can improve performance.

Founder of the Cloud TPU program here. If you'd like to experiment with TPU VMs for free and are willing to share your work with the world somehow (e.g. via publications or open-source projects), you can apply to participate in the TPU Research Cloud (TRC) program here: https://sites.research.google/trc/

Here's a quick overview of TPUs from last year's Google I/O keynote: https://www.youtube.com/watch?v=XFFrahd05OM&t=1565s

Yes, TPU VMs dramatically improve the Cloud TPU user experience. You now have direct access to the VM on each TPU host whether you are using JAX, PyTorch, or TensorFlow, which provides a lot more flexibility and control and can often improve performance.

Thanks for the very kind feedback! We've wanted to provide TPU VMs since the beginning of the Cloud TPU program, and I'm delighted that you're enjoying them. Many people across Google contributed to this launch.

We're definitely working on GKE integration, and we intend to provide more support for orchestration over time since individual Cloud TPU workloads are getting bigger and bigger. We'd love to hear any more detailed feedback you might have.

We only see code when customers open-source it or otherwise explicitly share it with us. We are directly in touch with several customers who are using the PyTorch / TPU integration, so we hear feedback from them, and we also run a variety of open-source PyTorch models on Cloud TPUs ourselves as we continue to improve the integration.

(I'm one of the Cloud TPU product leads)

We've seen multiple BERT-related PyTorch models training successfully on Cloud TPUs, including training at scale on large, distributed Cloud TPU Pod slices.

Would you consider filing a GitHub issue at https://github.com/pytorch/xla or emailing [email protected] to provide a bit more context about the specific issue you encountered?

Here's the current PyTorch/TPU troubleshooting guide, which provides information on how to collect and interpret metrics that are very helpful for debugging: https://github.com/pytorch/xla/blob/master/TROUBLESHOOTING.m...

Thanks!

When comparing different hardware configurations within or between public clouds, power measurements generally aren't available, whereas prices or realistic price estimates generally are.

Sure. The $2.48/hour per V100 GPU on GCP does not include the price of the CPU host; that is purely the price to rent a single accelerator. By contrast, a network-attached Cloud TPU v3 device includes both a CPU host and four connected TPU v3 chips that collectively deliver up to 420 teraflops. Furthermore, each individual V100 GPU on GCP has 16 GB of memory, whereas the Cloud TPU v3 device has 128 GB of HBM.

The best apples-to-apples performance-per-dollar comparison we have publicly available was published last fall, and it compared the performance and cost of using various Cloud TPU v2 Pod slice sizes with the performance and cost of using various numbers of V100 GPUs attached to a single GCP host:

https://cloud.google.com/blog/products/ai-machine-learning/n...

We went to great lengths to ensure that we trained exactly the same version of ResNet-50 to the same accuracy in the same way across all hardware configurations. The methodology predated MLPerf and is documented in full here:

https://github.com/tensorflow/tpu/blob/master/benchmarks/Res...

If you were going to do a similar performance-per-dollar comparison today, the simplest approach might be to try to get the code from NVIDIA's MLPerf 0.6 submissions running at scale on one or more major public clouds using the fastest-available networking technology that each cloud provides:

https://github.com/mlperf/training_results_v0.6/tree/master/...

It would be very interesting to see how distributed training performance using large-scale GPU clusters in public clouds compares with the published on-premise MLPerf performance numbers using exactly the same MLPerf code and methodology. With these measurements in hand, it would then be straightforward to make performance-per-dollar comparisons with Cloud TPU v3 Pod slices of various sizes.

Yes, Cloud TPU Pods are designed to train much larger models on much larger datasets. And, as you mention, if you are willing to adjust your model architectures and training algorithms to take full advantage of the hardware, you can sometimes achieve substantial gains.

Author of the blog post here.

As mentioned in other comments, I'd recommend doing a performance-per-dollar comparison in addition to looking at this pure performance comparison at maximum scale.

Author of the blog post here.

I'd recommend doing a performance-per-dollar comparison before drawing this conclusion.