jax-ml/scaling-book: A Jekyll Book About TPU Scaling, Not a Library
Home for "How To Scale Your Model", a short blog-style textbook about scaling LLMs on TPUs
At a glance
- What is it?
- How To Scale Your Model is a Jekyll and Distill-style textbook repository about parallelism and communication bottlenecks for LLMs on TPUs. It is documentation you build and read, not code you install, and that distinction decides whether it belongs in your stack.
- Who is it for?
- Adopt this if you are choosing a parallelism scheme for TPU training or inference and want the roofline and communication reasoning written down rather than inferred from scattered posts. Do not adopt it if you run on GPUs, if you need a dependency you can import, or if you want a stable citable artifact rather than a site that changes with every push to main.
- Can I use it commercially?
- Yes. MIT is a permissive licence: you can use, modify and sell software built on it, as long as you keep its copyright and licence notices.
- Is it still maintained?
- Yes. The repository last received commits 12 days ago.
- What is it written in?
- Mainly HTML, according to GitHub's language statistics.
Answers come from the project's GitHub data, last synced on September 15, 2026, and from our analysis. They are not legal advice.
DEEP OPEN-SOURCE ANALYSIS
The problem is scheme selection, not another JAX wrapper
The repository README states the aim plainly: to explain how TPUs work, how LLMs run at scale, and how to pick parallelism schemes during training and inference that avoid communication bottlenecks. That is a narrow problem, and it is a real one. Once a model no longer fits on one chip, the question stops being about model code and becomes about where tensors live and how much data moves between them. Most teams learn this by running into a wall, then reading forum threads that disagree with each other.
The audience is correspondingly narrow. The topics list names jax, llm-inference, llms, roofline and tpus. If you are writing PyTorch training loops on a single accelerator, nothing here addresses your bottleneck. The book is written by Jacob Austin, Sholto Douglas, Roy Frostig, Anselm Levskaya, Charlie Chen, Sharad Vikram, Federico Lebron, Peter Choy, Vinay Ramasesh and Albert Webson at Google DeepMind, with the README noting that many of the ideas were first derived by James Bradbury and Reiner Pope. That provenance matters for how you read it: this is the reasoning of people who build JAX and TPU systems, written for people about to make the same decisions.
What is actually in the repository: a site, a build script, a licence
The primary language is HTML, which tells you what you are cloning. This is a Jekyll site using a Distill-style theme, credited in the README to al-folio and the Distill team, not a Python or C++ package. There is no import path, no wheel, no versioned API. The homepage points at https://jax-ml.github.io/scaling-book/, and the GitHub Pages deployment is handled by a GitHub Action that runs automatically on new commits to the main branch. Reading the book requires no local build at all; the build instructions exist for people who want to edit it or work offline.
The one piece of tooling that is not about rendering is bin/convert_to_single_md.py. According to the README it combines all chapters into scaling-book-combined.md in the repository root, and it does five specific things: strips Jekyll frontmatter, converts {% include figure.liquid %} to standard markdown images, converts internal page links to anchor links, converts inline $$ math to $, and strips unsupported LaTeX commands while emitting warnings. That last behaviour is the interesting one. It means the combined file is a lossy derivative of the chapters, and the warnings are the only record of what was dropped. The README also gives the follow-on command, pandoc scaling-book-combined.md -o scaling-book.docx, for a Word version.
Licensing is MIT. The README does not attach conditions to reading or quoting the text, but it does state that contributing on GitHub requires signing a Google Contributor License Agreement, with the link https://cla.developers.google.com/clas. So the asymmetry is: MIT for the material, a CLA for patches back to the repository. That is a normal arrangement for a corporate-hosted project, and it is worth knowing before you plan a fork you intend to upstream.
Building the book locally, and why Ruby 3.4.5 is the first hurdle
The README gives a concrete dependency list: Ruby, ImageMagick and Jupyter. On macOS the stated commands are brew install imagemagick ruby followed by pip install jupyter. Then it warns that you should make sure the correct version of Ruby is found in PATH, and that you should have at least Ruby 3.4.5 installed. That version floor is the part most likely to bite. A system Ruby from an older distribution will not be enough, and the README does not spell out a version manager, so you are expected to know how to put a newer Ruby ahead of the system one.
Once that is settled, the sequence is git clone https://github.com/jax-ml/scaling-book.git, cd scaling-book, bundle install, then bundle exec jekyll serve. The README states that after a successful serve the book is available at http://127.0.0.1:4000/scaling-book. Note the path prefix: the site is served under /scaling-book, not at the root, which matches the GitHub Pages URL. There are no environment variables, no config keys and no secrets documented in the material, which is consistent with a static site. Nothing here needs a GPU, a TPU or a cloud account. The heaviest requirement is ImageMagick, presumably for asset processing, and the README does not explain which assets depend on it.
The combined-document path is lighter: python bin/convert_to_single_md.py, then optionally pandoc. If your goal is to read the book as one file or hand it to a colleague, that route skips Ruby entirely.
The conversion script is lossy and tells you so only in warnings
The most consequential limitation is not a bug, it is the design of the single-document export. Four of the five transformations are mechanical and safe: frontmatter removal, figure include rewriting, internal link rewriting to anchors, and inline $$ to $. The fifth, stripping unsupported LaTeX commands, removes content. The README says warnings are emitted, which means the script knows what it discarded, but it does not say the warnings are collected anywhere durable. If you pipe the output to a file and scroll past stderr, you can end up with a combined markdown file that silently lost notation from a book whose subject is quantitative reasoning about hardware.
There is a second, quieter loss. Internal page links become anchor links, which is correct inside a single file, but the reverse direction does not exist. If you cite a section from the combined markdown, your reader cannot follow it back to the live page without searching. For a textbook that the README asks you to cite in academic contexts, that is a friction point worth planning around.
And the export is a snapshot. The site deploys on every commit to main, and there are no retrieved releases, so there is no tag to pin your copy to. The BibTeX entry in the README gives the year as 2025 and the note field as a retrieval URL, which is honest about the fact that this is a living document rather than a fixed edition.
What this is not: a library, a course, or a GPU resource
The most useful comparison is not another book. It is the JAX documentation itself. The JAX docs describe what each API does and how to call it; this repository describes why a given sharding or parallelism choice costs what it costs on TPU interconnect, and the README frames the whole project around avoiding communication bottlenecks. If you need the signature of a collective operation, the API reference is the right place. If you need to decide whether a scheme is worth its communication, that is the question this book sets out to answer. They are complementary, and the book assumes you already know the APIs.
A second comparison is a vendor's own scaling guide. Those tend to be tied to one product line and are updated when the product changes. This one is tied to TPUs specifically, which is stated in the README and reflected in the topics. That is a boundary, not a flaw. The reasoning about rooflines transfers to other accelerators in outline, but the numbers and the topology do not, and the material will not tell you where the transfer stops.
What it is not, despite the repository shape, is software you adopt. There is no version to upgrade, no dependency graph to audit, no runtime to monitor. The maintenance question is inverted: the cost is in keeping your local build working and in re-reading sections when the underlying hardware generations move, not in patching a dependency.
Maintenance, contribution and the cost of a living document
There are no retrieved releases, so there is no changelog to diff between two points. The last push recorded is 2026-09-03, and the deployment action fires on main, which means the published site and the repository head are effectively the same thing at any moment. For a reader that is convenient. For anyone citing a specific claim in a paper or design doc, it is a hazard, because the sentence you quoted may be edited without any version marker. The mitigation is to generate scaling-book-combined.md at the commit you read and keep that file alongside your citation, rather than linking to the live site alone.
Contributing has its own cost. The README directs issues and questions to Giscus comments on the website itself or to GitHub discussions, and pull requests are welcome, but GitHub contributions require the Google CLA. If your organisation has restrictions on signing contributor agreements with third parties, that is a gate you will hit before your first patch, and the README does not describe an alternative path other than emailing jaaustin [at] google [dot] com.
On licence implications, MIT is permissive and the README does not add restrictions on reuse of the text. That is a statement about what the repository says, not legal advice. If you plan to redistribute the book inside a commercial training product, read the licence text itself rather than this summary.
The MIT licence covers the repository contents. It does not tell you anything about the TPU hardware the book describes, and the book's accuracy on that subject is bounded by when each chapter was last edited.
Who should clone this, and who should just read the site
Clone it if you are editing chapters, if you need an offline copy for an air-gapped environment, or if you want the combined markdown for a documentation pipeline. The build is three commands after the toolchain is in place, and the only real prerequisite is getting Ruby 3.4.5 or newer ahead of your system Ruby in PATH.
Read the hosted site and skip the clone if you only want the content. There is no reason to install ImageMagick and Jupyter to read prose, and the GitHub Action already publishes every commit.
Stay away if your stack is GPU-based, if you want something importable, or if you need a fixed edition with a version number to cite. The absence of releases is the clearest signal that this is a site first and an artifact second.
Before you rely on any specific number in the book, check the commit date against the TPU generation you are actually provisioning, and regenerate scaling-book-combined.md from that commit so your citation points at text that will not move. If the conversion warnings mention dropped LaTeX in the chapter you care about, read that chapter on the live site instead of trusting the combined file.
Editorial conclusion
Adopt this if you are choosing a parallelism scheme for TPU training or inference and want the roofline and communication reasoning written down rather than inferred from scattered posts. Do not adopt it if you run on GPUs, if you need a dependency you can import, or if you want a stable citable artifact rather than a site that changes with every push to main. Before relying on it, run the local build with Ruby 3.4.5 or newer, run python bin/convert_to_single_md.py and read the LaTeX stripping warnings, and check the last commit date against the TPU generation you actually rent.
Community notes