Sponsored Session: TorchTPU: Expanding TPU Programmabil... Kat Ko, Claudio Basile & Jana van Greunen

This talk introduces Torch TPU, a new platform that enhances PyTorch's compatibility with various hardware accelerators, particularly Google's TPUs. The speaker, Yana van Geenen from Meta, explains that the need for hardware optionality in machine learning workflows has become critical due to rising costs and computational demands. Collaborating with Google engineers, they have designed a stack that allows developers to easily port existing PyTorch models to TPUs with minimal changes required to their code. Kat, another speaker from Google, highlights the significant benefits of using TPUs, including their ability to handle diverse workloads efficiently. The session outlines the technical details of the implementation, including how PyTorch operations are translated into TPUs and the optimizations made to ensure strong performance. The speakers emphasize that this effort is not just experimental but has been validated against numerous models, showcasing impressive results in scalability and performance.

Hi everyone. I think we're seeing a last few people trickle in. Um but we're excited today to present about Torch TPU, which is bringing PyTorch to more hardware. Um so, I'll quickly introduce myself. I'm Yana van Geenen. I lead the PyTorch engineering effort at Meta. Um and I have presenting, I think most of the exciting details here, Kat um Claudio, sorry, and and Kat, um both from

Google, and we've been collaborating on this joint effort. So, you know, there maybe a few questions of why why are we looking at this in general? And what we've seen just because of, you know, various factors, whether it's cost or capacity crunch, it is very important these days. Welcome everyone. Uh to have hardware optionality when you're either training large models, running small experiments, or um especially running

inference at scale. And so, uh you know, as we've seen again, the different types of accelerators have different um performance points and optimization schemes. And so, it is very important to be able to take a model that you've developed and trained and actually port it to all the different um hardware uh platforms. The other thing that's changed is that um fundamentally, AI is now a cluster, shall

we say, business, um which is that everything really is distributed. And this makes actually going to different accelerators even more challenging, because not only is the compute side different, you also have a different cluster and communication setup. So, what we're hoping to bring is PyTorch as the portable foundation that can make this challenge easy for um our community and developers. And I think Kat is going to

explain a little bit more about um how this is going to work. Right. Yeah. Uh so, I'm Kat, uh eng lead at on Torch TPU from Google. And uh like Yana was mentioning, right, with hardware optionality as top of mind, we think TP TPUs are a great um option here. TPUs power AI at Google as well as across our cloud customers. And so, that means it's part

of a mature production hardened stack that runs a ton of workloads across Google's internal workloads. It also supports a broad variety of workloads across our cloud customers. And you get that's um world-class scale, uh Gemini being the best hero example, which is trained and served off TPUs. The challenge historically though with unlocking TPUs is that it comes with this usability tax. Um previously, developers would have had

to, you know, adopt a new framework like JAX, but, you know, considering that the broad majority of the ML community is now on PyTorch, right, we're all here today, um it was critical for us to find a path to unlock TPUs um for for you guys. Uh and really meet developers where they are at with their existing PyTorch workloads to port it easily over onto TPUs with

minimal friction, minimal code change. So, a year ago, we embarked on this journey with Yana, with the PyTorch core team, to do exactly that. Um and to, you know, we built this stack from the ground up uh with uh usability, portability in mind without sacrificing performance. And today, we're excited to announce that uh we now have our new PyTorch on TPU stack. Uh so, on the right,

you see we start with PyTorch aten ops. We lower those directly to stable HLO through your standard PyTorch to eager or compiled paths. Um that then gets passed to open XLA compiler, which then compiles an executable that can then run on TPUs. So, throughout this journey, we worked really closely with the PyTorch core team to make sure every step of this lowering path, every decision that we're

making, we are taking as native an approach as possible to really blend that PyTorch native usability with the power and scale of TPUs. Um so that today, you guys can now take your existing PyTorch workloads and run it on a TPU with minimal code change. So, what does that actually look like? If you want to get started on TPUs, all you have to do is take your

device string and replace it with TPUs, and you're done. No more additional mark steps, no more manual syncs, no lazy tensor surprises. Torch TPU works with the eager mode you already love. Uh and today, technically, you still have to add an additional import statement, but very soon, you won't even need that um like as demonstrated in the slide. So, you know, we're really trying to um make

it easy for you to port your workloads onto TPUs with minimal code change. We worked really hard to make this just a drop-in replacement. For more advanced workloads, um you know, we took the exact same approach here. We, you know, anytime we could have gone with maybe a more Google custom approach or aligned with PyTorch, we always chose to align with PyTorch. so, you know, for a

distributed case, for example, uh Torch TPU supports all your standard PyTorch distributed idioms. So, whether you're on FS whether you're using FSDP or your standard DDP, Torch TPU intercepts the collective at the process group layer and maps that directly onto a TPU optimized communication. So, if you have a distributed workload, you can port that over onto a TPU pod with minimal code change as well. So, this

isn't just a research project. We worked really hard over the last year to take this new stack and actually validate it against thousands of hugging face models, run it across a broad range of architecture types from LLMs to diffusion models, etc. Um we ran it across inference workloads, training benchmarking for quality, showcasing strong vanilla out-the-box performance numbers, and we were also able to demonstrate linear scaling up

to 256 chips. Um Claudio is going to dig more into this in his part of the talk, but what I really wanted to highlight in this slide was that it's not just a new stack that we have. We've actually worked really hard to battle test it um against a bunch of a broad variety of uh and prove to ourselves um that it's not you know, you don't

have to sacrifice performance to get that native developer experience. And we can't wait for you guys to try it. Um I'm just going to wrap my section with a quick demo of how you can take a uh Llama 3 8B training workload and run it on TPUs. Torch TPU lets you run standard PyTorch code on TPUs. First, import API from Torch TPU. And then set a device

variable to API.tpu_device. For multi-chip distributed training, use torch.distributed and call init_process_group with a backend set to TPU_DIST. In this example, we'll log the results to a Vertex project. And we'll fine-tune a Llama 3 8B model using Hugging Face Transformers. We can use a standard LoRA config. And then send our model to the TPUs by calling model.to(device). With Torch TPU, we can use existing PyTorch libraries for optimization

and training. And our training loop is the same as a typical PyTorch training loop, including the optimizer update steps. Finally, for this example, we'll log our metrics to Vertex so we can aggregate all the logs across multiple TPUs. We can launch the job using XPK. And in this case, it's on 256 Ironwood TPU chips. And in Vertex, we can see that the job has been launched using

a 4.4.8.2 TPU topology. Finally, in Vertex, we can see our parameters for the experiment. And the aggregated average loss. And this is all made possible by changing just a few lines of code with All right, with that, I'll hand it over to Claudio. Thank you. Thank you. Hi everyone. Um So, um the next part of this uh talk I wanted to discuss a little bit more about

the um the journey and also from internal point of view, how is how we build how did we build Torch TPU, how it works. Before doing that though, I want to just uh do a little bit of an overview on the TPUs. TPUs are quite different uh from GPUs in in many ways. And so, I thought maybe worth while spending a few minutes on it. Um in

the picture, you have a host. Um the host may be connected to multiple chips, four chips in this in this uh uh slide. Um the chips have a an ICI interconnect, uh which is really a chip-to-chip uh um optical link. And they are connected into a torus structure, which could be a 2D or 3D depending on the TPU generation. Those are bidirectional link, really, really fast fast

speed. Um the chip may have multiple chiplets, and and two in this case, like in Ironwood, for instance. And inside the chiplets, you will see uh tensor cores. The tensor core is really the the the uh muscle part of the TPU, which has a very large uh matrix multiply unit. Um in in Ironwood, you'll be able to do 256 by 256 matrix multiplications really, really fast. And

then you have also number of sparse cores. Those are additional units um that are designed for um irregular memory access, vector operations, and and all of those. and and so, by orchestrating this uh hardware, and one one one important thing I wanted to mention, this is mostly, especially the tensor core, is single-threaded. It's not like on a on a on a GPU, you're fundamentally launching a lot

of threads of executions to to to hide the the latency, but this is really the the maximum that the the acceleration really comes from that matrix multiplier and by using the sparse cores for sparse operations. So, this gives you a little bit of an idea how the TPU works and there is much more material on on on internet on the Google Cloud documentation about this. Um, when

we started as Scott was mentioning before, we wanted to take this as a as a PyTorch friendly approach as possible. And so, we started from PyTorch eager. As a matter of fact, we this started as a prototype just to validate the idea and and then we realized this this was the right path to go. And we support three different types of eager modes. The first one is

the one that you are probably familiar with. You know, it's for debugging is basically you may be familiar with a CUDA launch blocking environment variable which allows you to launch one at a time on the device and also to synchronously for the device to be ready before you dispatch next operation. This what we call strict eager is really what you are mostly familiar with with CUDA. You

launch operations one at a time, but you obviously are not waiting for them. The CPU and the GPU may run asynchronously and this is happening also on the on the TPU. What's more interesting interesting though is eager way that we call defer and fuse. When we build eager, we said, you know, is there a better way, a more performant way to do eager and that's that's what

we we did. And this is based on the technique where we are observing the stream of ops being dispatched by the PyTorch dispatcher and we may defer some of them up to certain points. We call them materialization points and then we can fuse these ops together and compile for maximum performance. And I'll show you later how this new eager mode performs with respect to your typical strict,

if you want, and the torch compiler. We also have a compilation cache that it's implemented in torch TPU. This compilation cache works on a single host so that if you have multiple jobs host, they don't have to recompile multiple times the same executable. It's all shared. And if you have multiple hosts and you can set it up with an environment variable that gives you a directory where

all these executables are persisted and shared across multiple hosts. Um, and so that also gives you a faster round trip times as you experiment and you run these jobs in production. Um, torch compile, we again followed the PyTorch approach. We use dynamo. the workload to an FX graph and then after that we go through a eager PyTorch torch TPU eager phase where we trace the graph, the

FX graph with a deferral method which basically allows us to accumulate all the ops, defer them and compile them into a single stablehlo program. And once that stablehlo program gets compiled, we get a cache key and we restore we store that cache key into the returned callable and that's how we are able to basically execute these torch compiled subsections of your workload. One thing you will see

here is that we did not use inductor for a number of reasons. First of all, it was much faster for us. Having built eager first, it was much faster for us to to reuse a lot of the code we had built to power torch compile. Again, we are still using dynamo, so the front end is the same. The other aspect is actually probably more profound, which is

XLA has been highly optimized for the last X number of years for TPU and it works really really well. One of the nice properties of the XLA compiler is that it can automatically overlap computation and communication. And this is critical if you're doing really large training workloads. And and so by using XLA, this came came for free. Okay. Now, having said that, we are also investigating an

inductor backend for Palace in the context of Helium, which I'll I'll mention more later down the road. Yes, kernels. Kernels are also very very critical. A lot of our our effort was to make sure that out of the box performance was great and that's that's what we we do normally and we tune it, but obviously if you are in a high performance workload where it's high throughput

or or low latency, you want to be able to use kernels. We started with Palace because that's really the kernel language that's available for the TPU and we support a seamless Palace integration. If you look at this picture here on the right hand side, you'll see this is PyTorch code. The call actual equals add vectors X and Y, that is a call to a kernel. X and

Y, those are torch tensors. Now, the kernel can be written as a Palace kernel here and it's basically mostly Palace with an annotation, a torch TPU specific annotation. Could be JAX code or coming up soon is Helium. And so in fact, we are highly committed to Helium and we want to support and make sure that this provides also a very valuable alternative as a DSL for for

PyTorch targeting TPUs. All right. Distributed, let me know on the on the time. I can talk for Yeah, I'm fine. >> [laughter] >> All right, distributed. So, the approach we took again was to can we provide operations at the lowest point so that everything else that's built on top of it will work out of the box, right? That was the premise. And so, what we did was

to implement the process group abstraction through the collectives, the TPU collectives. TPU collectives are, you know, they're different from the GPU. You don't have a nickel kind of library. Everything is in for collectives, everything is a stablehlo op, so they will go through the compiler and all of that. But as a user, as a PyTorch user, you don't see any of that because you're just using PyTorch

abstractions. As a as we did the implementation of process groups, a lot of things just started working out of the box including DDP and the tensor and fair scale, FSDP V2 and we we tried this on llama, various versions of llama, Gwen. There's a number of models that we run in a distributed way and it all worked. One thing I do want to emphasize that we do

support SPMD. We do support that. And that may be a a new part of a new news. You know, in the TPU typically, if you're using JAX, you know, you are really you follow the SPMD with sharding and Obviously, to meet the demands for PyTorch, we had to support SPMD and that's that's a big piece of it. We are in the process of improving performance, measuring performance,

scalability. We talked about 256 chips. We are actually doing more. We'll be doing more pod size, which is thousands, tens of thousands of chips. And then another aspect that we were also working on is really supporting an abstraction that's similar to the CUDA streams. I mentioned that the TPU is intrinsically single-threaded, especially for the tensor core then you have the sparse core chiplets that can be targeted

severally and so we want to expose them so that you get those CUDA streams. Results. So, >> So, that's we took a number of few thousands model from Hugging Face and we run them. Okay, we did a script. And we didn't do anything special. We just, you know, we took them out of the we downloaded and run them. About 70% of them run out of the box

with on CPU and TPU. We didn't have any issues. Some of them had some bugs, some unique bugs, some ops were missing. But fundamentally, you know, 70% success rate on thousands of models is is pretty decent. And this is data that is probably a month old. We probably by now have already closed the gap with a few ops, missing ops. Those are [snorts] some of the examples

of the models that we have validated. You'll see some transformer based models, some ResNets and and some diffusion models and so forth. And while we are doing this, we are also optimizing, measuring and optimizing performance on selections of models, Gemma, Llama, Gwen, which have business relevance. That chart shows you the number of ops that we have implemented. We I think now we are complete in terms of

the core retain ops, but we there are also more ops that may be needed. You know, we just looked at the what is that called, native functions YAML file. We looked at all the all the apps that have you know a CUDA implementations and said okay, maybe those are the apps that we should also implement. And I think by the end of the quarter we should be

in the 80% range. This is a little bit of a performance slide. I talked about the various this is various models the debug the strict and the deferred fused and compiled. You can see that this new eager mode really reduces the gap between the eager mode you're used to and the compiled mode that you're you're familiar with. And so that's something very interesting that you can benefit

with no flags no changes to to your This is a very important slide. We measure scaling on a training workload 8 to 256 and virtually the scaling was we expect that you know TPUs are known to scale very really really well with large number of chips and we confirm Okay, so if you're targeting performance obviously our goal is to make a model run out of the box

but if you're targeting performance you will have to take the hardware into account. If your head size dimension is 64 you're not utilizing the TPU at the maximum speed. So there are a few things that you will have to do if you want to have a high MMS view on the TPU and these are some some examples. And we will provide obviously we're going to provide tools

and documentations on on how to best utilize the hardware. There are some interesting problems that we are still tackling. This is a list of them. Bound of dynamism is a very interesting problem as well as a library of precompiled TPU TPU kernels is something that doesn't exist today and we're working on. And finally this is a summary of what you should expect in 2026 in no specific

order so not putting any ETA here opening up the GitHub repository so you guys can try yourself on on Google Cloud Helion VM Torch Titan and and so forth. If you want to know more here is a QR code and we're here for for questions. >> [applause] >> Yes. No no fully open source. Yeah yeah. No no no. We we are we are completely committed to the

PyTorch community and support that so it's going to be fully open source and we will accept contributions and all of that. Yeah. Yeah. When when into 2026 yes this year. When when is this going to be public right? Yes the plan is by the end of the Good. Thank you everyone.