第 1 讲 · 课程导论:为什么需要 ML 系统?· Course Introduction: Why ML Systems?

Advanced Topics in Machine Learning Systems (LLM Edition) · Lecture 01。个人学习讲义与配套实验,按主题整理并持续修订。

课程源码 ZIP ↓

个人整理的学习笔记,非官方讲义;课程来源与整理说明见课程介绍。配套示例未逐一运行验证。

本篇目录

对应 Maps to:CMU 15-779 Advanced Topics in Machine Learning Systems (LLM Edition)(Fall 2025,Zhihao Jia),Lecture 1 · Fri 8/29 · Course Introduction。 阅读材料:(必读)S. Keshav, How to Read a Paper, ACM SIGCOMM Computer Communication Review, 2007。本讲 §9 的示范评审用到第 2 讲的两篇必读论文:Abadi et al., TensorFlow: A System for Large-Scale Machine Learning(OSDI’16);Paszke et al., PyTorch: An Imperative Style, High-Performance Deep Learning Library(NeurIPS’19)。 依据说明:本讲以原课第 1 讲公开的 slides(01-course-introduction.pdf,38 页)为主线:ML 成功的驱动力、每美元算力的增长、规模定律、ML 硬件的并行性与异构性、硬件的快速演进、ML 系统的目标、MLSys 作为研究领域、监控摄像头练习、学 MLSys 的三个理由、课程地图、学习目标、与 CMU 其他课程的关系、上课形式、读论文与写评审、课程项目和成绩构成。slides 只给出了这些主题的标题和示意图;§3–§4 的硬件数字、§4 的 roofline 分析、§6 的设计空间搜索、§8 的”一行 PyTorch 沿栈下行”实验和 §9 的示范评审都是本讲义自己设计的,不代表课堂上讲了什么。 配套代码 Code:code/lecture01/(4 个脚本 + 4 个共享模块;01、04 只用标准库,02、03 用 PyTorch CPU 版;每个脚本 1–6 秒) 预计阅读时间:3–4 小时(含运行代码)


0. 本讲导航 · How to read this lecture

0.1 学习目标 Learning goals

学完这一讲,你应该能够:

  1. 说出 ML 成功的三个驱动力(模型、数据、算力),并解释为什么”把硬件变成有用的 FLOPs 的效率”本身就是一个值得研究的问题。
  2. 用具体数字描述 ML 硬件的三个特点:大规模并行(从集群到 CUDA core 的层次)、高度异构(存储层次的容量与带宽差几个数量级)、快速演进(每一代新增的硬件特性和”内存墙”)。
  3. 写出 roofline 模型,计算一个运算的算术强度和一块硬件的拐点,判断一个内核是计算受限、访存受限还是开销受限;并知道 roofline 只是上限。
  4. 说清 ML 系统的定义和目标,以及 MLSys 研究和普通 ML 研究、普通系统研究的区别(“协同设计”)。
  5. 画出这门课的课程地图(算法、计算图、并行、编译、内存、内核六层 + LLM 专题),知道每一层解决什么问题、对应哪几讲。
  6. 用 Keshav 三遍阅读法读一篇 MLSys 论文,按原课的格式写一份三段式评审,并知道课程项目的里程碑和要求。

0.2 先修与本讲范围 Prerequisites & scope

原课没有硬性先修,最好学过机器学习导论。本讲是导论,需要的背景都就地补上。

默认你已经会 Assumed本讲只给直觉,后面再展开 Previewed only
矩阵乘法;“训练就是最小化损失”的大意CUDA 编程模型、warp、共享内存(第 3–4 讲)
Python 基础(代码里不熟悉的写法都有注释)FlashAttention、warp 特化、mega-kernel(第 5–7 讲)
知道 GPU 是”很多核的并行处理器”(不知道也没关系,§3 会讲)编译器、图优化、超优化(第 8–11 讲);并行训练(第 12–15 讲);LLM 服务与微调(第 16–21 讲)

如果你对 C 语言的指针、浮点数格式、神经网络训练或计算机的内存层次完全陌生,可以先读 CMU11-868 的前置讲 P1–P4,本课不再重复。

0.3 本讲地图 Map

§1  开场:一行 PyTorch 背后发生了什么?
§2  ML 成功的驱动力:模型、数据、算力;规模定律;训练要多少 FLOPs        ← 代码 01(演示 5)
§3  ML 硬件(一):大规模并行、高度异构                                  ← 代码 01(演示 1、4)
§4  ML 硬件(二):快速演进与"内存墙";roofline 模型                      ← 代码 01(演示 2、3)、02
§5  ML 系统:定义、目标、研究领域
§6  课堂练习:在监控摄像头上部署目标检测(ML / 系统 / MLSys 三种做法)    ← 代码 04
§7  课程地图:六层软件栈 + LLM 专题;学习目标;与其他课程的关系
§8  把全课串起来:一行 PyTorch 的旅程                                   ← 代码 03
§9  怎样学这门课:读论文、写评审、做项目(附一份示范评审)
§10 小结   §11 常见误区   §12 练习   §13 延伸阅读   §14 下一讲预告

建议的阅读顺序:先读 §1–§3 并运行 01;读 §4 时运行 02(它测的是你自己的电脑);读 §6 时运行 04;读完 §7 的课程地图再运行 03,对照 §8 看”一行代码”怎样穿过每一层。§9 是方法论,可以在开始读第 2 讲的论文之前再回来读一遍。

0.4 记号 Notation(全课统一)

记号意思
FLOP、FLOPs、FLOP/s一次浮点运算(一次乘法或一次加法各算 1 次);运算次数;每秒运算次数(算力)
PpeakP_{\text{peak}}硬件的峰值算力(FLOP/s),本讲义一律取厂商标称的 dense 值
BB带宽(byte/s),不加说明时指显存(HBM)带宽
WW,QQ一个运算的 FLOPs;它要从显存读写的字节数
I=W/QI = W / Q算术强度 arithmetic intensity(FLOP/byte)
I⋆=Ppeak/BI^{\star} = P_{\text{peak}} / B拐点 ridge point(FLOP/byte)
NN,DD模型参数量;训练 token 数
TT,dd,dffd_{\text{ff}}token 数(序列长度 × 批大小);隐藏维度;前馈层中间维度
MFUModel FLOPs Utilization:实际有用 FLOP/s ÷ 峰值 FLOP/s
⏱在本机实测、每台机器会不同的数字;GPU 上的数字一律是估算

1. 开场:一行 PyTorch 背后发生了什么?· What happens behind one line of PyTorch?

先看一行再普通不过的代码:

y = model(x)

对写模型的人来说,这一行就是”把输入喂给模型、拿到输出”。但如果 model 是 Llama 3 405B,x 是一批几千个 token 的提示词,那么在这一行”返回”之前,下面这些事情都要发生:

  • 模型被拆成成千上万个算子(operator):矩阵乘法、加法、归一化、注意力……
  • 每个算子变成一个或几个内核(kernel)——在 GPU 上真正执行的程序;
  • 每个内核又被拆成几万个线程块和几百万个线程,分配到 GPU 上一百多个流式多处理器(SM)里;
  • 一块 GPU 的显存放不下模型,所以模型被切开放到几十、几百块 GPU 上,GPU 之间要不停地交换数据;
  • 每一层存储(寄存器、共享内存、L2、显存、主机内存、网络)的容量和速度相差几个数量级,数据要在它们之间搬来搬去。

原课主页把整门课的主题浓缩成一句话(转述):我们学习一个用高级语言写出来的 ML 模型,是怎样被分解成底层内核、并以分布式的方式在异构硬件加速器上执行的。 这一行 y = model(x) 和底层上亿个计算单元之间的所有软件,就叫 ML 系统(Machine Learning Systems)。

为了让这句话不那么抽象,本讲的代码 03 拿一个真实尺寸的 LLaMA 前馈块(Llama 3 8B 的维度)做了一次”沿栈下行”:它的 forward() 只有 5 行 Python,抓成计算图后是 12 个会产生内核的节点;在 H100 上估算,eager 模式(一个算子一个内核)处理 4096 个 token 要 1.78 ms,而”理论上不可能更快”的下界是 1.46 ms;处理 1 个 token(解码)时,eager 要 153 µs,下界是 109 µs,差距里将近三分之一是”启动内核”本身的开销。这门课的每一讲,都是在用某种办法缩小这个差距——§8 会把这个实验完整走一遍。


2. ML 成功的驱动力:模型、数据、算力 · What drives ML progress?

2.1 三驾马车

原课的第一个问题是:ML 这些年的成功,根本的驱动力是什么?slides 给的答案是一个三角形:模型(Model)、数据(Data)、算力 / 硬件(Compute / Hardware)。

  • 模型:Transformer、注意力、残差连接、归一化……这些结构让”堆更多参数”真的能变得更好。
  • 数据:互联网规模的文本、图像、代码,让大模型有东西可学。
  • 算力:没有足够的算力,前两者都无从谈起。GPT-3 这种规模的模型在十年前的硬件上根本训练不出来。

这三者是互相放大的:好的模型结构让”多给算力”变得划算(Transformer 比 RNN 好并行得多),更多的数据需要更多算力去消化,更多的算力又让更大的模型和更多的数据成为可能。

为什么一门系统课要从这里讲起? 因为三者里面,“算力”不是一个固定的数字,而是”硬件 × 软件效率”。同样一万块 GPU,训练时能用到峰值的 20% 还是 50%,等于平白多出或少掉一半的算力。ML 系统研究的经济价值就在这里——后面 §2.4 会给一个具体的数字。

2.2 每美元算力的指数增长与”专用化”

slides 引用了 Ray Kurzweil 在《奇点临近》(2005)里的一张图:每美元每秒能买到的计算量一百多年来一直指数增长,按他的外推,2023 年前后会超过人脑的算力。这张图的细节可以争论,但趋势本身是硬件行业的事实:计算越来越便宜。

不过,增长的来源变了。过去几十年,CPU 靠晶体管变小(摩尔定律)和频率提高”免费”变快;2005 年前后,频率提升因为功耗密度撞墙而停了下来(Dennard 缩放失效),单核性能的增长大大放缓。此后性能的增长主要来自两件事:

  1. 并行:一颗芯片上放更多的核——CPU 从单核到几十核,GPU 到上万个 CUDA core;
  2. 专用化(specialization):为某一类计算设计专门的电路——GPU 里专门做矩阵乘加的 Tensor Core、Google 的 TPU、各种 NPU。

这两件事直接塑造了 ML 系统要面对的硬件:并行度越来越高,硬件越来越”偏科”——矩阵乘法快得惊人,别的运算相对越来越慢。§4 会看到这个”偏科”有多严重。

2.3 规模定律:训练算力和推理算力都在扩大

slides 用一张”规模定律”(Scaling Law)的图把这条链串起来:模型的准确率和能力的提升 ← 训练与推理算力的扩大 ← 硬件的并行化和专用化。

  • 训练侧:Kaplan 等人(2020)发现语言模型的损失随参数量、数据量、算力呈幂律下降;Hoffmann 等人(2022,Chinchilla)进一步指出,给定算力预算,参数量和训练 token 数应该大致同比例增长(大约每个参数 20 个 token)。结论是:想要更好的模型,就要更多的训练算力。
  • 推理侧:近两年的推理模型(OpenAI o1 等)表明,让模型在回答前”多想一会儿”——生成更长的推理过程、采样多个答案再挑选——也能稳定地提升能力。这叫测试时算力扩展(test-time compute scaling)。于是推理也变成了一个”算力越多越好”的问题,而且推理要服务海量用户,总算力需求可能超过训练。

训练要多少算力? 一个稠密 Transformer 训练的总 FLOPs 可以用一条非常好用的公式估算:

Ctrain≈6NDC_{\text{train}} \approx 6 N D

直觉是:每个 token 前向时,每个参数参与一次乘法和一次加法(2N2N FLOPs);反向传播要对激活和参数各求一次梯度,大约是前向的两倍(4N4N);合计每个 token 6N6N,再乘以 token 数 DD。(这条公式忽略了注意力里和序列长度平方成正比的部分,上下文很长时会低估。)

代码 01 的演示 5 用它算了两个公开模型:

模型NNDD6ND6ND论文报告
GPT-3(2020)175B300B3.15×10233.15 \times 10^{23}3.14×10233.14 \times 10^{23}
Llama 3 405B(2024)405B15.6T3.79×10253.79 \times 10^{25}约 3.8×10253.8 \times 10^{25}

两个数都和论文报告吻合。四年里,单个旗舰模型的训练量涨了约 120 倍。

2.4 为什么”效率”值几十天、几千万美元

有了总 FLOPs,训练时间就是:

天数=CtrainGPU 数×Ppeak×MFU÷86400\text{天数} = \frac{C_{\text{train}}}{\text{GPU 数} \times P_{\text{peak}} \times \text{MFU}} \div 86400

MFU(Model FLOPs Utilization)是”真正用在模型计算上的算力 / 峰值算力”,大规模训练里常见的是 30%–50%。代码 01 的估算:用 16K 块 H100、MFU 40% 训练 Llama 3 405B,约需 68 天;MFU 每提高 10 个百分点,就能省下约 14 天——对一个上万块 GPU 的集群,这是几千万美元量级的电费和租金。

还有一个更硬的约束:显存。混合精度 Adam 训练每个参数大约要 16 字节(BF16 参数和梯度、FP32 主参数和两个 Adam 状态,第 13 讲细讲),Llama 3 405B 的训练状态约 6.5 TB,至少要 81 块 80 GB 的 H100 才放得下,这还没算激活值。所以对大模型来说,并行化(第 12–15 讲)不是”锦上添花的加速手段”,而是”能不能训练”的前提。

小结 §2:ML 的进步 = 模型 × 数据 × 算力,而”算力”= 硬件 × 系统效率。规模定律让算力需求在训练和推理两侧都持续膨胀;系统效率的每一个百分点都直接换算成时间和钱。


3. ML 硬件(一):大规模并行、高度异构 · Massively parallel, highly heterogeneous

3.1 从集群到 CUDA core:五层并行

slides 用一张图展示了 ML 硬件的层次。把它和代码 01 演示 1、4 的数字(H100 SXM 的厂商标称值)放在一起:

层次一个单元里有什么典型数量
集群 Cluster很多节点,靠高性能网络(InfiniBand / 以太网)相连1K–100K 块 GPU
节点 Node几颗 CPU + 主机内存(DRAM)+ 通常 8 块 GPU,GPU 之间用 NVLink 相连8 块 GPU
GPU上百个流式多处理器(SM)+ L2 缓存 + 显存(HBM)H100:132 个 SM,80 GB HBM
SMCUDA core(标量运算)+ Tensor Core(矩阵乘加)+ 寄存器 + 共享内存H100:128 个 CUDA core + 4 个 Tensor Core
合计一块 H100 = 16,896 个 CUDA core;1K 块 ≈ 1700 万个,100K 块 ≈ 17 亿个

slides 上写的”10M–1B CUDA cores”就是这么来的。程序员写的是一行 y = model(x),ML 系统要把它变成上亿个计算单元上的协同工作:怎样切分计算、怎样把数据喂到每个单元、怎样在它们之间同步。

为什么要有这么多层? 因为不同层次之间的”通信代价”天差地别。同一个 SM 里的线程可以通过共享内存在几十个时钟周期内交换数据;不同 SM 之间只能通过 L2 或显存;不同 GPU 之间要走 NVLink;不同节点之间要走网络。层次越远,通信越慢、越贵。所以几乎所有的并行策略都遵循同一个原则:通信越频繁的计算,放得越近。第 12–15 讲的张量并行放在节点内、流水线并行跨节点,就是这个原则的体现。

3.2 存储层次:异构首先是”存储的异构”

slides 图里的另一条线是存储:寄存器文件(Register Files)→ 共享内存(Shared Memory)→ 显存(Device Memory / HBM)→ 主机内存(DRAM)→ 网络。代码 01 的演示 1 给了一个直观的对比:

  • H100 每个 SM 有 228 KB 共享内存和 256 KB 寄存器,132 个 SM 加起来,全卡的片上存储大约 64 MB;
  • HBM 显存有 80 GB,是片上存储的约 1250 倍;
  • 但 HBM 的带宽(3.35 TB/s)远低于片上存储的总带宽(片上存储的带宽高一个数量级以上)。

于是出现了一个贯穿全课的矛盾:大的存储慢,快的存储小。数据大部分时间躺在又大又慢的 HBM 里,计算单元却只能直接操作又小又快的寄存器。一个内核快不快,往往不取决于它做了多少运算,而取决于它把数据在这些层次之间搬了多少次。

一个比喻:GPU 像一个有上万名厨师(计算单元)的厨房。每个厨师手边的砧板(寄存器)很小,每个灶台区有个小冰箱(共享内存),整个厨房共用一个大冷库(HBM),冷库在走廊尽头。厨师切菜很快,但每次去冷库取菜都要走很远。一个好的”菜谱”(内核)会一次把一批菜搬到小冰箱里、在砧板上连续做完几道工序再放回去,而不是每切一刀就跑一趟冷库。第 4 讲的分块矩阵乘法、第 5 讲的 FlashAttention、第 8–11 讲的算子融合,都是在设计这样的菜谱。

3.3 “异构”的另外两层意思

除了存储层次,“异构”还有两层意思:

  1. 计算单元的异构:同一块 GPU 上,Tensor Core 做矩阵乘法的速度是 CUDA core 做普通运算的十几倍(H100:BF16 Tensor Core 989 TFLOP/s,FP32 CUDA core 约 67 TFLOP/s)。能变成矩阵乘法的计算就快,变不成的就慢。
  2. 设备的异构:一个系统里同时有 CPU、GPU,可能还有 TPU、NPU、专用推理芯片,它们的编程模型、数值格式、内存都不一样。§6 的摄像头练习里就会遇到一块”只支持 int8、必须用编译过的引擎才能调用”的 NPU。

4. ML 硬件(二):快速演进与”内存墙” · Rapidly evolving hardware & the memory wall

4.1 每一代都有新东西:A100 → H100 → B200

slides 的第三张硬件图比较了几代 NVIDIA GPU 的结构。每一代不只是”更快”,而是多出了新的硬件特性,软件要改写才能用上:

代际新增的关键特性(和本课相关的)软件上的后果
A100(Ampere,2020)异步拷贝 cp.async(数据从显存直接搬进共享内存,不经过寄存器);BF16 / TF32 Tensor Core;结构化稀疏软件流水线:一边算这一块,一边搬下一块(第 6 讲)
H100(Hopper,2022)线程块簇(thread block cluster)与分布式共享内存(一个簇里的 SM 可以互相读共享内存);TMA(Tensor Memory Accelerator,专门搬张量的硬件单元);异步的 warpgroup 矩阵乘(wgmma);FP8warp 特化:一部分 warp 专门搬数据、一部分专门算(第 6 讲);FlashAttention-3
B200(Blackwell,2024)Tensor Memory(每个 SM 里专供 Tensor Core 使用的存储);FP4 / FP6;两颗裸片封装成一块卡又一轮内核重写

这张表说明了一件让软件工程师头疼的事:为上一代硬件精心调优的内核,在新一代硬件上往往不是最优的,甚至用不上新特性。 FlashAttention 从第一版到第三版,很大一部分改动就是在追硬件。手写的内核追不上硬件的演进速度——这是第 8–11 讲”让编译器自动生成和优化内核”的根本动机。

4.2 算力涨得比带宽快:内存墙

代码 01 演示 2 把四代 GPU 的标称值放在一起(全部取 dense 值;B200 是 HGX B200 的单卡数字):

GPUBF16 TFLOP/s最低精度最低精度 TFLOP/sHBM TB/sBF16 拐点最低精度拐点
V100(2017)125(FP16)FP161250.90139139
A100(2020)312INT86242.04153306
H100(2022)989FP819793.35295591
B200(2024)2250FP490008.002811125

从 V100 到 B200 的 7 年里:

  • 同精度(BF16/FP16)算力涨了 18 倍(年均 +51%),算上低精度涨了 72 倍(年均 +84%);
  • HBM 带宽只涨了 8.9 倍(年均 +37%),容量涨了 5.6 倍,NVLink 带宽涨了 6 倍。

“拐点”(峰值算力 ÷ 带宽,下一节解释)因此从 139 FLOP/byte 升到 281(BF16),算上 FP4 则到了 1125。要说明一个细节:拐点并不是每一代都单调上升——B200 的 BF16 拐点(281)略低于 H100(295),因为这一代 HBM 带宽一下涨了 2.4 倍,比 BF16 算力的 2.3 倍还多。但长期趋势和低精度下的趋势都很清楚:计算越来越便宜,搬数据相对越来越贵。 这个现象有一个名字:内存墙(memory wall)。卡间互联也一样(NVLink 只涨了 6 倍),所以还有一堵”通信墙”。

4.3 roofline 模型:判断瓶颈的最小工具

怎样把”内存墙”变成可以计算的东西?答案是 roofline 模型(Williams、Waterman、Patterson,2009)。它只看两个硬件数字和两个程序数字:

  • 硬件:峰值算力 PpeakP_{\text{peak}}(FLOP/s)和显存带宽 BB(byte/s);
  • 程序(一个内核):要做的运算 WW(FLOPs)和要从显存读写的字节数 QQ(bytes)。

一个内核至少要花”把运算做完的时间”和”把数据搬完的时间”中较长的那个:

t  ≥  max⁡ ⁣(WPpeak,  QB)t \;\ge\; \max\!\left(\frac{W}{P_{\text{peak}}},\; \frac{Q}{B}\right)

两边同除以 tt,就得到它能达到的算力上限:

Pattainable=min⁡ ⁣(Ppeak,  I⋅B),I=WQP_{\text{attainable}} = \min\!\left(P_{\text{peak}},\; I \cdot B\right), \qquad I = \frac{W}{Q}

II 叫算术强度(arithmetic intensity):每从显存搬 1 字节,做多少次运算。它只取决于算法和数据布局,与硬件无关。把 PattainableP_{\text{attainable}} 画成 II 的函数(两个坐标轴都取对数),就是一条”先斜后平”的折线——像一个屋顶,所以叫 roofline。斜坡和平顶的交点:

I⋆=PpeakBI^{\star} = \frac{P_{\text{peak}}}{B}

就是拐点(ridge point)。规则非常简单:

  • I<I⋆I < I^{\star}:访存受限(memory-bound)。加算力没用;要减少搬运——融合、分块、量化、复用。
  • I>I⋆I > I^{\star}:计算受限(compute-bound)。要用好 Tensor Core、降低精度、减少冗余计算。

几个典型运算的算术强度(BF16,2 字节/元素,按”理想复用”计算:每个输入只读一次):

  • 向量加法 z=x+yz = x + y:每个元素 1 次运算,搬 3 个元素(6 字节),I≈0.17I \approx 0.17。
  • 矩阵-向量乘(n×nn \times n 矩阵乘一个向量,LLM 解码 batch=1 的情形):每个权重 2 字节、只用一次,做 2 次运算,I≈1I \approx 1。
  • 矩阵乘法 m×km \times k 乘 k×nk \times n:W=2mnkW = 2mnk,Q=2(mk+kn+mn)Q = 2(mk + kn + mn)。方阵 m=n=km=n=k 时 I=n/3I = n/3,n=4096n = 4096 时约 1365。

代码 01 演示 3 算出了它们在各代 GPU 上最多能用到峰值的百分之几:

运算IIV100A100H100B200
向量加法0.170.12%0.11%0.06%0.06%
GELU 激活2.51.80%1.63%0.85%0.89%
GEMV n=4096n=4096(解码)1.00.72%0.65%0.34%0.36%
GEMM 64×4096×409664 \times 4096 \times 40966244.7%40.6%21.0%22.1%
GEMM 409634096^3(预填充)1365100%100%100%100%

三个结论:

  1. 逐元素运算和解码时的矩阵-向量乘,只能用到峰值的零头。 在 H100 上,Tensor Core 在这些运算期间几乎一直在等数据。
  2. 把多个 token 拼成一批,能让同一份权重被复用。 batch=64 时,同一个矩阵乘的上限从 0.34% 升到 21%。这就是 LLM 服务要做批处理(第 16 讲)的根本原因。
  3. 同一个访存受限的运算,在越新的卡上占峰值的比例越低(V100 上 0.12%,B200 上 0.06%;中间 B200 比 H100 略有回升,原因同 §4.2)。如果用 FP8 / FP4 的峰值来比,比例还要再低几倍。

roofline 故意忽略了什么? 它假设计算和访存可以完美重叠(取 max 而不是相加),假设每个数据只从显存读一次(理想复用),忽略了缓存、调度、启动开销和多卡通信。所以它给出的是上限,不是预测值。它的价值在于:拿到任何一篇论文的加速比,你都可以先问”按 roofline,理论上最多能快多少?论文做到了多少?剩下的差距在哪里?”

4.4 在你自己的电脑上画一条 roofline(代码 02)

roofline 是一个”任何处理器都适用”的模型。代码 02 在本机 CPU 上实测两个屋顶,再跑一组算术强度从 0.08 到 341 的运算,看实测点是不是贴着屋顶走。以下是在一台 Apple M4 Pro 上的结果(⏱,你的数字会不同):

测量结果 ⏱
内存带宽:copy / add / sum,单线程约 116 / 119 / 89 GB/s
内存带宽:同样三种运算,12 线程约 224 / 219 / 226 GB/s
矩阵乘法峰值(2048³,float32)约 3.2 TFLOP/s
本机拐点约 14 FLOP/byte(H100 的 BF16 拐点是 295)

为什么一块笔记本 CPU 的矩阵乘法能到 3 TFLOP/s?因为 PyTorch 在 macOS 上调用 Apple 的 Accelerate 库,它用到了芯片里专门做矩阵运算的协处理器——这正是 §2.2 说的”专用化”的缩小版:矩阵乘法有专用硬件,逐元素运算没有。

一组运算的实测结果(⏱):

运算II实测 GFLOP/s屋顶 GFLOP/s占屋顶
向量加法(12 线程)0.0817.618.893%
向量加法(1 线程)0.0810.018.853%
GEMV 4096×40960.5098.7112.788%
GEMM 2048×1×20480.5020.2112.618%
GEMM 2048×16×20487.9461177626%
GEMM 2048×256×204810232353234100%
GEMM 2048×2048×20483413205323499%

从这张表能读出 roofline 之外的四个重要现象:

  1. 要用满带宽,需要足够的并行度。 同一个向量加法,单线程只能达到屋顶的一半,12 个线程才接近屋顶:一个核心同时发出的访存请求不够多,喂不满内存系统。GPU 上完全一样——要有足够多的 warp 同时在等数据,才能把 HBM 带宽用满。第 3–4 讲会把这叫做占用率(occupancy)。
  2. roofline 是上限,不是预测值。 访存受限区的运算达到屋顶的 18% 到 93% 不等。离屋顶最远的是形状很”瘦”的矩阵乘法(K=1K=1,实际上是一个外积):数学库主要为常见的大方阵调优,罕见形状往往没人调——这是第 8–9 讲”自动生成、自动调优内核”的动机。
  3. 屋顶本身也是测出来的。 有几个点略超过 100%:我们的”计算屋顶”只是某一次测量的最好成绩,同一个运算两次测量也会差几个百分点。
  4. 还有第三种瓶颈:固定开销。 代码 02 的第三步测了一组很小的运算:
运算每次调用 ⏱占矩阵乘法峰值
1 个元素的加法0.41 µs0.0001%
4×4×4 矩阵乘法0.32 µs0.013%
64×64×64 矩阵乘法1.1 µs15%
256×256×256 矩阵乘法约 21 µs约 50%

一个 4×4×4 的矩阵乘法只有 128 次运算,却要花 0.3 µs——时间全花在 Python → C++ 分发、检查形状、分配输出这些”和张量多大无关”的步骤上。GPU 上的对应物是内核启动开销(几微秒)。LLM 解码时每一步只处理很少的 token,很多内核都小到被这个开销淹没——这正是第 7 讲 mega-kernel 要消灭的东西。

小结 §3–§4:ML 硬件有五层并行、六层存储;每一代都有需要软件改写才能用上的新特性;算力涨得比带宽快,越来越多的运算落在 roofline 的斜坡上。判断一个内核的瓶颈有三个候选:计算、访存、固定开销。roofline 只给上限。


5. ML 系统:定义、目标、研究领域 · What are ML systems?

5.1 定义与目标

把 §2–§4 合起来,原课第 1 讲给 ML 系统下的目标是(转述):在大规模并行、日益异构、快速演进的硬件平台上,高效地部署 ML 应用。

三个形容词各对应一类难题:

  • 大规模并行 → 怎样切分计算、怎样同步、怎样把通信藏在计算后面(第 3–7、12–15 讲);
  • 日益异构 → 怎样在不同的存储层次、计算单元、设备之间安排数据和计算(第 3–11 讲);
  • 快速演进 → 怎样让软件跟上硬件,最好是自动跟上(第 8–11 讲的编译与超优化)。

slides 把 ML 系统画在模型、数据、硬件三角形的中央:ML 系统是连接这三者的桥梁。它向上给模型开发者提供简单的编程接口(PyTorch 让你写 y = model(x)),向下把计算高效地映射到硬件上。

5.2 一张 ML 系统的”全家福”

为了让”ML 系统”这个词具体一点,这里按层列出一些你会在本课遇到的真实系统(不需要现在就记住):

层代表系统 / 库本课讲次
编程框架TensorFlow、PyTorch、JAX第 2 讲
厂商算子库cuBLAS、cuDNN、CUTLASS第 2、4 讲
内核 DSL 与编译器Triton、TVM、XLA、TorchInductor、TileLang第 8–9 讲
图优化与超优化TASO、PET、TENSAT、Mirage第 10–11 讲
分布式训练Megatron-LM、DeepSpeed(ZeRO)、PyTorch FSDP、Alpa、FlexFlow第 12–15 讲
LLM 推理服务vLLM、SGLang、SpecInfer、Mirage Persistent Kernel第 7、16–17 讲
微调与后训练PEFT / LoRA、QLoRA、各种 RL 训练框架第 18–19 讲

5.3 MLSys 作为一个研究领域

slides 指出,ML 系统已经成长为一个独立的研究领域:主要的 ML 会议(NeurIPS、ICML、ICLR)和系统会议(OSDI、SOSP、NSDI、ASPLOS、ISCA 等)都有大量 MLSys 论文和专门的研讨会,并且有了一个专门的会议 MLSys(2018 年以 SysML 的名字在斯坦福创办,后改名为 MLSys)。

MLSys 研究和普通的 ML 研究、普通的系统研究有什么不同? slides 的回答是:MLSys 用一种整体的(holistic)方法,把 ML、数据、系统和硬件的技术结合起来解决问题。具体说:

  • 普通的 ML 研究通常把硬件当成固定的黑盒,追求”同样的 FLOPs 下更准”或”同样的准确率下更少的 FLOPs”;
  • 普通的系统研究通常把工作负载当成固定的输入,追求”同样的程序跑得更快”;
  • MLSys 研究的特点是两边都可以动:可以为硬件改模型(例如 GQA 减少 KV 缓存的读取、MoE 让参数增长而计算不增长),也可以利用 ML 的特性改系统(例如”计算不必完全精确”让量化、推测解码、部分等价变换成为可能;“训练是迭代的、可预测的”让流水线调度和预取成为可能)。

最后这一点尤其值得强调:ML 计算有很多普通程序没有的特性——数值上可以容忍一定误差、计算图在训练时反复执行同样的结构、大部分时间在做少数几种算子。很多 MLSys 的经典工作,都是抓住了其中一个特性。


6. 课堂练习:在监控摄像头上部署目标检测 · Exercise: object detection on a camera

6.1 问题与三种做法

slides 里有一个练习:要在监控摄像头上部署一个目标检测模型,要求

  • 准确率 ≥ 90%
  • 时延 ≤ 10 ms
  • 内存 ≤ 100 MB

slides 对比了三种思路:

  1. 典型的 ML 做法:设计更准、更小的模型——剪枝、量化、蒸馏、低秩近似等。
  2. 典型的系统做法:为现有模型构建一个更快、更省内存的推理引擎,提高资源利用率和运行性能。
  3. MLSys 做法:模型、系统、硬件协同设计、协同优化——利用专用 AI 硬件,为这块硬件定制模型,再构建能把两者用起来的系统。

停下来想一想:为什么第三种”更好”?不就是把前两种加起来吗?关键在于”加起来”不是简单的并集:有些方案只有在两边同时改的时候才存在。

6.2 把三种做法变成三个”搜索范围不同的团队”(代码 04)

为了让这个论点可以检验,代码 04 构造了一个玩具设计空间。再强调一次:所有数字都是本讲义为演示而设定的,量级参考了常见的边缘设备,但不对应任何真实产品。

设计空间有六个维度:4 种模型尺寸(Det-S/M/L/XL,越大越准越慢)× 5 种数值格式(fp32、fp16、训练后量化 int8-PTQ、量化感知训练 int8-QAT、int4-QAT)× 3 种剪枝(不剪、结构化剪 30% 通道、非结构化剪 50% 权重)× 是否蒸馏 × 2 种推理引擎(框架直接逐算子执行的 naive;做了算子融合和内存规划的 optimized)× 2 种硬件(摄像头自带的 CPU;板载 NPU,int8 算力是 CPU 的 10 倍,但只支持 int8,而且必须用编译过的引擎才能调用)。一共 480 个组合,硬件能跑的有 312 个。

代价模型就是 §4 的 roofline 加固定开销:时延 = max(计算时间, 访存时间) + 层数 × 每层内核数 × 每个内核的开销;内存 = 权重 + 激活峰值 + 引擎本身。

三个团队的搜索范围:

  • ML 团队:可以改模型的一切(尺寸、量化、剪枝、蒸馏,包括需要重新训练的方法),但用框架自带的 naive 引擎、跑在 CPU 上。
  • 系统团队:模型是别人训练好的 Det-L(fp32 下 91.8%),不能重新训练,但可以做不需要训练的转换(fp16、训练后量化),可以自己写引擎、可以用任何硬件。
  • MLSys 团队:所有维度一起搜。

结果(输出完全确定):

团队搜索的方案数可行方案数最接近成功的方案卡在哪里
ML 团队1200Det-S + 蒸馏 + int8-QAT + 结构化剪枝 @CPU:85.5% / 9.7 ms准确率
系统团队70Det-L + int8-PTQ + optimized @NPU:89.8% / 8.0 ms准确率(差 0.2 个点)
MLSys 团队3128最好:Det-L + 蒸馏 + int8-QAT + optimized @NPU:91.9% / 8.0 ms / 28.7 MB—

每个团队失败的方式很能说明问题:

  • ML 团队被困在”要么准、要么快”之间。准确率达标的方案里最快的(Det-L + int8-QAT + 结构化剪枝)要 57 ms;能在 10 ms 内跑完的只有最小的 Det-S,而它怎么优化也到不了 90%。CPU 的算力和 naive 引擎每层三个内核的开销,是 ML 团队碰不到的天花板。
  • 系统团队找到了 NPU 这条快路(8.0 ms),但 NPU 只支持 int8,而不能重新训练的模型只能用训练后量化,掉了 2 个点,准确率差 0.2 个点没达标。不量化的话,fp16 在 CPU 上最快也要 150 ms。
  • MLSys 团队的 8 个可行方案全部用了 NPU、全部用了自己构建的 optimized 引擎、全部用了需要重新训练的 ML 技术(量化感知训练、蒸馏或剪枝)。NPU 的算力要靠”系统”才能用上;它只支持 int8,而 int8 要靠”ML”才能保住准确率——单独一方都做不到。可行方案的帕累托前沿上还有一个更快的选择:再加 30% 结构化剪枝,90.9% / 5.8 ms。

6.3 一个小插曲:纸面 FLOPs 减少 ≠ 时延减少

代码 04 最后单独比较了两种剪枝(Det-M、fp32、optimized 引擎、CPU):

剪枝方式纸面 GFLOPs实际时延权重
不剪6.0120.2 ms36.0 MB
结构化剪 30% 通道4.284.2 ms25.2 MB
非结构化剪 50% 权重3.0120.2 ms27.0 MB

非结构化剪枝把纸面 FLOPs 砍了一半,在稠密硬件上时延却一点没变:零散分布的零值没法让矩阵乘法跳过——硬件还是按稠密矩阵在算。它的存储也没有省一半:稀疏格式要额外存每个非零值的位置(这里假设 16 位下标),int8 时每个非零值占 1 + 2 = 3 字节,剪掉一半后(13.5 MB)反而比不剪(9.0 MB)更大。结构化剪枝则直接让矩阵变小,任何硬件都受益。

这就是”协同设计”在最小尺度上的含义:一个 ML 技术是否有用,取决于系统和硬件是否支持它。 NVIDIA 从 A100 开始在 Tensor Core 里支持 “2:4 结构化稀疏”(每 4 个元素里恰好 2 个是零),正是硬件为了让剪枝真正变快而做的让步(附加讲 A2)。

6.4 这个练习想教的三件事

  1. 搜索空间决定了你能找到什么。 三个团队用的是同一个代价模型,区别只在搜索范围。MLSys 的价值不在于某一个技巧,而在于把搜索空间扩大到跨越模型、系统、硬件的边界。
  2. “有效的优化”由瓶颈决定。 在 CPU 上,量化主要省的是计算;在 NPU 上,量化是”能不能用这块硬件”的门票。同一个技术在不同的瓶颈下价值完全不同。
  3. 既然要搜索,就让机器来搜。 这个玩具空间只有 480 个点,可以穷举;真实的空间(一个内核的分块参数、一个计算图的改写方式、一个模型的并行切分方式)动辄 101010^{10} 以上。这引出了本课的一条暗线:把专家的手工优化变成”搜索空间 + 代价模型 + 正确性验证”——第 9 讲的自动调优、第 10–11 讲的图替换与超优化、第 15 讲的自动并行,都是这个套路。

7. 课程地图:六层软件栈 + LLM 专题 · The course map

7.1 为什么要分层?为什么从下往上讲?

slides 用一张分层图描述了 ML 系统的全栈(自顶向下):ML 模型 → 算法优化 → 计算图级优化 → 并行化 / 分布式训练 → ML 编译 → 内存优化 → GPU / 内核编程 → 硬件。

分层是系统设计的老办法:每一层只关心自己的抽象,把复杂度隔离开。比如写计算图优化的人不需要知道一个矩阵乘法内核用了多少寄存器,只需要知道”这个算子的代价大约是多少”。但 MLSys 的有趣之处恰恰在于,最好的优化经常跨越层的边界:FlashAttention 同时改了算法(在线 softmax)和内核(分块、片上计算);Mirage 把计算图级和内核级放在同一个搜索空间里;推测解码同时改了算法和服务系统。

原课的讲授顺序是从下往上:先讲 GPU 和内核(第 3–7 讲),再讲编译(第 8–9 讲)、图优化(第 10–11 讲)、并行(第 12–15 讲),最后是 LLM 专题(第 16–21 讲)。理由是:上层的每一个决策(要不要融合、怎样切分、用什么精度),最终都要落到”内核在硬件上跑多快”上来评判;不先理解底层,上层的代价模型就是空中楼阁。

7.2 逐层看:每一层解决什么问题

① GPU / 内核编程(第 3–7 讲)——一块 GPU 怎样跑满?

slides 列出的主题:现代 GPU 架构、CUDA 编程、warp 特化、persistent kernel(mega-kernel)。核心问题是 §3–§4 说的两件事:让上万个计算单元都有活干(并行度、占用率),让数据尽量在片上存储里被复用(分块、合并访存、避免 bank 冲突)。第 5 讲用 FlashAttention 做案例;第 6 讲讲 Hopper 时代”一部分 warp 搬数据、一部分 warp 算”的 warp 特化;第 7 讲讲把整个模型编译成一个常驻 GPU 的大内核,消灭内核启动开销。

② ML 编译(第 8–9 讲)——内核能不能自动生成、自动调优?

slides 列出:基于分块(tile)的 ML 编译器、基于学习 / 搜索的优化。手写 CUDA 内核太难、太慢、跟不上硬件(§4.1)。Triton 让程序员”按块”而不是”按线程”写内核,把线程级的细节交给编译器;TVM / Ansor 定义一个巨大的”程序变体”空间,再用学习到的代价模型去搜索最快的那一个。

③ 计算图级优化(第 10–11 讲)——算子之间能不能重新组合?

slides 列出:代数变换、部分等价变换、超优化。一个计算图可以有很多种数学上等价的写法(例如把两个共享输入的矩阵乘法拼成一个更大的矩阵乘法),哪种写法最快取决于硬件。TASO 自动生成并形式化验证图替换规则;PET 允许”只在部分位置等价”的变换再自动修正;Mirage 把内核内部的结构也纳入搜索,能自动发现 FlashAttention 这类优化。

④ 并行化 / 分布式训练(第 12–15 讲)——一块卡放不下、算不完怎么办?

slides 列出:数据并行、模型并行、流水线并行、自动并行。§2.4 已经看到,大模型连训练状态都放不下一块卡。数据并行复制模型、切分数据;张量并行切分每个矩阵;流水线并行切分层;它们各有通信、显存、空泡的代价。自动并行(FlexFlow、Alpa)在巨大的组合空间里自动找最好的组合。

⑤ 内存优化(第 12–13 讲、附加讲 A1)——显存不够怎么办?

slides 列出:重计算(rematerialization)、零冗余(zero redundancy)。重计算是”用算力换显存”:反向传播时不保存所有中间结果,需要时再算一遍;ZeRO 是”用通信换显存”:数据并行里每块卡不再保存完整的优化器状态、梯度和参数,而是各存一份、用时再取。

⑥ 算法优化(第 18、21 讲、附加讲 A2)——能不能少算一点、少存一点?

slides 列出:剪枝(量化、稀疏、低秩近似)、微调、混合专家、检索增强。这一层直接改变计算本身:量化减少每个数的位数;LoRA 只训练一个低秩的增量;MoE 让每个 token 只激活一部分参数。§6 已经看到,这一层的技术是否有效,取决于下面几层是否支持。

⑦ LLM 专题(第 16–17、19–20 讲)

生成式 LLM 带来了一些专门的新问题:自回归解码让推理变成了访存受限的长循环;KV 缓存动态增长、管理困难;海量请求需要批处理和调度;RL 后训练把”生成”和”训练”捆在一起。第 16–17 讲讲推理服务(连续批处理、PagedAttention、RadixAttention、推测解码),第 19–20 讲是两次嘉宾讲座(后训练系统、开放模型训练)。

7.3 一条暗线:从手工优化到自动搜索

如果把第 9、10、11、15、17 讲放在一起看,会发现一个反复出现的模式(这也是主讲人 Zhihao Jia 的研究组一系列工作的共同特点):

讲次被自动化的专家技巧搜索空间代价模型正确性保证
9 AutoTVM / Ansor内核的分块、循环顺序、向量化程序变体学习到的代价模型 + 实测由构造保证
10 TASO / PET计算图的改写规则图替换序列算子级实测代价形式化验证 / 随机测试 + 修正
11 Mirage跨内核 / 线程块 / 线程的融合多层 µGraph硬件代价模型有限域上的随机测试(概率保证)
15 FlexFlow / Alpa并行切分策略SOAP 空间 / 两级层次空间执行模拟器 / 实测由构造保证
17 SpecInfer推测的候选(树)令牌树接受率验证机制保证输出分布不变

读这些论文时,可以统一问三个问题:空间多大?代价怎么估?正确性怎么保证?

7.4 学习目标与课程定位

slides 给出的三个学习目标(转述):理解现代 ML 框架的一般组成;学习面向生成式 AI 应用(LLM)的 MLSys 技术;实现你自己的 MLSys 项目。

slides 还区分了 CMU 的三门 MLSys 课:10-414/714 深度学习系统(秋季,侧重深度学习算法的设计与实现)、15-442/642 机器学习系统(春季,侧重 MLSys 的系统方面)、15-779 高级机器学习系统(侧重前沿 MLSys 技术,以论文阅读和讨论为主)。

在本仓库里,CMU11-868(LLM 系统)和本课重合最多,它以”自己实现每个组件”为主线;本课以”读懂并复现每篇论文的核心机制”为主线,两套讲义可以互相补充。


8. 把全课串起来:一行 PyTorch 的旅程(代码 03)· One line down the stack

现在可以完整地走一遍 §1 预告的实验了。代码 03 拿一个 LLaMA 式前馈块,维度取 Llama 3 8B(d=4096d = 4096,dff=14336d_{\text{ff}} = 14336),从上到下走五层。

8.1 第 1 层:程序员写的代码

def forward(self, x):
    var = x.pow(2).mean(dim=-1, keepdim=True)      # RMSNorm
    h = x * torch.rsqrt(var + self.eps)
    h = h * self.norm_weight
    a = F.silu(self.w_gate(h)) * self.w_up(h)        # SwiGLU MLP
    return x + self.w_down(a)                        # 残差

RMSNorm 故意用最基本的运算手写,这样下一层能看到每一个底层算子。

8.2 第 2 层:计算图

torch.fx.symbolic_trace 把 forward() 变成一张图:15 个节点,其中 1 个输入、1 个参数读取、1 个输出,剩下 12 个节点会产生内核——3 个矩阵乘法(w_gate、w_up、w_down),1 个归约(mean),8 个逐元素运算(pow、add、rsqrt、两个 mul、silu、mul、残差 add)。参数量 176.2M(BF16 下 352 MB)。

计算图是 ML 系统最核心的抽象(第 2 讲详细讲):有了它,系统就可以在执行之前”看到”整个计算,从而做全局的优化——融合、重排、切分、调度。

8.3 第 3 层:内核与代价(H100 估算)

eager 模式下,一个算子就是一个内核。代码用 §4.3 的 roofline 加 4 µs 的启动开销(一个常见的经验值)估算每个内核在 H100 上的耗时。分两个场景:

预填充(一次处理 4096 个 token):

内核FLOPs读写字节瓶颈估算时间
3 个矩阵乘法(各)4.8×10114.8 \times 10^{11}268 MB计算490 µs
mul_2(silu 输出 × up)5.9×1075.9 \times 10^{7}352 MB访存109 µs
silu2.3×1082.3 \times 10^{8}235 MB访存74 µs
RMSNorm 的 6 个内核、残差加法很少共约 336 MB访存 / 开销共约 128 µs
合计(12 个内核)1.73 GB1.78 ms

矩阵乘法占了 99.97% 的 FLOPs,却只占估算时间的 83%;剩下 17% 花在几乎不做运算的逐元素和归约内核上——它们全部访存受限。

解码(每步只处理 1 个 token):情况完全不同。三个矩阵乘法变成矩阵-向量乘,时间几乎全花在把 352 MB 权重从显存读一遍(每个 39 µs);其余 9 个内核小得可怜(几十字节到几十 KB),时间全部是启动开销。合计 153 µs 里,12 次启动开销占了 31%。

这两个场景对应 LLM 推理的两个阶段(第 5、16 讲):预填充是计算受限的,解码是访存受限加开销受限的。同一个模型、同一段代码,瓶颈可以完全不同——所以 LLM 系统经常要为两个阶段分别优化。

8.4 第 4 层:图优化——逐元素融合,以及理想下界

代码实现了一个最简单的图优化 pass:用并查集把”直接相连的逐元素 / 归约节点”合并成一个内核。融合后组内的中间结果留在寄存器和片上存储里,不再写回显存;只有从组外读进来的张量和要被组外使用的结果才算显存读写。结果是 6 个内核:

fused[pow+mean+add+rsqrt+mul+mul]   ← 整个 RMSNorm 一个内核
w_gate, w_up                        ← 矩阵乘法保持独立
fused[silu+mul]                     ← SwiGLU 的逐元素部分
w_down
fused[add]                          ← 残差
内核数预填充显存读写预填充估算解码估算
eager121.73 GB1.78 ms153 µs
逐元素融合61.33 GB1.64 ms(1.09×)129 µs(1.19×)
理想下界:整个块一个内核10.42 GB1.46 ms109 µs

“理想下界”假设整个块只用一个内核、所有中间结果都不落显存:只读输入和权重、只写输出。它不一定能做到——预填充时中间张量(4096 × 14336)远大于片上存储,需要像 FlashAttention 那样重新安排计算顺序才能接近它。但它告诉我们还剩多少空间:从 eager 到下界之间的差距,正是这门课要追的东西:

  • 把矩阵乘法后面的逐元素运算融进矩阵乘法内核的”尾巴”里(epilogue fusion)——第 8 讲;
  • 把 w_gate 和 w_up 这两个共享输入的矩阵乘法拼成一个——第 10 讲的图替换;
  • 把 RMSNorm 和后面的矩阵乘法融成一个内核——第 11 讲 Mirage 的经典例子;
  • 解码时把 12 次启动开销压到接近 0——第 7 讲的 mega-kernel。

8.5 第 5 层:真机验证——编译器真的生成了一个融合内核

估算终归是估算。代码的最后一层在本机 CPU 上用 torch.compile(PyTorch 2 的编译器,第 2、8 讲会讲它的结构)真的做一次融合。对 relu(x*2+1)*0.5 这个 4 个算子的链,编译器生成了一个叫 cpp_fused_add_mul_relu_0 的 C++ 内核,核心就是这样一个循环(节选):

for (x0 = 0; x0 < N; x0 += 4) {                     // 每次处理 4 个 float(向量化)
    auto tmp0 = Vectorized<float>::loadu(in_ptr0 + x0);  // 只读一次 x
    auto tmp3 = tmp0 * 2.0;                          // 4 个运算都在寄存器里完成
    auto tmp6 = tmp3 + 1.0;
    auto tmp7 = clamp_min(tmp6, 0);                  // relu
    auto tmp10 = tmp7 * 0.5;
    tmp10.store(out_ptr0 + x0);                      // 只写一次结果
}

eager 模式下这 4 个算子要各自读一遍、写一遍整个数组(8 次搬运),融合后只要读 1 次、写 1 次(2 次搬运),roofline 预测的加速比是 4.00×。实测(⏱,单线程,3200 万个 float32):

函数eager编译后加速比字节比预测编译后的有效带宽
relu(x*2+1)*0.59.9 ms2.4 ms4.19×4.00×约 114 GB/s(屋顶的约 99%)
silu(g)*u29.4 ms27.9 ms1.05×1.67×约 14 GB/s(屋顶的 13%)
RMSNorm(含按行归约)9.0 ms13.0 ms0.69×3.50×约 21 GB/s(屋顶的 18%)

(单线程内存带宽屋顶约 115 GB/s。同一台机器上运行三次,第一行的加速比分别是 4.16×、3.95×、4.19×;后两行每次都明显低于预测。)

第一行是一个漂亮的验证:访存受限的运算,省下多少字节就快多少——预测 4.00×,实测约 4×(3.95×–4.19×),编译后的内核跑满了带宽。

后两行同样重要,因为它们没有达到预测:

  • silu(g)*u 的 eager 版本本身就只有 23 GB/s,远低于带宽屋顶——在 CPU 单核上,silu 里的指数运算很贵,这个运算本身就是计算受限的,省字节帮不上忙。而在 H100 上,即使只看不用 Tensor Core 的 CUDA core,FP32 拐点也有约 20 FLOP/byte,silu 每搬 1 字节只做几次运算,它是访存受限的。一个运算是不是访存受限,取决于硬件的算力 / 带宽比。
  • RMSNorm 融合后反而变慢了:融合确实省下了字节,但编译器为”按行归约”生成的代码只跑到 21 GB/s,内核变成了受指令限制。roofline 只给上限,能不能达到,取决于内核写得好不好。 这正是为什么第 4 讲要学怎样写好归约,第 8–9 讲要学编译器怎样搜索更好的实现,第 11 讲 Mirage 要在搜索时同时考虑内核内部的结构。

(如果你的机器上没有 C++ 编译器,第 5 层会自动跳过,前四层的结论不受影响。)

小结 §8:一行 forward() → 15 个节点的计算图 → 12 个内核 → 每个内核有自己的瓶颈。预填充时矩阵乘法以外的”小运算”占了 17% 的时间,解码时启动开销占了 31%。融合是最基本的图优化;从 eager 到理想下界的差距,就是本课第 3–11 讲的主题。真机实验同时证明了 roofline 的威力(预测 4.00×、实测约 4×)和它的局限(瓶颈会转移,内核质量决定能不能达到上限)。


9. 怎样学这门课:读论文、写评审、做项目 · How to take this course

9.1 上课形式:讲授 + 论文讨论

原课每次 80 分钟:大约 50 分钟讲授一个主题的基本概念,30 分钟讨论这个主题下的两篇近期 MLSys 论文。学生在每次课前读论文、写评审,目的是(slides 转述):熟悉主题和论文,理解它们的长处和局限,学习并推广其中的想法。

为什么 MLSys 课要以读论文为主? 因为这个领域变得太快:本课第 6、7、11 讲的内容(warp 特化在 Hopper 上的用法、mega-kernel、Mirage)在两三年前还不存在,教科书跟不上。读论文是进入这个领域的唯一途径;而 MLSys 论文又有它特殊的读法——不只看想法,还要看实验是不是真的支持结论。

9.2 Keshav 三遍阅读法

原课第 1 讲的必读材料是 S. Keshav 的《How to Read a Paper》。它的核心是:不要从第一个字读到最后一个字,而是读三遍,每一遍的目标不同。

第一遍:五到十分钟,鸟瞰全文。

  • 仔细读标题、摘要、引言;
  • 只读各节和小节的标题,跳过正文;
  • 读结论;
  • 扫一眼参考文献,看有没有你读过的。

读完第一遍,你应该能回答五个 “C”:Category(这是什么类型的论文:新系统?测量研究?对已有系统的分析?)、Context(和哪些论文相关、建立在什么理论或系统之上)、Correctness(它的假设看起来成立吗)、Contributions(主要贡献是什么)、Clarity(写得清楚吗)。很多时候第一遍之后就可以决定”不用再读了”——这完全正常。

第二遍:一小时左右,抓住内容,但跳过细节。

  • 仔细看图、表和示意图——尤其是坐标轴的标签、误差棒、基线是谁;
  • 把没读过的重要参考文献记下来;
  • 跳过证明和实现细节。

读完第二遍,你应该能向别人复述这篇论文的主要论点和支持它的证据。

第三遍:对初学者可能要四五个小时,“在脑子里重新实现一遍”。

  • 假设你是作者,在同样的假设下自己把这个工作做一遍,再和论文比较;
  • 质疑每一个假设,找出论文没说的东西和隐藏的缺陷;
  • 想想”如果是我,会怎样做下一步”。

对本课来说,每一讲的配套代码就是一种”第三遍”:用几百行 Python 把一篇论文的核心机制重新实现一遍,你会立刻发现哪些地方是论文一句带过、实际上很微妙的。

9.3 读 MLSys 论文时要问的问题

slides 给了五个通用问题(转述):这篇论文要解决什么问题?为什么这个问题重要又困难?以前的工作为什么解决不了?新意在哪里?结果好不好?

针对 MLSys 论文,本讲义再补充一份检查清单,它们大多来自 §4 和 §8 的经验:

  1. 瓶颈是什么? 论文针对的是计算、访存、通信、开销还是空泡?用 roofline 或 α-β 模型估一下:理论上最多能快多少?
  2. 基线强不强? 对比的是最新版本、调优过的系统,还是一个很弱的实现?
  3. 端到端还是微基准? 一个内核快了 3 倍,端到端可能只快 5%(§8 里矩阵乘法以外的部分只占 17% 的时间)。
  4. 在什么硬件上测的? 换一代 GPU,瓶颈可能完全不同(§4.2)。
  5. 工作负载有没有代表性? 批大小、序列长度、模型尺寸是不是真实场景里的取值?
  6. 有没有消融实验? 加速来自论文的核心想法,还是来自一些顺带的工程优化?
  7. 正确性 / 准确率有没有代价? 如果是近似方法(量化、剪枝、部分等价变换),准确率掉了多少?如果是精确方法,怎样保证精确?
  8. 开销藏在哪里? 搜索 / 编译 / 调优要多久?额外的显存从哪里来?

9.4 评审的格式与一份示范评审

原课要求的评审有三段(转述):

  1. 用自己的话概括第一篇论文(不要照抄论文的句子);
  2. 用自己的话概括第二篇论文;
  3. 写两篇之间的联系,例如:比较它们的优劣;用一篇的想法去解决另一篇的问题;或者提出一个结合两者的新想法。

原课后来补充说明:不需要写很长,抓住要点即可。

下面是一份本讲义写的示范评审,对象是第 2 讲的两篇必读论文。建议你先自己读这两篇论文、写完评审,再回来对照。

TensorFlow(OSDI’16):这篇论文要解决的是”怎样用一个统一的编程模型,既支持研究人员快速试验新模型,又能在从手机到数千台机器的异构环境里高效地训练和部署”。它的核心设计是把整个计算——包括运算和模型的可变状态(变量)以及更新状态的操作——都表示成一张数据流图,先构图、后执行(延迟执行)。因为系统在执行前能看到完整的图,它就可以做全局的优化,把子图放到不同设备上,并在跨设备的边上自动插入发送 / 接收操作;以前参数服务器架构里写死在系统中的”参数更新规则”,也变成了图里普通的操作,用户可以自己定义。论文的局限也来自同一个设计:先构图再执行让调试和动态控制流变得别扭,用户写的是”构建图的代码”而不是”计算本身”。

PyTorch(NeurIPS’19):这篇论文的出发点和 TensorFlow 相反:它认为可用性应该放在第一位,并论证”即时执行”(eager / define-by-run)的命令式编程不一定要牺牲性能。PyTorch 让每个运算在 Python 里立即执行,模型就是普通的 Python 程序,自动微分通过运算时记录的”磁带”反向回放实现。为了拿回性能,它把控制流(留在 Python 里)和数据流(交给 C++ 核心和 GPU)分开:Python 只负责按顺序发射内核,内核在 GPU 上异步执行,所以只要内核足够大,Python 的开销就能被藏起来;再配合缓存显存分配器、引用计数立即释放内存、多进程数据加载等工程设计。论文报告它在常见基准上的速度与当时最快的框架相差不大。

两篇的联系:两篇论文代表了 ML 框架设计的两极——“先看见整个计算再优化”(TensorFlow)对”立即执行、保持灵活”(PyTorch)。TensorFlow 的优势正是 PyTorch 的短板:没有全局的图,就难以做跨算子的融合、内存规划和自动切分;而 PyTorch 的优势(易用、易调试、支持任意 Python 控制流)正是 TensorFlow 1.x 的痛点。有意思的是两者后来走向了融合:TensorFlow 2 默认采用即时执行,再用 tf.function 把需要加速的部分转成图;PyTorch 2 用 torch.compile(TorchDynamo)在不改变用户写法的前提下,从即时执行的 Python 代码里”抓出”计算图交给编译器。本讲 §8.5 用的正是后者——一个可以讨论的问题是:如果用户代码里有依赖数据的控制流(例如 if x.sum() > 0),这种”事后抓图”的方法应该怎样处理?

注意这份评审的写法:每一段都先说”要解决什么问题”,再说”核心想法”,最后说”代价或局限”;第三段不是简单地罗列异同,而是指出两者的互补关系和后来的发展,并留了一个开放问题。

9.5 课程项目

原课的成绩构成是:课程项目 50%,论文评审 30%,课堂参与(展示、Piazza)20%。项目由 1–3 人的小组完成,里程碑如下(细节见原课 Logistics 页,附加讲 A3 会展开):

里程碑要求
一页开题(5%)题目和成员;引言(核心想法);问题(要解决的研究问题,清楚地定义);现状(已有方法及其局限,你的方法和它们的关系);实施计划(怎样实现和评估,基于哪个框架、用什么数据)
中期检查(不计分)到 office hour 讨论或发一份简短的进展说明
期末展示(15%)8–10 分钟的录播:概述、问题、相关工作、方法、评估,时长也计分
期末报告(30%)最多 8 页(不含参考文献),MLSys 会议格式:引言、问题、相关工作、系统总览图、方法、评估、结论;提交代码

原课提供了一份项目点子清单,也鼓励学生带自己研究相关的题目。几个选题的经验(本讲义的建议):

  • 从瓶颈出发:找一个你能测到的瓶颈(某个内核没跑满带宽、某种并行策略空泡太大、某类请求排队太久),而不是从”我想用某个技术”出发;
  • 新硬件特性 + 旧问题:每一代硬件的新特性(§4.1)都会让一批旧问题有新的解法;
  • 被忽视的工作负载:长上下文、小批量、MoE、多模态、RL 后训练……主流系统往往只为最常见的场景调优(§4.4 的”瘦”矩阵乘法就是一个小例子);
  • 先估算再动手:开题前用 roofline 估一下”理想情况下能快多少”,如果上限只有 5%,换个题目。

10. 小结 · Summary

  1. ML 的进步 = 模型 × 数据 × 算力,而算力 = 硬件 × 系统效率。 规模定律让训练和推理两侧的算力需求持续膨胀(GPT-3 到 Llama 3 405B,单模型训练量涨了约 120 倍);MFU 每提高 10 个百分点,一次万卡训练就能省下约两周。
  2. ML 硬件大规模并行、高度异构、快速演进。 一块 H100 有 132 个 SM、近 1.7 万个 CUDA core;片上存储约 64 MB,HBM 有 80 GB;每一代都有需要软件改写才能用上的新特性。
  3. 算力涨得比带宽快(内存墙)。 V100 到 B200,BF16 算力涨了 18 倍、带宽涨了 8.9 倍;越来越多的运算落在 roofline 的斜坡上。
  4. roofline 模型:t≥max⁡(W/Ppeak,Q/B)t \ge \max(W/P_{\text{peak}}, Q/B),算术强度 I=W/QI = W/Q,拐点 I⋆=Ppeak/BI^{\star} = P_{\text{peak}}/B。瓶颈有三个候选:计算、访存、固定开销。roofline 只给上限。
  5. ML 系统的目标是在大规模并行、日益异构、快速演进的硬件上高效地部署 ML 应用;MLSys 研究的特点是跨层协同设计——摄像头练习里,只有同时动模型、系统和硬件的团队找到了可行方案。
  6. 课程地图:内核(3–7)→ 编译(8–9)→ 图优化(10–11)→ 并行与内存(12–15)→ LLM 专题与算法优化(16–21)。一条暗线是”把手工优化变成搜索空间 + 代价模型 + 正确性验证”。
  7. 怎样学:用三遍阅读法读论文,用五个问题 + MLSys 检查清单审视论文,按三段式写评审;用”先估算再动手”的习惯选课程项目。

11. 常见误区 · Common misconceptions

  1. “FLOPs 越少就越快。” 只对计算受限的运算成立。§8 里矩阵乘法以外的运算只有 0.03% 的 FLOPs,却占了 17% 的时间;§6.3 里非结构化剪枝砍掉一半 FLOPs,时延纹丝不动。
  2. “峰值算力就是能达到的算力。” 厂商标称的是理想条件下的峰值(而且宣传页常写稀疏加速后翻倍的数字)。大规模训练的 MFU 通常只有 30%–50%;访存受限的运算只能用到峰值的千分之几。
  3. “新一代 GPU 会让所有运算都同比例变快。” 算力和带宽涨幅不同:访存受限的运算只跟着带宽涨,计算受限的运算跟着算力涨。同一个程序的瓶颈可能在换卡之后转移。
  4. “roofline 能预测一个内核的耗时。” 它只给上限(下界时间)。§4.4 里同样访存受限的运算,实测从屋顶的 18% 到 93% 不等;§8.5 里融合后的 RMSNorm 只跑到 18%。
  5. “融合总是更快。” 融合省下了显存读写,但如果融合后的内核写得不好(例如归约没有向量化),或者运算本身是计算受限的,就可能不快甚至更慢(§8.5)。
  6. “一个运算是不是访存受限,只看算法就知道。” 算术强度只看算法,但”是否访存受限”要和硬件的拐点比较。同一个 silu,在 H100 上访存受限,在 CPU 单核上计算受限。
  7. “MLSys 就是工程,没有研究问题。” MLSys 的很多核心贡献是新的抽象(数据流图、PagedAttention 的分页)、新的搜索空间(SOAP、µGraph)和新的正确性保证(形式化验证的图替换、推测解码的无偏性),而不只是”把代码写快”。
  8. “读论文就是从头读到尾。” 三遍阅读法的第一遍只要五到十分钟,大部分论文读完第一遍就够了;把时间留给真正需要第三遍的那几篇。

12. 练习 · Exercises(含参考答案)

练习 1(拐点):A100 的 FP16 Tensor Core 峰值是 312 TFLOP/s,HBM 带宽 2.039 TB/s。(a) 计算它的拐点。(b) 一个 FP16 的向量加法 z=x+yz = x + y 在 A100 上最多能用到峰值的百分之几?

答:(a) 312/2.039≈153312 / 2.039 \approx 153 FLOP/byte。(b) 每个元素 1 次运算,读 2 个、写 1 个 FP16(6 字节),I=1/6I = 1/6;Pattainable=(1/6)×2.039 TB/s≈0.34P_{\text{attainable}} = (1/6) \times 2.039 \text{ TB/s} \approx 0.34 TFLOP/s,占峰值约 0.11%(和代码 01 演示 3 一致)。

练习 2(批处理与算术强度):解码时,一个 d×dd \times d 的 BF16 权重矩阵乘以 bb 个 token 的激活(形状 b×db \times d)。(a) 写出算术强度 I(b)I(b)。(b) d=4096d = 4096 时,bb 至少要多大,这个矩阵乘法在 H100 上才可能是计算受限的?

答:(a) W=2bd2W = 2bd^2,Q=2(d2+2bd)Q = 2(d^2 + 2bd)(权重读一次,输入读一次,输出写一次),I(b)=bd2d2+2bd=bdd+2bI(b) = \dfrac{bd^2}{d^2 + 2bd} = \dfrac{bd}{d + 2b}。当 b≪db \ll d 时 I≈bI \approx b。(b) 需要 I(b)≥295I(b) \ge 295:4096b/(4096+2b)≥295.3⇒b≥295.3×4096/(4096−590.6)≈345.14096b / (4096 + 2b) \ge 295.3 \Rightarrow b \ge 295.3 \times 4096 / (4096 - 590.6) \approx 345.1,即 b≥346b \ge 346。也就是说,一批至少要三百多个 token 同时解码,才能让权重读取不再是瓶颈——这是第 16 讲连续批处理的量化动机。

练习 3(训练时间):用 2048 块 H100 训练一个 70B 参数的稠密模型,数据 15T token,MFU 45%。估算需要多少天。

答:C=6×70×109×15×1012=6.3×1024C = 6 \times 70 \times 10^9 \times 15 \times 10^{12} = 6.3 \times 10^{24} FLOPs。每秒有效算力 =2048×989.4×1012×0.45≈9.12×1017= 2048 \times 989.4 \times 10^{12} \times 0.45 \approx 9.12 \times 10^{17} FLOP/s。时间 ≈6.9×106\approx 6.9 \times 10^6 秒 ≈80\approx 80 天。(可以用 perf_model.training_days(6.3e24, 2048, H100, 0.45) 核对。)

练习 4(解码的时延下界):一个 8B 参数的 BF16 模型在单块 H100 上以 batch=1 解码。忽略 KV 缓存和启动开销,每生成一个 token 至少要多久?每秒最多生成多少个 token?

答:每步至少要把全部权重读一遍:8×109×28 \times 10^9 \times 2 字节 =16= 16 GB,16/3.35≈4.816 / 3.35 \approx 4.8 ms,即每秒最多约 210 个 token。注意这和算力几乎无关(每步只有 2×8×109=1.6×10102 \times 8 \times 10^9 = 1.6 \times 10^{10} FLOPs,在 989 TFLOP/s 下只要 16 µs)。所以降低解码时延的手段是减少要读的字节(量化)、让一次读取服务更多 token(批处理、推测解码),而不是加算力。

练习 5(融合的字节账):对 y = silu(g) * u(gg、uu 都是 nn 个 BF16 元素),(a) eager 模式下两个内核一共读写多少字节?(b) 融合成一个内核后呢?(c) 如果它是访存受限的,融合的理论加速比是多少?(d) 代码 03 在 CPU 上实测只有约 1.05×,为什么?

答:(a) silu:读 gg 写 ss,4n4n 字节;mul:读 ss、uu 写 yy,6n6n 字节;合计 10n10n。(b) 读 gg、uu 写 yy:6n6n。(c) 10/6≈1.67×10/6 \approx 1.67\times。(d) 在 CPU 单核上,silu 里的指数运算很贵,这个运算本身是计算受限的(eager 版本的有效带宽只有约 23 GB/s,远低于约 115 GB/s 的屋顶),省字节帮不上忙。

练习 6(改一改协同设计的玩具):修改 04_codesign_camera.py:(a) 让 NPU 也支持 fp16(算力设为 1000 GOP/s)。系统团队现在能找到可行方案吗?(b) 把训练后量化的准确率损失从 2.0 改成 1.0,又会怎样?(c) 这说明了这个练习的结论有什么局限?

答:(a) 系统团队可以把 Det-L 转成 fp16(不需要重新训练),在 NPU 上计算时间 15/1000=1515 / 1000 = 15 ms,仍然超过 10 ms,还是找不到(自己运行确认)。(b) Det-L + int8-PTQ @NPU 的准确率变成 90.8%,系统团队就找到可行方案了。(c) 结论依赖于具体数字:如果硬件足够强,或者某种不需要训练的转换足够好,单靠一方也能成功。练习想说明的不是”系统团队一定失败”,而是搜索空间越大,能找到的方案越多;真实项目里,跨越模型 / 系统 / 硬件边界往往是唯一的出路,但不是每次都是。

练习 7(读论文):用 Keshav 的第一遍方法(不超过十分钟)读《How to Read a Paper》本身,回答五个 C。

答(参考):Category:方法论 / 经验分享型的短文,不是研究论文。Context:面向研究生和新研究者,建立在作者自己多年的读论文经验上,也提到了做文献调研的方法。Correctness:它的”假设”是”读论文的时间有限、大部分论文不值得精读”,对大多数研究者成立。Contributions:提出三遍阅读法,并给出文献调研的流程。Clarity:非常清楚,只有两三页。

练习 8(写评审):读完第 2 讲的两篇必读论文(TensorFlow、PyTorch)后,按原课的三段式格式写一份自己的评审,然后和 §9.4 的示范评审对照:你的第三段提出了什么不同的联系?

答:开放题。一些可以写的方向:两者的自动微分实现方式(图上的符号微分 vs 运行时的磁带)对内存优化(例如重计算)有什么影响;两者对分布式训练的支持方式(图里的发送 / 接收操作 vs 外部的通信库);如果把 TensorFlow 的”变量也在图里”的设计用到 PyTorch 的编译器里,能带来什么新的优化机会。


13. 延伸阅读 · Further reading

  • S. Keshav, How to Read a Paper(2007):本讲必读,两三页,值得每学期重读一次。
  • S. Williams, A. Waterman, D. Patterson, Roofline: An Insightful Visual Performance Model for Multicore Architectures(CACM 2009):roofline 模型的原始论文。
  • J. Kaplan et al., Scaling Laws for Neural Language Models(2020);J. Hoffmann et al., Training Compute-Optimal Large Language Models(Chinchilla, 2022):规模定律。
  • A. Ratner et al., MLSys: The New Frontier of Machine Learning Systems(2019):MLSys 会议创办时的白皮书,描述了这个领域关心的问题。
  • NVIDIA Hopper 与 Blackwell 架构白皮书:TMA、线程块簇、分布式共享内存、Tensor Memory 的一手资料,第 6–7 讲会用到。
  • H. He, Making Deep Learning Go Brrrr From First Principles(博客,2022):用”计算、访存、开销”三类瓶颈解释深度学习性能,和本讲 §4.4、§8 的视角完全一致,强烈推荐。
  • 本仓库 CMU11-868 第 1 讲:从 LLM 的角度估算”训练一个 100B 模型要多少资源”,与本讲的 §2 互补。

14. 下一讲预告 · Next lecture

第 2 讲《ML 系统 101:基础 + TensorFlow / PyTorch》会从本讲 §8 的”计算图”出发,系统地介绍一个深度学习系统的标准部件:张量与算子、计算图、自动微分、计算图优化(Conv + BatchNorm 折叠)、并行化概览、ML 编译概览、内存优化(重计算与零冗余),最后精读两篇必读论文,比较 TensorFlow 的”延迟执行”和 PyTorch 的”即时执行”各让哪些优化变得容易或困难。配套代码会实现一个极小的计算图 IR minigraph.py,在第 10、11、15 讲继续使用。

在读第 2 讲之前,建议你先用三遍阅读法的前两遍读完 TensorFlow 和 PyTorch 两篇论文,并试着写一份评审(练习 8)。

← 返回课程目录