Conclusions and Further Reading

Part 11 of How To Scale Your Model (Part 10: JAX | ...)

Thank you for reading! Here we'll include a few more references for further study.

Thank you for reading this set of essays and congratulations on making it all the way to the end. Before we conclude, a few acknowledgments:

Acknowledgments

This document represents a significant collective investment from many people at Google DeepMind, who we’d like to briefly acknowledge!

• James Bradbury, Reiner Pope, and Blake Hechtman originally derived many of the ideas in this manuscript, and were early to understanding the systems view of the Transformer.
• Sholto Douglas wrote the first version of this doc and is responsible for kicking off the project. He is more than anyone responsible for the overall narrative of this doc.
• Jacob Austin led the work of transforming this first version from rough notes into a more polished and comprehensive artifact. He did much of the work of editing, formatting, and releasing this document, and coordinated contributions from other authors.
• Most of the figures and animations were made by Anselm Levskaya and Charlie Chen.
• Charlie Chen wrote the inference section and drew many of the inference figures.
• Roy Frostig helped with publication, editing, and many other steps of the journey.

We’d also like to thank many others gave critical feedback throughout the process, in particular Zak Stone, Nikhil Sethi, Caitlin Stanton, Alex Dimitriev, Sridhar Lakshmanamurthy, Albert Magyar, Diwakar Gupta, Jeff Dean, Corry Wang, Matt Johnson, Peter Hawkins, and many others. Thanks to Ruiqi Gao for help with the HTML formatting.

Thank you all!

Further Reading

There is a bunch of related writing, including the following:

• Writing TPU Kernels with Pallas: increasingly, TPU programming involves writing custom kernels in Pallas. This series discusses how to write kernels and many lower level TPU details that aren’t mentioned here.
• How to Optimize a CUDA Matmul Kernel for cuBLAS-like Performance: a Worklog: while GPU and CUDA specific, this is an excellent blog post showing how to optimize a matmul kernel in CUDA. This might be a good deep dive into how TPUs and GPUs are different.
• Distributed arrays and automatic parallelization: this is a really nice guide to parallelism APIs in JAX and is a good way to learn how to actually implement some of the ideas we’ve discussed here.
• Rafi Witten’s High Performance LLMs 2024 Class: our former colleague Rafi gave a great course on TPU performance engineering and the slides are all on GitHub. This covers a bunch of things in more depth than we do here.
• [2211.05102] Efficiently Scaling Transformer Inference: a detailed paper on the mathematics of Transformer inference. This is the inspiration for a lot of this document.

There remains a lot of room for comprehensive writing in this area, so we hope this manuscript encourages more of it! We also believe that this is a fruitful area to study and research. In many cases, it can be done even without having many hardware accelerators on hand.

Feedback

Please leave comments or questions so that we can improve this further. You can reach our corresponding author, Jacob Austin, at jaaustin [at] google [dot] com, or suggest edits by posting issues, pull requests, or discussions on GitHub.