WEBVTT

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

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

s1
00:00:15.280 --> 00:00:21.119
So, put your model somewhere.

s2
00:00:21.119 --> 00:00:22.640
Put your agent.

s3
00:00:22.640 --> 00:00:24.640
You're hoping for the best.

s4
00:00:24.640 --> 00:00:26.800
You're waiting for something to crash.

s5
00:00:26.800 --> 00:00:27.920
Everything looks good.

s6
00:00:27.920 --> 00:00:30.320
Everything's fine.

s7
00:00:30.320 --> 00:00:46.399
and you see this and that's the problem right in these type of bugs there is no crash there's no warning no error and there's high confidence that's not a quality issue

s8
00:00:46.399 --> 00:00:58.800
because uh right this is something that you don't really know what and how and why um welcome to this talk my name is Ival this is aaf and together

s9
00:00:58.800 --> 00:01:05.199
We're going to take you on a journey of how we ended up fixing those type of bugs.

s10
00:01:05.199 --> 00:01:06.320
A little bit about us.

s11
00:01:06.320 --> 00:01:10.880
We work at AI21, which is an AI research lab.

s12
00:01:10.880 --> 00:01:23.119
We started as a foundation model company, most famously known for Jamba, which is a hybrid architecture between transformers and mamba, which is an SSM state.

s13
00:01:23.119 --> 00:01:32.479
And while we were doing those, while we were training those models, while we were shipping those models into production and had users and we had a lot of workload,

s14
00:01:32.479 --> 00:01:36.880
we got into several interesting bugs.

s15
00:01:36.880 --> 00:01:43.280
And these are the bugs that I think are the hardest to deal with because this is not a quality problem.

s16
00:01:43.280 --> 00:01:53.280
It's not something you can take your research team and try to optimize or solve or make the model be better at something.

s17
00:01:53.280 --> 00:01:54.799
This is an engineering problem.

s18
00:01:54.799 --> 00:02:02.640
This is an issue where that there is high confidence but the output is bad.

s19
00:02:02.640 --> 00:02:19.920
So let's break let's dive deep to the first case what we call the imposter request where just to set up the scene what are we talking about we are talking about how we during the training of our jamba model more specifically we did gpo

s20
00:02:19.920 --> 00:02:31.599
which is a type of RL training and again this is a hybrid model layers of mamba and attention and just to make sure we're all aligned what is the type of a request

s21
00:02:31.599 --> 00:02:34.640
So life lives life of a request.

s22
00:02:34.640 --> 00:02:44.080
So we start with the prompt token tokenization and then in the forward pass we're doing both prefill and then decode.

s23
00:02:44.080 --> 00:02:50.879
After that we finish the forward pass detoenization to go back to text.

s24
00:02:51.120 --> 00:02:59.280
And the thing about the crime here is that it's bad on so many levels but mainly on these three.

s25
00:02:59.280 --> 00:03:01.840
This is what we call the one in thousand gibberish.

s26
00:03:01.840 --> 00:03:15.680
It's not something that will happen in the first 500 or 900 requests, but it will happen in the 10,00 which is rare enough to duplicate it easily but not

s27
00:03:15.680 --> 00:03:15.920
right.

s28
00:03:15.920 --> 00:03:18.959
It's too common to ship it.

s29
00:03:18.959 --> 00:03:19.760
Sorry.

s30
00:03:19.760 --> 00:03:28.400
H also it only happened in VLM not in other not in any other infra framework and it's something which is laid on set.

s31
00:03:28.400 --> 00:03:32.560
It's not something that will happen if you have only few requests.

s32
00:03:32.560 --> 00:03:34.640
You need some sort of workload.

s33
00:03:34.640 --> 00:03:36.000
So it's rare.

s34
00:03:36.000 --> 00:03:39.120
It's light on set and it's very engine specific.

s35
00:03:39.120 --> 00:03:44.000
It's a very very hard task and we had to bring one of our best detectives to handle that.

s36
00:03:44.000 --> 00:03:47.599
So I'll give it to SF to explain how.

s37
00:03:47.599 --> 00:03:48.159
All right.

s38
00:03:48.159 --> 00:03:49.280
Hey guys, thank you.

s39
00:03:49.280 --> 00:03:50.400
Uh thanks you.

s40
00:03:50.400 --> 00:03:55.599
So we're going to start with um trying to reproduce something that was very difficult to reproduce.

s41
00:03:55.599 --> 00:04:04.239
Um and basically um VLM has got a lot of um a lot of flags, a lot of the CLI flags and a lot of knobs you guys can turn and tweak.

s42
00:04:04.239 --> 00:04:12.959
Um and one of the things that helped us understand how to even reproduce it because when we try to reproduce it on the first time just sending prompts here, prompts there a few batches

s43
00:04:12.959 --> 00:04:16.799
um it didn't really help us manage to get gibberish back from our model.

s44
00:04:16.799 --> 00:04:19.840
The model responded back just fine.

s45
00:04:19.840 --> 00:04:25.759
So what we did was is we tried to make it happen in a very very short um amount of time.

s46
00:04:25.759 --> 00:04:29.280
So we'll get a a quick feedback loop when we try to deb debug it.

s47
00:04:29.280 --> 00:04:39.199
So what we did was is we took one of the one of the um most default and most common um um flag that the VLM allows you to play with which is GPU memory utilization

s48
00:04:39.199 --> 00:04:48.479
which basically allows you to to choose how much memory how much GPU memory you want to allocate for your for your weight for your activations and for your KV cache and so on.

s49
00:04:48.479 --> 00:04:52.000
and we reduced it from 90% to 20%.

s50
00:04:52.000 --> 00:05:05.280
And once we did that and then we started uh running a lot of requests uh simultaneously um all of a sudden request number let's say 8 854 suddenly returned gibberish

s51
00:05:05.280 --> 00:05:10.240
and when we did that we uh we sampled all of the batches with temperature zero.

s52
00:05:10.240 --> 00:05:18.000
So we'll be able to deterministically and constantly get the same request to get to return to return and respond with gibberish.

s53
00:05:18.000 --> 00:05:27.680
Um so um like Ival said it happened only in VLM and um we used um in order to uh another way to reproduce it and to understand where the issue really came from.

s54
00:05:27.680 --> 00:05:45.120
We used um uh hugging faces transformers as a baseline since transformers is a very had a very um vanilla and uh plain implementation of our mamba kernels as opposed to VLM which uh all the kernels and all the engine have gone through a lot of changes and modifications

s55
00:05:45.120 --> 00:05:48.400
to support a lot of um cool features that VLM supports.

s56
00:05:48.400 --> 00:05:55.520
So we use transformers as our baseline to understand whether or not there is an issue with our um with our inference or not with the model or not.

s57
00:05:55.520 --> 00:06:11.440
So what we did was we took VLM and we sent uh all of our prompts through VLM and we generated a response all the responses we got and now in our hands we have the response along with the log props because because in VLM you're you able to get your log props out and inspect them.

s58
00:06:11.440 --> 00:06:23.520
Then what we did was we took the full sequence the prompt and the generation and we uh uh passed it over to to hugging faces forward pass but all we did was run just the prefill

s59
00:06:23.520 --> 00:06:34.880
um and then we uh uh samp we took this the logit out of the prefill um um response we ran it through softmax and then we were able to u to compare the divergence

s60
00:06:34.880 --> 00:06:38.000
uh in the distributions of our tokens.

s61
00:06:38.000 --> 00:06:41.440
Um that's a short uh pseudo code of how that look like.

s62
00:06:41.440 --> 00:06:43.840
You can see here that we uh take up the prompt.

s63
00:06:43.840 --> 00:06:46.800
We we run it through gen uh BLM's generate.

s64
00:06:46.800 --> 00:06:50.800
We get the uh the response back along with the uh with the log props.

s65
00:06:50.800 --> 00:06:54.000
We pass it over to to hugging faces forward pass.

s66
00:06:54.000 --> 00:06:56.080
We only run it with prefill.

s67
00:06:56.080 --> 00:07:02.240
Um we've created some function called compute log log props which um runs this is the softmax.

s68
00:07:02.240 --> 00:07:09.599
Um then you calculate the difference between them and then you'll be able to tell the divergence between every one of the tokens log props.

s69
00:07:09.599 --> 00:07:10.000
All right.

s70
00:07:10.000 --> 00:07:19.759
So now that we have the tools in our hand to understand where the issue could maybe come from um we started to look at different suspects in VLM's engine.

s71
00:07:19.759 --> 00:07:24.400
So the first thing we looked at was um the CUDA prefill kernel of Mamba.

s72
00:07:24.400 --> 00:07:30.639
Um we looked at it we inspected all of the all the math that's being done here that's being done there.

s73
00:07:30.639 --> 00:07:37.120
Um, and we looked at the tensor in, and the tensor's out before before we called the prefill and after.

s74
00:07:37.120 --> 00:07:38.960
Everything looks just fine.

s75
00:07:38.960 --> 00:07:45.360
Second thing we did was running Nvidia's compute sanitizer tool to really see if we have any out of bound memory.

s76
00:07:45.360 --> 00:07:49.919
Um, any other memory uh bugs or issues.

s77
00:07:49.919 --> 00:07:51.759
Looked okay to me.

s78
00:07:51.759 --> 00:07:56.240
Um, then what we did was we tried to isolate between the decode kernels and the pre-filled kernels.

s79
00:07:56.240 --> 00:07:59.360
Now, we saw that the pre-filled kernels were working just fine.

s80
00:07:59.360 --> 00:08:03.680
So we tried to not call the decode kernels because in Mamba you're able to do that.

s81
00:08:03.680 --> 00:08:12.639
Um so what we did was um we moved all of our calls and all of our computations to go through the pre through the pre-fill kernel and there you have it.

s82
00:08:12.639 --> 00:08:14.879
The gibbish all of a sudden kind of vanished.

s83
00:08:14.879 --> 00:08:19.120
So we were like okay it's got to be the decode kernels.

s84
00:08:19.120 --> 00:08:25.199
But you know how it is in software you get excited too quickly and then you figure out it's not what happened.

s85
00:08:25.199 --> 00:08:42.080
So what we did was uh we tried to start playing with the with VLM's engine and we kind of needed to go and you know lift the hood up and see what we can do to maybe get a bit better understanding and maybe you know get our hands dirty because BLM didn't really give us more um tools

s86
00:08:42.080 --> 00:08:45.920
to really um debug our kernel and our and our forward pass.

s87
00:08:45.920 --> 00:08:55.839
So once so once uh once a tensor once the request gets all the way to your forward pass and before it goes into your prefill and decode kernels you don't really have any identity.

s88
00:08:55.839 --> 00:08:59.440
You can't really tell what prompt uh is currently being processed.

s89
00:08:59.440 --> 00:09:03.120
It's all just tensors and and numbers and matrices.

s90
00:09:03.120 --> 00:09:12.800
So what we did was is we added to the request ID to some class called forward context that uh that we propagated all the way down to uh to Mamba's

s91
00:09:12.800 --> 00:09:17.040
forward pass just before the the prefill and the decode kernels were called.

s92
00:09:17.040 --> 00:09:28.000
And there we just managed to uh you know have a simple uh if condition with a request ID, the one that gave us gibberish and put a break point there and then we were able to

s93
00:09:28.000 --> 00:09:33.440
to to infer and to really um inspect all the metadata that comes along with it.

s94
00:09:33.440 --> 00:09:43.519
And the second we did that, we saw that the request was um for the first time when it went through the uh through the forward pass, it's actually doing decode before prefill.

s95
00:09:43.519 --> 00:09:59.200
the scheduler um decided that it's um that that this request should be doing decode before prefill and as Uval said earlier in our in the life cycle of a prompt a prompt should first be uh going through prefill

s96
00:09:59.200 --> 00:10:10.880
and then decode and what happens was is that when when in Mamba um you run a request first uh with a with decode first after a long a lot of other requests

s97
00:10:10.880 --> 00:10:21.440
were already computed the state was already kind of um overused and we were using uh the data and the computations of stale requests, requests that came before it.

s98
00:10:21.440 --> 00:10:28.880
So now we were actually running decode on on previous requests and that kind of generated gibberish for us.

s99
00:10:28.880 --> 00:10:34.880
So the kernels weren't doing the wrong thing, they were called at the wrong time for the wrong requests.

s100
00:10:34.880 --> 00:10:38.079
And now and why did it matter only for Mamba?

s101
00:10:38.079 --> 00:10:47.040
The reason is was is that in uh in attention uh when you write the tokens KV you write it you write the tokens cavies before you actually you read it.

s102
00:10:47.040 --> 00:10:58.320
So even if you have stale data it's being overwritten but for mamba as uh as I said when uh when you first go through the decode kernels you first read the state and then you compute over it.

s103
00:10:58.320 --> 00:11:07.440
So what happens was is you just use over you use stale data um um when you do the when you do a decode and the fix was relatively simple.

s104
00:11:07.440 --> 00:11:23.920
Well, we just needed to make sure that what we do is that when a request first when the when theuler first classifies a request, it's got to it's got to make sure that um that um that it sets that that if a that if that if you sees a request that's whose tokens were never been computed.

s105
00:11:23.920 --> 00:11:30.720
Um we and and they're zero to mark them as to mark them as uh as prefill as as prefill.

s106
00:11:30.720 --> 00:11:38.800
So the so when they get to the forward pass, they'll actually just be um uh used for prefill and not decode and not u and not chunked.

s107
00:11:38.800 --> 00:11:41.519
You can see that it was merged after some time.

s108
00:11:41.519 --> 00:11:45.519
Um and that really leads us to and then we thought everything was fixed, right?

s109
00:11:45.519 --> 00:11:47.760
We thought everything was fixed and there you have it.

s110
00:11:47.760 --> 00:11:49.279
No more issues.

s111
00:11:49.279 --> 00:12:01.360
But that but that was almost the case because after a little bit of time it gets us to case number two which um surfaced another issue that we've faced in our in our RL and our inference.

s112
00:12:01.360 --> 00:12:10.959
Um so we ran RL and and when while we were running um our our trainings our post trainings um and we looked at our evaluations and all of our benchmarks

s113
00:12:10.959 --> 00:12:15.279
we saw that we had some log prop spikes between the rollout and the FSTP step.

s114
00:12:15.279 --> 00:12:19.040
Uh so and that was before any weight app update.

s115
00:12:19.040 --> 00:12:23.519
So some weights uh same inputs and the two log props should be identical.

s116
00:12:23.519 --> 00:12:25.040
Now they weren't.

s117
00:12:25.040 --> 00:12:33.920
Um we saw that every 12 step cons constantly um there was a log prop spike and that was kind of weird.

s118
00:12:33.920 --> 00:12:36.399
Now what would you guys do right?

s119
00:12:36.399 --> 00:12:39.360
What can what what's the what's the first thing to do here?

s120
00:12:39.360 --> 00:12:44.560
So we wanted to find some lever that changes how things fail and not just how much they fail.

s121
00:12:44.560 --> 00:12:53.920
We want to see how much um we want to tweak some knobs that that don't just tell us, hey, this error um this error is very very bad.

s122
00:12:53.920 --> 00:12:57.839
This error happens uh this many times or or [clears throat] and so and so on.

s123
00:12:57.839 --> 00:13:07.839
We wanted to tweak some knobs that kind of tell us that once we tweak that knob, um we understand how it's wired to uh anything in the in in VLM's

s124
00:13:07.839 --> 00:13:08.399
engine.

s125
00:13:08.399 --> 00:13:12.959
And so we'll be able to specifically go and and debug that specific part.

s126
00:13:12.959 --> 00:13:18.079
Um that's some uh some some cool meme that you all wanted to put in.

s127
00:13:18.079 --> 00:13:26.320
Um so what we did was is we decided to um to increase the the the amount of rollouts per prompt.

s128
00:13:26.320 --> 00:13:35.839
Uh since we saw uh in our default RL engine we have um eight rollouts per prompt and we saw that it happened con deterministically every 12 steps.

s129
00:13:35.839 --> 00:13:40.240
We decided okay let's try to tweak it up a bit and uh and increase the amount of rollouts per prompt.

s130
00:13:40.240 --> 00:13:45.839
So we started doubling it from say eight to 16 to 64 32 and 128.

s131
00:13:45.839 --> 00:14:03.360
And you can see here that it's almost um almost um there's a pattern here that that the more we increased it the closer um it happened because what we wanted to achieve here we wanted to try to reproduce the issue as fast as possible so we'd have a faster debug uh debug loop feedback loop.

s132
00:14:03.360 --> 00:14:13.360
So when we when we when we ran it on 128 rollouts per prompt, it happened immediately on step one and we didn't have to wait for step 12 and step 24 and so on.

s133
00:14:13.360 --> 00:14:20.240
Now you might think, okay, so you guys played with the with the GPU memory utilization before you you tweaked it, you decreased it.

s134
00:14:20.240 --> 00:14:25.040
It looks like, you know, when you test on pressure, um, it really surface things up.

s135
00:14:25.040 --> 00:14:27.360
So So we we thought that as well.

s136
00:14:27.360 --> 00:14:34.959
And when we reduced the GPU memory from zero from 0.9 to 0.2, two, it actually caused the issue to go away.

s137
00:14:34.959 --> 00:14:37.519
So, we actually pulled the wrong lever here.

s138
00:14:37.519 --> 00:14:47.920
Um, and the reason is is because we noticed that Mamba kernels um used um in 32 unsigned in 32 index pattern uh pointer.

s139
00:14:47.920 --> 00:14:55.760
So, once the offset went past some you know 4 billion um uh numbers, it wrapped around instead of throwing an error.

s140
00:14:55.760 --> 00:15:04.480
Um so uh when we shrank the GPU memory, VLM allocated a small state buffer and um and the cache index never got large to hit that slot.

s141
00:15:04.480 --> 00:15:10.079
So we were so we were just not reaching um far enough for the buffer to trigger an overflow.

s142
00:15:10.079 --> 00:15:13.519
So again the fix was rather simple.

s143
00:15:13.519 --> 00:15:26.320
All we needed to do was just change one word, one one data type variable um from u in32 to size t which basically means for most modern architectures hardware

s144
00:15:26.320 --> 00:15:34.320
architectures size t would mean to uh it would be now changed to unsigned 64 uh bit and that's a very large number.

s145
00:15:34.320 --> 00:15:39.279
We didn't we never reached that number and that overflow now never happened.

s146
00:15:39.279 --> 00:15:45.760
So what we can see here is that we had um two scenes and one criminal.

s147
00:15:45.760 --> 00:15:48.639
Um both kind of you know they had similar symptoms.

s148
00:15:48.639 --> 00:15:55.519
Both had silent gibberish and and silent log prop spikes which also kind of pro sometimes generated gibberish.

s149
00:15:55.519 --> 00:15:58.320
They were both around the mamba state cache.

s150
00:15:58.320 --> 00:16:08.399
Um they were both surfaced by memory pressure whether it was for worse or for the best and uh both found via log props um forensics.

s151
00:16:08.480 --> 00:16:18.560
um stateful inference inference systems don't fail loudly they lie to you confidently I mean obviously sometimes you get crash you get out of bounds errors you get other you know

s152
00:16:18.560 --> 00:16:29.360
exceptions and so on but sometimes there are some errors that don't surface up and you don't get a trace log you don't get anything you have to go and dig and understand why things happen

s153
00:16:29.360 --> 00:16:41.600
um so if there some takeaways to take from this um presentation is build a log props comparison script if you need to compare your quality you need to compare it to understand whether you modelize the issues or not.

s154
00:16:41.600 --> 00:16:48.880
Um log props comparison script with a baseline of some other inference framework that you have or built is always great.

s155
00:16:48.880 --> 00:16:52.000
Um reproducing underression constrain memory.

s156
00:16:52.000 --> 00:16:58.959
Uh crank the scale up, play with other knobs that the inference framework gives you and really try to understand where the issue comes from.

s157
00:16:58.959 --> 00:17:05.280
Uh look for what moves um the failure shape, the timing, the space and the location.

s158
00:17:05.280 --> 00:17:15.360
And when things don't really have identity uh thread identity through and what it also I want you to take from this and don't be afraid to even you know for complex

s159
00:17:15.360 --> 00:17:28.079
uh systems like VLM or any other um complex framework don't be afraid to go dig in the code get your hands dirty um sometimes you know model languages LLMs are

s160
00:17:28.079 --> 00:17:35.360
um they might tell you how things work but you know without you seeing it in your own eyes getting your hands dirty you won't get full understanding of what's going on.

s161
00:17:35.360 --> 00:17:36.960
Um, thank you.

s162
00:17:36.960 --> 00:17:38.400
You guys can add us on LinkedIn.

s163
00:17:38.400 --> 00:17:43.520
Scan the QR code to read the actual blog that we've published with this finding.

s164
00:17:43.520 --> 00:17:46.352
Um, yeah, that's it.

s165
00:17:46.352 --> 00:17:48.352
[applause]