{"data":{"slug":"jax-ml-scaling-book","name":"scaling-book","tagline":"Guide on scaling LLMs on TPUs","github_url":"https://github.com/jax-ml/scaling-book","owner":"jax-ml","repo":"scaling-book","owner_avatar_url":"https://avatars.githubusercontent.com/u/58486408?v=4","primary_language":"HTML","stars":1441,"forks":205,"topics":["jax","llm-inference","llms","roofline","tpus"],"archived":false,"github_pushed_at":"2026-09-19T16:14:09+00:00","maintenance_label":"Very active","stars_delta_30d":150,"url":"https://www.graphcanon.com/tools/jax-ml-scaling-book","markdown_url":"https://www.graphcanon.com/tools/jax-ml-scaling-book.md","api_url":"https://www.graphcanon.com/api/graphcanon/tools/jax-ml-scaling-book","graph_url":"https://www.graphcanon.com/api/graphcanon/graph?tool=jax-ml-scaling-book","description":"Home for \"How To Scale Your Model\", a short blog-style textbook about scaling LLMs on TPUs","homepage_url":"https://jax-ml.github.io/scaling-book/","license":"MIT","open_issues":10,"watchers":11,"ai_summary":"Provides insights and methodologies for scaling machine learning models particularly focusing on the use of tensor processing units (TPUs).","readme_excerpt":"# How To Scale Your Model\n\nThis book aims to demystify the art of scaling LLMs on TPUs. We try to explain how TPUs work, how LLMs actually run at scale, and how to pick parallelism schemes during training and inference that avoid communication bottlenecks. The book is available at https://jax-ml.github.io/scaling-book.\n\n### Acknowledgments\n\nThis book was 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. Many of the ideas were first derived by James Bradbury and Reiner Pope.\n\nThe website uses a Distill-style Jekyll theme created by https://github.com/alshedivat/al-folio and the Distill team. Thank you!\n\n### Running Locally\n\nTo build this repo locally, you will need Ruby, ImageMagick, and Jupyter installed, which for MacOS can be installed with Homebrew using\n\n```\nbrew install imagemagick ruby\npip install jupyter\n```\n\nAfter this is installed, you should make sure the correct version of Ruby is found in PATH. You should have at least Ruby 3.4.5 installed. Then you ought to be able to run\n\n```\ngit clone https://github.com/jax-ml/scaling-book.git\ncd scaling-book\nbundle install\nbundle exec jekyll serve\n```\n\nOnce you have run jekyll serve successfully, the book will be available at `http://127.0.0.1:4000/scaling-book`.\n\nThe Github Pages deployment is handled by a GitHub Action that runs automatically on new commits to the main branch.\n\n### Generating a Single Document\n\nTo combine all chapters into a single markdown file with standardized formatting:\n\n```bash\npython bin/convert_to_single_md.py\n```\n\nThis generates `scaling-book-combined.md` in the repository root. The script:\n- Strips Jekyll frontmatter\n- Converts `{% include figure.liquid %}` to standard markdown images\n- Converts internal page links to anchor links\n- Converts inline `$$` math to `$`\n- Strips unsupported LaTeX commands (with warnings)\n\nTo convert the combined markdown to a Word document:\n\n```bash\npandoc scaling-book-combined.md -o scaling-book.docx\n```\n\n### Contributing and Contact\n\nIf you see any issues or have questions, please leave a comment on the website itself (powered by Giscus) or in the GitHub discussion. Feel free to send a PR if you want to contribute. You can also email jaaustin [at] google [dot] com.\n\nTo contribute on GitHub you will need to sign a Google \"Contributor License Agreement\" (CLA). You can do that here: https://cla.developers.google.com/clas.\n\n### Citation\n\nFor attribution in academic contexts, please cite this work as\n\n```Austin et al., \"How to Scale Your Model\", Google DeepMind, online, 2025.```\n\nBibTeX citation\n\n```\n@article{scaling-book,\n  title = {How to Scale Your Model},\n  author = {Austin, Jacob and Douglas, Sholto and Frostig, Roy and Levskaya, Anselm and Chen, Charlie and Vikram, Sharad and Lebron, Federico and Choy, Peter and Ramasesh, Vinay and Webson, Albert and Pope, Reiner},\n  publisher = {Google DeepMind},\n  howpublished = {Online},\n  note = {Retrieved from https://jax-ml.github.io/scaling-book/},\n  year = {2025}\n}\n```\n\n\n\n*This book was originally called \"How To Scale Your Dragon\", after the Dreamworks film, hence the dragon imagery.*","github_created_at":"2025-02-03T01:38:25+00:00","created_at":"2026-07-11T11:43:48.777731+00:00","updated_at":"2026-09-20T06:33:54.803201+00:00","categories":[{"slug":"inference-serving","name":"Inference & Serving","url":"https://www.graphcanon.com/categories/inference-serving","markdown_url":"https://www.graphcanon.com/categories/inference-serving.md","api_url":"https://www.graphcanon.com/api/graphcanon/categories/inference-serving"},{"slug":"llm-frameworks","name":"LLM Frameworks","url":"https://www.graphcanon.com/categories/llm-frameworks","markdown_url":"https://www.graphcanon.com/categories/llm-frameworks.md","api_url":"https://www.graphcanon.com/api/graphcanon/categories/llm-frameworks"}],"tags":[{"slug":"jax","name":"jax"},{"slug":"llm-inference","name":"llm-inference"},{"slug":"llms","name":"llms"},{"slug":"roofline","name":"roofline"},{"slug":"tpus","name":"tpus"}],"trust":{"provenance":{"is_fork":false,"github_id":926297959,"owner_type":"Organization","methodology":"github_public_v1","parent_repo":null,"near_duplicate_slugs":[]},"computed_at":"2026-09-19T18:31:21.769Z","maintenance":{"label":"Very active","score":96,"methodology":"github_public_v1","releases_90d":0,"days_since_push":0,"last_release_at":null,"stars_delta_30d":150,"open_issues_delta_30d":3},"security_summary":{"status":"no_lockfile","scanner":null,"low_count":0,"high_count":0,"last_scan_at":"2026-07-11T11:43:50.067Z","medium_count":0,"scan_profile":"none","critical_count":0}},"capability_facts":{"mcp":{"source":"repo_scan","observed_at":"2026-09-19T18:31:22.866Z","server_manifest":false},"scan":{"source":"repo_scan","observed_at":"2026-09-19T18:31:22.866Z"},"languages":{"value":["html","javascript"],"source":"github.language+package.json","observed_at":"2026-09-19T18:31:22.866Z"},"license_spdx":{"value":"MIT","source":"github.license","observed_at":"2026-09-19T18:31:22.866Z"}},"decision_facts":{"hosting":null,"pricing":null,"requirements":{"min_ram_gb":null,"requires_docker":false},"constraints":{"min_ram_gb":null,"requires_docker":false},"when_to_use":["You are working specifically with machine learning models that leverage Tensor Processing Units (TPUs) for performance and are looking to understand optimization techniques.","If you require detailed methodologies focused on scaling LLMs using JAX on TPUs, this resource offers applicable insights.","For users encountering scaling challenges in model inference and serving phases, especially within the context of TPU architectures."],"when_not_to_use":["Do not use if your project focuses exclusively on GPU scaling or other hardware not aligned with Tensor Processing Units (TPUs).","If you are looking for a general approach to any framework's scalability without emphasis on TPUs.","This resource is unsuitable if you need information about model training phases, as it emphasizes inference and serving phases."],"source":"enrich:decision_facts","observed_at":"2026-07-16T20:18:06.560Z"},"constraint_facets":{"min_ram_gb":null,"requires_docker":false},"decision_summary":[{"label":"Adopt for","value":"scaling-book"}]}}