jax-ml/scaling-book resource

Health v2 · maintenance only

Flags: no_releases

Adoption not part of the score

1371 stars · 194 forks observed · 2026-08-28

What it is AI-extracted, prompt v1, taxonomy v1, 2026-08-30, confidence not recorded

An open-source, blog-style online textbook called 'How To Scale Your Model' that explains how TPUs and GPUs work and how to scale LLM training and inference efficiently. It is built as a Jekyll site with interactive Jupyter-based content, published via GitHub Pages.

Use cases

• learn how TPUs work and how LLMs run at scale
• understand roofline analysis for accelerator performance
• choose parallelism schemes for LLM training and inference
• estimate the cost and time to train or serve a large Transformer model
• understand collective communication operations like AllGather
• learn how to avoid communication bottlenecks when scaling models

When to choose

• you want a free, in-depth systems-level explanation of LLM scaling on TPUs/GPUs
• you need to reason about parallelism strategies (data, tensor, pipeline, etc.) for large models
• you want to ballpark training/inference costs and memory requirements
• you have basic LLM and JAX familiarity and want to go deeper on hardware efficiency

When to avoid

• you need hands-on training code or a runnable framework rather than a textbook
• you are a complete beginner with no understanding of Transformers or LLM training
• you need vendor-specific GPU (CUDA) optimization guides rather than TPU-focused systems theory

Facets

learning-resource · maturity active

documentation developer-tools large-language-models gpu-computing tutorials performance machine-learning python ruby tpu llm-training llm-inference roofline-analysis parallelism jax textbook distill-style web-server

Member repositories

For agents

markdown · JSON · MCP: product_card(name="jax-ml/scaling-book")