We are here to talk to you about PyTorch and how it's a great tool
for you to scale your deep learning models in a cheap and a fast and efficient way.
So when we're scaling neural networks up, what that means is
we are adding more and more layers
to your neural network.
And what this does is it allows your model to capture more nuance in your data
and allows your model to perform more complex tasks.
Well, all of this doesn't come for free.
The model will start to require more memory and more compute.
Let's see the llama model, for example.
It has four variants
ranging from 7 billion parameters
to 70 billion parameters.
The smallest variant, 7 billion,
was trained using 2 trillion tokens
and it required around
roughly 200,000 GPU hours.
That's a long time.
So PyTorch, along with offering
the modular building blocks to build your neural networks,
it also offers some utilities to distribute to this training workload.
And let's take a look at one of them
called Distributed Data Parallel (DDP).
So when you're training your model
with DDP, there are three steps that need to happen.
First, the forward pass where you take the data
and you pass it through the model,
computes the loss
which is then back propagated
through the model, which then gives you the gradients.
The third step where we update the model's weights is preceded
by a synchronization, where all the computed gradients
from each of these replicas are communicated with each other.
Now, the hallmark of DDP and really all distributed training
is an overlap of
the computation and communication.
Essentially what
that means is simultaneously we are doing the back propagation
while communicating all of the calculated gradients.
This saves us time and keeps the GPU running
at a near 100% utilization.
At a very high level,
what that looks like:
You first divide the model into buckets.
Each replica calculates
the gradients of the first bucket
and while it is calculating
the gradients of the second bucket,
these first bucket gradients are synchronized.
And this is happening simultaneously
with the computation of the gradients in the second bucket.
And similarly, as each bucket is being calculated,
the preceding buckets gradients are being communicated.
This ensures that all the GPUs are running
at full utilization and you're not having any idle workers.
They're all working very hard to train your model.
And this is the
case where the model fits in one GPU.
For example, the 7 billion model can fit on
on any cloud GPUs today.
But when you start scaling that up, like, for example, 70 billion
or even the 30 billion models, it's very difficult to fit them in one GPU.
So in that paradigm, in those regimes,
the DDP model does not work.
Now I'll call upon Raghu to talk to us about what in PyTorch allows us
to train models which are larger than what your single GPU can accommodate.
All right, so you've heard
all about distributed data parallel, and how do you scale a small model
on very large number of GPUs and reduce your training times.
So now let's look at what happens
when we have a model that does not fit in a single GPU.
And that is where FSDP comes to your rescue.
FSDP stands for "Fully Sharded Data Parallel"
and what fully sharded data parallel does is it
takes the model, breaks it down into what are called
units, and all of these units are
then sharded across GPUs.
So you can take that--
Think of it as shredding the model and then each GPU owns
small portions of this model.
So that's where shards are coming in.
Now, what
happens after that is pretty much very much like what DDP is doing.
Think of instead of a model, you are thinking of it as a unit.
And during the initial unit
construction, from the shards,
you are going to do an AllGather.
And this happens in the forward pass.
So you're gathering the unit, you're computing on top of it.
And this happens across all the units.
So as soon as your unit is computed, you lose that unit's memory
and then you go to the next unit and so on.
And that's the forward pass.
Once the entire forward pass is
computed, your loss is
computed, and you go--
So this is your forward pass.
And now you're going to do an AllGather
in the backward pass.
And during the backward pass, what is happening is you are computing
again, very much like here, you are gathering the unit, you're computing
the back propagation, but just in the reverse fashion.
Once you have computed the gradients,
then those are again,
very much like what you have done for your DDP.
Those are synchronized across all the GPUs that are responsible for holding
that particular portion of the model.
So once you are synchronized, that completes your entire step
and then you continue doing all of this and then FSDP, very much like DDP,
you are going to utilize overlap in a significant way
because imagine in DDP there was only one single synchronization step, in FSDP,
I have more opportunities for doing overlap and keeping those GPUs
continuously busy while you are doing your computation.
And that is how you achieve scale of these models.
So what are the typical ways in which people train these large models?
Most of the world knows about them as HPC systems,
very large scale systems with very high speed interconnects,
state of the art GPUs and servers and training these models.
So what happens if I have "not a HPC system" that, you know, I have a good node.
So let's say that instead of this GPU, we call this a node.
Typically nodes have eight,
maybe 16 GPUs in them.
And many of these, the typical HPC system
has an HPC interconnect things like
the interconnect
where people
may have heard the term InfiniBand.
So very high speed interconnect between nodes;
and people may have heard the term Ethernet.
Ethernet is far more
common in many environments, whether it is cloud
or on-prem environments, whereas InfiniBand is less common.
So there is a common misconception that these larger models
will always need InfiniBand, even though that is true
for some of the larger models and it is faster.
But with Ethernet also, you can train these models in an efficient manner.
And that is where IBM helped, worked with the PyTorch community,
and introduced this concept of rate limiter.
What rate limiter is doing, is really balancing off the trade off
between the overlap in terms of communication and computation
by managing the memory in the GPU better.
And that is how you reduce the amount
of communication per
GPU computation time step. And you are increasing the amount
of compute while you're keeping that communication constant.
And that is how you can get away with training on Ethernet
while achieving similar kinds of benefits as InfiniBand.
So Raghu, we saw quite a few things about how PyTorch helps
in scaling up your deep learning workloads and you AI workloads.
We saw FSDP, we saw the rate limiter API.
And they do a pretty good job.
Absolutely.
And we also saw how
interconnects play a role on how do you scale up and so on, right?
But what happens inside a single node between that CPU and GPU?
Tell me a little bit more about it.
Yeah, you know,
it's an interesting question you raise because there is
some more efficiency gains that we can we can opt in.
And let's take a look at, you know, what's what's happening on the board.
So the reason-- one of the reasons --why PyTorch is
so popular is because of its programing paradigm
called "eager mode".
And eager mode allows developers
to have a very Pythonic approach to their programs.
It allows dynamism, like, you know, you can do your if-else blocks.
You're going to have for-loops within that.
It's very flexible.
Yeah.
And we already spoke about CPU and GPU.
Yeah.
So I'm going to draw a small little schematic over here,
which sort of illustrates what's happening in
eager mode.
Let's say this is your CPU queue
and this is your GPU
queue. And queue is essentially
the order of instructions that each chip is launching.
Mm hmm.
Now, taking a step back, your AI programs,
your AI models are essentially a sequence of operations.
Small AI models, small deep networks have lesser instructions comparatively.
And larger models have many more.
Mm hmm. Mm hmm.
So every line of code that I write in
in my Python program basically translates to going to the GPU.
Is that essentially what's happening? Yes.
Almost every line is an independent instruction
that your CPU is deriving from the program.
And so we can probably think of your model,
the execution of your model, as a sequence of instructions over here.
Mm hmm.
And because you're using a hardware accelerator like a GPU,
you're going to queue up those instructions
onto the relevant operations specific to the back end that are using.
Okay.
So let's say you have CUDA operations
on an Nvidia GPU.
So this is kind of how eager mode works.
This is the appropriate paradigm that allows you to iterate
or interact with your program in a very, you know, 1-to-1 way.
Okay.
So what is happening between,
you know, these times when-- what happens there?
And that's, you know, if you've hit the
nail on the head,
these empty spaces over here
are actually times when the GPU is not working.
These are
idle.
And, you know, that's what we do
not like about our GPU is because they're a very expensive resource.
We don't like them to be idle.
We want them to be working all through,
like you want to maximize the utilization over there.
And these get more and more as you scale.
These are the model size, right? Exactly.
Like as you as your models get larger, the skew starts getting much longer.
And as the queue starts getting longer, the number of these idle spaces
start increasing
quite a lot.
Yeah.
And these essentially translate to a lot of costs, unnecessary costs.
You're burning GPU hours without actually getting any bang for your buck.
Okay, so how do we address this?
So here we have an interesting tradeoff.
It poses a pretty interesting engineering challenge.
We want that you got more, like, flexibility.
It makes programing fun. Mm hmm.
But we also want that efficiency. Mm hmm.
So earlier last last year
at the PyTorch conference, we announced 2.0.
So PyTorch 2.0 packs in a very interesting new paradigm.
So it's essentially a new compiler
called TorchDynamo.
And this still allows you
to program your PyTorch code, as you always have.
It's still got that flexibility, that interactivity.
But what it does differently is instead of having
separate instructions
queued up like in eager mode,
your program is essentially converted to a graph of operations.
It's almost a static graph.
And so all of these instructions are sort of merged together.
Mm hmm.
And they form like it's a 1-to-m.
Mm hmm.
And sometimes you might have graph breaks, You know,
you might want to come back to eager mode paradigm.
So your entire program has been translated
to two instructions over here instead of many different instructions.
And on the GPU,
again, you're queuing up only two instructions
instead of all of n instructions over here.
Yeah.
And quite often you won't even have this break.
So your entire program, your entire model, is this one block of instructions in it
which executes seamlessly on the GPU.
That is pretty cool.
And how do you actually achieve this from a programing standpoint?
Like what does, say me, as a developer,
what do I need to do to get this?
One line of code.
That's wonderful.
All of this
goodness in one API called torch.compile.
So that's like super easy for me as a developer to get up to.
What kind of speed ups have you seen with torch.compile?
Let's say just about training it.
Oh, torch.compile
primarily targets training workloads
and depending on the architecture and the model, it varies.
But we have seen speed ups of at least 40 times. Wow.
So that is an order of magnitude.
Many orders of magnitude of on display and eager mode.
Okay, wonderful.
And at IBM, what we have been doing is to really work with the community
and to look at torch.compile from an inference lens.
So what are the beauty of compile?
I feel, is that it not only works for training and giving
that kind of speed up, but it actually works for inferencing as well.
The same kind of gaps that you're seeing, we see that being provided to.
Are you saying the same fundamental principle applies?
I understand, but it's not exclusive to this training.
It's not exclusive to training.
So I think torch.compile is the new paradigm that is going to
change efficiency, completely across both training and inferencing.
So fantastic times for PyTorch.
And for the users. Yes, I know.
And what else are we doing to get PyTorch to be more
community
oriented, to grow that option of larger models?
Yeah, I mean, it's no secret like we're living right now
in the age of language models, large language models which are
just awe inspiring and what simple algorithms are able to do.
Mm hmm.
Well, not so simple, but you're probably familiar with Llama.
It's a language model from Meta. Yeah. Yeah.
And a lot of that code is open sourced.
You know, the model itself, the way it's there, available to the community
in a very permissive license.
Moreover,
we also have a lot of scripts available on how you can be using these downstream.
Okay.
There is one GitHub repository that you should check out.
Okay. It's called llama recipes.
So it's under the facebookresearch/llama-recipes on GitHub.
Oh, we'll add a link to to that in the video.A
And you'll notice
quite interestingly that they also use FSDP there too.
Okay, especially for fine tuning llama you know.
Yeah.
You want to adapt all of that knowledge and that model
to a particular domain or to a particular topic of interest.
Yeah. So you want to fine tune it.
Mm hmm. And of course, like we already saw how
[crosstalk] really works and everything.
Yeah. Yeah.
And how FSDP helps resolve a lot of that with an easy to use API.
And you can see that in practice
at the llama recipes repository.
Wonderful! So I think everything is coming together
with Python in the models, the way we train them, tune them,
and I'm sure you know, very soon it'll be around inferencing tools.
So this is fantastic news for the PyTorch community.
Yeah, I'm really looking forward to what the community comes up with
and what they're, how they start using this.
Absolutely!
Probably had a like a, probably a new step in this age.
Absolutely.
Looking forward to that. Thank you, Suraj.
Thanks Raghu.
Thanks for watching.
And if you like the video, please click on like and subscribe to our channel.
And if you want to learn more about PyTorch, about the concepts
we spoke about in this video, and even Llama,
do check out the links below in the video description.