这是 DeepMind Scaling Book 系列的第 11 部分。
总结 (Conclusions)
How To Scale Your Model Part 11 (Part 10: JAX | Part 12: GPUs)
我们已经走过了漫长的道路，从简单的矩阵乘法 Roofline 模型到在成千上万个 TPU 上训练和部署最先进的 LLM。
感谢您阅读这一系列文章，并祝贺您坚持到了最后。 在我们结束之前，还有一些致谢：
致谢 (Acknowledgments)
这份文档代表了 Google DeepMind 许多人的重大集体投资，我们想简要地感谢他们！
- James Bradbury, Reiner Pope, 和 Blake Hechtman 最初推导了这份手稿中的许多想法，并且是早期理解 Transformer 系统观点的人。
- Sholto Douglas 撰写了这份文档的第一个版本，并负责启动该项目。
- Jacob Austin 领导了将第一个版本从粗略笔记转化为更完善和全面的工件的工作。
- Charlie Chen 撰写了推理部分并绘制了许多推理图表。
- Roy Frostig 协助出版、编辑和旅程中的许多其他步骤。
延伸阅读 (Further Reading)
如果您想通过阅读原始论文来深入了解，这里有一些很好的起点：
基础模型 (Foundational Models)
- Attention Is All You Need: Transformer 的原著。
- The Llama 3 Herd of Models: LLaMA 3 论文。
并行与系统 (Parallelism and Systems)
- GSPMD: General and Scalable Parallelization for ML Automating Sharding on TPU Pods: JAX 的自动分片背后的技术。
- GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism: 流水线并行。
- Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism: 张量并行的经典论文。
- ZeRO: Memory Optimizations Toward Training Trillion Parameter Models: FSDP 的基础。
硬件 (Hardware)
- In-Datacenter Performance Analysis of a Tensor Processing Unit: 原始 TPU 论文。
- TPU Deep Dive: 深入了解 TPU 架构。
- Domain specific architectures for AI inference: AI 推理的硬件和模型深入研究。
- Writing TPU Kernels with Pallas: 越来越多的 TPU 编程涉及在 Pallas 中编写自定义内核。
其他资源
- Huggingface Ultra-Scale Playbook: 类似于本书的 GPU 模拟，更深入地讨论了 PyTorch 如何实施并行技术。
- Stas Bekman’s ML Engineering Handbook: 极其使用的 ML 基础设施指南，涵盖了如何与云提供商谈判、集群管理以及 GPU 吞吐量的经验测量等本书未涉及的主题。
我们希望这本“书”对你解开大模型的神秘面纱有所帮助。快乐 Scaling！
再次提醒，如果你对 GPU 感兴趣，请查看第 12 节！
