结论与延伸阅读

Part 11 of How To Scale Your Model (第10部分：JAX | 第12部分：GPU)

感谢阅读！这里我们将包含一些进一步学习的参考资料。

感谢阅读这组文章，恭喜你一直读到了最后。 在我们结束之前，有一些致谢：

致谢

本文档代表了 Google DeepMind 许多人的重大集体投入，我们想简要致谢！

• James Bradbury、Reiner Pope 和 Blake Hechtman 最初推导了本手稿中的许多想法，并且很早就理解了 Transformer 的系统视角。
• Sholto Douglas 编写了本文档的第一版，并负责启动了这个项目。他比任何人都更负责本文档的整体叙事。
• Jacob Austin 领导了将这个第一版从粗略笔记转变为更精炼和全面的成品的工作。他做了大量的编辑、格式化和发布这份文档的工作，并协调了其他作者的贡献。
• 大部分图形和动画由 Anselm Levskaya 和 Charlie Chen 制作。
• Charlie Chen 编写了推理部分并绘制了许多推理图形。
• Roy Frostig 帮助完成了出版、编辑和旅程中的许多其他步骤。

我们还要感谢许多在整个过程中提供关键反馈的其他人，特别是 Zak Stone、Nikhil Sethi、Caitlin Stanton、Alek Dimitriev、Sridhar Lakshmanamurthy、Albert Magyar、Diwakar Gupta、Jeff Dean、Corry Wang、Matt Johnson、Peter Hawkins 等人。感谢 Ruiqi Gao 帮助完成 HTML 格式化。

感谢大家！

在你离开之前，你可能也会喜欢阅读关于 NVIDIA GPU 的新第12章！

延伸阅读

有很多相关的文章，包括以下内容：

• TPU Deep Dive：一个精彩的深入探讨 TPU 架构的文章，与本书精神相同。
• Domain specific architectures for AI inference：一个硬件和模型深度解析，与本书精神相同。
• A Domain-Specific Supercomputer for Training Deep Neural Networks：原始的 TPU 论文之一，其中有很多关于 Google TPU 项目的精彩细节，本书未涵盖。
• Making Deep Learning Go Brrrr From First Principles：一个更关注 GPU 和 PyTorch 的 LLM roofline 和性能工程教程。
• Writing TPU Kernels with Pallas：越来越多地，TPU 编程涉及用 Pallas 编写自定义内核。本系列讨论如何编写内核以及这里没有提到的许多低级 TPU 细节。
• How to Optimize a CUDA Matmul Kernel for cuBLAS-like Performance: a Worklog：虽然是 GPU 和 CUDA 特定的，这是一篇优秀的博客文章，展示了如何在 CUDA 中优化 matmul 内核。这可能是深入了解 TPU 和 GPU 差异的好方式。
• Distributed arrays and automatic parallelization：这是 JAX 中并行化 API 的一个非常好的指南，是学习如何实际实现我们这里讨论的一些想法的好方法。
• Rafi Witten’s High Performance LLMs 2024 Class：我们的前同事 Rafi 开设了一个关于 TPU 性能工程的优秀课程，幻灯片都在 GitHub 上。这比我们这里更深入地涵盖了很多内容。
• [2211.05102] Efficiently Scaling Transformer Inference：一篇关于 Transformer 推理数学的详细论文。这是本文档很多内容的灵感来源。
• Huggingface Ultra-Scale Playbook：某种程度上是本书的 GPU 对应物，更深入地讨论 PyTorch 如何在训练期间实现并行化技术和内存节省技术。
• Transformer Inference Arithmetic：一个包含本书许多相同想法和一些优秀插图的博客。
• Stanford CS336 Slides and Videos：一个涵盖 LLM 训练和服务许多细节的精彩斯坦福课程，有一些有用的练习。作业 1 和 2 特别相关。
• Stas Bekman’s ML Engineering Handbook：一个高度实用的 ML 基础设施指南，涵盖本书未涉及的主题，如如何与云提供商谈判、集群管理和 GPU 吞吐量的经验测量。

这个领域仍然有很大的空间进行全面的写作，所以 我们希望这份手稿能鼓励更多这样的工作！我们也相信 这是一个值得研究和探索的富有成效的领域。在许多情况下，即使 手头没有很多硬件加速器也可以进行这项工作。

反馈

请留下评论或问题，以便我们进一步改进。你可以通过以下方式联系我们的通讯作者 Jacob Austin： jacobaustin123 [at] gmail [dot] com，或者在 GitHub 上发布 issue、pull request 或讨论来建议编辑。