Serving LLaMA 3-70B on TPUs | How To Scale Your Model

[jacobaustin123]

Serving LLAMA on TPUs!

[samacqua]

"Notably, at all batch sizes greater than 2k, FLOPs is always smaller than our KV loading time in this regime." Is this a typo, where it should be 200 instead of 2k?

[jacobaustin123]

Rather, it meant to say "Notably, at all sequence lengths greater than 2048, we spend more time on KV cache loading than we do on FLOPs! So while we can improve our hardware utilization by increasing batch size, at long context lengths KV loading always dominates the total step time." I've gone ahead and updated the doc.

[gitnicos]

I think Question 3 answer has mistaken wording:
If we were to run on a 4x4, we’d still be fine ICI-wise, so our latency would drop to 17 / 2 = 8.5ms, but our throughput would remain the same.
instead the throughput is twice as big but per chip will remain the same

[jacobaustin123]

Fixed, not says "but our throughput per chip would remain"

[LeonEricsson]

A couple of places throughout this section and the last use $Y > F/2550$ to determine ICI bounds for TP. This should be $Y > F/2200$ as we are using the v5e. 2550 came from v5p.

[jacobaustin123]

Thanks, just pushed a fixed for this. They're close enough I sometimes forget which is which.