[MUSIC PLAYING]
DEBI CABRERA: We're back at Google Cloud Next.
And I'm here with Alex.
And he's going to talk to us about his session, Accelerating
AI Inference Workloads with Google Cloud GPUs and TPUs.
Thank you for being here with us, Alex.
ALEX SPIRIDONOV: Thank you so much for having me.
DEBI CABRERA: Wonderful.
So can you introduce yourself, and tell us
a little bit about what you do at Google,
and also a brief synopsis of your talk.
ALEX SPIRIDONOV: Of course.
So my name is Alex Spiridonov.
I'm a group product manager at Google.
I lead our AI Accelerators Software Team.
So our team creates the software for Cloud TPUs and GPUs.
And I lead the Product Management Team, working closely
with engineers and our many amazing collaborators
across Google in the industry.
Now, the focus of my talk today is
going to be on inference, which is one of the key workloads
that we see with AI today.
Once you have trained your large language
model or your generative AI model
or once you have taken an open model and fine tuned it,
then you need to deploy it for inference,
which is also called serving.
And that is a really key workload.
And what our customers really care about
is cost efficiency at scale.
Essentially, they want to take those large language
models or GenAI models, they want
to deploy it to their millions of customers around the world,
and they want to serve them with as high performance as possible
with as low cost as possible.
So balancing that high performance
and cost optimizations becomes really, really important.
And so this week, we're really excited to announce
some major advancements in inference
on Google Cloud for both TPUs and GPUs.
We announced a new high performance inference
engine that is optimized for large language models.
It's called JetStream.
And it provides high performance,
cost-efficient, scalable serving or inference for LLMs on XLA
devices, starting with Cloud TPUs.
And best of all, it's open source.
So it's all out there in GitHub and we invite developers
and we invite community to build with us.
We also announced MaxDiffusion.
Now, whereas JetStream is for large language models,
MaxDiffusion is for diffusion models
or text-to-image, text-to-video generative AI models.
That's another really important workloads
that a lot of our customers are interested in.
And MaxDiffusion is a collection of high performance,
open source, reference model implementations that
are optimized to be really, really
efficient on our accelerators.
So our developers can take one of those open models,
THEY can customize it, and they can deploy it for inference,
and get that great performance and cost efficiency.
And then lastly, we're sharing some exciting updates
to our GPU product portfolio for inference,
announcing both additions to the A3 VM family,
as well as some ongoing partnerships and collaboration
that we have with NVIDIA to bring
the next-generation Blackwell platform from NVIDIA to our AI
hyper computer architecture.
DEBI CABRERA: Wow, so many exciting launches.
OK, so going back to inference, that balance
of cost and efficiency.
What would you say are key considerations
to keep in mind when selecting GPUs and TPUs?
ALEX SPIRIDONOV: Absolutely.
So the important thing to take away here
is that Google Cloud is the only place on the planet that
offers our customers a choice of two
proven, high-performance platforms for inference.
And that's Cloud TPUs and Cloud GPUs.
Those are both purpose-built AI accelerators
that are really excellent for serving models of all shapes
and sizes, whether that's LLMs, generative AI
models like diffusion models, or even ranking and recommendation
models.
And what becomes really important
is those combined hardware and software performance
optimizations.
So let's take TPUs for example.
As I mentioned, we built JetStream,
which is our new high-performance LLM inference
engine for XLA devices.
And what makes JetStream so fast and so cost-efficient
is that combined power of the hardware and software
optimizations, because we have the TPU hardware, which
is really optimized for high-performance,
high-throughput, low latency, and cost efficiency.
And then we build this very advanced software stack
that really takes full advantage of those hardware capabilities,
with key performance optimization patterns
like continuous batching, efficient attention algorithms,
quantization, the list goes on and on.
And we support popular open models,
like Google's Gemma, and also open models like LLaMA.
And we have a long roadmap of future performance
optimizations.
We're really building this for the long term.
And we're inviting the community to build with us.
So that's how we achieve peak performance and cost
efficiency on the TPU side.
On the GPU side, Google Cloud has an amazing collaboration
with NVIDIA to bring the latest NVIDIA
GPUs, like the NVIDIA H100s and the future HGX B200
GPUs to Google Cloud.
And not only are we bringing the latest hardware innovations
from NVIDIA to Google Cloud, we're
also optimizing and supercharging
our software stack to run really well on those NVIDIA's GPUs.
So, for example, we collaborated closely with NVIDIA
to optimize our JAX ML framework to be really, really
efficient on NVIDIA GPUs.
JAX is a really powerful ML framework
because it is Python-based and is compiler-oriented,
which means that it's both easy to use
and it's also really fast.
And now, thanks to our great collaboration with NVIDIA,
it is also highly optimized for both large-scale training
and inference on NVIDIA GPUs.
So in summary, it's really about giving our customers a choice.
Whether you want the latest and greatest
NVIDIA GPUs and the software stack
that's optimized to take full advantage,
whether that's NVIDIA's own software stack or the JAX ML
framework, or if you want the end-to-end hardware
software optimizations and TPUs, all of those options
are available to our customers.
And they're only available on Google Cloud.
DEBI CABRERA: Wow.
I love what you're saying about meeting customers
where they are.
We're giving them the options.
Here, we have three amazing options for them to choose from.
Wonderful.
OK, so in terms of the unique demands of generative models,
like large language models, what are some recommendations
for TPU versus GPU deployment?
ALEX SPIRIDONOV: Yeah, that's a great question.
And, again, we really think of it
as an "and" and as a compliment.
So it's not a matter of choosing one versus the other.
It's really about giving our customers a choice
and giving them the opportunity to choose from all
of those great options.
Now, to your specific question about,
what is it about those workloads that makes them
so unique and so challenging?
Large language model inference and generative AI inference
is really different from traditional web
serving in the era where Google was born
or even from serving maybe traditional deep neural
networks.
And the reason why it's different
is that it's just a more complex process.
There are more complex inference tasks that are involved.
So let's take large language models for example.
There's actually two stages in which inference happens.
There is the prefilled stage where you send your prompt.
And then there is the decode stage or the generation stage.
And that's where the model actually makes the predictions.
It generates the output and response
to the user's prompt or the user's query.
And those two stages are actually quite different.
And they have different computational demands.
The first one is more, we could say, compute bound
and the other one is more memory bound or memory bandwidth bound.
And so the bottom line is that it really
takes very careful optimization and very careful
tuning of the software stack to really
be able to perform all of those stages in the prediction process
very, very efficiently.
But the good news for our developers
is that we have already done that work for them.
We have optimized JetStream and we are partnering with NVIDIA
to optimize the GPU software stack to run really,
really well on TPUs and GPUs on Google Cloud.
And not only that, we also provide reference models.
So we provide "out of the box" open source implementations that
are optimized for inference both on TPUs and GPUs.
That makes it really easy for developers
to get started because they can take one of those open models
in a framework of their choice, whether that's JAX or PyTorch.
And take that model, they can customize it.
That's all open source.
And then very easily, they can deploy it for inference
on an AI accelerator of their choice.
So we're investing not only in industry-leading performance
and cost efficiency, we also want the developer experience
to be as seamless as possible.
And we want our software and our tools that we're building
and that we're building with our partners
to really fit seamlessly into that developer
workflow on Google Cloud.
DEBI CABRERA: I mean, it sounds like, really, there's
nothing to sacrifice here.
It sounds like, I want to work with this today.
It's so much flexibility, and options,
and we're meeting customers and devs where they are.
It really sounds like a wonderful slew of launches
here that people can use.
ALEX SPIRIDONOV: Absolutely, we're
really excited to put this in the hands of our customers
and the hands of our developers.
It is all available today.
It's all in open source on GitHub.
You can go to our JetStream repository
to get started with LLM inference.
We have high-performance implementations
in JAX and PyTorch available to make
getting started really easy.
Similarly, with GPUs, all of that software
and all of those reference models
are also available for our developers to use.
And we invite the community to build with us.
We are intentionally open sourcing software packages
and software capabilities, like JetStream,
like MaxDiffusion, like MaxText, which
is the MaxDiffusion for large language models.
It's those reference implementations of LLMs.
We are intentionally putting them out there for our community
to embrace, to adopt, and to build with us, because inference
is an incredibly dynamic space.
It is very exciting because there's just so much innovation
that is happening.
We really want to invite the community
to build with us, to contribute the latest and greatest
algorithms, performance optimizations built in the open,
so that together, we can advance the state of the art.
And together, we can build the best platform
for anyone looking to get started with LLMs and GenAI.
So whether you're a researcher pushing
the frontiers of natural language understanding or text
generation, or if you're a developer looking
to integrate the latest cutting-edge LLM and GenAI
capabilities into your applications,
we want to give you the foundation
that you need to succeed.
DEBI CABRERA: Amazing.
Well, you heard the gentleman, go ahead and try this
out yourselves.
And if folks are going to go try it out themselves, what
are some common pitfalls or challenges
that one may face when using inference optimization?
And how can they be avoided?
ALEX SPIRIDONOV: That is an excellent question.
So here, it's really important to take a step back and think
about the entire user journey or think about the entire life
cycle of those models, because this initial performance
optimization and tuning is really, really important.
But then comes probably the most important step,
which is actually deploying those models to production.
That's how you actually serve those models,
you serve queries and live traffic from your users.
And that's how you make those AI-powered applications actually
available to your end users.
And there is a set of challenges and there's
a set of considerations that are important there
to minimize cost.
The short answer that I would give here
is that we have amazing integration with Google
Kubernetes Engine or GKE.
GKE is an orchestration layer that makes it really easy
and seamless for our users to manage all of this compute.
So whether you're running a small model
using just a few accelerators or if you're
serving the largest model that can use dozens or maybe even
hundreds of accelerators, GKE makes it very, very easy
and seamless for you to orchestrate and manage
all of that capacity.
But not only that, it also helps you with cost management,
because it does auto-scaling.
Let's say, you build an amazing new AI-powered app and it makes
the front page of "The New York Times,"
suddenly you have all of those users.
So GKE will automatically scale up.
It will provision more GPUs and more GPUs
for you so that you can meet all of that demand
without missing a beat.
And then let's say the demand quiets down,
then it will automatically provision down your capacity, so
that you're not paying for those accelerators
that you're not using.
So you can think of it as auto-scaling up and down,
even scaling down to zero if necessary.
And this is just one of the advanced capabilities of GKE.
There are many, many others that I would encourage all
of our developers to check out.
But GKE really provides amazing support and integration
for all of Google Cloud accelerators.
And it is one of the key tools and the key capabilities
that we give our users to make their cost
management as easy and as seamless as possible.
DEBI CABRERA: Wonderful.
Again, this pattern of you make it easy for people
to get in there and just get started.
So in 60 seconds or less, actually,
if someone wants to get started, can you tell our audience,
how can they easily get started?
ALEX SPIRIDONOV: Great.
So the easiest way to get started
is you can go to our GitHub repository.
So you can either check out our JetStream repository
on github.com/google/jetstream.
We also have a PyTorch version, so jetstream/pytorch.
For our MaxDiffusion, it's similar,
github.com/google/maxdiffusion.
MaxText also has a repository.
We also published a blog post earlier today,
so you don't need to memorize all of those links.
You can just check out our blog post on the wonderful Google
Cloud Blog where you can see all of those announcements,
you can see performance benchmarks.
You can easily click through and go to the links
and check out the software for yourself today.
Please reach out with your questions, make the pull
requests.
Send us your feedback.
We really want to build this with the community in mind.
We want to build this with our users and for our users.
And we invite all of you to build with us.
DEBI CABRERA: Amazing.
So make sure to check out the links in the description.
As Alex said, we'll have everything you need in there.
Thank you so much for chatting with us today.
And make sure to check out the full talk,
and that also will be linked in the description.
ALEX SPIRIDONOV: Thank you for having me.
[MUSIC PLAYING]