Arun's blog

Notes on Local AI - How do i compute my flops

What are flops?

Table of Contents

What are flops?

flops

Yes, you got it right. It is Floating-Point Operations Per Second. Technically it can be Floating-Point Instructions Per Second (FLIPS). Or FLOPs, or you know, FLIP FLOPS. Doesn't really matter. I wish it had a cooler name.

I love it when people just say FLOP. Floating-Point Operation Per..... COMPLETE IT! Did your GPU just die without getting the last token? Although 1 FLOP means just one floating-point operation like an add or multiply. We all have agreed silently to it. No one questions it.

How do I compute my flops? (Matrix Multiplication)

Anyways, how do I compute my FLOPs? Let's talk about matrix multiplication:

Say we have A = [M, K] and B = [K, N] — the matrix multiplication of A and B results in C = [M, N].

What exactly is matrix multiplication though? It's just taking a row from A, a column from B, doing a multiply of 2 elements, and then adding all the results.

AI SLOP STARTS

Matrix Multiplication Formula
When multiplying an [M, K] matrix A by a [K, N] matrix B, the resulting [M, N] matrix C is calculated using the following summation equation:

C_ij = Σ (A_ik * B_kj) (for k = 1 to K)

Where:

Visual 2x2 Matrix Example

[ a  b ]   [ e  f ]   [ (ae + bg)  (af + bh) ]
[ c  d ] x [ g  h ] = [ (ce + dg)  (cf + dh) ]

AI SLOP ENDS

So in the nice visual example above, for every pair of elements—say a and e—we do 1 multiply and 1 add. So 2 FLOPs (argh!). If you sum up all of the operations, it will come out as 12.

You might also see 2 * M * K * N widely used to compute FLOPs. That 2K is actually 2K - 1 (since adding K numbers takes K - 1 additions). But generally, the matmul in LLMs has big Ks like 8192, so the -1 is negligible and rounded off for easier computation.

Exact FLOPs = M * N * (2K - 1)

So in this 2x2 case: 2 * 2 * (4 - 1) = 12.

Full Model FLOPs: Forward vs. Backward Pass

Now that we have covered how to calculate Matrix Multiplication FLOPs, let's look at how FLOPs for a full model—both forward and backward pass—are calculated.

In the land of inference, we only care about a forward pass. For training, we have a forward pass, a backward pass, and an optimizer step. The optimizer step is a relatively low-FLOP module compared to the forward and backward passes. So we just ignore the FLOPs from the optimizer in the bigger picture as a rounding error.

Inference has 2 parts—prefill and decode—but for the purpose of this blog, let's stick to only forward and backward pass terminology and not worry about inference prefill, decode, or training specifics.

The Forward Pass

Let's talk about the forward pass. Assuming you know what that is: if not, check out my transformers blog for a primer on the architecture.

Yes, Transformers have LayerNorms, Softmax, RoPE, and residuals. But Self-Attention and the MLP are the main FLOP-heavy modules. The rest of them become rounding errors.

Define Your Shapes

Before you compute FLOPs, you should do one thing for sure. It's probably something you learnt as a toddler. SHAPES! Define the shapes or you WILL lose your mind.

Let's go:

  1. Sequence length (S) - how many tokens in one sequence.
  2. Batch size (B) - One sequence only? Noob! Shove in more.
  3. Hidden dimension (D) - usually 4096 or 8192 (can't wait for these numbers to get outdated).

We pass B * S tokens (let's call total tokens T = B * S) into the model, so our input matrix shape is [B*S, D]. From our earlier matmul formula [M, K] @ [K, N], this means M = B*S.

Attention Block

We have Q, K, V projections first. So the [B*S, D] input gets multiplied with a [D, D] weight matrix.

Using our very trusted 2 * M * K * N for FLOPs calculation, we get:

There is a nuance if you are using GQA here. Let's stick to standard multi-head attention so my brain doesn't hurt much.

Next, we have the attention magic: (Q @ K^T) @ V

Put those together and it comes out to: 4*B*S²*D

S² is the painful part here. Longer sequences quadratically increase FLOPs. If your sequence length is 128k, it's over.

And finally, the output projection: 2*B*S*D² (same size as one projection from above).

So total Attention FLOPs is: 8*B*S*D² + 4*B*S²*D

MLP Block

Let's stick to a classic Transformer here. No MoEs. More brain hurt if we do that.

Typically, the projections here go from D to 4D (4D is called the intermediate size, but let's stick to using D and not introduce I here).

Total MLP FLOPs is: 16*B*S*D²

Whole Transformer Block - 1 Layer - Total FLOPs

Hmmmmmmmm, isn't there a shortcut? The forward pass is supposed to be 2 * N * T (where N = params and T = tokens).

Let's add up the heavy matmuls and ignore the S² attention one for a second:

Now let's think: how many parameters (N) are in one Transformer layer?

Look at those two numbers side-by-side:

So it's literally:

Forward Pass FLOPs = 2 * (B*S) * (12D²) = 2 * T * N

Lowkey proud that I derived this ;)

So you have an 8B parameter model and 1 token? Forward pass is 2 * 8B * 1 = 16 GFLOPs. Done. Go home.

Now the Backward Pass

Great, we can price out inference. What happens when we actually have to train the thing?

We need to compute gradients during the backward pass to learn something. If you remember your calculus, for a linear layer Y = X @ W, we need to calculate 2 things:

For every single matmul in the forward pass, we need to do 2 matmuls in the backward pass.

Backward Pass FLOPs = 4 * N * T

Simplified for brain hurt—although it's different if you are using activation checkpointing to save memory. Then you have to recompute a forward pass again during the backward pass. That adds another 2NT, making the backward pass 6NT by itself. Memory isn't free (it's expensive)—but you can pay for it with FLOPs, not just $.

The Legendary 6NT & MFU

So that is IT! Without activation checkpointing, your cheat sheet for a full training step is:

Whenever a frontier lab says "we trained a model with 10^25 FLOPs," it means they did 6NT. Not count every single FLOP. Or maybe they did count it. Who knows!

One more thing you pray for: when you take that total FLOPs number and divide it by the time it took your GPU to run it, you get actual FLOP/s—which gives you a magical percentage called MFU (Model FLOPs Utilization). You pray for more MFU. (Keep praying, or your VC will not like you.)