Two Bugs That Hid in Plain Sight: A vLLM Debugging Detective Story — Asaf Gardin & Yuval Belfer

AI Engineer · 18 min · 165 sentences · from YouTube's caption track

Each timecode opens YouTube at the start of that sentence. Line anchors (#s42) are the cue ids in the WebVTT, and every line carries its start and end seconds. All transcripts has every talk, and the whole corpus as one file.

00:15So, put your model somewhere.
00:21Put your agent.
00:22You're hoping for the best.
00:24You're waiting for something to crash.
00:26Everything looks good.
00:27Everything's fine.
00:30and 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
00:46because 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
00:58We're going to take you on a journey of how we ended up fixing those type of bugs.
01:05A little bit about us.
01:06We work at AI21, which is an AI research lab.
01:10We 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.
01:23And 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,
01:32we got into several interesting bugs.
01:36And these are the bugs that I think are the hardest to deal with because this is not a quality problem.
01:43It's not something you can take your research team and try to optimize or solve or make the model be better at something.
01:53This is an engineering problem.
01:54This is an issue where that there is high confidence but the output is bad.
02:02So 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
02:19which 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
02:31So life lives life of a request.
02:34So we start with the prompt token tokenization and then in the forward pass we're doing both prefill and then decode.
02:44After that we finish the forward pass detoenization to go back to text.
02:51And the thing about the crime here is that it's bad on so many levels but mainly on these three.
02:59This is what we call the one in thousand gibberish.
03:01It'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
03:15right.
03:15It's too common to ship it.
03:18Sorry.
03:19H also it only happened in VLM not in other not in any other infra framework and it's something which is laid on set.
03:28It's not something that will happen if you have only few requests.
03:32You need some sort of workload.
03:34So it's rare.
03:36It's light on set and it's very engine specific.
03:39It's a very very hard task and we had to bring one of our best detectives to handle that.
03:44So I'll give it to SF to explain how.
03:47All right.
03:48Hey guys, thank you.
03:49Uh thanks you.
03:50So we're going to start with um trying to reproduce something that was very difficult to reproduce.
03:55Um 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.
04:04Um 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
04:12um it didn't really help us manage to get gibberish back from our model.
04:16The model responded back just fine.
04:19So what we did was is we tried to make it happen in a very very short um amount of time.
04:25So we'll get a a quick feedback loop when we try to deb debug it.
04:29So 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
04:39which 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.
04:48and we reduced it from 90% to 20%.
04:52And 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
05:05and when we did that we uh we sampled all of the batches with temperature zero.
05:10So we'll be able to deterministically and constantly get the same request to get to return to return and respond with gibberish.
05:18Um 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.
05:27We 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
05:45to support a lot of um cool features that VLM supports.
05:48So 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.
05:55So 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.
06:11Then 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
06:23um 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
06:34uh in the distributions of our tokens.
06:38Um that's a short uh pseudo code of how that look like.
06:41You can see here that we uh take up the prompt.
06:43We we run it through gen uh BLM's generate.
06:46We get the uh the response back along with the uh with the log props.
06:50We pass it over to to hugging faces forward pass.
06:54We only run it with prefill.
06:56Um we've created some function called compute log log props which um runs this is the softmax.
07:02Um 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.
07:09All right.
07:10So 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.
07:19So the first thing we looked at was um the CUDA prefill kernel of Mamba.
07:24Um we looked at it we inspected all of the all the math that's being done here that's being done there.
07:30Um, and we looked at the tensor in, and the tensor's out before before we called the prefill and after.
07:37Everything looks just fine.
07:38Second thing we did was running Nvidia's compute sanitizer tool to really see if we have any out of bound memory.
07:45Um, any other memory uh bugs or issues.
07:49Looked okay to me.
07:51Um, then what we did was we tried to isolate between the decode kernels and the pre-filled kernels.
07:56Now, we saw that the pre-filled kernels were working just fine.
07:59So we tried to not call the decode kernels because in Mamba you're able to do that.
08:03Um 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.
08:12The gibbish all of a sudden kind of vanished.
08:14So we were like okay it's got to be the decode kernels.
08:19But you know how it is in software you get excited too quickly and then you figure out it's not what happened.
08:25So 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
08:42to really um debug our kernel and our and our forward pass.
08:45So 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.
08:55You can't really tell what prompt uh is currently being processed.
08:59It's all just tensors and and numbers and matrices.
09:03So 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
09:12forward pass just before the the prefill and the decode kernels were called.
09:17And 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
09:28to to infer and to really um inspect all the metadata that comes along with it.
09:33And 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.
09:43the 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
09:59and 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
10:10were 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.
10:21So now we were actually running decode on on previous requests and that kind of generated gibberish for us.
10:28So the kernels weren't doing the wrong thing, they were called at the wrong time for the wrong requests.
10:34And now and why did it matter only for Mamba?
10:38The 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.
10:47So 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.
10:58So 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.
11:07Well, 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.
11:23Um we and and they're zero to mark them as to mark them as uh as prefill as as prefill.
11:30So 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.
11:38You can see that it was merged after some time.
11:41Um and that really leads us to and then we thought everything was fixed, right?
11:45We thought everything was fixed and there you have it.
11:47No more issues.
11:49But 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.
12:01Um 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
12:10we saw that we had some log prop spikes between the rollout and the FSTP step.
12:15Uh so and that was before any weight app update.
12:19So some weights uh same inputs and the two log props should be identical.
12:23Now they weren't.
12:25Um we saw that every 12 step cons constantly um there was a log prop spike and that was kind of weird.
12:33Now what would you guys do right?
12:36What can what what's the what's the first thing to do here?
12:39So we wanted to find some lever that changes how things fail and not just how much they fail.
12:44We 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.
12:53This error happens uh this many times or or [clears throat] and so and so on.
12:57We 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
13:07engine.
13:08And so we'll be able to specifically go and and debug that specific part.
13:12Um that's some uh some some cool meme that you all wanted to put in.
13:18Um so what we did was is we decided to um to increase the the the amount of rollouts per prompt.
13:26Uh 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.
13:35We decided okay let's try to tweak it up a bit and uh and increase the amount of rollouts per prompt.
13:40So we started doubling it from say eight to 16 to 64 32 and 128.
13:45And 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.
14:03So 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.
14:13Now you might think, okay, so you guys played with the with the GPU memory utilization before you you tweaked it, you decreased it.
14:20It looks like, you know, when you test on pressure, um, it really surface things up.
14:25So So we we thought that as well.
14:27And when we reduced the GPU memory from zero from 0.9 to 0.2, two, it actually caused the issue to go away.
14:34So, we actually pulled the wrong lever here.
14:37Um, and the reason is is because we noticed that Mamba kernels um used um in 32 unsigned in 32 index pattern uh pointer.
14:47So, once the offset went past some you know 4 billion um uh numbers, it wrapped around instead of throwing an error.
14:55Um 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.
15:04So we were so we were just not reaching um far enough for the buffer to trigger an overflow.
15:10So again the fix was rather simple.
15:13All 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
15:26architectures size t would mean to uh it would be now changed to unsigned 64 uh bit and that's a very large number.
15:34We didn't we never reached that number and that overflow now never happened.
15:39So what we can see here is that we had um two scenes and one criminal.
15:45Um both kind of you know they had similar symptoms.
15:48Both had silent gibberish and and silent log prop spikes which also kind of pro sometimes generated gibberish.
15:55They were both around the mamba state cache.
15:58Um 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.
16:08um 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
16:18exceptions 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
16:29um 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.
16:41Um log props comparison script with a baseline of some other inference framework that you have or built is always great.
16:48Um reproducing underression constrain memory.
16:52Uh crank the scale up, play with other knobs that the inference framework gives you and really try to understand where the issue comes from.
16:58Uh look for what moves um the failure shape, the timing, the space and the location.
17:05And 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
17:15uh 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
17:28um 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.
17:35Um, thank you.
17:36You guys can add us on LinkedIn.
17:38Scan the QR code to read the actual blog that we've published with this finding.
17:43Um, yeah, that's it.
17:46[applause]