Who uses Google TPUs for inference in production?

#1

Post by arthurdelerue » Mon, Mar 11, 2024, 4:22 PM UTC

I am really puzzled by TPUs. I've been reading everywhere that TPUs are powerful and a great alternative to NVIDIA.

I have been playing with TPUs for a couple of months now, and to be honest I don't understand how can people use them in production for inference:

- almost no resources online showing how to run modern generative models like Mistral, Yi 34B, etc. on TPUs - poor compatibility between JAX and Pytorch - very hard to understand the memory consumption of the TPU chips (no nvidia-smi equivalent) - rotating IP addresses on TPU VMs - almost impossible to get my hands on a TPU v5

Is it only me? Or did I miss something?

I totally understand that TPUs can be useful for training though.

Re: Who uses Google TPUs for inference in production?

#2

Post by hiddencost » Mon, Mar 11, 2024, 6:40 PM UTC

Google is using them in prod. I think they're so hungry for chips internally that cloud isn't getting much support in selling them.

Re: Who uses Google TPUs for inference in production?

#3

Post by emadm » Mon, Mar 11, 2024, 6:44 PM UTC

https://pytorch.org/blog/high-performance-llama-2/

Re: Who uses Google TPUs for inference in production?

#4

Post by amelius » Mon, Mar 11, 2024, 7:16 PM UTC

hiddencost wrote: post #2

Google is using them in prod. I think they're so hungry for chips internally that cloud isn't getting much support in selling them.

Or maybe they are just using nVidia. Who knows ...

Re: Who uses Google TPUs for inference in production?

#5

Post by eklitzke » Mon, Mar 11, 2024, 7:16 PM UTC

amelius wrote: post #4

hiddencost wrote: post #2

Google is using them in prod. I think they're so hungry for chips internally that cloud isn't getting much support in selling them.

Or maybe they are just using nVidia. Who knows ...

Lots of people know.

Re: Who uses Google TPUs for inference in production?

#6

Post by htrp » Mon, Mar 11, 2024, 7:26 PM UTC

We've previously tried and almost always regretted the decision. I think the tech stack needs another 12-18 months to mature (doesn't help that almost all work ex Google is being done in torch).

Re: Who uses Google TPUs for inference in production?

#7

Post by nivekney » Mon, Mar 11, 2024, 7:27 PM UTC

eklitzke wrote: post #5

amelius wrote: post #4

Earlier quoted context omitted.

Or maybe they are just using nVidia. Who knows ...

Lots of people know.

marks as solved

Re: Who uses Google TPUs for inference in production?

#8

Post by kccqzy » Mon, Mar 11, 2024, 7:41 PM UTC

Apparently Midjourney uses it. GCP put out a press release a while ago: https://www.prnewswire.com/news-releases/midjourney-selects-...

Re: Who uses Google TPUs for inference in production?

#9

Post by danjl » Mon, Mar 11, 2024, 7:51 PM UTC

hiddencost wrote: post #2

Google is using them in prod. I think they're so hungry for chips internally that cloud isn't getting much support in selling them.

I would guess that Google's vertexAI managed solution uses TPUs. Also Google uses them internally to train and infer for all their research products.

Re: Who uses Google TPUs for inference in production?

#10

Post by vineyardmike » Mon, Mar 11, 2024, 7:51 PM UTC

amelius wrote: post #4

hiddencost wrote: post #2

Google is using them in prod. I think they're so hungry for chips internally that cloud isn't getting much support in selling them.

Or maybe they are just using nVidia. Who knows ...

Beyond the fact that this is hardly a secret, there’s lots of other signs.

1. They have bought far less from NVidia than other hyper scalers, and they literally can’t vomit without saying “AI”. They have to be running those models on something. They have purchased huge amounts of chips from fabs, and what else would that be?

2. They have said they use them. Should be pretty obvious here.

3. They maintain a whole software stack for them, they design the chips, etc. Then they don’t really try to sell the TPU. Why else would they do this?