Weight Folding, CUDA Streams, and the Bug That Made My Model Speak Backwards — Filip Makraduli
AAI Engineer
컴퓨터/소프트웨어AI/미래기술
스크립트
00:00:00.
00:00:13Hello everyone, thank you for coming and
00:00:18I'll start the talk now. So
00:00:22this talk is around a paper
00:00:27that I did which is very simple.
00:00:33The proposition is very clear.
00:00:36It's basically two lines of
00:00:39algebra that make the RMS norm
00:00:42layer in transformers cheaper,
00:00:45quicker, and kind of improve it
00:00:48as like a layer in the transformer
00:00:51architecture. Similar to how
00:00:54layer norm once used to be the standard
00:00:57and then it was substituted by RMS norm.
00:01:00This follows along this way of thinking.
00:01:04And I got the chance to kind of meet
00:01:09some people from the open source world
00:01:11and I co-authored this paper
00:01:14together with Nils Graf who was the
00:01:17kind of the creator of this.
00:01:21And the work follows from there.
00:01:23So this is presented on archive.
00:01:26You can have a look, read it, test it out.
00:01:29There is a repo as well.
00:01:31And the concept, let's say the idea
00:01:36and the way of thinking, it's easiest
00:01:38to explain with maybe flash attention.
00:01:40So in a similar way of how flash attention
00:01:45kind of waits until there is a multiplication
00:01:48and tries to limit this communications
00:01:52between memory so that the whole process
00:01:54is faster.
00:01:55This is kind of a similar thought along those lines.
00:01:58And it does certain improvements
00:02:01that make the RMS norm process much quicker
00:02:08and in effect improve the whole transformer.
00:02:15And one question is, okay, why RMS norm?
00:02:19Since that layer does almost none of the math.
00:02:24And that's true.
00:02:26So the share of the kind of math portion
00:02:29if you look at it is quite small.
00:02:31However, the clock time or wall time,
00:02:34as they say, is quite big.
00:02:36And for example, in one decode step,
00:02:40so right when like inferences perform,
00:02:43the RMS norm can be started like 33 times.
00:02:47Of course, it depends on the model and so on.
00:02:50In the paper, you have the specific models
00:02:52and how this was tested.
00:02:54And the question is how this can be improved
00:02:59and how this weight for the matrix multiplication
00:03:03can be kind of avoided.
00:03:06And the reason why this is slow
00:03:09is because the GPUs are not slow or bad at math,
00:03:15but they're bad at everything else around the actual math.
00:03:20So that means starting the work, the actual work.
00:03:24So for example, starting the process as it happens
00:03:30in some of the experiments 33 times,
00:03:32that takes a long time.
00:03:35And for example, fusing each normalization
00:03:40into the matrix multiplication can help avoid this.
00:03:44Also doing weight folding
00:03:46can help in kind of moving data between memory
00:03:50and that's a process that's also slow for GPUs.
00:03:54And also weighting.
00:03:56So for example, deferring the division
00:04:00that's done in the RMS norm layer
00:04:02is also a way to avoid this waiting step.
00:04:06So basically what this paper does
00:04:09is it improves all these three aspects
00:04:13by doing a few algebraic tricks
00:04:16in the way RMS norm is computed.
00:04:19That's it.
00:04:21And math-wise, these are the tricks.
00:04:25It's mainly around the first two propositions.
00:04:29One is weightless normalization.
00:04:31You can see that here.
00:04:33And deferred normalization.
00:04:35So that's the second one.
00:04:37And now in more newer architectures,
00:04:39there is a situation where RMS can appear twice.
00:04:44For example, in Gemma 4 this happens.
00:04:48So cancelling the pre-normalization also works.
00:04:52And all of this is algebraically proven in the paper.
00:04:58And the first proposition is this
00:05:00where kind of the gain and the weight fold
00:05:04fold into one matrix.
00:05:06W you can see here with an asterisk.
00:05:09And that is computed offline,
00:05:12similar to how maybe in flash attention
00:05:14you compute some stuff on the side
00:05:16so that there is no communication
00:05:18between memory all the time.
00:05:20So this is one step that's kind of done,
00:05:24this weight folding.
00:05:26And the other step is deferring the scalar divide
00:05:32of the matmul so that they can be done in parallel.
00:05:35So in a normal case, you would have to compute once,
00:05:38then wait and compute again.
00:05:41In this case, the idea is to kind of split this
00:05:44so that it can be parallelized.
00:05:48And the third one, which is kind of a version of this,
00:05:51is that there is kind of, if there are two,
00:05:56because this is scale invariant,
00:05:58one of them can be dropped and this still works.
00:06:01And this is applicable to newer models
00:06:04that can have this architecture and implementation.
00:06:10So in order to make this happen in real life,
00:06:13especially this proposition number two.
00:06:16So for this one, for example, it's easy.
00:06:20There is a repo called transformer tricks.
00:06:22You can just apply this to any model and it works.
00:06:25But in order to do this, there is some kernel work.
00:06:29So it's not as straightforward to do.
00:06:31So in order for me to do that,
00:06:35I was implementing this and I came out with this experiment once.
00:06:42So it looks okay in general where it's like,
00:06:46okay, the prompt is the transformer architecture,
00:06:48revolutionized NLP because,
00:06:51and then there is some kind of expected output.
00:06:54But in the output I got,
00:06:56I saw this repetition and one step lag.
00:06:59As you can see here, the word because appears again.
00:07:01And there was something happening with the GPU streams
00:07:07and I was trying to figure out what was happening.
00:07:10And I was getting this one step lag and kind of outputs
00:07:15that were from the past in a way.
00:07:18And in debugging all of this,
00:07:21I realized that in the process of building something like this.
00:07:26So as I explained the proposition to or deferring this to operations,
00:07:32in CUDA, you can do two things.
00:07:35You can do like tensor cores that do one part of the matrix multiplication
00:07:39and you can do CUDA cores that kind of run stuff like element-wise operations,
00:07:44reductions, square roots, and so on.
00:07:47So the idea was to do this in parallel and get the benefit
00:07:52of what I was explaining in the paper to actually test out this concept.
00:07:58So this is how it was supposed to look like.
00:08:00So there is, if you do things sequentially,
00:08:03there is this idle waiting time when the vector unit computes the RMS and scaling,
00:08:10and then there is a matrix multiplication.
00:08:12So the idea was, okay, with flash norm, which is the technique in the paper,
00:08:16you're supposed to do those both in parallel.
00:08:19So the matrix unit computes the matmul and the vector unit computes the RMS.
00:08:23So in that way you save time.
00:08:26However, you cannot just do this in Python.
00:08:28You have to go a bit lower.
00:08:30And I did that with CUDA codes like this.
00:08:35And this looked in general okay at that time.
00:08:42However, I realized that I did something slightly wrong.
00:08:47And that thing was that the join in the end where you're supposed to join the two streams was implicit in my case.
00:08:57And when I tested this out, the unit test worked.
00:09:01The quality seemed similar like perplexity testing and so on because it's just like similar generation.
00:09:08But over long generation I was able to see this problem.
00:09:11So I had no idea what this was.
00:09:14And the reason was that when I was doing this implicit join, basically one of the streams hadn't finished the work.
00:09:24So I got race conditions that kind of read the past from the unfinished matrix multiplication.
00:09:31So the idea that I had to fix this was around the fact that I had to be explicit about the join.
00:09:40And wait until one of the operations is finished so that I'm certain that when I join I'm not reading from the past.
00:09:48So that was the realization in this exploration of CUDA streams.
00:09:54And this is how I had things done.
00:09:57So the join was implicit.
00:10:00So the post scale read like an old buffer value.
00:10:06And how this is fixed is with this where basically you need to mark the end of the matrix multiplication.
00:10:14Then mark the end of the RMS.
00:10:17And then post scale wait for the first stream.
00:10:21And then wait for the second stream.
00:10:24And that fixed the bug and made kind of the paper work and the model speak forward instead of backwards.
00:10:33And that was the cool maybe academic perspective.
00:10:38But I also wanted to try things, right?
00:10:40Deploy this, test it out, see how I can make it work in maybe a more production setting.
00:10:47And you can also read the paper and see all the tests.
00:10:50Some are done, most are done around llama models, but like this works for other architectures as well.
00:10:56So what you can do for this specific paper is, for example, the weight folding that I explained, the proposition one, you can just do it with some code in the repo.
00:11:10That's like flash, you say flashify and it does that.
00:11:13However, with this second thing that I mentioned, you need to do a bit of kernel work if you want to do that.
00:11:19Like I explained in my example.
00:11:22And these are some results that are based on llama models and there are different kind of details that you can have a look at as well.
00:11:29Like what happens if you do only deferred normalization, what happens if you do a full fused kernel.
00:11:36So there are a lot of experiments of going lower here to test all the prepositions.
00:11:41And these have been our results in different, let's say, levels of scrutiny and detail.
00:11:48But even the simple one with like weight folding shows some improvement.
00:11:54And this also works with like the day-to-day tools that you use in a model.
00:11:59So it's not like you have to reinvent the wheel or, you know, do things from scratch.
00:12:05So it works with torch compile because it's kind of like a new checkpoint and that's it.
00:12:13Flash attention does similar tricks at a different layer.
00:12:16And also it works with quantized models.
00:12:18So it's totally cool to actually apply this and you can get a model that has this cool new normalization layer.
00:12:27And where you can get this details and codes to actually run this is this transformer tricks repo.
00:12:34So it has different algebra tricks like I explained as well as this paper that I mentioned.
00:12:41And also there is the GitHub, not the GitHub, but the HuggingFace model repo where I've done this with some models.
00:12:50And you can have a HuggingFace link to that model and test it out.
00:12:54And what you also can do with this HuggingFace models is to deploy them in production.
00:13:00So when I was thinking about doing this, I realized that, okay, now that let's say the science is done and there is a link to a HuggingFace model,
00:13:11Superlink's inference engine was a cool way to actually deploy any HuggingFace model.
00:13:19And we've done this at hackathons where people would bring like a custom HuggingFace model or checkpoint that they have with their fine-tuned stuff.
00:13:27And you can test out like, even if you have some version of this algebraic tricks that you want to improve a model and test your own research ideas,
00:13:37you can actually try that out and have a deployed version on a cluster of this model and not have to worry about this glue code around deploying models.
00:13:48So that's pretty cool. And the point is that if you have the full cluster open source and the model inference open source,
00:13:58you can actually test out this kind of maybe more novel research ideas where if you want to do kernel manipulation or flash norm and things like that,
00:14:12it's much more difficult to do this at the rented end point where you don't own the inference.
00:14:18So it's, you want something that's portable and flexible to actually allow you to do this stuff,
00:14:23but it's also production ready enough so that you can test things out at scale.
00:14:27And you can, for example, use site to combine this with other models.
00:14:33Like, as you can see in the top left, there is, you can have this flashified models with different other models to do agentic tasks if you want,
00:14:43and kind of do that end-to-end bigger use case.
00:14:46And the way site works is this production cluster helps you deploy the models.
00:14:52So you can have a look at size repo as well for more details on this.
00:14:57And also there is a smarter queuing mechanism that helps you, especially if you work with smaller models,
00:15:03because when doing the flash norm stuff, I worked with, like, smaller llama models and also with small agents from HikingFace.
00:15:11So having a way to deploy smaller models that can also work on, like, the same GPU so that you don't have to spend your money on GPU costs,
00:15:24but actually kind of switch models around, especially smaller models, it was quite useful.
00:15:30And you can also control the model configs through an API as well as the cluster,
00:15:35which is also pretty convenient without having, like, an infra guy supporting you in your open source research.
00:15:40So that's cool as well.
00:15:43And you own your cloud, which is useful if you want open weights, open models, open source.
00:15:50And there is also, like, a catalog that Cy has of different models, not just the ones I mentioned,
00:15:57but you can have a look.
00:15:59There's also re-ranking embedding models if you're building something along those lines.
00:16:03And with that, I'm kind of finishing this story of my research journey where I co-authored this paper around the technique that improves the transformer,
00:16:15but also found a way kind of to bring this to, let's say, production and test it out and find a way to play around with this open source models.
00:16:24And feel free to contact me on LinkedIn, maybe if you have any questions or contributions.
00:16:30A lot of this stuff that I've mentioned, like, some of them are PRs on, like, VLLM or on Hugging Face.
00:16:36You might find them all around.
00:16:38You can also see the, check out the paper.
00:16:40That's the archive link that you have there.
00:16:43And you also have the Cy repo and my LinkedIn.
00:16:47So, thank you very much for attending.
00:16:58And you can catch me for questions.
00:16:59We'll be here.
00:17:00Close by.
00:17:13you
커뮤니티 글
아직 글이 없습니다. 이 영상에 대한 첫 번째 글을 작성해 보세요!
이 영상에 대해 글쓰기