TPU v4 provides exaFLOPS-scale ML with efficiency gains | Hacker News Reader

TPU v4 provides exaFLOPS-scale ML with efficiency gains

cloud.google.com

112 pointszekrioca38 comments

I am very impressed with what Google has done for the state of machine learning infrastructure. I'm looking forward to future models based on OpenXLA which can run between Nvidia, Apple Silicon, and google's TPU's. My main limiter to using TPU more often is model compatibility. The TPU hardware is clearly the very best, just not always cost-effective for those of us who are starved of available engineering hours. OpenXLA may fix this if it lives up to its promise.

That said, it's also incredible how fast things move in this space:

> Midjourney, one of the leading text-to-image AI startups, have been using Cloud TPU v4 to train their state-of-the-art model, coincidentally also called “version four”.

Midjourney is already on v5 as of the date of publication of this press release.

> model compatibility

JAX will compile your model with XLA today, letting you run it on Nvidia, Apple Silicon (CPU for now), google's TPU, etc.

Overall, XLA really shines on TPUs, is doing very well on GPUs and sadly lags behind on CPUs (reaching about single core performance).

As a bonus, JAX is basically Numpy++: you can use it to build scientific/numerical software and run it easily on GPU and TPU!

yeah we trained v5 on TPUs and continue to train on them.

I refuse to care about them until they sell them on PCIe Cards.

The lock-in is bad enough when dealing with niche hardware on-prem, i certainly won't deal with niche hardware in the cloud.

can't agree you, until then, I have no interests in its TPUx

This is part of why I’m hesitant to use ARM on AWS, its not easy enough to just buy a representative equivalent for onprem

It is easy especially now.

(With Altra availability ramping up)

https://www.ipi.wiki/products/com-hpc-ampere-altra?variant=4...

https://tenstorrent.com/

,,Midjourney, one of the leading text-to-image AI startups, have been using Cloud TPU v4 to train their state-of-the-art model, coincidentally also called “version four''

This sounds quite bad in a press release when Midjourney is at v5. Why did they move away?

Sounds like they are just out of date, it is possible MJv5 was also trained on TPUs.

We didn't, v5 was trained on TPUs too

They're so non confrontational. Their performance comparisons are against "CPU". Just come out and say it, even if it's not apples to apples. If the 3D-torus interconnect is so much better, just say how it compares to NVidia's latest and greatest. It's cool that midjourney committed to building on TPU, but I have a hard time betting my company on a technology that's so guarded that they won't even post a benchmark against their main competitor.

The paper compares it to A100.

https://arxiv.org/pdf/2304.01433.pdf

Nvidia's last gen...

That's addressed in section 7.4 of the paper:

> The Google tradition is to write retrospective papers after systems are deployed and running production apps. Both TPU v4s and A100s deployed in 2020 and both use 7nm technology. The newer, 700W H100 was not available at AWS, Azure, or Google Cloud in 2022. The appropriate H100 match would be a successor to TPU v4 deployed in a similar time frame and technology (e.g., in 2023 and 4 nm).

That's not really addressing the point though, is it. That tradition at Google evolved when they had technology that was well ahead of their competitors, which they currently are not and therefore this paper is more like a marketing document than a recruiting resource. It's inevitable that people will want comparisons before buying TPU access and giving them a mapreduce style paper that comes out years after the technology is mature isn't going to satisfy them.

Have you tried to get an h100? Its not really possible right now

This is very impressive technology and engineering.

However, I remain a bit skeptical of the business case for TPUs for 3 core reasons:

1) 100000x lower unit production volume than GPUs means higher unit costs

2) Slow iteration cycle - these TPUv4 were launched in 2020. Maybe Google publishes one gen behind, but that would still be a 2-3 year iteration cycle from v3 to v4.

3) Constant multiple advantage over GPUs - maybe 5-10x compute advantage over off the shelf GPU, and that number isn't increasing with each generation.

It's cool to get that 5-10x performance over GPUs, but that's 4.5yrs of Moore's Law, and might already be offset today due to unit cost advantages.

If the TPU architecture did something to allow fundamentally faster transistor density scaling, it's advantage over GPUs would increase each year and become unbeatable. But based on the TPUv3 to TPUv4 perf improvement over 3 years, it doesn't seem so.

Apple's competing approach seems a bit more promising from a business perspective. The M1 unifies memory reducing the time commitment required to move data and switch between CPU and GPU processing. This allows advances in GPUs to continue scaling independently, while decreasing the user experience cost of using GPUs.

Apple's version also seems to scale from 8GB RAM to 128GB meaning the same fundamental process can be used at high volume, achieving a low unit cost.

Are there other interesting hardware for ML approaches out there?

> Are there other interesting hardware for ML approaches out there?

Google also has Coral, which is a non-cloud mobile-focused TPU that you can buy and plug in (USB or PCIe).

https://coral.ai/products/

Naturally, this is orders of magnitude less powerful than the kind of TPU being discussed here but I just like the idea of local compute.

Just received my mSATA Coral TPU in the mail. I'd been waiting 11 months for it after backordering on Digikey. Perhaps this speaks to the parent comment's concerns over unit volume and iteration cycle? Hopefully that will improve in the future and modules like these will become widespread.

> If the TPU architecture did something to allow fundamentally faster transistor density scaling, it's advantage over GPUs would increase each year and become unbeatable.

It is completely unreasonable to expect something like that.

>100000x lower unit production volume than GPUs

This is obviously an exaggeration, I wonder what the actual ratio is between TPU proudction and e.g. A100 production.

Probably more close to 1000x? I see a fairly large number of TPU pod these days and I don't think A100 is as prevalent as high end consumer GPU, which is typically measured in millions, not billions.

> 100000x lower unit production volume than GPUs means higher unit costs

Two points. Nvidia's RTX 3000 series (3060 Ti, 3080, and many other flavors) ships 6 or more flavors per generation. The related silicon has names like the GA102, GA103, GA104, GA106, and GA107. So only 1/6th of the consumer market for Nvidia silicon can be amortized over any single design.

I wouldn't be at all surprised to see Google making the TPUs by the million. I found a vague reference to 9 exaflops and single facilities (one of many) costing $4 billion to $8 billion.

So I wouldn't assume that the consumer GPU market/number of silicon designs is 100,000 times larger than the TPUv4 market.

> Slow iteration cycle

True. Then again generations make much less difference than they used to. Gone are the days where even after a multiple generations that average performance increases by 2x. Sure nvidia's 4000 series claims 2x ... on raytracing. But normal game performance seems to be more like 15%. Sure various trickery like DLSS helps, but similar tricks are increasing the performance of older cards as well. Similarly apple's a14 -> a15 -> a16 (or m1 -> m2 if you prefer) chips have had modest performance increases and mostly have increases in perf/watt.

> 4.5yrs of Moore's Law

It's dead Jim.

I believe that Nvidia uses the same chip from A16 up to the A100, and maybe for some of the Quadro chips. That easily puts the unit count into the several millions.

The picture in this article shows 8 racks with (according to the paper on arxiv) has 16 TPU sleds each, 4 TPUs per sled. that's only 512 chips. According to the paper, it is one of eight in a 4096-chip supercomputer. If you give them 10-100 of those around the world, you get 40,000-400,000 chips. That's enough for reasonable scale. Nvidia should still have 100x (or more) their scale.

They are not all the same chips; some use different processes, for one. I assume that it would eat into any large volume efficiency claims. I am sure there are more A100s than TPUv4 chips out there, but I would hesitate to say the scale is over 100x.

> 3) Constant multiple advantage over GPUs - maybe 5-10x compute advantage over off the shelf GPU, and that number isn't increasing with each generation.

The advantage of this hardware isn’t just the raw compute capacity. It’s the massive IO bandwidth.

GPUs are fast, but they’re not designed for massive all-to-all communication across large networks. That’s where modules like this shine.

Why do you need all-to-all? With deep networks you could just put different layers on different GPUs, then run the forward and backward passes asynchronously. They would only need to communicate in a chain-like manner.

Just as a quick example (there are lots of other examples), I currently am working on a model that does significant tensor-level parallelism. The input to the model is huge, so we replicate the weights for the encoder part of the model over 12 A100s, and then run an all-reduce over the outputs before then distributing that reduced tensor to another set of GPUs for the next set of operations. In this setup we need significant bandwidth not only between every device in a node but also across nodes (infiniband).

Isn’t the simpler example data parallel training? That’s a reduce and broadcast.

You’re correct, I was going to add on to my answer that this is combined with DDP for a “3D” style parallelism that could specifically benefit from the TPU’s network topology, but by the time I got done fixing all the typos (and still missing a few) from writing it on my phone I completely forgot :)

To be honest, I think regulators should have stepped in to stop the acquisition of Melannox (Infiniband) by NVidia. There are basically no non-NVidia options for Infiniband now. (Infiniband being about half an order of magnitude ahead of Ethernet at any one time, with lower overhead and lower cost per unit bandwidth. Often the basis for high bandwidth Ethernet cards as well. And ostensibly an open standard.)

I agree that I’m uncomfortable with the consolidation of so much of the HPC stack under Nvidia (and would be 10x more uncomfortable if the ARM acquisition went through), but I don’t think there was really a monopoly there. I think the real “market” is networking, not just infiniband, and in that sense there is a ton of competition, just no competition that hits the performance characteristics that Nvidia/Melanox achieve.

Biggest advantage is the high bandwidth interconnect. It’s very painful training large models on multiple GPUs. TPUs help a lot with that, although the software side of things is pretty brutal.

> 2) Slow iteration cycle - these TPUv4 were launched in 2020. Maybe Google publishes one gen behind, but that would still be a 2-3 year iteration cycle from v3 to v4.

Is that substantially different from the Nvidia release cadence? According to the Wikipedia announcement dates, their last three generations were H100 in March 2022, A100 in May 2020, V100 in March 2017.

As of now, Apple's version is uncompetitive for training. It is decent for inference, which is good but doesn't cover the training usecase which is the main usecase for TPUv4 and Nvidia GPUs like A100. A100 provides 312 TFlops, but M1 GPU only 5.1 TFlops, which is 2 orders of magnitude slower.

This blog [1] does make the argument that M1 Max has about 8x lower performance than a consumer grade Nvidia GPU 2080, but also uses 8x less watts, so it's possible Apple can make a product with similar performance/watt. However, I would say that they are farther away than TPUv4 for now, because the product doesn't exist.

This is really cool!

Reply on news.ycombinator.com