Cutting LLM Memory by 84%: A Deep Dive into Fused Kernels

The Logit Bottleneck

To get us began, let’s put some extra numbers on the logit bottleneck. We think about an enter matrix X with form [NxD], a weight matrix W with form [DxV] and a logit matrix Y=X@W with form [NxV]. Within the context of an LLM, N could be the sequence size multiplied by the batch dimension (i.e. the overall variety of tokens within the batch), D the scale of the hidden state and V the vocabulary dimension. 

For a Llama3 8B mannequin, we might have a context window of 8192 tokens, a hidden state with 4096 dimensions and a vocabulary dimension of 128,256 tokens. Utilizing a modest batch dimension of 8, we get N = 8192x8 = 65,536.

This leads to the Y matrix having form [NxV]=[65,536x128,256], or roughly 8.4 billion parts. In bfloat16, this is able to take up 16.8GB of reminiscence. Nonetheless, if we observe greatest practices and use float32 for the loss calculation to make sure numerical stability, the necessities double to 33.6GB.

To place this quantity in perspective, we might additionally want round 16GB of reminiscence to carry the weights of Llama3 8B in reminiscence in bfloat16. One most GPUs, this leaves no area for the large overhead of the optimiser states (e.g. Adam’s moments) and different activations, ensuing within the notorious PyTorch OOM error.

Usually, this drawback is handled by utilizing:

• Gradient accumulation: Use a smaller batch dimension and accumulate gradients over a number of batches between every optimiser step, emulating a bigger batch dimension whereas holding much less information in reminiscence.
•  Activation checkpointing: PyTorch shops all intermediate activations for reuse within the backward move, checkpointing clears these activations and recomputes them on-the-fly through the backward move. This results in massive reminiscence financial savings however will increase coaching time for the reason that variety of required ahead passes is doubled.
• Micro-batching the loss: As a substitute of computing the loss over the N dimension directly, we are able to slice it and accumulate the loss over smaller chunks with dimension n < N. Now, we solely maintain a slice of dimension [n, V] in reminiscence at a time.
• Blended precision coaching: Utilizing half precision throughout coaching gives 2x reminiscence discount and vital speedups on Tensor Cores.

Whereas these options appear enticing, all of them have vital drawbacks: gradient accumulation and activation checkpointing decelerate coaching, blended precision could be unstable and micro-batching requires (gradual) PyTorch stage iteration and despite the fact that n is chosen to be smaller than N, the vocabulary dimension stays big compared.

Extra importantly, these options don’t handle the issue we’ve handled repeatedly all through this collection: information motion. Certainly, we’re nonetheless losing time by writing billions of logits to VRAM solely to learn them again milliseconds later.

The Kernel Resolution

As we’ll see in a minute, the ahead and backward move of the cross-entropy loss contain dot merchandise, matrix multiplication and a softmax. As we discovered on this collection, these are all operations that may be tiled effectively. In different phrases, we are able to carry out them iteratively whereas solely holding a small piece of the inputs in reminiscence at any time.

Moreover, cross-entropy is usually preceded by a matrix multiplication: the linear projection from the hidden state into the vocabulary area. This can be a nice alternative for operator fusion: fusing a number of operation inside a single kernel, leading to massive speedups and potential reminiscence features.

Within the following sections, we’ll check out the right way to derive and effectively fuse the ahead and backward passes by a kernel combining a linear layer with cross-entropy.

As talked about within the final article, Triton kernels don’t natively register in PyTorch’s autograd. Due to this fact we have to derive the gradient ourselves, an exquisite event to brush up on some calculus 😉

The maths behind Fused Linear Cross-Entropy

Definition and Ahead Move

On this part, we derive the mathematical expression for our Fused Linear Cross-Entropy layer to see the way it naturally lends itself to tiling.

For 2 discrete chance distributions p and q, cross-entropy is outlined as:

In our context, p is the one-hot vector representing the goal token, whereas q is the mannequin’s distribution over the vocabulary. We get hold of q by making use of a softmax to the logits l, themselves the outputs of the previous linear layer.

Since p is constructive for a single goal token y, the summation collapses. We will then substitute the numerically secure softmax (as mentioned within the last article) to derive the ultimate expression:

By substituting the logits l with the linear layer x . w, we see that the ahead move boils down to a few main portions:

1.  The goal logit x . w_y.
2. The log-sum-exp (LSE) of all dot merchandise.
3. The worldwide most logit used for numerical stability.

Because of the net softmax algorithm, we are able to compute these portions with out ever materialising the total vocabulary in reminiscence. As a substitute of an O(V) reminiscence bottleneck, we iterate over the hidden dimension D and the vocabulary V in small tiles (D_block and V_block). This transforms the calculation into an O(1) register drawback.

To parallelise this successfully, we launch one GPU program per row of the enter matrix. Every program independently executes the next steps:

1. Pre-compute the goal logit: Carry out a tiled dot product between the present row of X and the column of W related to token Y.
2. On-line discount: Iterate by the hidden and vocabulary blocks to:
    1. Monitor the operating most (m)
    2. Replace the operating sum of exponentials (d) utilizing the net softmax method:

Now that we’ve a greater understanding of the ahead move, let’s check out the derivation of the backward move.

Backward Move

Notation

To derive our gradients effectively, we’ll use Einstein notation and the Kronecker delta.

In Einstein notation, repeated indices are implicitly summed over. For instance, an ordinary matrix multiplication Y = X@W simplifies from a verbose summation to a clear index pairing:

The Kronecker delta (δ_ij) is used alongside this notation to deal with id logic. It is the same as 1 if i=j and 0 in any other case. As we’ll see, that is notably helpful for collapsing indices throughout differentiation.

Matrix Multiplication

On this part, we derive the back-propagated gradients for matrix multiplication. We assume the existence of an upstream gradient ℓ. 

To find out the way it back-propagates by matrix multiplication, we use the apply the chain rule to the inputs x and the load matrix w. Right here y represents the multiplication’s outputs:

We begin by deriving the partial derivatives of y with respect to x, following these steps:

1. Categorical y by way of x and w.
2. Discover that w is a continuing with respect to the spinoff of x, so we are able to pull it out of the spinoff.
3. Categorical the truth that the partial spinoff of x_ik with respect to x_mn is 1 solely when i=m and okay=n utilizing the Kronecker delta.
4. Discover that ẟ_kn enforces okay=n, due to this fact w_kj * ẟ_kn reduces to w_nj.

Then, we think about the total expression and procure the gradient. We derive the final step by noticing as soon as once more that 1/y_ij * ẟ_im reduces to 1/y_mj.

Nonetheless, matrix notation is conceptually nearer to our Triton kernel, due to this fact, we rewrite this expression as a matrix multiplication by utilizing the id X_ij = [X^T]_ji:

We observe the very same steps to derive the gradient with respect to W:

Then, the back-propagated gradient follows:

Which is equal to the matrix notation:

Cross-Entropy

On this part, we’ll deal with cross-entropy utilized to discrete chance distributions. Contemplating a tensor of j logits, with a label y, the cross-entropy is computed as follows:

The place x_y corresponds to the logit related to the label.
As soon as once more, we have an interest within the partial spinoff of any output i with respect to any enter okay. Due to the normalising issue, each aspect i impacts the worth of each different aspect, due to this fact, the partial spinoff is obtained by defining the operate piecewise relying on the worth of i:

Summing each circumstances, we get hold of the gradient:

And in matrix notation:

The place y_{one scorching} is a vector of zeros with the entry akin to the label set to 1. This consequence tells us that the gradient is just the distinction between the prediction and the bottom reality.

Fused Linear Cross-Entropy

Combining the linear projection with cross-entropy in a single expression, we get:

Because of the chain rule, deriving the gradient of this expression boils right down to multiplying the gradients we computed beforehand:

The place x and y check with the inputs and outputs to the linear layer respectively and w to the related weight matrix.

Notice: in a batched setting, we’ll want to cut back the W gradients over the batch dimension. Usually, we use a sum or imply discount.

Kernel Implementation

With the speculation established, we are able to implement the fused kernel in Triton. Since cross-entropy is usually the ultimate layer in a language mannequin, we are able to mix the ahead and backward passes right into a single kernel. This fusion provides two benefits: it minimises the overhead of a number of kernel launches and considerably improves information locality by retaining intermediate values on-chip.

We are going to analyse the kernel step-by-step from the angle of a single program occasion, which, in our parallelisation technique, handles one particular row of the enter matrix.

1. Setup and Goal Logit Pre-computation

The preliminary part entails customary Triton setup:

• Program Identification: We use tl.program_id to find out which row of the enter matrix the present program is chargeable for.
• Parameter Initialisation: We outline tiles utilizing D_BLOCK and V_BLOCK and initialise the operating most (m) and sum (d) required for the net softmax algorithm.
• Pointer Arithmetic: We calculate the bottom reminiscence addresses for our tensors. Pointers for X (enter) and dX (gradient) are offset utilizing the row stride so every program accesses its distinctive token vector. Conversely, the W (weight) pointer stays on the base handle as a result of each program should ultimately iterate by your complete vocabulary area.
• Masking and Early Exit: We outline an ignore_index (defaulting to -100). If a program encounters this label (e.g. for padding tokens), it terminates early with a lack of 0 to avoid wasting cycles.

2. Computing the Goal Logit

Earlier than the primary loop, we should isolate the goal logit x . w_y. We iterate over the hidden dimension D in D_BLOCK chunks, performing a dot product between the enter row X and the precise column of W akin to the ground-truth label Y.

As a result of W is a 2D matrix, calculating the pointers for these particular column tiles requires exact stride manipulation. The illustration under helps visualising how we “bounce” by reminiscence to extract solely the required weights for the goal token.

As soon as the tiles are loaded, we solid them to float32 to make sure numerical stability and add their dot product to an accumulator variable earlier than transferring to the following iteration.

Right here’s the code to date:

Subsequent, we execute the ahead move, which processes the vocabulary area in two nested levels:

1. Tiled Logit Computation: We compute the logits for a V_BLOCK at a time. That is achieved by iterating over vocabulary dimension V (outer loop) and the hidden dimension D (inside loop). Inside the inside loop, we load a tile of X and a block of W, accumulating their partial dot merchandise right into a high-precision register.
2. On-line Softmax Replace: As soon as the total dot product for a logit tile is finalised, we don’t retailer it to VRAM. As a substitute, we instantly replace our operating statistics: the utmost worth m and the operating sum of exponentials d utilizing the net softmax method. By doing this “on the fly”, we be sure that we solely ever maintain a small V_BLOCK of logits within the GPU’s registers at any given second.

Following these iterations, the ultimate values of m and d are used to reconstruct the LSE. The ultimate scalar loss for the row is then computed by subtracting the goal logit (x . w_y) from this LSE worth.

Right here’s a visible illustration of the ahead move:

Right here’s the code for the ahead move:

We are actually right down to the final a part of the kernel: the backward move. Our purpose is to compute the gradients with respect to X and W utilizing the expression we derived earlier:

To stay memory-efficient, we as soon as once more course of the vocabulary in tiles utilizing a two-staged method:

1. Recomputing Normalised Possibilities (P): As a result of we didn’t retailer the total logit matrix through the ahead move, we should recompute the activations for every tile. By reusing the Log-Sum-Exp calculated within the ahead move, we are able to normalise these activations on-the-fly. Subtracting the ground-truth label Y from the goal logit inside this tile provides us an area chunk of the gradient logit, P.
   2. Gradient Accumulation: With a tile of P in hand, we calculate the partial gradients. For dX, we carry out a dot product with blocks of W^T; for dW, we multiply by tiles of X^T. To soundly mixture these values throughout your complete batch, we use Triton’s tl.atomic_add.
   This operation acts as a thread-safe +=, guaranteeing that completely different packages updating the identical weight gradient don’t overwrite each other.

Listed below are some extra particulars on the implementation:

• The Stride Swap: When computing P . W_T, we don’t truly have to bodily transpose the large W matrix in reminiscence. As a substitute, we invert the shapes and strides in W’s block pointer to learn the rows of W as columns of W^T. This leads to a “free” transpose that saves each time and VRAM.
• Numerical Precision: It’s value noting that whereas X and W is likely to be in bfloat16, the buildup of dW and dX through atomic_add is normally carried out in float32 to forestall the buildup of tiny rounding errors throughout hundreds of rows.
• Rivalry Notice: Whereas atomic_add is important for dW (as a result of each program updates the identical weights), dX is personal to every program, which means there may be zero rivalry between program IDs for that particular tensor.
• Atomic Add Masking: atomic_add doesn’t assist block pointers. Due to this fact, we implement the pointer and masks logic for dW explicitly.

The next determine is a illustration of the backward move for one iteration of the outer loop (i.e. one block alongside V and all blocks alongside D):

Right here’s the total code for the backward move:

This concludes the implementation of our kernel! The complete code together with the kernel and benchmark script is out there here.

Reminiscence Benchmark

Lastly, we examine our kernel with the PyTorch baseline utilizing hyperparameters impressed from Llama3 and an A100 GPU. Particularly, we think about a sequence size of S=16,384, a batch dimension of B=1 and an embedding dimension of D=4096; the vocabulary dimension is about to V=128,256.

As anticipated, the PyTorch baseline allocates a large intermediate tensor to retailer the activations, leading to a peak reminiscence utilization of 36.02GB. As compared, our Triton kernel reduces the height reminiscence utilization by 84% by allocating solely 5.04GB utilizing D_BLOCK=64 and V_BLOCK=64!

Utilizing even smaller block sizes would enable for additional reminiscence features at the price of effectivity.

Atomic Limitations and Manufacturing Scaling

On this article, we targeted on the technical and mathematical instinct behind fused Linear Cross-Entropy kernels. We used atomic operations like tl.atomic_add to maintain the code minimal and readable. Nonetheless, whereas our kernel efficiently slashed reminiscence utilization by a staggering 86%, the Triton kernel is considerably slower than native PyTorch.

Sadly, the identical atomic operations which make this kernel simpler to write down and comprehend come at the price of a large site visitors jam since hundreds of threads attempt to modify the identical reminiscence handle directly. Usually, tl.atomic_add is performant when rivalry is low. In our present implementation, we’ve:

1. Excessive Rivalry: For the load gradient, each single program within the batch (as much as 16,384 in our check) is making an attempt to replace the identical reminiscence tiles concurrently. The {hardware} should serialise these updates, forcing hundreds of threads to attend in line.
2. Numerical Non-associativity: In computer systems, floating-point addition is non-associative. Rounding errors can accumulate otherwise relying on the order of operations, which is why correctness exams would possibly move on a T4 however fail on an A100, the latter has extra streaming multiprocessors (SMs) performing extra concurrent, non-deterministic additions.

Notice on Precision: On Ampere and newer architectures, the TF32 format can additional contribute to those discrepancies. For strict numerical parity, one ought to set allow_tf32=False or use larger precision varieties through the accumulation steps.

Path to Manufacturing

To maneuver past this academic implementation and towards a production-ready kernel (I like to recommend trying on the Liger-Kernel implementation), one may implement a number of optimisations:

• Changing dX Atomics: Since every program “owns” its row of X, we are able to use easy register accumulation adopted by a tl.retailer, eliminating atomics for the enter gradients fully.
• A devoted dW Kernel: To optimise the computation of dW, manufacturing kernels typically use a special grid technique the place every program handles a block of W and iterates by the batch dimension, accumulating gradients regionally earlier than a single international write.
• Micro-batching: Superior implementations, similar to these within the Liger-Kernel library, course of the sequence by blocks alongside the N dimension, making the reminiscence scaling fixed within the sequence size slightly than linear. This allows the use a lot bigger batch sizes at a diminished reminiscence value.

Conclusion

This concludes our deep dive into fused linear cross-entropy kernels. Thanks for studying throughout, and I hope this text gave you each the instinct and the sensible understanding wanted to construct on these concepts and discover them additional.

Should you discovered this convenient, think about sharing the article; it genuinely helps assist the effort and time that goes into producing this work. And as all the time, be at liberty to contact me when you’ve got questions, ideas, or concepts for follow-ups.

Till subsequent time! 👋

Sources

1. Introducing Meta Llama 3: The most capable openly available LLM to date
2. LigerKernel (lecture)
3. LigerKernel Linear Cross-Entropy Implementation
4. Unsloth Implementation (cross-entropy only)