Two Bugs That Hid in Plain Sight: A vLLM Debugging Detective Story — Asaf Gardin & Yuval Belfer
AAI Engineer
Computing/SoftwareInternet Technology
Transcript
00:00:00So, you put your model somewhere, put your agent, you're hoping for the best,
00:00:24you're waiting for something to crash, everything looks good, everything's fine.
00:00:29And you see this.
00:00:33And that's the problem, right, in these type of bugs.
00:00:37There is no crash, there's no warning, no error, and there's high confidence.
00:00:43That's not a quality issue.
00:00:46'Cause, right, this is something that you don't really know what and why.
00:00:53Welcome to this talk. My name is Yuval. This is Asaf.
00:00:58Together, we're gonna take you on a journey of how we ended up fixing those type of bugs.
00:01:04A little bit about us.
00:01:06We work at AI21, which is an AI research lab.
00: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.
00:01:22And 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 workload, we got into several interesting bugs.
00:01:36And these are the bugs that I think are the hardest to deal with.
00:01:41Because this is not a quality problem.
00: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.
00:01:53So, this is an engineering problem. This is an issue where there is high confidence, but the output is bad.
00:02:02So, 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?
00:02:12We are talking about how we, during the training of our Jamba model, more specifically, we did GRPO, which is a type of RL training.
00:02:22And again, this is a hybrid model.
00:02:24Layers of Mamba and attention.
00:02:27And just to make sure we're all aligned, what is the type of a request?
00:02:31So, life of a request.
00:02:34So, we start with prompt, tokenization, and then in the forward pass, we're doing both prefill and then decode.
00:02:44After that, we finish the forward pass, detokenization to go back to text.
00:02:51And the thing about the crime here is that it's bad on so many levels, but mainly on these three.
00:02:59This is what we call the one in thousand gibberish.
00:03:01It's not something that will happen in the first 500 or 900 requests, but it will happen in the one thousand.
00:03:09Which is rare enough to duplicate it easily, but not, right, it's too common to ship it.
00:03:18Sorry.
00:03:19Also, it only happened in VLLM, not in other, not in any other inference framework.
00:03:26And it's something which is late on set.
00:03:28It's not something that will happen if you have only few requests.
00:03:32You need some sort of workload.
00:03:34So, it's rare.
00:03:35It's late on set.
00:03:36And it's very engine specific.
00:03:38It's a very, very hard task, and we had to bring one of our best detectives to handle that.
00:03:43So, I'll give it to Asaf to explain how.
00:03:47All right.
00:03:48Hi, guys.
00:03:48Thank you, Yuval.
00:03:49Thanks, Yuval.
00:03:50So, we're going to start with trying to reproduce something that was very difficult to reproduce.
00:03:55And basically, VLLM has got a lot of flags, a lot of the CLI flags, and a lot of knobs you guys can turn and tweak.
00:04:04And 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, it didn't really help us manage to get gibberish back from our model.
00:04:17The model responded back just fine.
00:04:19So, what we did was is we tried to make it happen in a very, very short amount of time.
00:04:25So, we'll get a quick feedback loop when we try to debug it.
00:04:29So, what we did was is we took one of the most default and most common flags that the VLLM allows you to play with, which is GPU memory utilization, which basically allows you to choose how much GPU memory you want to allocate for your weight, for your activations, and for your KV cache, and so on.
00:04:48And we reduced it from 90% to 20%.
00:04:51And once we did that, and then we started running a lot of requests simultaneously, all of a sudden, request number, let's say 854, suddenly returned gibberish.
00:05:05And when we did that, we sampled all of the batches with temperature zero, so we'll be able to deterministically and constantly get the same request to return and respond with gibberish.
00:05:17So, like Yuval said, it happened only in VLLM.
00:05:20And we used, in order to, another way to reproduce it and to understand where the issue really came from, we used Hugging Faces Transformers as a baseline, since Transformers is a very, had a very vanilla and plain implementation of our Mamba kernels, as opposed to VLLM, which all the kernels and all the engine have gone through a lot of changes and modifications to support a lot of cool features.
00:05:46So, we use Transformers as our baseline to understand whether or not there is an issue with our inference or not, with the model or not.
00:05:55So, what we did was we took VLLM and we sent all of our prompts through VLLM and we generated a response, all the responses.
00:06:02We got, now in our hands we have the response along with the log probs, because in VLLM you're able to get your log probs out and inspect them.
00:06:11Then what we did was we took the full sequence, the prompt and the generation, and we passed it over to Hugging Faces forward pass.
00:06:20But all we did was run just the pre-fill, and then we took the logit out of the pre-fill response, we ran it through Softmax, and then we were able to compare the divergence in the distributions of our tokens.
00:06:38That's a short pseudocode of how that looked like.
00:06:41You can see here that we take up the prompt, we run it through VLLMs generate, we get the response back along with the log probs, we pass it over to Hugging Faces forward pass, we only run it with pre-fill.
00:06:55We created some function called compute log probs, which runs the Softmax, then you calculate the difference between them, and then you'll be able to tell the divergence between every one of the tokens log probs.
00:07:09All right, so now that we have the tools in our hand to understand where the issue could maybe come from, we started to look at different suspects in VLLMs engine.
00:07:19So the first thing we looked at was the CUDA pre-fill kernel of Mamba.
00:07:24We looked at it, we inspected all of the math that's being done there, and we looked at the tensor in, the tensors out before we call the pre-fill and after.
00:07:36Everything looks just fine.
00:07:38Second thing we did was running NVIDIA's compute sanitizer tool to really see if we have any out-of-bound memory, any other memory bugs or issues.
00:07:48It looked okay to me.
00:07:52Then what we did was we tried to isolate between the decode kernels and the pre-fill kernels.
00:07:56Now, we saw that the pre-fill kernels were working just fine, so we tried to not call the decode kernels, because in Mamba you're able to do that.
00:08:03So what we did was we moved all of our calls and all of our computations to go through the pre-fill kernel, and there you have it.
00:08:12The gibberish all of a sudden kind of vanished.
00:08:14So we were like, okay, it's got to be the decode kernels.
00:08:18But you know how it is in software.
00:08:20You're getting excited too quickly, and then you figure out it's not what happened.
00:08:24So what we did was we tried to start playing 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 better understanding and maybe, you know, get our hands dirty.
00:08:38Because VLM didn't really give us more tools to really debug our kernel and our forward pass.
00:08:45So once a tensor, once the request gets all the way to your forward pass, and before it goes into your pre-fill and decode kernels, they don't really have any identity.
00:08:55You can't really tell what prompt is currently being processed.
00:08:59It's all just tensors and numbers and matrices.
00:09:02So what we did was is we added the request ID to some class called forward context that we propagated all the way down to Mamba's forward pass just before the pre-fill and the decode kernels were called.
00:09:16And there we just managed to, you know, have a simple if condition with a request ID, the one that gave us gibberish, and put a break point there.
00:09:26And then we were able to infer and to really inspect all the metadata that comes along with it.
00:09:33And the second we did that, we saw that the request was, for the first time, when it went through the forward pass, it's actually doing decode before pre-fill.
00:09:43The scheduler decided that this request should be doing decode before pre-fill.
00:09:50And as Yuval said earlier, in the lifecycle of a prompt, a prompt should first be going through pre-fill and then decode.
00:10:00And what happens was is that when in Mamba you run a request with decode first, after a lot of other requests were already computed, the state was already kind of overused.
00:10:15And we were using the data and the computations of stale requests, requests that came before it.
00:10:21So now we were actually running decode on previous requests.
00:10:25And that kind of generated gibberish for us.
00:10:28So the kernels weren't doing the wrong thing.
00:10:30They were called at the wrong time for the wrong requests.
00:10:34And now -- and why did it matter only for Mamba?
00:10:37The reason is was is that in a tension, when you write the tokens KV -- you write the tokens KVs before you actually -- you read it.
00:10:47So even if you have stale data, it's being overwritten.
00:10:50But for Mamba, as I said, when you first go through the decode kernels, you first read the state, and then you compute over it.
00:10:58So what happens was is you just use stale data when you do a decode.
00:11:05And the fix was relatively simple.
00:11:07We just needed to make sure that what we do is that when a request first -- when the scheduler first classifies a request, it's got to make sure that it sets -- that if you see the request whose tokens were never been computed, and they're zero,
00:11:26to mark them as pre-fill, so when they get to the forward pass, they'll actually just be used for pre-fill and not decode, and not chunked.
00:11:38You can see that it was merged after some time.
00:11:41And that really leads us -- and then we thought everything was fixed, right?
00:11:45We thought everything was fixed, and there you have it.
00:11:47No more issues.
00:11:48But that was almost the case, because after a little bit of time, it gets us to case number two,
00:11:54which surfaced another issue that we've faced in our RL, in our inference.
00:12:01So we ran RL, and while we were running our trainings, our post trainings, and we looked at our evaluations and all of our benchmarks,
00:12:10we thought that we had some logprop spikes between the rollout and the FSDP step.
00:12:15And that was before any wait update.
00:12:18So some waits, same inputs, and the two logprops should be identical.
00:12:23Now they weren't.
00:12:24We saw that every 12-step constantly, there was a logprop spike, and that was kind of weird.
00:12:33Now what would you guys do, right?
00:12:36What's the first thing to do here?
00:12:39So we wanted to find some lever that changes how things fail and not just how much they fail.
00:12:44We want to see how much -- we want to tweak some knobs that don't just tell us, hey, this error -- this error is very, very bad -- this error happens this many times, or -- and so on.
00:12:57And so we wanted to tweak some knobs that kind of tell us that once we tweak that knob, we understand how it's wired to anything in VLM's engine.
00:13:08And so we'll be able to specifically go and debug that specific part.
00:13:14That's some cool meme that you all wanted to put in.
00:13:19So what we did was is we decided to increase the amount of rollouts per prompt.
00:13:26Since we saw in our default RL engine we have eight rollouts per prompt, and we saw that it happened deterministically every 12 steps, we decided, okay, let's try to tweak it up a bit and increase the amount of rollouts per prompt.
00:13:40So we started doubling it from 8 to 16 to 64, 32, and 128.
00:13:45And you can see here that it's almost -- almost -- there's a pattern here.
00:13:51The more we increased it, the closer 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 loop, feedback loop.
00:14:05So 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 24 and so on.
00:14:12Now, you might think, okay, so you guys played with the GPU memory utilization before.
00:14:18You tweaked it, you decreased it.
00:14:20It looks like, you know, when you test on pressure, it really surfaced things up.
00:14:24So we thought that as well.
00:14:26And when we reduced the GPU memory from 0.9 to 0.2, it actually caused the issue to go away.
00:14:35So we actually pulled the wrong lever here.
00:14:38And the reason is, is because we noticed that Mamba kernels used unsigned int 32 index pointer.
00:14:48So once the offset went past some, you know, 4 billion numbers, it wrapped around instead of throwing an error.
00:14:55So when we shrank the GPU memory, BLM allocated a small state buffer, and the cache index never got large to hit that slot.
00:15:04So we were just not reaching far enough for the buffer to trigger an overflow.
00:15:09So, again, the fix was rather simple.
00:15:13All we needed to do was just change one word, one data type variable from UN32 to size T, which basically means for most modern architectures, hardware architectures,
00:15:27size T would mean that it would be now changed to unsigned 64 a bit.
00:15:32And that's a very large number.
00:15:34We never reached that number, and that overflow now never happened.
00:15:38So what we can see here is that we had two scenes and one criminal.
00:15:45Both kind of, you know, they had similar symptoms.
00:15:48Both had silent gibberish and silent log prob spikes, which also kind of sometimes generated gibberish.
00:15:55They were both around the Mamba state cache.
00:15:57They were both surfaced by memory pressure, whether it was for worse or for the best.
00:16:03And both found via log prob forensics.
00:16:09Stateful inference systems don't feel loudly.
00:16:12They lie to you confidently.
00:16:13I mean, obviously, sometimes you get crashed.
00:16:15You get out-of-bounds errors.
00:16:16You get other, you know, exceptions and so on.
00:16:19But sometimes there are some errors that don't surface up, and you don't get a trace log.
00:16:24You don't get anything.
00:16:25You have to go and dig and understand why things happen.
00:16:29So if there are some takeaways to take from this presentation is build the log props comparison script.
00:16:36If you need to compare your quality, you need to compare it to understand whether you modelize the issues or not.
00:16:41Log props comparison script with a baseline of some other inference framework that you have or built is always great.
00:16:48Reproducing under pressure and constrain memory, crank the scale up, play with other knobs that the inference framework gives you and really try to understand where the issue comes from.
00:16:59Look for what moves the failure shape, the timing, the space, and the location.
00:17:05And when things don't really have identity, thread identity through.
00:17:10And what also I want you to take from this, and don't be afraid to even, you know, for complex systems like VLLM or any other complex framework, don't be afraid to go dig in the code, get your hands dirty.
00:17:23Sometimes, you know, model languages, LLMs are, 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.
00:17:34Thank you, you guys can add us on LinkedIn, scan the QR code to read the actual blog that we've published with this finding.
00:17:43Yeah, that's it.
00:18:04We'll see you next time.
Community Posts
No posts yet. Be the first to write about this video!
Write about this video