All About Transformer Inference | How To Scale Your Model

[jacobaustin123]

Serving big models!

[clu0]

In the first figure for sharding KV comparing MHA and MQA, I think the locations of the reduce-scatter and all-gather are reversed.

[rohanmahajan1993]

Had question and was trying to understand more.

1. "For prefill, since we’re doing self-attention, so this simplifies to
   This is great because it means the arithmetic intensity of attention during prefill is
   Θ(T). That means it’s quite easy to be compute-bound for attention. As long as our batch size and sequence length are both fairly large, we’ll be fine!"

why do we need our batch size also to be fairly large, it seems like the equation is defined by token length; so isn't that sufficient

"Why is this, conceptually? Mainly, we’re compute-bound in linear portions of the model because the parameters (the memory bandwidth-heavy components) are reused for many batch items. However, every batch item has its own KV cache, so a bigger batch size means more KV caches. We will almost always be memory bound here unless the architecture is adjusted aggressively."

another question why the batch size effects this being memory bound; doesn't batch size cancel out here in the equations you're showing; I get the point that t=1 for generation, and since we have to fetch the kv cache stuff for s, we're still doing kind of half the memory work for that but we lose the s*t multiplication for the flopers inference.

[jacobaustin123]

1. I think this should basically just say "as long as our sequence length is fairly large". I think before this was support to say "if our batch size is large our MLPs are good, and if our sequence length is long, our attention is good".

2. Could you restate this slightly? The idea is, for our MLPs, we load one matrix for B entries, so the FLOPs / byte loaded increases with batch size (since our activations are small). For attention, as we increase B, we load B times more KV caches, so our intensity doesn't improve.

[rohanmahajan1993]

On 2, it was totally my bad. I think initially I misread and thought the text was implying that increasing batch size decreases the intensity of the attention but in reality the point it was that it doesn't change the attention part.

[mmcclean-aws]

I have a question on the Generation section and the calculations around time for HBM loading. It seems to only consider the weight loading time and not also the KV cache loading time as well. Does this need to be added to the calculation ?

[jacobaustin123]

Do you mean under "Distributing Inference Over Multiple Accelerators"? I think in general here we're looking at the matmul vs. ICI comms overlap, which is unrelated to the KV caches. Those will always be memory-bound and more model parallelism will always improve things

[mmcclean-aws]

Am referring to this formula. The time for HBM comms is only considering model weights and not KV cache loading time. For long sequences the KV cache can be larger than the model weights so should also be taken into consideration

[jacobaustin123]

Think about what we're trying to calculate here. We want to know how much model parallelism we can do before ici communication during our matmuls takes longer than weight loading. At that point, we hit a limit and more model parallelism won't reduce the time our matmuls take. It will reduce the time KV loading takes, so if all you care about is latency, you can do even more.

So I guess the point is, this is the rule for how much model parallelism you can do before you start hitting diminishing returns on throughput at a fixed batch size. But you can go beyond it if you're willing to take a throughpout hit.

[Lhongpei]

About the part of "exact number will differ on the kind of quantization and hardware".If we use int8 to store parameters but bf16 for computation, why is the result of critical 120？ According to the following formulas:

$$ \frac{2BDF}{2BD + 2DF + 2BF} \geq \frac{\text{TPU FLOPs/s}}{\text{HBM Bandwidth}} = \frac{1.97E+14}{8.20E+11} = 240 $$

For int8, the communication will decrease to half, then the critical B will be 480?

[jacobaustin123]

Not quite. If we simplify, we end up with 2B > 240, or B > 120.

[aasthavar]

For the formula:

FFW params = d_model^2 * ffw_multiplier * 3 * n_layers

I am unsure where the *3* came from? Shouldn't it be just *2*?

For one layer, given:

• Input Projection:
  • W_in.shape = [d_model, d_mlp]
  • b_in.shape = [d_mlp]
• GeLU
  • no weights
• Output Projection:
  • W_out.shape = [d_mlp, d_model]
  • b_out.shape = [d_model]

Total params for one layer:

params = (d_model * d_mlp) + (d_mlp * d_model) 
       = (d_model^2) * ffw_multiplier * 2

Total params across all layers:

total_params = params * n_layers 
             = (d_model^2) * ffw_multiplier * 2 * n_layers

[jacobaustin123]

So this is a little confusing in this book because we assume what you describe in Section 5 for simplicity.

However, GeGeLU, the gated variant described in https://arxiv.org/abs/2002.05202v1 has two up-projection matrices whose output is pointwise multiplied. There’s a note on this in Section 4.

[aasthavar]

Shouldn't there be positional embedding matrix (of shape = [max_seq_len, d_model]) as well, along with ffw+vocab+attn params?

[jacobaustin123]

Sometimes. Many models today use some kind of parameter free positional embedding, or where the number of params is very small, eg ROPE, T5-style relative positional embedding, etc.

Either way, the number is so small it essentially doesn’t matter.

[VijayGKR]

For the newly defined model we're using in the worked problems, should the number of heads be 16 instead of 32, or qk dim be 128 instead of 256? Right now, d_qkv * n_heads is 8192, instead of matching d_model at 4096.

[jacobaustin123]

Probably yes, maybe worth updating. Still, it's not a strict rule that d_qkv * n_heads = d_model. The main thing that matters is the number n_kv heads

[andrewkchan]

Note on ICI bounds for generation: during training we want to be compute-bound, so our rooflines look at when our ICI comms take longer than our FLOPs. However, during generation, if we’re memory bandwidth bound by parameter loading, we can increase model sharding beyond this point and improve latency at a minimal throughput cost.

I don't understand why there would be any throughput cost when increasing model sharding to improve latency in generation while memory-bound.

1. If we're holding batch size constant, then doesn't improving latency strictly improve throughput?
2. Otherwise, IIUC increasing the number of accelerators we shard over yields more HBM, which means we can increase batch size. As long as we are memory-bound, we can actually therefore improve throughput while holding latency constant, no?

[jacobaustin123]

I agree that increasing the number of chips in theory should leave throughput constant if we don't change the batch size. If we're memory-bound on 8 chips, and then we increase to 16 and we're still memory-bound, we're using 2x more chips for 2x more throughput (since the weight loading time is halved), so totally constant.

In practice we tend to suffer some downside from trying to do more model sharding, like some ICI overhead that isn't totally free, since it's rare for us to be fully memory-bound. So I think you have the right idea, it's just that in practice it's likely to have some overhead.

[andrewkchan]

I see. In that case I'm confused about what definition of decode throughput we're using. I was assuming throughput = batch_size / latency, where a decrease in latency with batch size constant must improve throughput. Maybe worth clarifying?

[jacobaustin123]

I see, I typically mean batch size / latency / num chips, i.e. tokens/sec/chip. I'll clarify, but I think that's the quantity that matters.

[zhiyisun]

Minor typo in the first "Takeaway"

Takeaway: Transformer matmuls are compute-bound iff the per-replica token batch size is greater than

[jacobaustin123]

This was intended, cause iff is sometimes used in math to mean "if and only if". It's a kind of niche usage though.

[mayun8172]

In the multi query attention full algorithm,

The step 7 is confused to me

6. O[B_Z, S, K_Y, M] = Q[B_Z, K_Y, M, H] *_H K[B_Z, S, K_Y, H]
7. O[B_Z, S, K, M] = Softmax_S(O[B_Z, S, K_Y])

Do we should softmax along Y instead ?

7. O[B_Z, S, K, M] = Softmax_Y(O[B_Z, S, K_Y, M])