Latest / Agentic Conversations (formally mlops.community) / Handling Multi-Terabyte LLM Checkpoints // Simon Karasik // #228

Transcript

Handling Multi-Terabyte LLM Checkpoints // Simon Karasik // #228 · Agentic Conversations (formally mlops.community)

0:00
My name is Simon Simon Karasik. I'm a machine learning engineer
0:05
at Nebula CI and broken on large language models at the moment.
0:11
So never CI is a cloud company where we're making a NI specific
0:17
cloud. But there is an arrogant team is
0:20
working on alms and I'm coming from this team and yeah, that's
0:24
what I do now. And about coffee, you know, I
0:27
tried to quit coffee some time ago.
0:30
I didn't drink coffee for months.
0:33
At first it was challenging but then became fun.
0:36
I realised that tea is great. You can have tea with milk, tea
0:40
without milk, a green tea, black tea.
0:44
What is happening in the loves community?
0:47
We are back for another podcast edition.
0:49
As per usual I'm your host Dimitrios and today we're
0:52
talking with Simon. But before we get into all of
0:54
that, I've got a little bone to pick with the Australians out
0:57
there. For the life of me, I cannot
1:00
understand why you call lettuce salad.
1:04
Salad is so much more elaborate than just lettuce, so stop doing
1:09
it on injustice and make sure if you're gonna call something
1:13
salad, you at least have some dressing or the very least oil
1:19
and vinegar on that please. Alright, that's it with my bone
1:22
to pick. Now we got into a deep, deep
1:25
conversation with Simon all about checkpoints and
1:28
checkpoints when it comes to training large language models.
1:32
As you will hear, he is currently training a 300 billion
1:37
parameter model which sounds absurdly large.
1:41
I'm very excited to see what comes out of that.
1:44
But he talks about how checkpoints become so big when
1:50
you are training at that level. The checkpoint management in
1:53
itself is something that you have to think about and scaling
1:57
laws for what you're going to do with these four terabyte
2:01
checkpoints becomes an issue. And how you look at the
2:06
management of OHH, Do we need one per week?
2:10
Do we need one per hour? Do we need checkpoints that we
2:13
had a month ago? How do you manage all of that?
2:17
He breaks down eloquently. I really appreciated it.
2:20
And he's got like the coolest job in the world because he's
2:24
working at Nebulous AI and they are specifically a cloud
2:30
provider that is dedicated to AI workloads.
2:34
So what does that mean? Well, that means that they have
2:37
tools for training and inference.
2:40
They also have a whole lot of GPU availability, which
2:43
everybody needs if you're going to be doing any kind of AI these
2:47
days, it seems like. And there's also like a UX
2:50
that's built for the engineer and developers, which is top
2:53
notch so you don't have this mess of services.
2:57
There's not that paradox of choice of which service should I
3:01
use. They really think through that.
3:03
And what Simon told us is he's been dog fooding the whole LM
3:07
generation process and as you will see he got into some of the
3:13
nitty gritty when it comes to networking, checkpointing.
3:17
I feel like an expert after talking to him about
3:19
checkpointing and As for like a little icing on the cake we
3:24
could call it. He talked about the differences
3:27
that he has seen over his years between Sikit learn models and
3:32
training those and training 1 billion parameter large language
3:37
model, quote unquote large. I guess we could call that small
3:40
language model these days and then training something like
3:43
this 300 billion parameter large language model.
3:46
So hope you all enjoy and a huge shout out to Nebulus for
3:50
sponsoring this episode. If you like it and you think it
3:54
is valuable, we would love to see you.
3:57
Leave some feedback in Spotify. Drop a star, Leave a review.
4:03
Comments on YouTube are always fun.
4:05
Subscribe. You know what to do, I'll see
4:08
you all on the other side. Or maybe in San Francisco on
4:11
June 25th for the in person AI Quality Conference we're
4:14
organising. Whenever I think about
4:23
checkpoints, I think about Checkpoint Charlie, because I
4:27
love Berlin and I go to Berlin quite a bit.
4:31
So I have that association, I say the word checkpoint and then
4:34
next thing that comes up is, Charlie, we are going to be
4:38
talking all about checkpoints today.
4:40
But before we get into that, I want to talk to you a little bit
4:43
about how you came to be, where you're at.
4:46
I know you've done some work with traditional ML and you are
4:50
now training large language models.
4:52
So I want to hear about the juxtaposition there and like how
4:56
things differ in that regard. But maybe you can talk to us
4:59
about what you were doing with ads before you started training
5:03
large language models. Yeah, sure.
5:05
Now first of all, I also used in Berlin recently and also has
5:09
association with Checkpoint Charlie and it was one of my
5:12
ideas to have this like a logo for some resentation.
5:15
I don't know and I did it good to find a good image, yes, so
5:20
now we're working with Welsh language models.
5:22
But before I worked quite a lot in advertisement in machine
5:26
learning. For advertisement I worked at
5:28
Yandex. It's kind of Russian Google.
5:30
Some people call it like this. And I was focused on ML
5:34
infrastructure there and I was making some different ranking
5:39
systems like click prediction, like if you see some bottle
5:44
without advertised, will the user click on this bottle or
5:48
will not click? And I was also working on
5:51
infrastructure because it turned out that infrastructure is
5:54
really big, is really complex. It's not just like let's deploy
5:58
this model, it takes like maybe 20 different steps, different
6:02
teams, different systems for the model to get trained, to collect
6:06
the data to a pushed production. And it was like, wow, Dutch.
6:12
Ouch. We'll break that down a little
6:14
bit more. You can't just say that and then
6:15
keep me hanging, OK? Why was it so big?
6:17
What? What?
6:17
What did that look like? Paint the picture a little bit
6:19
more. Because if you want to create a
6:22
great advertisement system that is accurate, it's like
6:25
personalised. You have to collect lots of data
6:28
and process it. Because it was like I was making
6:31
some deep learning model like some recommender model and it
6:35
doesn't just go to production. Is this model is a feature to
6:38
enhance the model that is another model.
6:41
It's like a big combination and like OK well I need to train
6:44
this Dee lyric model in Pytorch but it's not regular pytorch,
6:48
it's like Pytorch especially for recommender systems and it has
6:52
some special addons developed by my my colleagues.
6:57
The sound it comes eventually date then we have some date as
7:01
it's coming every hour but it's not actually every hour because
7:05
it takes time for the data to actually come from your click on
7:09
the online through the data store and then get in trained
7:14
and pushed production and that's where lots of lots of pieces
7:18
like you said well developed for years her hand and even like
7:22
we've had this like a service that is serving like 10 or 20
7:27
deep learning models for different teams and also was a
7:32
big big like infrastructure and insurance.
7:37
And I remember like I do Samsung, you know, like I reduce
7:41
some load, some save some percentage of CPUs, like let's
7:46
remove this model and save like one percentage of CPUs.
7:50
But it ends U like saving as 1000 of CPUs.
7:55
It was like wow and I have never seen the CPUs.
7:58
I just see this dashboard like does that mean Ccus we consume?
8:02
But it's not like all data centre like to enter computers.
8:07
I don't know. But I Did you get a medal or
8:12
anything? Did you get a raise after you
8:15
did that? Yeah, does.
8:16
Anyone notice I guess. Sure, like everybody's noticed
8:20
because if he if he release some CPUs, that means now we can use
8:24
it for something more useful. But you can't just remove some
8:29
model from production. It takes time to check that
8:31
nobody needs this model, like it doesn't make any profit.
8:35
Then you can remove this model and you'll use this space for
8:39
something else. Damn, that's crazy what you just
8:42
said. How there was basically like
8:45
dead models out there in the wild.
8:48
Still, but there wasn't proper cleaning that had happened.
8:53
So they were still just like hanging out and you recognised
8:58
that and you were able to save CPU costs and just like open up
9:03
CPUs for newer things and potentially better things to
9:07
come out. Or at least new experiments that
9:10
are going to be not dead in the wild.
9:14
And what did the process look like for like that Garbage
9:18
collection? Actually it was a kind of
9:21
another project. It was another project I tried
9:24
to connect all the pieces together because we had this
9:27
runtime like inference. Then we had some data art, we
9:31
had training pipelines and they were not really interconnected.
9:35
And happens like maybe you had some model in run time, you
9:38
removed this model, but the process that was training this
9:41
model was still active, it was still consuming some CPUs every
9:45
day. It was another project to
9:48
actually connect all the pieces together so that if you delete a
9:52
model from runtime you get wording.
9:54
Hey man is still trying this model.
9:56
Maybe you should stop. It wow, that is so wild to think
10:00
about that you had so many models out there and being
10:04
trained and just being in production that it's almost like
10:09
people lost count of them or there wasn't like clear owners
10:12
of it. Is that safe to say?
10:14
It's some kind of historical reasons.
10:16
You know, like you have a system that is active for like 10 years
10:19
or more. Somebody created a model like 10
10:22
years ago and deployed it. It makes the job, it works well.
10:28
Eventually the system gets a factored changes changes and
10:31
then you have some blind of config that you have no idea how
10:36
it comes here. Who made this model?
10:38
The person who made this model is already, I didn't know in
10:41
another company for five years. And nobody knows you want you
10:46
don't want to crash the roduction.
10:48
So you're really making it step by step.
10:51
Yeah, I can understand that. You're being very cautious about
10:57
how you go about collecting those models or getting them out
11:00
of production because it is a non zero chance that it could
11:07
actually be very valuable and there could be a lot of money
11:10
being made. And then you went and you said
11:13
alright, enough of the deep learning recommender systems,
11:17
I'm going to go to the source. I like large language models now
11:20
or what happened? Give me the breakdown on how you
11:24
started working with large language models and training
11:26
large language models. Yeah, it was just an
11:29
opportunity. I saw this like job.
11:32
I decided this is really interesting.
11:34
I see that how large language models are changed in the world.
11:39
I even remember when I was working back in advertisement,
11:42
my team lead was so much impressed, he said.
11:46
No matter what we do, LMS will replace everything.
11:51
Ohh, no. He was one of those.
11:52
Guys, yeah, he was kind of fatalist.
11:54
But still, I it was impressed. If you know, like when you saw
11:59
that I'm doing this super tough job, NASA can replace me.
12:04
So you go to judge and say, hey, could you write this code for
12:08
me? And it does the job.
12:10
You write, you're writing an article you like spending one
12:13
hour. Then you ask CHC and it makes it
12:17
better. The better than this was made
12:19
making point our and no it's mind blowing.
12:24
Yeah, yeah. So that made you realise there's
12:28
a lot of potential here and what happened?
12:33
Yes. So I joined this team and we're
12:37
starting making this language models.
12:41
I was Williams the first join maybe like like peeps join and
12:48
what's impressive, what was impressive for me because you
12:51
know back in the days when they work in ads, it was like really
12:56
large team, like there was people for everything.
12:58
Like you can talk to this guy, to this guy, and ideally.
13:02
Have two. There was a million models in
13:04
production. Yeah, and they haven't really
13:06
had an idea of how it all works at such a big scale.
13:10
You joined a new team within Yandex that was working on large
13:14
language. Models No I I joined Nebulous.
13:16
I joined Nebulous. So and now neighbours is working
13:19
on this LNG projects. OK, and it's a startup, so you
13:24
don't have someone to get you coffee in the morning and
13:26
someone to write your reports for you.
13:28
You got to do everything now on your own.
13:30
Yeah, sounds like that. So it's like neighbours is a
13:33
cloud company. It's mostly cloud, cloud
13:36
developing, cloud engineering, but we still have this team of
13:39
large language models. It's like we are like a small
13:44
island of machine learning actually machine learning guys
13:47
around big infrastructure guys and like I have to come U with
13:51
ideas, my team has to come U with idea how can we make this
13:57
machine learning big machine learning training like out of
14:01
better like cloud services. For sure you have Kubernetes,
14:05
you have storage, you have some CI, but it's this a building
14:10
blocks what you need to do a real big pipeline.
14:14
Ohh I. See.
14:15
Was it tough? It's an interesting journey.
14:18
So it's all, it's almost like your dog fooding the LM journey
14:23
that your customers are going on and you're doing it yourself so
14:28
that when customers say, oh we want to train your large
14:31
language model, we want to train a large language model, then you
14:36
understand what they're about to go through.
14:38
Yeah, exactly. So the food is one of the terms
14:42
we're now using and it's actually now you know, like
14:45
we're one step ahead of our customers.
14:49
So whatever it happens, first we might find some problem, we fix
14:55
it and we know that now our cloud is 100% production rate
15:00
because we train such big model and we know what can go wrong
15:05
and what can go well. Yeah, yeah.
15:08
And talk to me about the models that you've trained so far.
15:12
Yeah. So we started like with a kind
15:16
of testing. So while we're one of our first
15:20
goal was to verify that we can train a good model.
15:23
So we tried to reproduce a LAMA 7 billion parameters in the way
15:27
we produce this data collection pipeline, training, pipelines.
15:31
We had this kind of budget, we limited as LAMA and made sure
15:35
that we can do a good job and then we decided to go bigger.
15:40
So now we are trading a 300 billion parameters model.
15:45
It's like huge. It's still in progress.
15:48
So I maybe can't share lots of details, but I hope you will
15:52
hear about it soon. But it's not like 300 billion
15:57
parameters, it's like more than 1000 GPUs to train, just like
16:02
insane, you know? And I have never seen this virus
16:06
like I have never been to data centre, but I can imagine like a
16:09
big big room full of computers just doing some matrix
16:14
multiplications. Yeah, yeah.
16:17
That's where you can go in the winter if it gets too cold.
16:21
OK. That big room with all those
16:23
GPUs, Wow. O that is impressive.
16:26
Now I want to talk about checkpoints and what that
16:32
experience has been like, like what the learnings you have seen
16:35
while training have been. I think for the uninitiated, can
16:41
I give you my understanding of checkpoints?
16:44
And then you can. Yeah, sure.
16:45
Correct me and tell me where it's wrong.
16:47
I look at checkpoints like when you're playing video games and
16:50
you get to a certain part in the game that the next time you die
16:55
you go back to that part. So you don't go all the way back
16:58
to the beginning. You just go back to whatever the
17:02
level 4, level 5 wherever you get to the beginning of level 4
17:06
and then you face the final boss of level 4.
17:11
But is that a good way of putting it?
17:14
Or tell me where I'm I'm off? Yeah, it's, you know, it's
17:17
really good comparison. It's really how it happens.
17:20
But you know, like if what can happen, like you wake up one
17:25
morning and see like and your loss plots, that's your loss has
17:30
exploded. You know like ohh wow, we need
17:34
to rollback. I hope we have a checkpoint
17:36
because if you don't have a checkpoint, you're in trouble.
17:39
You need to redo like days of work.
17:42
But if you haven't checkpoint, you still have maybe some hours
17:45
to retrain. How often are you checkpointing?
17:49
Maybe like once every hour, maybe two hours.
17:53
Nice. So just to make sure like if
17:55
something happens, like we don't have to kind of redo many steps,
18:00
because on one hand, you know, like checkpoints and it's also
18:03
not free because you want, if I want to save checkpoint, I still
18:06
have to freeze the training like for 4:50 minutes, maybe for some
18:11
minutes to save the checkpoint. So if I do checkpoints too
18:15
frequently, I make the training slower.
18:18
But if I don't do checkpoints, I'm in the risk.
18:22
You're taking a risk. So you found the sweet spot is
18:26
once an hour, once every two hours.
18:28
Yeah, it just sounds like some reasonable, some reasonable
18:31
amount of times it's OK if something happens, maybe we are
18:34
OK to redo one hour of training. Yeah, and your checkpointing
18:41
everything like or. How do you calculate the
18:43
checkpoint size I? Guess.
18:45
Yeah, sure. So like you know it's quite an
18:49
impressive because if you have a large language model or hugging
18:53
face, it's like often it's like inference model in a way it's
18:58
already quantized and Justice you are ready to push it into a
19:01
GPU. And it's like you have a
19:04
perimeter and often it's like like 2 bytes per parameter and
19:10
it's like you have Lama 7 billion parameters and like 14
19:15
storage, 10 gigabytes. But if you train a model like
19:19
this, you need much much more space because you have
19:22
optimizer, you have parameters, it's like 3 numbers.
19:27
Like for every bit on you have you need to have like 3 billions
19:30
of numbers to save and like. Each number is like a floating
19:36
point number. It's like 4 bytes and it's like
19:39
OK, you have flower lamma 7 billion on hugging face.
19:43
This is just like short 10 gigabytes.
19:45
But to train, Lama says you need to have 70 gigabytes of space
19:52
like Seventh or something like that.
19:54
And it's not all OK, we save 70 gigabytes, but inside in memory
20:00
you have much more. You have like gradients, you
20:03
have all this kind of stuff that are allocated when you train.
20:08
You don't save them, but it is allocated on GPUs.
20:12
So if I wanted to train Samsung, it's like, OK, no problem, It's
20:16
like certain gigabytes. No, it's not sorted.
20:19
It's like hundreds of gigabytes. Is that I need to be allocated
20:23
in GPU memory? And that's presumably why things
20:28
get very expensive. Very quickly, yeah, sure.
20:30
Because like you need to have lots of just memories to
20:34
allocate the model. Yeah.
20:35
And you need to have the confidence that you're doing the
20:38
right thing, because I can imagine a lot of fear that
20:43
people have when they're training models is they're
20:46
training something and it's going to come out in 20 days or
20:50
14 days and it's going to be absolute shit and not very
20:55
useful. So how are you making sure to
20:58
cover your back on that? Yeah, sure.
21:01
So first of all, we start with small experiments.
21:05
So we start with like 1 billion model, 3 billion model,
21:09
something that can be trained like in one day, in three days,
21:12
not unlike on 100 GPUs, but maybe on 8 or 16 GPUs.
21:17
Just to make sure that we are doing the right job.
21:20
Because we know there are papers about scaling clothes about like
21:25
if you make your model 10 times bigger, what will be the result
21:29
O knowing that we can say OK, if it's a small model, it's OK big
21:33
model also should be OK. So this is kind of the 1st way
21:37
you make sure that everything is all right.
21:40
Also what we do training, we have a lots of plots, we do, we
21:46
have weights and biases. It's like we have 100 dashboards
21:50
there because we have laws, we have gradients, we have
21:53
gradients on this layer, on that layer, and this really helps
21:57
because if something goes wrong, we can always go to this
22:01
particular plot and see what happens there.
22:05
Hmm, OK. But it feels like there's
22:09
different types of checkpoints that you can do, right?
22:12
Like there's basic checkpoints and maybe there's more in-depth
22:17
checkpoints. In my off on that one.
22:21
Actually I guess as there is only one type of checkpoint we
22:25
can do so we just save as a model with optimizer state and
22:30
basically it says everything we need to have to restore our
22:34
training. O If something crashed, we still
22:36
can restore some training courses checkpoint and we'll
22:40
still have all the same loss which will continue right to
22:44
where it stopped. O when it is a smaller model
22:49
like you were talking about, or checkpoints any different.
22:55
Yeah, so there is a difference, and the difference is in size O.
22:59
If you train a model like 300 billion parameters, each
23:03
checkpoint is like 3 1/2 terabytes, just the same like
23:08
you see three half terabytes. You know, my MacBook has like
23:12
only 400 gigabytes. Oh my God.
23:16
It's like more than like the hard drive of my laptop and it
23:21
needs to be saved like every hour like with many replicas and
23:25
it doesn't fit into one computer like if you even if you have
23:30
this super powerful GPU virtual machine, it has one terabyte of
23:38
memory and and checkpoint district terabytes.
23:42
So you need to have like several virtual machines to save this
23:46
checkpoint and to load it. And it's really a lot O if you
23:51
have a small model, like relatively small like 1 billion
23:54
parameters like some years ago it was the highest you can get.
24:00
I remember like 1 bird was released, like Bird large was
24:04
300 million and it was like a lot, nobody could run it.
24:10
It was really heavy and nowadays we don't even count.
24:14
So this. One, we don't even bad an eye.
24:16
That's so true. Yeah.
24:18
So yeah, if you have a small checkpoint, it's not a problem
24:21
to save it. Maybe it's 50 gigabytes, maybe
24:23
70. It's still not so much like our
24:27
machines that we on which we run training are much bigger and you
24:31
can save 70 gigabytes to disc. It will take some time, but it's
24:36
still it's OK. It's not a big issue, but now
24:39
when you have terabytes you need to say need to save, need to do
24:44
it in parallel. You need to distribute the
24:46
checkpoint, amount of costs, amount of rental machines,
24:52
Causeway. If you try to save St rebates
24:55
from one virtual machine, it will take really long.
25:00
Yeah. So we do some parts to make it
25:04
faster. Ohh, yeah, tell me about those.
25:07
Yeah, sure. So first of all, we do it in
25:09
parallel on many machines. We slit our checkpoints.
25:13
So our checkpoint into parts O if you have training and
25:17
retraining on 8 machines, each machine will save a part of this
25:22
checkpoint so that yeah, OK, in total it's like 3 terabytes, but
25:27
for every machine it just like like 300 gigabytes.
25:32
It's still quite a lot, but it's not so much.
25:36
Era. It reminds me of torrents, like
25:39
when you tore into a music file or something and you get it back
25:44
from many different pieces of that music file.
25:48
Yeah, kind of sounds like that, yeah.
25:51
And we also realised, OK, like we have this machines and we
25:58
have network to the storage like you know, regular Internet
26:01
network and our machines are connected via super fast
26:06
InfiniBand network. It's like a GPU, GPU network,
26:09
because if you're training model on many machines, you need to
26:12
communicate a lot. That's why there is a super fast
26:16
network and OK we have a network between hosts that is much
26:22
faster than the network to the storage.
26:24
O When we save our model our checkpoint, it is saved like a
26:29
small bits from every cost. Then we'll load small bits.
26:35
So when we load the checkpoint, every host has only a small part
26:41
of the model. It's not enough, but in total
26:44
every host has everything and then we use our superfast
26:48
network to exchange the parts. How often are you needing to use
26:54
the checkpoints? Is it like once a week, once a
26:58
day? Yeah, it it depends, you know,
27:03
if everything goes well, we don't want it to ever rollback.
27:07
Yeah, ideal world. We have a training we just
27:10
started. We wait one month.
27:12
So it depends, you know, actually we use checkpoints
27:17
always because on top of what we train, we have another
27:22
background process that takes this checkpoint and runs some
27:26
evaluations, some validation because we don't want to stop
27:30
our main training. But we still want to collect all
27:33
these benchmarks like MMU and other.
27:37
So like we have a background process that reads every new
27:40
checkpoint and tries to validate it.
27:43
But also tricky because if our checkpoint is that big like
27:47
Twitter bytes, this another process also has lots of job to
27:51
do. Just the singing about these
27:55
gigantic checkpoints that you have and how you can even work
27:59
with them. So you you basically.
28:01
The ideal scenario is you set it and forget it, but the realistic
28:06
scenario is like you're going in and you're trying to figure
28:09
things out every once in a while on how it looks.
28:14
Yeah, got it. So first of all, let's read.
28:20
Some realise that it's important to save checkpoints properly in
28:25
a way we don't need. So we save it every hour, but we
28:31
don't need to store checkpoints from every hour, but we need to
28:35
store checkpoints from yesterday and from last week.
28:38
You know, kind of this exponential schedule because you
28:43
know, like if you store like it every hour, every hour, it's
28:46
like 3 terabytes. You can quick quickly reach 1
28:49
petabyte of disc and actually we use Kubernetes to orchestrate
28:54
our training also was a quite a funny and tough way to get it
29:00
working. Turkish trade this 100 machines
29:03
working altogether. Well, because I've heard a lot
29:06
of people like to use SLURM, right?
29:08
Was there a reason that you chose Kubernetes instead?
29:11
So actually no. I've heard about slum only
29:15
recently when I was talking to my colleagues who work with
29:19
clients and they said like we need slower because like clients
29:23
lose use slower And I'm like what's this?
29:28
You know, I can you know, like it's happens sometimes you have
29:30
some technology that is there around like for years, but you
29:34
have never touched it and then you realise that you just needed
29:37
to use right Google query search writings.
29:42
Yeah. So we use this because I don't
29:45
know it's it feels like natural because I don't know like I had
29:50
experienced with Kubernetes, my colleagues have experienced with
29:53
Kubernetes. Yeah, it's initial way to go.
29:56
Yeah yeah, I know the Databricks team who trained DBX, they came
30:04
on last week and they were talking about how they made that
30:08
choice too. And they weren't with
30:09
Kubernetes, but I have heard a lot of the researchers enjoy
30:14
slur much more than Kubernetes. Yeah, I agree.
30:18
I read the bit about slower, so it seems to rise from MPI from
30:24
this kind of more scientific computations.
30:26
And maybe the guys who are doing research, they come from this
30:29
background and these are used to MPI, but those guys have coming
30:35
from more kind of engineering infrastructure as I'm more
30:39
familiar with Kubernetes. Yeah, yeah.
30:43
Are there any lessons that you learned when it comes to storage
30:46
choices? It was very impressive for me
30:49
when I was trying to understand how does storage work in cloud.
30:54
Because you know, like if I have my laptop, it's all clear.
30:56
I have my laptop, I have SSD or just save it Button cloud you
31:03
have like data centre, you have virtual machines, you have
31:07
storage. Where does storage live?
31:10
If you have like, you know, like 3 and you say just three, let's
31:14
save it. What actually happens?
31:17
And turns out to be a really big, big job, because first of
31:23
all, what was the mind blowing for me?
31:25
You know you have like a virtual machine and you have a disc on
31:28
this virtual machine, but it's not there.
31:31
It's not if you have a real time machine, disc is not really
31:34
connected. This machine is right, another
31:37
machines that have these discs and that's why they're called
31:44
network lists. Yeah, cloud uses natural discs,
31:48
Network discs, sorry. Because like you know, you have
31:51
a real machine and you want you want to go to another real time
31:55
machine with most CPUs with more memory, but you still have the
31:59
disc. You can't just take the disc,
32:01
put it into another machine. You can't move like all your
32:04
data from one machine to another.
32:07
That's why Cloud uses this natural discs.
32:11
And OK, it was the first aside for me how hard it is and
32:17
regarding that choice of storage, so we had several
32:20
alternatives that are considered and we tried.
32:23
So we had this file storage like NFS that as tree like storage
32:30
and we also considered so to do data processing we are using a
32:36
product called Tractor AI. It's a product built on top of
32:41
white thesaurus. It's a kind of big data system
32:45
that was previously built-in Yandex and now it's like a
32:49
cooperation between neighbours and Yandex.
32:51
It's open source products. So we're contributing,
32:54
contributing to it and we have this structure that is built on
32:58
top of this white thesaurus. And so like we also have this
33:02
option. Because you're talking about
33:04
storage. You're not just talking about
33:05
these gigantic TB storage of checkpoints.
33:09
You're also talking about all the data breakdown, everything
33:11
that you're thinking about when it comes to storage.
33:14
I care about security and scalability because it could be
33:19
that you have some storage. It works just well when you
33:22
touch it to 1 virtual machine, but then like you have like 10
33:26
virtual machines and them all want to do something but storage
33:30
is not scalable. And that says all men like you
33:33
had this one GB of bandwidth but now like you all have this one
33:39
GB like not not each one GB but one GB for all.
33:44
So also what's important when I'm coming to choose a storage
33:49
is what features does IT support?
33:52
Because, you know, like we all worked with S3, it looks like
33:57
just great. It's just easy.
33:59
Let's push something there. But then when they worked with
34:04
it, I faced this old issue that story is not a file system.
34:12
Because like we we had, we tried to store checkpoints in S3 and
34:17
we used this kind of tools that kind of as three file system
34:23
that make you feel that you have a bucket and a street but it
34:27
looks like a disc, it looks like a mount and usual machine.
34:31
It worked just well. But then our checkpoints and
34:34
library was doing some move like let's move the checkpoint from
34:39
here to there and with three you can't do it because three is
34:44
like let's copy the data and delete the old data.
34:48
That's how we do move. And there's we just observed
34:52
that our training became very slow.
34:55
So it feels like you we're hitting, it's not necessarily
35:00
the scale scaling laws that we were talking about from like the
35:03
papers and the training. But you were hitting different
35:05
types of scaling laws where you didn't necessarily need to think
35:09
about these things like when you have smaller amount of data or
35:13
you have smaller checkpoints but all of a sudden S3 wasn't
35:18
possible for you. So you went with.
35:20
Did you end up choosing NFS? So at the moment we use this
35:28
tractor AI. OK.
35:30
Yeah, that it was. It was an experiment and it went
35:33
really well. We like what we have.
35:36
And another motivation to use Structure I was the fact that we
35:40
already use it to do data processing.
35:42
We already stores the data, we process the data.
35:45
We know it is a reliable system and most of the time we use the
35:50
system. So it's kind of natural to use
35:53
it in one more place and like have everything there.
35:57
It's easier for us to manage like this.
36:00
There's the and besides storage in hitting scaling laws with
36:07
storage, what are some other pieces that have become like
36:12
paramount in your setup? Yeah.
36:16
First of all, it's governesses, so they use it a lot.
36:20
Like to have some problem. OK, let's deploy a demon set.
36:22
We have some problem, Let's do this.
36:24
Let's deploy that. It became a tool we used a lot.
36:28
It's personal tool to master. And UH, we even use some tools
36:36
on top of governesses to orchestrate our training.
36:39
Yeah, so we use Argo. Argo is a kind of orchestration
36:44
framework and we found it really useful because what happens with
36:49
model training? It's not like you just have one
36:54
port, another port, they're all connected.
36:57
If you're training model on 100 machines, they all have to be
37:01
coordinated and if one machine files, everything should fail
37:06
and so we used Argo to implement this kind of coordination.
37:12
OK, it it wasn't Argo workflows then yeah, it.
37:15
Wasn't Argo workflows? No, it was interesting one about
37:19
Argo workflows. I so I was just at Cube Crown a
37:22
few weeks ago and it feels like everybody was wearing an Argo
37:25
shirt and it's gotten a lot of traction these days.
37:29
It is what Cube Flow is built on top of.
37:34
Yeah. So I think a lot of people have
37:36
gone and they said OK, we wanna get a little bit further down
37:40
and have more knobs and have I think Argo is matured nicer than
37:45
Cube Flow and so people are going towards Argo more than
37:50
cubes though these days. But that's that's interesting
37:53
that you chose it too and you saw the value in it.
37:56
What other add-ons where you're using with Kubernetes?
38:01
Because I know there's especially you talked about how
38:04
networking was so important, like where you're doing special
38:07
stuff with networking where you're using a specific flavour
38:10
of Kubernetes. Was it your own stood up?
38:13
Was it the cloud version? What does that look like?
38:16
Yeah. So we use a standard coordinates
38:19
as it's available in our cloud, but we for sure we have some
38:23
additional stuff as we deploy to coordinate this.
38:27
First of all, it is GPU operator because GPUs are really complex
38:33
and need to monitor them. Well, you need to check them.
38:35
So we deploy GPO operator that's making sure that everything will
38:38
work very smooth. And for example, if you operator
38:42
chefs your GPU, and it sees that GPU is an unbroken for some
38:47
reason, it's still too hot for example, it will stop it.
38:52
It will just tell combinators, Hey conferences, this note is
38:56
not good. Let's don't use it.
38:59
Relax for a little bit. Take a little break.
39:02
Yeah, as a product that you found useful is we deployed our
39:07
custom a piece of software that runs this network testing that
39:13
it it cheques that all the nodes that we're using to train are
39:17
well connected. They have the right speed, just
39:20
in case you know it can happen. If you train a model, it trains
39:24
altogether. And if something is if OneNote
39:28
slow for some reason, it looks like all training is slow.
39:33
So you really need to have some paranoid mentoring to make sure
39:38
that OK, this is not the slow or this is not this is OK.
39:43
So we do this testing to make sure that everything is well
39:46
connected, works smoothly. So what are you doing in those
39:50
cases? Because it's basically like like
39:52
the army, you're only as fast as your weakest point, or you're
39:56
only as fast as your slowest point and you're monitoring the
40:00
Kubernetes nodes and so let's say that one of the nodes for
40:04
some reason is going very, very slow.
40:08
What do you do? Do you stop training and then
40:09
try and debug? Or you do you go into it and try
40:12
and just update on the fly? So because you know one great
40:19
feature of cloud is you don't rent actual machines, you don't
40:24
run specific GPUs, you run kind of resource.
40:28
But it's it's a concept, it's idea.
40:30
So if we see that something doesn't work and for example
40:33
this GPU is not working, we can always disable it and we will be
40:38
given another. So OK, we we automatically stop
40:42
the training, we have the process that does monitoring, it
40:45
stops the training like this, it removes this broken GPU and we
40:49
are given another note automatically and we continue
40:53
there and then we can talk to our cloud team.
40:56
They will debug this, they will fix it and later they will know
41:00
what goes wrong, they will fix it and we will not have this
41:04
kind of issue like this anymore. O Most of the time it's not even
41:09
on your side, it's just that the GPU itself for some reason isn't
41:13
working that well. Yeah, it could be the reason
41:16
because you know it's it's super hard piece of hardware.
41:20
It's not like you have a machine, it has eight GPUs.
41:24
They're connected, they're connected inside, they're
41:27
connected to each other, they're connected to other machines.
41:30
It's super hard system and there is always something that can go
41:33
wrong, but because it's cloud it's easier to just OK replace
41:39
it and do something later. Yeah.
41:43
And so you're monitoring the system and the coordinated
41:48
clusters, I imagine. Are you using something like a
41:52
Prometheus on that? And then you're also monitoring
41:56
the training and are there other things that you're monitoring?
42:01
Yeah, sure. So first of all we use one
42:05
weights and biases to collect our like trail clocks, I mean
42:10
like loss and et cetera. And what is the biases already
42:14
does pretty good job. It also select some statistics
42:18
about the realtor machine, it shows sanju temperature.
42:23
I don't know like the amount of memory consumed.
42:27
So we have quite a lot of information coming from weights
42:30
and biases. Also there are lots of cloud
42:34
based mentoring. So because when you create a
42:37
realtor machine in cloud, it already has some monitorings in
42:41
it, we just go to another tab in UI and we see the temperature,
42:47
you know we CPU. Load.
42:49
Et cetera. Also we collect logs of training
42:55
and became also quite a challenging ohh task because
43:02
it's not like it's interesting how machine learning is getting
43:06
closer to kind of microservices, because I used to do micro
43:11
services before and I use some logging systems with
43:15
microservices and it's all distributed but they're
43:20
connected. The same about training.
43:22
You have a training, it's running on 100 machines, but
43:25
it's all interconnected and you need to collect logs from every
43:29
machine to eventually have a picture of what went wrong.
43:35
Because if training fails, just every note, every machine has a
43:40
bunch of logs of load. Looks like hey, something went
43:44
wrong and failed and failed and failed and failed and you need
43:48
to have a way to find this one specific machine that's had a
43:53
real error because others fail because that one failed.
43:57
Yeah, because it's all linked. So if one goes down, everyone
44:02
goes down. Yeah.
44:03
So that it's it's funny you mentioned that because I do
44:06
remember there was this blog post that came out, I want to
44:11
say it was like 2020 or early 2021 from Open AI and they
44:16
talked about how they scaled their Kubernetes cluster to over
44:21
6000 nodes. And I remember when that came
44:25
out and this was pre Open AI being, you know what it is
44:29
today. And I was so fascinated by that,
44:33
by how big of training jobs they needed and the type of problems
44:39
that they were looking at. And so this was as they were
44:42
probably training like GT2 and they were scaling it up to 6000
44:47
nodes and getting to that point. So I I understand things can get
44:53
very complex and very quickly. And so you're looking at the
44:57
training job, So you're monitoring the training job.
44:59
You're also monitoring the GPUs, making sure they're working and
45:03
then making sure that you have a way to debug if things go
45:08
offline. And then you're also monitoring
45:10
just the system in general and the Kubernetes cluster, making
45:14
sure that that is working and there's every node is firing how
45:18
it should. Yeah, exactly.
45:21
So we try to monitor everything. So there's.
45:25
Two last questions that I want to get to for you.
45:29
First one is around pre training versus fine tuning and how much
45:33
of like this conversation that we're having with checkpoints
45:37
changes as we start looking at fine tuning.
45:41
Against the difference is first of all if you do fine tuning,
45:46
you don't need so many GPUs if you do retraining in to have
45:51
like 1000 Jews. But if you then want to fine
45:54
tune this model, you have just 1000 of rows of data, 5000, I
46:00
don't know, not many, not billions froze and you don't
46:03
need some engineers. You cater will with just two
46:07
machines. But the problem is the model is
46:11
still as big as it was. It was like 3 terabytes.
46:15
It's still treacherous bytes, but now you need to fix this
46:18
street terabytes not into 100 machines into two machines or
46:23
four machines into all these three terabytes for two
46:26
machines. And creates a different kind of
46:31
workload for the cloud, different kind of requirements.
46:34
Because you know, like if you want to load each point from 100
46:38
machines, you want your storage to be scalable in terms of
46:42
consumers. You want your storage to be able
46:45
to work with 100 clients at the same time, but now you don't
46:50
need so many. You don't have so many
46:52
consumers. You just want to load the three
46:55
terabytes from 2 machines. But now you need this network to
46:59
be super fast. This, this is this particular
47:03
network from storage to this one machine to be super fast as the
47:06
rise will wait, I don't know like one hour waiting for your
47:10
data to get into your machine. But weren't you saying that the
47:17
networking that you need for pre training is also it needs to be
47:22
super fast or there you can have it be a little less?
47:26
So when we do pre training because we have many machines
47:30
and we kind of split our checkpoints.
47:33
So in a way that as I mentioned every machine is loading and
47:38
saving only a small part of checkpoint it's saving like
47:42
starts gigabytes out of St terabytes.
47:45
So for every particular machine you don't need to save and load
47:49
so much data. You just need to make sure that
47:51
it works, it works and like. If you have some good OK speed
47:57
of the network you will finish it.
48:00
But now you have all the same amount of data that you need to
48:05
read from one host and that's like what if hardware, more data
48:11
that need to be transferred over the same network but over one
48:14
network channel. And it's kind of different
48:18
design from the storage perspective, from the networking
48:21
perspective. And it's funny also network,
48:24
it's, it's interesting kind of how complex networking is
48:29
because I remember when I was doing some experiments and then
48:33
Nitric guys from network are coming to me and say, hey, what
48:37
are you doing? Like we see our mentoring just
48:40
explode it and I'm say ohh, well I'm testing this and say like,
48:45
OK, let us tune our network so that it will be able to scale
48:50
better to do the thing. And they ask me different
48:55
questions like how does your network do we go from here to
49:00
there or from here to there? Because this is different layers
49:03
of network and they require different parts of the network
49:07
to work. I'm like, I don't know, let's
49:10
discuss. That is fascinating.
49:13
Networking is an art 100%. You you got to a point where you
49:19
were like, all right, well we can get some optimization here
49:24
if I meet with the networking team and I see how our training
49:29
or and this was in the training or this was in the fine tuning
49:32
phase. It was and trained when we were
49:34
trying to hold the checkpoint from 100 machines.
49:37
Now it's a big load, right at the same time, the same moment.
49:42
So it was that moment. Good stress test.
49:45
Yeah, the networking chief. Yeah.
49:48
Yeah, I'm sure the first person to catch that was just like,
49:51
what is going on here? Who is doing what are what are
49:55
people doing here? We gotta go talk to Simon.
49:57
Yeah, it was exactly the sync and I was like it because one of
50:01
the colleagues who are doing network, she's sitting in my
50:04
office like right near me. So she came to me and said,
50:07
like, hey, what are you doing? Ohh yeah.
50:12
We gotta have a conversation. Come to my.
50:14
Come to my ohh it says that's great.
50:19
So that's the differences, big differences that we need to keep
50:22
in mind when it comes to fine tuning verse pre training.
50:25
What about when it comes to like the traditional ML training with
50:30
Sikit learn model versus now when you're training LMS's?
50:35
We've gone over how wildly different it is when it comes to
50:39
checkpoints and sizes, but are there things that you want to
50:41
highlight that you keep in mind now, or that you learned from
50:47
making that jump from like traditional ML and working with
50:50
psychic Learn and then going into like more deep learning and
50:54
LS? Maybe it's the first lesson I
50:57
have is Unless you're good recycling turn, you don't need
51:02
by church. Unless they are good with Piper,
51:05
she don't need like super fancy deploying models on the edge.
51:09
Plus, because it's not like, OK, let's just use GPU, I feel it's
51:15
like a much bigger shift of complexity that you get.
51:20
So unless you're good with something simpler, I believe we
51:25
should use something simpler. All these else shift to
51:28
something bigger. But it's interesting how it
51:31
happened with large language models.
51:33
Causeway, I talked to my friend and he told like you know,
51:39
previously we had some problem but we didn't have data and we
51:42
couldn't do anything. But with large language models
51:46
we can have zero data, but we can still solve the problem in a
51:51
way that now many problems that he used to solve recycling could
51:56
be solved just with prompting of charge.
51:59
PT maybe, yeah, but it requires some extra complexity from the
52:04
PC side. So what you're saying is like
52:10
there were use cases that we wanted to tackle, but we didn't
52:14
have the data for and now we don't need the data because it's
52:18
already in the model and we can just go and prompt.
52:20
It yeah, it's like super impressive how it changed
52:24
because it technology that wasn't there like a year ago,
52:28
two years ago it came and it changed everything.
52:33
But I do like this notion of keep it simple.
52:37
It is so funny that we need to keep repeating it because it's
52:40
so easy to overcomplicate and wanna jump to calling ourselves
52:45
like it's just like we get to wear these medals of all right
52:49
yeah I've trained a a model that is I used 1000 GPU's to train
52:54
this model or whatever it is you know bigger is better type thing
52:58
but if you don't need to then why why do you do it right.
53:04
Yeah, but you know, like, I guess just some people like to
53:09
do cool things, they like to do hard things and this is what
53:13
masturbates them to to work. This is the driver.
53:16
OK, they can do it simple, but they don't like it.
53:19
But they can do it complex and they enjoy it.
53:24
There's extra. Well, dude, this has been super
53:29
fun to talk with you. I appreciate this.
53:31
You gave me a sobering look at what it takes to actually train
53:37
these models, and you are encountering some of the
53:40
challenges like day in and day out.
53:43
And I appreciate you sharing these challenges with us because
53:46
now hopefully anybody out there that wants to do this knows what
53:52
it takes. Yeah, it's really.
53:57
Thank you. Hold up.
53:59
Wait a minute. We gotta.
54:00
Talk real fast. Because I am so excited about
54:02
the MLB's Community Conference that is happening on June 25th
54:07
in San Francisco. It is our first in person
54:09
conference ever. Honestly, I'm shaking in my
54:12
boots because it's something that I've wanted to do for ages.
54:15
We've been doing the online version of this and hopefully
54:19
I've gained enough of your trust for you to be able to say that.
54:22
I know when this guy has a conference, it's going to be
54:26
quality. Funny enough, we are doing it.
54:29
The whole theme is about AI quality.
54:32
I teamed up with my buddy Moe at Colina who knows a thing or two
54:37
about AI quality, and we are going to have some of the most
54:42
impressive speakers that you could think of.
54:45
I'm not going to list them all here because it would probably
54:47
take the next two to 5 minutes, but just know we've got the CTO
54:52
of Crews coming to give a little keynote.
54:55
We've got the CEO of you.com come in, we've got Chip, we've
55:00
got Linux, We've got the whole crew that you would expect.
55:04
And I am going to be doing all kinds of extracurricular
55:08
activities that will be fun and maybe a little bit cringe you.
55:14
You may see, hear or see me playing the guitar.
55:17
Just come. It's going to be an awesome
55:19
time. We'd love to have you there.
55:20
And that is again June 25th in San Francisco.
55:25
See you all there.