WEBVTT

NOTE Sentence-level transcript of https://www.youtube.com/watch?v=tNAV259UhCM

NOTE One cue per sentence. Cue ids are the line anchors on /transcripts/tNAV259UhCM.html. A cue ends where the next begins, or 2 s after its last word.

s1
00:00:01.309 --> 00:00:03.309
[music]

s2
00:00:13.120 --> 00:00:14.639
Hello everyone.

s3
00:00:14.639 --> 00:00:16.720
Thank you for joining me here.

s4
00:00:16.720 --> 00:00:24.560
Um yeah, today I'm going to talk about ray actors, vision tokens, and the gill.

s5
00:00:24.640 --> 00:00:31.039
um engineering a multimodal SFT data pipeline that keeps the GPUs busy.

s6
00:00:31.039 --> 00:00:38.000
So there's a lot of talk about you know how to optimize training paradigms.

s7
00:00:38.000 --> 00:00:44.879
Um but I want to give a different perspective on the whole ecosystem of what it takes to get a model trained.

s8
00:00:44.879 --> 00:00:52.079
Um and so the question I'm going to be asking here is does your training loop really depend on the GPUs right now?

s9
00:00:52.079 --> 00:00:53.920
Is it blocked on the GPUs?

s10
00:00:53.920 --> 00:01:04.320
And I think we're going to learn that we're going to see in this experiment that the hardest part of optimizing multimodal train throughput is not necessarily kernel optimization.

s11
00:01:04.320 --> 00:01:09.600
It is making sure that the G GPUs are constantly fed with data.

s12
00:01:09.600 --> 00:01:13.200
So let's go over the overview of what this topic what this talk will be about.

s13
00:01:13.200 --> 00:01:23.840
Anyway, so we will establish a baseline multimodal training pipeline um and then we will encounter bottlenecks and then address them one at a time.

s14
00:01:23.840 --> 00:01:35.920
So loosely they are one concurrency how fast can you process one image before feeding it to the model two pre-fetching how can we make sure that we have data on demand

s15
00:01:35.920 --> 00:01:53.200
three um how can we minimize data copies across different processes and four f additional framework optimizations um so just tweaking your bells and whistles um yeah so let's set up our premise

s16
00:01:53.200 --> 00:02:02.079
We all know that GPUs incur expensive costs and are generally want to maximize their utilization while we have them at hand.

s17
00:02:02.079 --> 00:02:15.520
Multimodal processing is different than text only um because there's a lot of CPU inensive operations that have to happen before you actually feed the model um the data.

s18
00:02:15.520 --> 00:02:31.920
So whether that is um you know loading resizing or normalizing images using like a library like pill for audio you do a lot of loading and cropping you can use like numpy arrays and then videos are some sort of combination of both

s19
00:02:31.920 --> 00:02:42.879
um and then it's also true that usually your training clusters have a lot more CPUs than GPUs at hand yet the GPU is what is starving in this pipeline.

s20
00:02:42.879 --> 00:02:53.360
So today we're going to introduce this equation called like weight time ratio which is generally I formulated as like the time it takes to prep your data and transfer it to the GPUs.

s21
00:02:53.360 --> 00:02:59.360
Um and then we divide all of that by the time it takes to do both of those and the train as well.

s22
00:02:59.360 --> 00:03:04.080
So basically the denominator is like your total pipeline.

s23
00:03:04.080 --> 00:03:13.599
Um and we're going to see that during our run we will go from a utilization of 15% to utilization of 90%.

s24
00:03:13.599 --> 00:03:29.120
So our baseline has like 85% of its spends 85% of its time um backed on data acquisition whereas after we do these optimizations less than 10 and so for the scope of this talk

s25
00:03:29.120 --> 00:03:30.640
GPUs are not the goal.

s26
00:03:30.640 --> 00:03:35.120
It's making sure that the data is never the bottleneck is the goal.

s27
00:03:35.120 --> 00:03:40.400
And so here's a little bit of background as to like what our stack will look like.

s28
00:03:40.400 --> 00:03:53.519
um we kind of take like a Quen 3VL based model um kind of try to train it on the paradigms of chat VQA and like multi-term VLM uh conversations.

s29
00:03:53.519 --> 00:04:08.080
So to recap before the GPUs do any math um we have to load the data we have to transform the data and then we have to we have to transport it to the GPUs.

s30
00:04:08.720 --> 00:04:12.480
Okay, this is a this is a overview of what our data actually looks like.

s31
00:04:12.480 --> 00:04:15.200
So, we will be using JSON L files.

s32
00:04:15.200 --> 00:04:21.519
Um, and then images are linked relatively uh to where the example lives.

s33
00:04:21.519 --> 00:04:25.280
And then we're going to assume that our data is just stored in S3.

s34
00:04:25.280 --> 00:04:35.600
Um, in this example, for example, it seems like that image 33471.jpeg seems to be a picture of like a blue bus.

s35
00:04:35.600 --> 00:04:37.280
Um, yeah.

s36
00:04:37.280 --> 00:04:39.120
So this is our system layout.

s37
00:04:39.120 --> 00:04:46.160
On the left we have a data generator which is responsible for the data loading, data transforming and data transporting.

s38
00:04:46.160 --> 00:04:51.919
Um and the that's like labeled as like fan out to the mega megatron data parallel ranks.

s39
00:04:51.919 --> 00:04:58.240
Um so the parallel ranks receive partitions of the data as decided by the coordinator.

s40
00:04:58.240 --> 00:05:03.040
Um it collates it performs forward pass and backwards pass on this data.

s41
00:05:03.040 --> 00:05:06.880
So, Megatron takes care of optimizing the model weights across the DP ranks.

s42
00:05:06.880 --> 00:05:15.440
And we're going to assume that it's like a it's kind of it is kind of working as a black box.

s43
00:05:15.440 --> 00:05:19.199
Um, so let's start with the naive baseline path.

s44
00:05:19.199 --> 00:05:23.680
So for each example, we fetch the images from object store.

s45
00:05:23.680 --> 00:05:31.919
Um, we kind of like we tokenize it, we resize it, we write resize the images, normalize them, and then feed them to the model.

s46
00:05:31.919 --> 00:05:46.560
Um yeah and then when we do that we can see that there's serial operations in the red which is namely getting all the images one by one and then processing them one by one.

s47
00:05:46.560 --> 00:05:59.919
So if your if your trainer is asking for a batch size of like 256 then you have to serially process all these images and then we can see that our train time is somewhere between 15 to 20 seconds.

s48
00:05:59.919 --> 00:06:05.199
uh however the data the time spent to actually fetch the data is more than a 100 seconds.

s49
00:06:05.199 --> 00:06:09.840
So that's why we result in only 15% utilization of our GPUs.

s50
00:06:09.840 --> 00:06:20.639
Um and then you can even see a profiling flame graph that we spend most of the time doing S3 loads and then you can see a long a little tail at the end about like loading the actual image on the pill

s51
00:06:20.639 --> 00:06:22.560
uh library.

s52
00:06:22.560 --> 00:06:26.319
So okay, solution one, let's try to add some concurrency to this.

s53
00:06:26.319 --> 00:06:35.440
Instead of those red singular uh you know serial processes, uh let's try to parallelize it by kind of having like a worker pool.

s54
00:06:35.440 --> 00:06:41.680
So when you request 256 examples, maybe let's just process all of those at once.

s55
00:06:41.680 --> 00:06:48.240
Um yeah, so for example, for example, we have two costs.

s56
00:06:48.240 --> 00:06:50.880
One is fetching the images, two is processing.

s57
00:06:50.880 --> 00:06:55.600
So that's like decoding, vision pre-processing, and then tokenizing your text.

s58
00:06:55.600 --> 00:06:58.960
Um, samples are very are independent.

s59
00:06:58.960 --> 00:07:02.639
So this just means it's like embarrassingly parallelizable.

s60
00:07:02.639 --> 00:07:06.560
Um, and so yeah, you have a design decisions here.

s61
00:07:06.560 --> 00:07:09.280
What forms of concurrency do you want?

s62
00:07:09.280 --> 00:07:18.000
So there's asynchronous programming, uh, which is suitable for like IO work such as making SQL queries or like performing network requests.

s63
00:07:18.000 --> 00:07:27.680
There's threads uh which you can use for like CPU intensive uh workloads that relinquish the Python gill like the global interpreter lock.

s64
00:07:27.680 --> 00:07:32.560
Um so this it's not straightforward what operations actually release the lock.

s65
00:07:32.560 --> 00:07:44.319
So you have to do a little more of your investigation but common uses are like using hugging face transformers tokenizers using the pill transforms API or like some pandas operations

s66
00:07:44.319 --> 00:07:54.720
processes um are more heavyweight but also provide total isolation between your workload so you can leverage multiple cores on your machine as you finish your work.

s67
00:07:54.720 --> 00:08:03.120
Um, yeah, doing a deep dive on the right solution would warrant its own talk, but like there's plenty of online resources that can help you decide for your use case.

s68
00:08:03.120 --> 00:08:07.199
What's important for our solution is that we need to match our bottleneck.

s69
00:08:07.199 --> 00:08:10.319
Not simply going with a lever that's worked for us before.

s70
00:08:10.319 --> 00:08:20.560
Um, so we'll choose async IO for fetching images here and then use ray actors as our processors uh as our process multipprocessing library.

s71
00:08:20.560 --> 00:08:30.400
Um, ray actors are essentially the same as like general processes, but they offer us APIs to communicate across workers that just make it so that we don't have to reinvent the wheel.

s72
00:08:30.400 --> 00:08:34.240
So, yeah, your process is only as slow as your slowest link.

s73
00:08:34.240 --> 00:08:37.599
Um, and we can see that our weight time ratio has gone down.

s74
00:08:37.599 --> 00:08:49.680
Um, on the second example, which is in the pink, uh, you can see that the time spent fetching data per batch has gone basically, yeah, around like around like 40 40 seconds.

s75
00:08:49.680 --> 00:08:53.519
Um, so this is this gives us a 25% speed up.

s76
00:08:53.519 --> 00:09:01.920
However, you yeah that's the as you can see the there's a huge gap that starts generating especially as we're going through many many steps.

s77
00:09:01.920 --> 00:09:12.959
Um, so can we somehow limit our bottleneck of waiting on the trainer to request before provisioning our worker pool to fetch the data?

s78
00:09:12.959 --> 00:09:16.480
So can we just like preemptively fetch our data?

s79
00:09:16.480 --> 00:09:28.399
Um so we came up with solution two which was that we have a background thread that fills up a queue ahead of time and it's always stays ahead of the consumer which means

s80
00:09:28.399 --> 00:09:39.680
that you your producer is producing examples at a multiplier of your tra of your consumer and so the trainer is never bottlenecked on actually requesting the data.

s81
00:09:39.680 --> 00:09:43.600
When it requests data it's it is available on demand.

s82
00:09:43.600 --> 00:09:45.440
There's a caveat to prefetching.

s83
00:09:45.440 --> 00:09:56.560
It makes your model and data loader checkpointing and resumption like a little intricate of a process and it can also suffer from a cold start period in which your queue is actually getting filled up

s84
00:09:56.560 --> 00:09:58.720
uh so it can be ready for training.

s85
00:09:58.720 --> 00:10:07.360
Um but uh if you work around those those details I think pre-fetching is a is a great idea.

s86
00:10:07.839 --> 00:10:16.720
Um yeah, so I think making sure your producer is always ahead of your consumer will be will pay great dividends for you here.

s87
00:10:16.720 --> 00:10:26.800
Um unsurprisingly, we achieved close to 0 seconds of data fetching and processing per batch because we've already fetched data and processed it way before the trainer even needs it.

s88
00:10:26.800 --> 00:10:33.120
Um and so we've effectively solved both the data load and the transform bottlenecks.

s89
00:10:33.120 --> 00:10:36.000
So after these two mitigations, we thought we were done.

s90
00:10:36.000 --> 00:10:51.519
throughput up GPUs fed however we observe that the data transport costs still linger in fact they've gone up so every sample transfers a large amount of like image arrays across process boundaries multiple times

s91
00:10:51.519 --> 00:11:03.040
every sample is pickled on the worker copied to the coordinator and then repickled onto the trainer per sample per step we have a mem copy that nobody really asked for

s92
00:11:03.040 --> 00:11:08.399
so what h how this manifest is that your data generator can often oo.

s93
00:11:08.399 --> 00:11:12.000
Um, and so multimodal tensors are huge.

s94
00:11:12.000 --> 00:11:13.760
They're megabytes per sample.

s95
00:11:13.760 --> 00:11:23.760
Um, and so by default, every sample is returned by value, which means this that these expensive arrays are copied over from process to process.

s96
00:11:23.760 --> 00:11:27.440
So a fix for this, you can come up with an object store.

s97
00:11:27.440 --> 00:11:32.000
Um, this is a nice convenience method of that array provides you.

s98
00:11:32.000 --> 00:11:36.079
Each sample carries a object reference which is just a pointer to the data.

s99
00:11:36.079 --> 00:11:38.320
Um and then the driver just holds the pointer.

s100
00:11:38.320 --> 00:11:41.279
The tensor never round trips to the head node.

s101
00:11:41.279 --> 00:11:50.399
Um and the Megatron data parallel ranks they can just directly call the race store to fetch the data as they need.

s102
00:11:50.880 --> 00:11:58.079
So with every choice of your parallelism and concurrency solution, you always have to think about the commu added communication costs.

s103
00:11:58.079 --> 00:12:00.800
So plan wisely around that.

s104
00:12:00.800 --> 00:12:01.519
Okay.

s105
00:12:01.519 --> 00:12:01.760
Yeah.

s106
00:12:01.760 --> 00:12:14.959
And we were able to eliminate an extraneous like put object uh kind of call that our prefetcher solution had that we did not have to deal with by using the ray store.

s107
00:12:14.959 --> 00:12:22.320
And we can see that our communication costs go down and the ratio of the time spent waiting also goes down.

s108
00:12:22.320 --> 00:12:25.440
So to recap, we took a serial baseline.

s109
00:12:25.440 --> 00:12:28.320
We added concurrency for data loading and processing.

s110
00:12:28.320 --> 00:12:34.800
We added pre-fetching to stay ahead of training requests and we added the race store to minimize data transfer.

s111
00:12:34.800 --> 00:12:39.760
We went from a baseline weight ratio of 85% to 20%.

s112
00:12:39.760 --> 00:12:51.519
And you can see that the saw to uh pattern that you see in the top graph is alleviated uh to a large extent which means that our utilization of the GPUs has gone up.

s113
00:12:51.600 --> 00:13:00.639
So when we're scaling up and now that we have like a good multimodal training pipeline, let's try to scale it up to like a production run.

s114
00:13:00.639 --> 00:13:07.040
While doing that, we observed that our weight time ratio for the same pipeline went up when all we did was scale some numbers.

s115
00:13:07.040 --> 00:13:09.839
So we just added four times the data parallel rings.

s116
00:13:09.839 --> 00:13:15.839
We added like 200 data streams and then batch size is x by4.

s117
00:13:15.839 --> 00:13:29.040
Um so it turns out that since the data generator and the worker pool are colllocated for fast data transfer what happens is that there's a bottleneck on the network interface card

s118
00:13:29.040 --> 00:13:31.360
of that of that node.

s119
00:13:31.360 --> 00:13:44.399
Uh which means that while we're packing all of this working and processing into specific nodes that their network throughput is not able to keep up with the demand of the training pipeline.

s120
00:13:45.200 --> 00:13:56.240
Um so at this point we did a little bit of digging into the frameworks that we are using and we saw that by default Ray schedules workers on this it prefers to schedule

s121
00:13:56.240 --> 00:14:04.480
workers on the same nodes and colllocate workers which means that it's easy to bottleneck one machine before you move to another machine.

s122
00:14:04.480 --> 00:14:16.399
So what we did was consider a possibility where we disagregate this whole process where workers are spread across multiple nodes.

s123
00:14:16.399 --> 00:14:27.199
Um so they can do their own image processing and storing and in a lazy just in time manner are retrieved and used by our DP rings.

s124
00:14:27.519 --> 00:14:36.560
So yeah to re and then so these are the two parameter two flags that we kind of tuned one is a zero copy retrieval and second is a spread

s125
00:14:36.560 --> 00:14:38.160
scheduling mechanism.

s126
00:14:38.160 --> 00:14:48.240
Um neither of these actually helped in a small scale experiment but together they had um a synergistic effect in a large scale experiment.

s127
00:14:48.240 --> 00:14:55.279
Um yeah so config can lose a config can do poorly at a small scale but can do really good at a larger scale.

s128
00:14:55.279 --> 00:15:01.120
So you should always think about how to ablate across your your scaling ladders.

s129
00:15:01.120 --> 00:15:09.120
And these were the two flags that ended up making a huge difference uh in how how fast our data processing pipeline was.

s130
00:15:09.120 --> 00:15:22.079
And so now this is the final pipeline that we've accomplished a concurrent eager um and a data access pattern aware system that shares the training across loads and enables that GAPU

s131
00:15:22.079 --> 00:15:26.480
stay warm while providing smoothness and stability in training.

s132
00:15:26.480 --> 00:15:31.600
Overall we achieved a 50% additional throughput at scale just from these two flags.

s133
00:15:31.600 --> 00:15:37.120
It's the same knob opposite verdict set by the bottleneck.

s134
00:15:37.120 --> 00:15:41.839
So these are the learnings we get from this talk if nothing else.

s135
00:15:41.839 --> 00:15:46.639
Um GPUs are expensive but sometimes the biggest battles happen outside of the GPU.

s136
00:15:46.639 --> 00:15:50.560
So first ask yourself are you CPU bound or GPU bound?

s137
00:15:50.560 --> 00:15:53.040
Are you IO bound or computebound?

s138
00:15:53.040 --> 00:15:56.399
And then solving one bottleneck often exposes another.

s139
00:15:56.399 --> 00:16:04.399
So first it was throughput for us then it became weight time then it became memory then it became retrieval cost.

s140
00:16:04.399 --> 00:16:06.639
Um, classic moves are not the end.

s141
00:16:06.639 --> 00:16:12.000
You know, parallelism and prefetching matter and I think a lot of content has like covered this.

s142
00:16:12.000 --> 00:16:20.240
Uh, but so do the semantics of the frameworks that you're using and the and the constraints that you're dealing with and scale always changes the answer.

s143
00:16:20.240 --> 00:16:33.680
Tuning parameters like zero copy waste um and spreading workers become important only at larger scales and they may not manifest as obvious improvements at smaller levels.

s144
00:16:33.680 --> 00:16:37.600
Overall, profiling your baseline helps you move the needle.

s145
00:16:37.600 --> 00:16:40.639
Um, yeah, thank you everyone for attending the talk.

s146
00:16:40.639 --> 00:16:44.000
Uh, if you have any questions, I'm happy to answer them.

s147
00:16:44.000 --> 00:16:46.000
Thanks.