Weight Folding, CUDA Streams, and the Bug That Made My Model Speak Backwards — Filip Makraduli

EnglishالعربيةFrançaisहिन्दीBahasa Indonesia日本語한국어PortuguêsРусский中文
AAI Engineer
Computing/SoftwareInternet Technology

Transcript

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

Key Takeaway

Algebraic weight folding and parallelized kernel execution optimize RMS norm layers in transformer models to eliminate redundant memory bottlenecks and accelerate inference.

Highlights

  • Algebraic optimizations make the RMS norm layer in transformer architectures faster and cheaper by reducing memory communication overhead.

  • During a single decode step in model inference, the RMS norm layer can execute up to 33 times.

  • Weight folding combines the gain and weight matrices into a single offline-computed matrix to eliminate redundant memory operations.

  • Deferred normalization splits the scalar division of matrix multiplication so the operations run in parallel across tensor and vector cores.

  • Implicit joins in CUDA streams cause race conditions and data corruption by reading buffer values from unfinished matrix multiplications.

Timeline

Optimizing Transformer RMS Norm Layers

  • Two lines of algebra improve the speed and cost efficiency of the RMS norm layer in transformers.
  • Clock wall time for RMS norm is high during inference because the layer executes up to 33 times in a single decode step.
  • GPUs spend more time initiating work and moving data between memory than performing actual mathematical computations.

Transformer models spend a disproportionate amount of inference wall time executing normalization layers repeatedly rather than performing core mathematical operations. Hardware limitations dictate that memory communication and process startup overhead constrain overall execution speed. Algebraic adjustments target this inefficiency by altering how normalization steps interact with memory.

Core Algebraic Propositions for Flash Norm

  • Weight folding merges gain and weight into a single matrix calculated offline.
  • Deferred normalization separates scalar division from matrix multiplication to allow parallel processing.
  • Implicit stream joins in custom CUDA kernels create race conditions that read outdated buffer values from unfinished calculations.

Three core mathematical propositions streamline transformer normalization. Weight folding removes redundant run-time memory transfers, while deferred normalization allows tensor cores and vector cores to execute matrix multiplication and square roots simultaneously. Writing custom CUDA kernels for these parallel operations requires explicit stream synchronization to prevent data race conditions and backward model generation.

Production Deployment and Open Source Tools

  • The transformer tricks repository provides code to apply weight folding and kernel optimizations to existing models.
  • Production clusters and inference engines like Cy allow developers to deploy custom HuggingFace models without complex infrastructure code.
  • Smaller open-source models share GPU resources through optimized queuing mechanisms to reduce infrastructure costs.

Research insights transition into production environments using open-source model repositories and specialized inference engines. Custom checkpoints containing algebraic tricks integrate with standard tools like torch compile and quantized models. Production clusters manage smaller models on shared GPUs to facilitate scalable testing and deployment of novel architecture modifications.

Community Posts

No posts yet. Be the first to write about this video!

Write about this video