Google Transcript: Co-Designing TPUs, Distributed Systems, and Frontier AI Models
Google Dev Labs Fireside Chat with Jeff Dean and Bill Jia on Scaling AI Systems, September 20, 2026
Introduction and the Early Days of Google
Host: So, welcome everybody. Thank you so much for joining us for all your enthusiasm. This is our seventh Dev Labs, and every single time we do it, it gets bigger and better. And this is by far the best one yet. And we got so much interest, and thank you for your feedback. So the way that this is going to work is we're going to talk about a couple of questions, talk about the evolution of TPUs. This is Jeff, who doesn't need any introduction at all. He was employee number 30 at Google. So let's start with that story, and we've got Bill here too, who's going to also be asking questions, but we also want him to be able to give us a little bit of tidbits in terms of execution of how AI infrastructure, the vision, and how it's getting executed now and also in the next couple of years, if that's okay. All right, let's do it. So Jeff, tell us a little bit more. You were employee number 30 at Google.
Jeff Dean: Yep. Yep.
Host: 1999. Most of the folks here weren't born yet, or they might have been, or they were in diapers. What was that like?
Jeff Dean: No, it was a lot of fun. I mean, I think when I joined, we were all in this tiny little office in downtown Palo Alto, wedged above what's now a T-Mobile store, sort of roughly that kind of scale of office. And we were trying to build a really high-quality search product. And I'd gotten some exposure to that in a prior job of information retrieval, particularly on the web, using the graph structure of the web to augment the existing information on pages and so on. And so Google seemed like a natural place to go try that in a less researchy environment, more of a people using my work directly, which was super fun. It was super high energy, and we were growing our traffic like 7% every week or something. And if you do 1.07 to the 52nd, you know you're like melting every year and trying to avoid falling over on Tuesdays at noon. So that was a lot of fun. And as the company has continued to grow, we've done more and more things, which is always exciting to see.
Host: And then what happened in 2001, if we fast-forward over?
Jeff Dean: So to avoid melting, one of the things we were doing was constantly rewriting the whole search and indexing and query serving system to be more efficient and use different data structures and more compact representations of things. And we got to a point in 2001 where we realized that the way we were sort of scaling both the index size and the amount of capacity we had for serving the index was we would make more and more partitions of the index. We called them shards to have a bigger index. And then every one of those shards, you would make more and more copies of it to provide capacity for serving more queries simultaneously. And we eventually got to the point where, instead of 7 shards and 10 copies each, we had like 60 shards and 20 copies of each in every data center. And we eventually did the math and realized that we could put, instead of having the index on disk, we could put the entire index of the web in memory on the 1,200 machines that would otherwise be 20 copies of the 60 partitions. And that was really fun. We sort of sped up all of our search system by like 4x in 3 days, which was pretty fun, yeah.
The Genesis of the TPU: From Speech Recognition to Custom ASICs
Host: And then in 2013, you did some back-of-the-envelope calculations.
Jeff Dean: Yeah. Yeah. So I started to work on deep learning training infrastructure and training deep learning models because it seemed like the right abstraction. I'd actually gotten exposed to neural nets as an undergrad in 1990, and I did an undergrad thesis on parallel training of neural nets because I just felt like, oh, this is the right thing. We just need to get 32 wimpy processors instead of one, and then we'll be able to train amazing models. Turns out we needed like a million times more compute power, not 32. But eventually, through Moore's Law improvements and so on, we started to have that starting maybe 2008 or 2009. And then in 2011, we started a project to really try to scale up and do distributed training of models for speech, for vision, and for language on—at the time, we had CPUs in our data centers, lots of CPUs. So we used about 16,000 CPU cores to train a large vision model, and we got these amazing improvements. And we were doing that for speech models as well, and the models were amazing improvements in quality. So we got like the equivalent of the last 20 years of speech research in terms of improvements in word error rate in like the two months that we tried to build a deep acoustic speech model. And so that was like, oh my goodness, this is going to be so good. But how do we serve it?
And at the time, I did some calculations where I said, okay, suppose 100 million people start talking to their phone for a few minutes a day, then how will we serve that? And turns out we would have needed to double the number of computers Google had just to roll out the better speech recognition model in one tiny, obscure part of Google. And that seemed a little excessive, not to mention impractical. So we decided that specialized hardware was the way to go. And that's sort of the genesis of the TPU family: we wanted to build specialized accelerators that basically were really, really good at low-precision linear algebra and nothing else. And then that's how you can serve the speech models and vision models and so on much more efficiently. So it turned out it was like 30 to 80 times better performance per watt than contemporary CPUs and GPUs of the day in 2015 when we had the chip available. And that was the first iteration.
Bill Jia: Yeah, TPU v1, which was really targeting inference, right? And then subsequent generations have really mostly targeted training and inference as well. And then more recently, we've even started to split those lines apart again because training and inference are somewhat a bit different.
Host: And when you published the paper in 2017, it was accepted. It's the most widely cited paper in the 50-year history—
Jeff Dean: Yeah, of ISCA, yeah. The paper about TPU v1 is a many-authored thing; there's probably like 35 authors or something. But that's nice considering it's relatively recent in the 50-year history of computer architecture.
Evolution of TPU Architecture, Software Stacks, and Optical Switching
Bill Jia: That's really just amazing that Google was considered as the top-notch software company, but then started to push on the hardware roadmap. That's really just amazing. And doing hardware is one thing, right? Because to make the TPU fantastic, there's a lot of board design, chipset design, TPU itself design. But also, we also need to make the compiler and framework work really well. Early on, you actually pushed on XLA, which is our TPU compiler, and also pushed on TensorFlow. If you can introduce what the early journey looked like pushing on this software stack for TPU.
Jeff Dean: Yeah. I mean, I think what you want as a machine learning developer or researcher is you just want to be able to express your high-level idea and then have it sort of magically run on a large-scale system without you having to think too much about all the performance things you might need to do under the covers to really get good performance. And so that means having a good compiler, having really good interconnects, having frameworks that express the abstraction of: am I talking to 4 chips on my local machine, or am I talking to 10,000 chips in a large distributed setup? You'd like to not have to think about that too much, or do different things for those two settings.
Host: Talk about the iterations of different TPUs, and where do you see it all going?
Jeff Dean: Yeah. So after TPU v1, which is really just a PCIe card that fit in a slot—and we actually bought a whole bunch of those without knowing how we would use them. I cornered our CFO at the time and said, we need to buy a lot of these because we're going to use them; we're not sure what for, but can we please buy a lot? So we did. And then TPU v2 was kind of the first system that we designed with training in mind, with not just a single chip, but a whole supercomputer-like system of many chips interconnected with a 2D torus. Over the generations, we've introduced liquid cooling, starting in TPU v3. So it's always exciting when you have pipes going to the surface of your chips in your system; leaks are not a good thing.
And then TPU v4: as you scaled up the pod size—TPU v2 and v3 had fixed wiring between 256 chips and then 1,024 chips—as you start to get larger and larger systems, you start to have failures in individual chips or boards or systems. You need to be more fault-tolerant because you still want to be able to use the whole topology of a much larger pod, even if something's broken. And so TPU v4 introduced this notion of an optical reconfigurable network between racks of machines.
Host: And so you can sort of assemble like Lego bricks from all over the data center floor, right?
Jeff Dean: With funny little switches that have little micro-adjustable mirrors to make it seem as though these eight racks are right next to each other, even though they're scattered across many aisles in the data center. And then that's been super useful. TPU v5 and v6 and v7 all introduce lower and lower precision computational formats, which has been super useful. You get much higher performance from that, although it's always weird to compare FLOPS when some FLOPS are 16-bit, some are 8-bit, and some are 4-bit. If you can make use of them in your machine learning algorithms, the lower-precision ones are a really great way to get even more performance.
Pathways: Orchestrating Irregularity and Massive Distributed Training
Host: And so the first TPU was 2015; 2017 was when the paper came out. It's just legendary now. And then 2018, Pathways. Let's talk about Pathways a little bit and why that's so important.
Jeff Dean: Yeah. So we started to build Pathways really to handle training of much more irregular and much sparser models, and also as the underlying system infrastructure for it to enable us to have a single programming model, a single process that would be driving a large distributed training system. It's much nicer as a programmer if your single Python process just looks like it has 8,000 devices attached to it, and you're like, great, I just want to do an AllReduce across all 8,000 devices, and it will do that without you having to sort of do anything special.
The underlying software in both the XLA compiler does some of that, and then at the slightly higher level, Pathways sort of orchestrates data movement between lots of different chips. So you have chips within the same TPU pod, and if you need to communicate there, it'll use the ICI high-speed links that are in the TPU pod. But if you then need to have something that spans multiple pods, or you need to talk from this chip on this pod to this chip on this pod, then Pathways will sort of orchestrate that transfer using the data center network, or even a wide area network if you're cobbling together a training job of pods in Oklahoma, Texas, and Iowa or something. And that kind of just gives you this nice abstraction of: I just have a whole bunch of compute, and I'm going to allow the system to figure out how best to use it.
Co-Designing Gemini Models with Hardware and Infrastructure
Bill Jia: Yeah, it's amazing. So Jeff, very early on, you were involved a lot in networking design for TPU, TPU itself design, then the compiler and framework, and of course Pathways, which actually orchestrates the entire data movement, the networking traffic reconfiguration, everything. So that's all the AI infrastructure stack. Now you are very deeply involved in designing the Gemini models, the current Gemini models, and future Gemini models as well. Then the Gemini models can really be co-designed with the entire infrastructure stack. So from your point of view, what kind of benefits have we already collected, and what kind of future thoughts do you have so that we can push the models' cutting-edge boundary?
Jeff Dean: I think the whole set of TPU generations has really benefited from the fact that we are big users of the ML compute ourselves, and we have lots of people trying to push the boundaries of new research ideas that might stress the hardware in ways that existing workloads don't. And so by having everything under one Alphabet roof, we actually can have a lot of interaction between the future TPU generation designers, the software designers of the compilers and the infrastructure, and then the machine learning researchers who can say, hey, this kind of thing seems to be working at small scale, we think that's going to be really important for large-scale future training runs in a year or two years, and we want to make sure our hardware does that.
Because as a computer architect, if you're doing this in isolation, just trying to guess where the fast-moving AI field is going in the two- to six-year timeframe that your particular chip you're designing now has to be relevant in, it's super hard. The more insight you can get of, yeah, we think this thing will work—and the more iteration. It's not just like we think we need to do this; it's often like the hardware people say, well, that's hard, but we could do this, would that be useful? And you have this kind of back-and-forth of: yes, that would actually enable us to do something very close to this, or maybe even better. And that's super important for making effective co-design.
Reliability at Scale: Goodput and Fault Tolerance across 100,000 TPUs
Bill Jia: Yesterday, I was on stage to give the opening. I was talking about, as we start to scale up Gemini models—remember, I remember every single month we sit in the same meetings to talk about reliability. So I'm on the infra side, and Jeff represents Google DeepMind. We say, hey, you know what, the Gemini models become bigger and bigger, then we use tons of TPUs to train the models. Eventually, we use 100,000 TPUs to train the models. That's the one metric we are looking at together, which is called goodput. Goodput.
Then initially, our goodput was not good. I remember what we said together: badput dominates goodput. Initially, it was like 60%. As we scale up more and more TPUs to pre-train that model, it becomes very complicated. We need bigger scale, we need more data, we need more parallelism. The goodput was like in the beginning 50%, 60%, 70%, but it's just not good enough, right? But long story short, now we use massive-scale training infrastructure, train much more complicated models, and goodput can be as high as 95%, even 98%.
Jeff Dean: Yeah, and I think that's actually a combination of a whole bunch of things: better operational practices, better QA testing of the chips when they're deployed, software things that can deal better with failures and continue to make forward progress even if some part of the system is down. All these things together matter a lot and make a big difference.
Bill Jia: Yeah, I think that's another benefit if we own the model, we own the data, and we also own the entire hardware, networking, and all the software stack and the operations team.
Jeff Dean: You can say, hey, it's super important, because if you think about an individual tray of TPUs, if you don't think that's part of a larger system, you're like, okay, fine, we can repair that in a week when we sweep by here. But if it's part of a pod that is being used and there's broken trays in it, then you want to go fix them immediately, because the blast radius of the thing you need to repair is actually much bigger than it might seem.
Host: So, when you're talking about 100,000 chips, reliability really isn't a software patch problem anymore. Would you say it's more of a system design? Is it more related to physics?
Jeff Dean: I think reliability is really a property of a whole system, and there's many different things you want to make robust. One thing you can do is make robust wholes out of unreliable parts. This is actually from the very beginning of Google. We bought El Cheapo consumer PCs to serve our search traffic, and we made robust software systems on top of that that would enable us to handle individual machine failures and still give the functionality that whatever system on it was supposed to be providing. Maybe it's a distributed file system, so you replicate data across multiple machines, and you always have some that are available even if some of the replicas of a piece of data are down.
And I think you can do the same thing in large-scale AI training systems. If you have a 20-pod default configuration, but one of them is down, you can make progress with 19 of the pods while you're repairing the 20th.
Operational Coordination and Mitigating Silent Data Corruption
Host: And Bill, you kind of touched on this a little bit yesterday with your keynote about OCS and Jupiter. Do you want to kind of talk a little bit about that and how that plays into reliability?
Bill Jia: Yeah, I think there's a long journey. The first thing is we look at the interrupts every single 10,000 chips, every single part, how many interrupts per day we have. We have to minimize that, because if that interrupts too many, as we scale up the training cluster, with way too many interrupts, that's not very good, right? So one thing we did is improve hardware reliability. Then not only we do that, but also yesterday I was mentioning that even before pre-training starts, software actually scans the entire fleet, every single component: what's your vital sign? If the vital sign is somewhat problematic, we actually either fix it or exclude it before training even starts. That's the first thing.
The second thing: training gets started. But when training gets started, even if the hardware shows very strong vital signs, training lasts for weeks; it still can have problems. How can we identify which specific components and servers have problems, then fix that, exclude it, replace it, and fix that? That's the second thing.
Now also, as Jeff mentioned, there's a lot of operational coordination. If the data center is maintaining the power of this entire data center row, and this row is actually doing the training, that's of course not very good. So we actually coordinate a lot with data center setups, Site Reliability Engineers, and all the ML researchers and engineers. That's a lot of coordination; we have to be in sync.
And I think Jeff also leads a lot on the Gemini model design and software design side. Sometimes things fail, but the training actually combines data parallelism and model parallelism together. If that's within one data replica, maybe the additional replica just keeps going, right? You just average weights automatically, so you don't have to be bothered by this one replica going down. So it's a lot of very comprehensive improvements. Anything to add there?
Jeff Dean: No. I mean, I think there's all kinds of things that can happen, and it's very hard to predict every possible way something might fail. So building robust things that can detect failures, even though you're not necessarily sure what caused the failure, but you can detect it. Occasionally, you have chips that, maybe when the temperature goes up, they start to have unreliability problems, like they add 2 plus 2 and get 5 or something.
Bill Jia: That's the most nasty one, a silent data corruption error. That's horrible, that's horrible.
Jeff Dean: Yeah, and that can happen in the chip itself, it can happen in a flaky network link, it can happen lots of places. We're using a lot of the same principles we used in the very early days of Google where, because we bought these El Cheapo consumer PCs, they not only didn't have ECC, they didn't have parity in their memory. And so if you use a lot of computers to do something and none of them have parity, then you are going to get random bit flips in things. A lot of the computations we did, you wanted to be robust to bit flips. One of the ways you could do that was: fine, I'm processing a billion web pages, if I drop one of them, it's probably not the end of the world. So you could just ignore that by checksumming in software above the hardware layer to make things robust, even if a particular machine or maybe the network link was not reliable.
Open Source Strategy: JAX, StableHLO, and PyTorch on TPU
Host: Now, Google is a big believer in open source, so I'm going to shift gears a little bit to talk about open source. I know we have a lot of academics in the audience, we have a lot of our partners as well. So let's talk a little bit about Google's AI infrastructure strategy as it relates to making sure that it's not a gated resource. So all these things that we talk about, we want to make sure that we're sharing it with the world. Let's talk about JAX and StableHLO and that evolution of open source as a whole, and how is that so relevant today to ensure that this doesn't become a gated resource.
Jeff Dean: Yeah, I think for us it's super important that we interact with the broader ecosystem. Open sourcing things so that people can see the source, they can modify it, they can help us improve it, and have it be a collaborative project that spans many organizations rather than just something we hold to ourselves is super important. That's why we open sourced TensorFlow, why we open sourced JAX, why we open sourced the StableHLO representation for the XLA compiler. And we'll open source more things in the future. We've been big contributors in general to open source, to the Linux kernel and so on, for many years. Usually we're one of the larger organizations that contributes to collective open source efforts because we believe in that, and we think the whole ecosystem benefits when everyone works together.
Host: It's a perfect segue for you to talk about PyTorch and your experience at Meta.
Bill Jia: Yeah, I was talking with Jeff right before this fireside chat. Jeff has been leading a lot of AI infrastructure open source strategy. As Jeff mentioned, in the past, Google open sourced TensorFlow and JAX, and Kubernetes was also created and open sourced by Google, and Android, and a lot of awesome stuff. Now as we push the TPU into the community and emphasize TPU on Google Cloud, we are doubling down on this open source strategy.
As a matter of fact, not only do we open source JAX core, we also build up a lot of upper-level libraries, like how can we do RL—we call it TuneX, which is actually built upon JAX core. We open source that. How to do checkpointing, how to do inference stuff: we actually build a lot of upper-level libraries and frameworks on top of JAX, and we are open sourcing that because we want people using that and leveraging this open source strategy to use TPU directly. That's one thing.
And of course, we want to meet where the customer is, right? So we also want to embrace whatever mature products exist in the open source community. PyTorch is very mature; PyTorch is very widely used in the community. You know what? Let's embrace PyTorch on TPU. We actually have this project called Torch-TPU. Right now, it is in private preview. We want to do the public preview next quarter, just in a couple of months, and then in Q4, we want a public offering on GitHub, and everybody can use it. And as I mentioned yesterday, if we use Torch-TPU, it's just several simple lines of code change: the backend device is TPU, and then hopefully all the training and all the inference run on TPU. We actually support both eager mode and compiled mode as well. That's on the framework level.
We also engage—I'm sure in the audience we have a lot of users using vLLM and SGLang. We also talk with vLLM and SGLang to make sure this kind of upper-level open source inference framework can also be used on TPU as well. So a lot of work right now: we open source our own stuff, and we also embrace mature open source products in the community.
Automating Hardware Design with Machine Learning
Host: And we do have a lot of first-year and second-year PhD students in the audience too. How many of you are PhD students? What are some of the things that they should be focusing on? Some of them are also undergrads too, just so you know.
Jeff Dean: Okay. It's always exciting to be at that stage in your career, because I think you can find things you really enjoy and think are important, and figure out a way to push some area forward in terms of making an impact in the world. I think that's just a fun time.
Host: Is there one physical bottleneck that you're thinking about just from an infrastructure perspective?
Jeff Dean: There's so many. I'm pretty bullish on more and more specialized hardware because I think that's the way you really get much more efficient systems. And we now have workloads where a handful of workloads are going to be a lot of the compute in the world, right? And so if you think about that, that is crying out for specialization. And the problem with specialization is if what you want to do changes in the future, then the thing you've lovingly crafted in hardware for maybe two years is no longer perhaps as relevant.
So I think to really make specialization work well, you need to automate much more of the hardware design process. Right now, the way you do hardware is you get a large team of people. Some of them take the high-level spec and build low-level RTL from it. Then because you've hand-translated this, you end up with another team of people that are going to verify that the first team of people did the right thing. And then you have another set of people who lovingly lay out the chip. That just seems like, if you can make a lot more automated loops that are searchable by reinforcement learning or other evolutionary techniques, and you can make those loops run fast enough—which current EDA tools typically are not designed to run very quickly—then you have the opportunity to really have a much more automated exploration loop of the design process, and perhaps to compress the design cycle tremendously.
If you could, for example, design a new chip with 10 people in 3 months instead of 150 people and 2 years, you'll see much more specialized hardware in the world. And you will have to bet much less on your future: what compute do I want to do 2 to 6 years from now becomes more of a 3 to 6 months to 4 years from now, and that's a much easier bet to make.
Inference Workloads, Agentic Systems, and Tool Bottlenecks
Bill Jia: If we go back one or two years, at that time, a lot of industry, Google included, plus a lot of frontier labs, the major focus was on how to make that model fantastic. There's a lot of focus on making the pre-training and post-training work, so a lot of hardware strategy was on the training side. But now, as large models are maturing, a lot of traffic starts to go into the agent world and on the inference side. So Jeff, if you can share any thoughts from a hardware design point of view: how can we design hardware to focus on and emphasize inference traffic?
Jeff Dean: Inference and training are somewhat different. For inference, you basically have a model, and you just want to serve a lot of requests for it. So you want to move as little information that doesn't change as possible. The things that don't change are the model weights, and the things that change are the requests and the KV cache and so on. So you need to design a system that is really efficient at minimizing this data movement there.
And then I also think, as you see not just a prompt and then a response, but much more independent—an agent does some stuff, it decides to call a tool, the tool runs, the tool results come back, you put that together into the model's context and let it figure out what to do next—we're going to realize that all of our tools are way too slow, because they were designed for human-speed iteration. If you're compiling code as part of your tool, you're really sad your compiler is slow. If you make your inference hardware super fast, it can generate code maybe even faster than you can compile it, and certainly faster than you can run it.
So one of the things we've been doing is making some of our internal tools faster to make them better. It turns out you can actually translate from one programming language to another quite effectively, because you really just have an entire, fully specified spec of what you want. You have the whole program written in an interpreted language like Python, and you just want the exact equivalent program in Go, Rust, C++, or whatever. And an agent can actually do that quite well, which is quite different than the normal kind of coding interaction we have with an agent, which is like, please make me a web server, and it has to fill in all kinds of details with all kinds of assumptions that may not be the ones you wanted. But with a fully specified tool, it can actually do a really good job of that: it can run all the unit tests, translate all the unit tests, run those in the new system, and verify behavior side-by-side to make sure it behaves the same way.
Bill Jia: We actually have a live example at Google internally we call Project Tern. Because a lot of models were built on TensorFlow, but we're migrating to JAX, we translate all the TensorFlow models and migrate to JAX, and do all the unit tests and code tests automatically.
Host: Now, there was a certain paper in NeurIPS that was rejected that you all had submitted that said very little impact. It was a distillation paper.
Jeff Dean: Oh yeah. This is like a motivational story; it's supposed to be an inspiration. Turns out distillation is important. This is a paper that my colleagues Geoff Hinton, Oriol Vinyals, and I submitted about how to take one shape model and use distillation to distill it into a student model.
Bill Jia: Which year was it?
Jeff Dean: It was 2015 or 2016. We were originally thinking of it in the context of training a big ensemble of different kinds of specialized models for vision. That is one of the sets of experiments we had in the paper: you have 20,000 vision classes, but you train a special one for animals, another one for cars, and another one for whatever, and then you just distill that into a single model that's good at all those things. And it got rejected, but that's fine; we put it on arXiv. People read it anyway.
Host: Don't be discouraged.
Jeff Dean: Exactly.
Q&A: Commodity Hardware vs. Specialized Supercomputers
Host: Okay, so we have a few minutes left. We're going to take a couple of questions. There's a few folks running around with microphones. There we go, one over here. Maybe we can start here. If you want to stand up and introduce yourself and give us a little one-liner about your background.
Audience Member (Ying): My name is Ying. I used to be in Google Brain for a decade. Good to see you again. My question for Bill and Jeff is: we talked about commodity software and commodity compute from the early days of Google; that's how Google scaled in the early 2000s. Nowadays, the AI supercomputer becomes more and more specialized, more like the original supercomputing framework. So how do you see the two different design patterns for hardware and machine learning systems? Do you think in the future we should move more towards commodity hardware, or should we continue to build up supercomputing machines? Thank you.
Jeff Dean: The reason we were able to scale with commodity hardware for search is that search is this very nice problem where, if you break it down, you have almost no communication between machines, and you have a lot of work to do on a single machine. So you don't need anything exotic in terms of the interconnect. We actually had 100-megabit Ethernet on the machines, and then we shared a 1-gigabit uplink per rack across 40 machines, so you actually had like a 4-to-1 oversubscription. But that's fine because you're sending like "Palo Alto restaurants" to each machine, and then you're getting back like a little snippet that says there's these 10 results. Whereas training really does push the boundaries: you need a whole bunch of connectivity.
You'd like to be able to train on a single chip—that would be the best—but that just takes too long or doesn't fit. And so you end up having to partition the problem across many chips. The way you partition it, no matter how you slice it—model parallelism or data parallelism—you usually end up with quite a lot of communication. And so that's why you have these more exotic interconnects between machines. You want the maximum performance per chip so that you don't have to spread it out over as many chips, but you still have to spread it out over a lot. That's why we're using liquid cooling to make each chip as performant as possible in that setting.
I think inference is going to be more like: you could have specialized hardware, but that is not as exotic, particularly for small models. But as you start to get larger models, then you start to have communication needs that are starting to look pretty exotic relative to say Ethernet commodity hardware and commodity networking, but it's definitely more mainstream and commodity than training. Does that make sense?
Q&A: Open Source Contributions, CUDA Parity, and AI Agents for Developers
Host: I think we have another question right here.
Audience Member (Andra): Jeff, my name is Andra from Uber. I have a question regarding the ecosystem of JAX and OpenXLA, especially given this open source community people have built. I see this is pretty much an era of AI infra akin to Android and iOS. What do you see in terms of how young engineers can contribute to the kind of library functions of OpenXLA? Because for most users, they are heavy CUDA users, or they build libraries or functions based on CUDA libraries. For OpenXLA, what do you see young engineers can contribute to?
Bill Jia: First of all, OpenXLA and JAX, and on top of JAX, the JAX stack has been entirely open sourced. That's one bucket of work. Another bucket of work is Torch-TPU. Right now internally we're collaborating with Meta; we're developing, but eventually we're going to put it onto GitHub and let all the community researchers and engineers contribute. That's our two parallel efforts. Any young engineers in the community who want to contribute to the open source repository are more than welcome to work with Google. We can discuss how we can co-develop this open source repository.
In addition, I think eventually a lot of users will be using either the JAX stack or the PyTorch stack for TPU or GPU. Right now, if people in the community have any issues using it, a lot of people call Google or call Nvidia because they have experienced engineers, or they put questions onto community discussion boards hoping somebody understands the question and has similar experience to answer it. But in my opinion, that's all very slow. Where the community can help contribute on open source strategy is: imagine if we can have a JAX agent and a PyTorch agent that run on GPU or TPU—it doesn't matter. If everybody in the community contributes data sources to help train that agent to be really powerful, then imagine if I have a TPU co-worker agent. When I'm running training, post-training, or inference, I don't have to ask Google or go to a discussion board; I directly use this agent to help me. If the agent can answer a large portion of questions, that's fantastic. But this needs all the community to contribute, because we need a lot of data points and experiences.
Jeff Dean: Maybe I'll give a meta-point about open source contributions. There's lots of ways for people to contribute. It's really good to interact with some of the owners of different repositories to try to identify: I'm thinking of doing this, would that be useful, or do you have ideas about things I should work on? Before you just chuck over 5,000 lines of code to someone and say, "Hey, here it is," I think it's a really good idea to get some buy-in or guidance about what would be useful. If you're eager to contribute but don't necessarily have a particular thing you want to work on, there's often a long list of things that people have in mind that would be useful.
Audience Member (Andra): Thank you for the answer. The one hesitation I have interacted with is that many functions or module kernels still heavily rely on CUDA. Do you see the trend of implementing something equivalent using JAX, OpenXLA, or Torch-TPU so that we have some kind of generalization compatibility compared to CUDA functions? Would that be a good direction?
Bill Jia: Yes, I think eventually. At Google internally, when we look at the internal JAX stack and also PyTorch stack on TPU, we look at all the CUDA functions, and we look at our functions as well. We are making sure we are at least on par, if not more. So I think that is a necessity we have to nail down. And in that layer, we actually want to open source low-level TPU SDK work.
Jeff Dean: I also think people want to think at a much higher level of abstraction. Thinking in terms of JAX or PyTorch expressions is what you want to do, not necessarily: how can I parallelize this with some kernel code, be it a TPU kernel language or a GPU kernel language, in order to get maximum performance? Hopefully the compiler and underlying system can give you that, and you can just think in beautiful abstractions like matrix multiply.
Q&A: Scaling Networks, 3D Tori, and Synchronous Training
Host: I think we have time for one last question. Just one last question, go ahead.
Audience Member (John): Hi, John from NYU. Working on networking for many years. I was very impressed by the talk from Bill Jia and the other fellow regarding scalability here. I know these TPUs are connected in a 3D torus. Compared with Ultra Ethernet Consortium and others looking into fat-tree interconnects, from a scalability point of view, when you scale up TPUs or GPUs up to tens of thousands or even one million, is this 3D torus still able to handle all-to-all communication at that large scale?
Jeff Dean: Our largest pods for TPU v5e or TPU v4, I believe, are 9,600 and something chips. So that's the scale at which we do 3D tori. Beyond that, we have Pathways as a software abstraction on top of many of those 3D-torus-connected pods, and that uses a data center network, which is whatever fabric you have in your data center, or even a multi-metro training setup where maybe you have 5 pods in this building in Oklahoma and 8 more in this building in Iowa, with a high-speed WAN link between them. That seems to have worked pretty well for us for scaling up.
You do sometimes want to map the computation you're doing so that you have data-parallel and model-parallel dimensions. You might typically want the model-parallel aspect to be within a single pod or a slice of a pod, and then have data-parallel replicas across those pods or slices of pods. That seems to work pretty well for us. The nice thing about torus-based networks is they're super easy to connect locally, as opposed to needing very complicated wiring in your data center floor.
Host: Except for the funny mirrored optical switches.
Jeff Dean: Anything to add, Bill?
Bill Jia: Yeah, that's very much true. We have ICI within the rack, within the cube. Then we use the OCS to scale it up to 9,600. Then going up, we have our datacenter fabric, orchestrated by the Pathways software. Going beyond 100,000 chips, we use cloud data center networking to connect everything together. This kind of scalability is working pretty well for us.
Jeff Dean: And I would say even at that scale, we're able to do fully synchronous training, which is nice from the ML reproducibility and interpretation standpoint. At some point, asynchronous training will come back, but so far we've been able to push synchronous training quite far, and we'll see what happens.
Host: That's a great way to end. We're out of time. Thank you so much, Jeff. Thank you, Bill, for your time. We have some really cool swag outside by the patio for you, which includes some famous Jeff Dean memes on a pen and a Dev Labs t-shirt. Thanks so much, everybody.