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 contemplate 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 can be the sequence size multiplied by the batch measurement (i.e. the entire variety of tokens within the batch), D the scale of the hidden state and V the vocabulary measurement. 

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

This ends in the Y matrix having form [NxV]=[65,536x128,256], or roughly 8.4 billion components. In bfloat16, this might take up 16.8GB of reminiscence. Nonetheless, if we comply with finest 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 house 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 measurement and accumulate gradients over a number of batches between every optimiser step, emulating a bigger batch measurement whereas holding much less knowledge in reminiscence.
•  Activation checkpointing: PyTorch shops all intermediate activations for reuse within the backward cross, checkpointing clears these activations and recomputes them on-the-fly throughout the backward cross. This results in giant 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 without delay, we will slice it and accumulate the loss over smaller chunks with measurement n < N. Now, we solely maintain a slice of measurement [n, V] in reminiscence at a time.
• Combined precision coaching: Utilizing half precision throughout coaching supplies 2x reminiscence discount and important speedups on Tensor Cores.

Whereas these options appear engaging, all of them have important drawbacks: gradient accumulation and activation checkpointing decelerate coaching, combined precision could be unstable and micro-batching requires (gradual) PyTorch stage iteration and although n is chosen to be smaller than N, the vocabulary measurement stays enormous as compared.

Extra importantly, these options don’t deal with the issue we have now handled repeatedly all through this sequence: knowledge 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 cross of the cross-entropy loss contain dot merchandise, matrix multiplication and a softmax. As we discovered on this sequence, these are all operations that may be tiled effectively. In different phrases, we will 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 house. It is a nice alternative for operator fusion: fusing a number of operation inside a single kernel, leading to giant speedups and potential reminiscence good points.

Within the following sections, we’ll check out the right way to derive and effectively fuse the ahead and backward passes by means of 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. Subsequently we have to derive the gradient ourselves, a beautiful event to brush up on some calculus 😉

The mathematics 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 acquire 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 are able to 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 cross boils down to a few major 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.

Due to the web softmax algorithm, we will 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 means of the hidden and vocabulary blocks to:
    1. Monitor the working most (m)
    2. Replace the working sum of exponentials (d) utilizing the web softmax components:

Now that we have now a greater understanding of the ahead cross, let’s check out the derivation of the backward cross.

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, a regular 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 identification logic. It is the same as 1 if i=j and 0 in any other case. As we’ll see, that is significantly 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 means of matrix multiplication, we use the apply the chain rule to the inputs x and the burden 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. Specific y when it comes to x and w.
2. Discover that w is a continuing with respect to the spinoff of x, so we will pull it out of the spinoff.
3. Specific the truth that the partial spinoff of x_ik with respect to x_mn is 1 solely when i=m and ok=n utilizing the Kronecker delta.
4. Discover that ẟ_kn enforces ok=n, subsequently w_kj * ẟ_kn reduces to w_nj.

Then, we contemplate the total expression and acquire 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, subsequently, we rewrite this expression as a matrix multiplication by utilizing the identification X_ij = [X^T]_ji:

We comply with 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 concentrate on 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 ok. Due to the normalising issue, each factor i impacts the worth of each different factor, subsequently, the partial spinoff is obtained by defining the perform piecewise relying on the worth of i:

Summing each instances, we acquire the gradient:

And in matrix notation:

The place y_{one scorching} is a vector of zeros with the entry similar to the label set to at least one. This outcome tells us that the gradient is solely the distinction between the prediction and the bottom fact.

Fused Linear Cross-Entropy

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

Due to the chain rule, deriving the gradient of this expression boils all the way down to multiplying the gradients we computed beforehand:

The place x and y seek advice from the inputs and outputs to the linear layer respectively and w to the related weight matrix.

Word: 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 will implement the fused kernel in Triton. Since cross-entropy is usually the ultimate layer in a language mannequin, we will 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 knowledge locality by maintaining 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 section entails customary Triton setup:

• Program Identification: We use tl.program_id to find out which row of the enter matrix the present program is accountable for.
• Parameter Initialisation: We outline tiles utilizing D_BLOCK and V_BLOCK and initialise the working most (m) and sum (d) required for the web 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 deal with as a result of each program should ultimately iterate by means of the complete vocabulary house.
• 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 principle 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 particular column of W similar 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 means of reminiscence to extract solely the mandatory weights for the goal token.

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

Right here’s the code to date:

Subsequent, we execute the ahead cross, which processes the vocabulary house 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 (internal loop). Inside the internal 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 working statistics: the utmost worth m and the working sum of exponentials d utilizing the web softmax components. 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 cross:

Right here’s the code for the ahead cross:

We are actually all the way down to the final a part of the kernel: the backward cross. 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 throughout the ahead cross, we should recompute the activations for every tile. By reusing the Log-Sum-Exp calculated within the ahead cross, we will normalise these activations on-the-fly. Subtracting the ground-truth label Y from the goal logit inside this tile provides us a neighborhood 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 the complete batch, we use Triton’s tl.atomic_add.
   This operation acts as a thread-safe +=, guaranteeing that completely different applications updating the identical weight gradient don’t overwrite each other.

Listed here are some further particulars on the implementation:

• The Stride Swap: When computing P . W_T, we don’t really 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 ends in 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 stop the buildup of tiny rounding errors throughout hundreds of rows.
• Rivalry Word: 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. Subsequently, we implement the pointer and masks logic for dW explicitly.

The next determine is a illustration of the backward cross 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 cross:

This concludes the implementation of our kernel! The total code together with the kernel and benchmark script is on the market here.

Reminiscence Benchmark

Lastly, we evaluate our kernel with the PyTorch baseline utilizing hyperparameters impressed from Llama3 and an A100 GPU. Particularly, we contemplate a sequence size of S=16,384, a batch measurement of B=1 and an embedding dimension of D=4096; the vocabulary measurement is ready 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. 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 permit for additional reminiscence good points 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 put in writing and comprehend come at the price of a large site visitors jam since hundreds of threads attempt to modify the identical reminiscence deal with without delay. Usually, tl.atomic_add is performant when rivalry is low. In our present implementation, we have now:

1. Excessive Rivalry: For the burden gradient, each single program within the batch (as much as 16,384 in our take a look at) is attempting 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 in another way relying on the order of operations, which is why correctness exams may cross on a T4 however fail on an A100, the latter has extra streaming multiprocessors (SMs) performing extra concurrent, non-deterministic additions.

Word 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 throughout 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 might implement a number of optimisations:

• Changing dX Atomics: Since every program “owns” its row of X, we will use easy register accumulation adopted by a tl.retailer, eliminating atomics for the enter gradients solely.
• A devoted dW Kernel: To optimise the computation of dW, manufacturing kernels typically use a unique grid technique the place every program handles a block of W and iterates by means of the batch dimension, accumulating gradients domestically earlier than a single world write.
• Micro-batching: Superior implementations, reminiscent of 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 quite than linear. This permits the use a lot bigger batch sizes at a lowered reminiscence value.

Conclusion

This concludes our deep dive into fused linear cross-entropy kernels. Thanks for studying all through, 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 handy, contemplate sharing the article; it genuinely helps assist the effort and time that goes into producing this work. And as at all times, be happy to contact me you probably have 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)