如何让模型高效扩展：TPU 上大语言模型的系统视角

1. 0 导论 导论 训练大语言模型常常给人一种炼金术般的感觉，但理解和优化模型性能并非一定如此。本书旨在揭开语言模型扩展科学的神秘面纱：TPU（以及 GPU）如何工作、它们如何彼此通信，LLM 如何在真实硬件上运行，以及如何在训练和推理期间并行化模型，使其在超大规模下仍能高效运行。如果你曾想过“训练这个 LLM 应该有多贵”“我自己部署这个模型需要多少内存”或“什么是 AllGather”，我们希望本书能对你有所帮助。
2. 1 第 1 章 Roofline 模型详解 在硬件上运行算法时，我们受到三项因素的限制：计算机执行数学运算的速度（OPs/秒）、用于搬运数据的可用带宽（字节/秒），以及用于存储数据的总内存容量（字节）。这些“Roofline”约束让我们能够为给定计算所需的时间确定上界和下界。
3. 2 第 2 章 如何理解 TPU 本章全面介绍 TPU 如何工作、它们如何通过网络连接以支持多芯片训练与推理，以及这些因素如何影响我们常用算法的性能。其中甚至也有一些对 GPU 用户很有价值的内容！
4. 3 第 3 章 分片矩阵及其乘法 训练大型机器学习模型时，我们必须把参数或输入拆分（即“分片”）到许多加速器上。由于 LLM 主要由矩阵乘法构成，理解这个问题归根结底就是理解：当矩阵被拆分到不同设备上时，应当如何进行矩阵乘法。本章基于 TPU 通信原语的成本，建立一套简单的分片矩阵乘法理论。
5. 4 第 4 章 你需要掌握的全部 Transformer 数学 本章将快速回顾 Transformer 架构，重点介绍如何计算 FLOPs、字节数以及其他值得关注的量。
6. 5 第 5 章 如何并行化 Transformer 训练 本章讨论 LLM 训练期间使用的四种主要并行方案：数据并行、完全分片数据并行（FSDP）、张量并行和流水线并行。对于每一种方案，我们都会计算系统从什么时候开始受到通信瓶颈的限制。
7. 6 第 6 章 在 TPU 上训练 LLaMA 3 让我们运用上一章学到的知识，仔细研究如何在 TPU v5p 上训练 LLaMA 3 模型。这些模型有多大？不同配置下的训练成本有多高？它们如何分片？我们将通过一系列粗略估算，具体说明前面各章的结论如何映射到真实模型上。
8. 7 第 7 章 Transformer 推理详解 对 Transformer 执行推理可能与训练大不相同，部分原因在于推理增加了一个需要考虑的新因素：延迟。本章将从使用模型采样单个新 token 开始，一路讲到如何把大型 Transformer 高效扩展到多个加速器切片上，并将其作为推理引擎的一部分运行。
9. 8 第 8 章 在 TPU 上部署 LLaMA 3-70B 服务 让我们仔细研究如何在 TPU v5e 上部署 LLaMA 3-70B 模型。按照 Roofline 估算，部署不同模型的成本是多少？它们的 KV 缓存有多大？应该使用什么批大小？推理期间，参数和激活值如何分片？我们将通过一系列粗略估算，推导生产环境中的延迟与吞吐量。
10. 9 第 9 章 如何对 TPU 程序进行性能剖析 到目前为止，本系列一直是纯理论的：基于硬件 Roofline 进行粗略估算。这种理解能带你走得很远，但大量优化最终取决于实践细节：XLA 编译器如何工作，以及当它未能得到理想结果时，如何使用 JAX/TensorBoard Profiler 等性能剖析工具判断应该做什么。本章将讨论这些内容。
11. 10 第 10 章 使用 JAX 为 TPU 编程 如何使用 JAX 高效地为 TPU 编程！本章的大量内容取自 JAX 的 shard_map JEP 文档（https://jax.readthedocs.io/en/latest/jep/14273-shard-map.html）。你可以在 Google Colab（https://colab.sandbox.google.com/）上使用免费的 TPU 运行本章代码示例。
12. 11 第 11 章 总结与延伸阅读 感谢阅读！本章还将提供一些参考资料，供你进一步学习。
13. 12 第 12 章 如何理解 GPU 我们在 Google 热爱 TPU，但 GPU 同样出色。本章将深入探索 GPU 的世界：每块芯片如何工作、它们如何通过网络连接，以及这些特性对 LLM 意味着什么，尤其是与 TPU 相比时。尽管 NVIDIA、AMD、Intel 等厂商提供了众多 GPU 架构，本章将重点关注 NVIDIA GPU。本章建立在第 2 章（https://jax-ml.github.io/scaling-book/tpus/）和第 5 章（https://jax-ml.github.io/scaling-book/training）的基础上，建议先阅读这两章。

原书署名

作者与章节贡献者

第 0—11 章的共同署名如下；第 10 章与第 12 章的特有署名另列于后。

参考文献

1. [transformers] Vaswani et al.. Attention Is All You Need . 2017.
2. [tpu_paper] Jouppi et al.. TPU v4: An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddings . 2023.
3. [glu] Shazeer. GLU Variants Improve Transformer . 2020.
4. [mqa] Shazeer. Fast Transformer Decoding: One Write-Head Is All You Need . 2019.
5. [gmqa] Ainslie et al.. GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints . 2023.
6. [moe] Shazeer et al.. Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer . 2017.
7. [zero] Rajbhandari et al.. ZeRO: Memory Optimizations Toward Training Trillion Parameter Models . 2019.
8. [megatron] Shoeybi et al.. Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism . 2019.
9. [DeepSeek3] DeepSeek-AI et al.. DeepSeek-V3 Technical Report . 2024.
10. [llama3] Grattafiori et al.. The Llama 3 Herd of Models . 2024.
11. [esti] Pope et al.. Efficiently Scaling Transformer Inference . 2022.
12. [paged] Kwon et al.. Efficient Memory Management for Large Language Model Serving with PagedAttention . 2023.
13. [spec1] Leviathan et al.. Fast Inference from Transformers via Speculative Decoding . 2022.
14. [spec2] Chen et al.. Accelerating Large Language Model Decoding with Speculative Sampling . 2023.
15. [eagle] Li et al.. EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty . 2024.
16. [medusa] Cai et al.. Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads . 2024.