PyTorch (einops), resource accounting (FLOPs, memory, arithmetic intensity)
- 日期
- Wed April 1
- 讲师
- Percy
- 课程资料
- 10 项
- 内容产物
- 4 项
课程简介
CS336 Lecture 2：PyTorch (einops) 与资源核算
学习目标
本讲的核心目标是建立"资源核算"(resource accounting)的思维方式：在训练模型之前，先精确估算所需的计算量(FLOPs)和内存占用。具体包括：
- 理解张量(Tensor)存储的基本原理，以及不同精度(fp32、fp16、bf16、fp8、fp4)对内存和计算速度的影响
- 掌握 einops 库的三种核心操作：einsum、reduce、rearrange
- 推导训练深度网络的计算量公式：6 × token数 × 参数量
- 理解算术强度(arithmetic intensity)与 Roofline 模型，判断操作是内存受限还是计算受限
- 掌握梯度累积(gradient accumulation)与激活检查点(activation checkpointing)两种内存优化技巧
课程主线
本讲属于系统层面(systems)的内容，关注"给定有限资源下如何训练最优模型"。课程从最底层的张量存储开始，逐步构建到完整的训练循环。
实际案例：Marin 项目
讲师提到 Marin 项目的 1e23 FLOPs 训练已完成，结果与 IsoFLOPs 曲线的预测损失值相差在 0.05 以内。这展示了资源核算的实际价值。
1. 张量与数据类型
1.1 张量是基础构建块
在深度学习中，所有数据——参数、梯度、优化器状态、激活值——都是张量(Tensor)。例如 DeepSeek 3.2 模型的参数就是一组不同形状的张量。
每个张量具有：
- 秩(rank)：维度数量。例如
torch.zeros(4, 8)是秩 2 的矩阵 - 数据类型(dtype)：决定每个元素占用的字节数
1.2 浮点数精度详析
float32 (fp32)——默认精度
- 组成：1 位符号位 + 8 位指数位 + 23 位尾数
- 每个元素 4 字节
- 在深度学习中使用 fp32 往往过于浪费
float16 (fp16)——半精度
- 组成：1 位符号位 + 5 位指数位 + 10 位尾数
- 每个元素 2 字节
- 问题：动态范围差，
1e-8下溢为 0，训练中易出现 NaN
bfloat16 (bf16)——Google Brain 2018 年开发
- 组成：1 位符号位 + 8 位指数位 + 7 位尾数
- 与 fp32 相同的动态范围，但分辨率更差
- 深度学习训练过程粗糙，高分辨率并非必需，因此 bf16 是很好的折中
fp8——2022 年标准化
- 两种变体：E4M3（范围 [-448, 448]）和 E5M2（范围 [-57344, 57344]）
- 由 NVIDIA Transformer Engine 支持
fp4——NVIDIA 2025 年开发
- 仅 4 比特，所有可能值从 -6 到 6
- 实际应用时使用分块缩放(block scaling)，每块可整体放大或缩小
- Nemotron 3 Super 使用 fp4 训练
混合精度训练
- 通用规则：参数、激活值、梯度用 bf16，优化器状态用 fp32
- PyTorch 提供 AMP 库自动处理，在安全处（如矩阵乘法）使用低精度，在危险处（如指数运算）保持 fp32
1.3 内存计算基础
内存占用 = 元素个数 × 每元素字节数
示例：4×8 的 fp32 张量占用 128 字节。GPT-3 前馈层中一个矩阵约 2.3 GB。
关键操作：PyTorch 默认张量在 CPU 上，需要显式 .to(device) 移到 GPU。
2. einops：命名维度的张量操作
2.1 动机
传统 PyTorch 代码中的 x.transpose(-2, -1) 容易混淆维度。einops 用命名维度替代索引，代码更清晰、不易出错。
einsum——带记账的广义矩阵乘法
z = einsum(x, y, "seq1 hidden, hidden seq2 -> seq1 seq2")- 命名出现在输出中的维度保留
- 未出现在输出的维度（如
hidden）自动求和
reduce——sum/mean/max/min 的推广
y = reduce(x, "... hidden -> ...", "sum")...代表任意数量的批量维度- 右侧未出现的维度被聚合
rearrange——重塑与分解
# 将 total_hidden 分解为 heads × hidden1
x = rearrange(x, "... (heads hidden1) -> ... heads hidden1", heads=2)
# 处理后合并
x = rearrange(x, "... heads hidden2 -> ... (heads hidden2)")- 括号表示维度的合并或分解
- 分解时需指定其中一个维度的值
2.2 实际优势
在反向传播中，einsum 使维度推导变得清晰：
- 通过形状匹配确定索引顺序，无需记忆转置规则
3. 计算量核算：FLOPs
3.1 术语澄清
- FLOPs（小写 s）：浮点运算次数，衡量已完成的计算量
- FLOP/s 或 FLOPS（大写 S）：每秒浮点运算次数，衡量硬件速度
注意 H100 规格表中的 1979 teraFLOP/s 包含稀疏性(sparsity)，稠密计算需除以 2。
3.2 矩阵乘法 FLOPs
对于 B×D 与 D×K 的矩阵乘法：
FLOPs = 2 × B × D × K原因：对每个 (i, j, k) 三元组，做一次乘法和一次加法。
直观理解：线性层 FLOPs = 2 × token数 × 参数量。这正是 Transformer 6ND 公式的雏形。
3.3 其他操作
- 逐元素操作（如加法）：FLOPs = 矩阵大小
- 对于足够大的矩阵，矩阵乘法主导所有其他操作
4. 实际性能与 MFU
4.1 测量实际性能
torch.cuda.synchronize() # 等待 GPU 异步任务完成
time = benchmark(lambda: x @ w) # 多次运行取平均
actual_flop_per_sec = num_flops / time4.2 模型 FLOPs 利用率 (MFU)
MFU = 实际 FLOP/s ÷ 规格书 FLOP/s- MFU ≥ 0.5 是良好水平
- 纯矩阵乘法可达 0.8
- 若 MFU 约 0.1，说明需要优化
5. 算术强度与 Roofline 模型
5.1 为什么 MFU 不到 1？
需要从加速器架构角度理解：
- HBM（高带宽内存）→ 加速器芯片（计算核心）
- 每步计算需：从内存发送输入 → 计算 → 发回输出
- 时间取决于：计算速度（FLOP/s）和内存带宽（bytes/s）
5.2 关键比量
加速器强度：每传输 1 字节可做多少计算
H100: 1979e12/2 FLOP/s ÷ 3.35e12 bytes/s ≈ 295 FLOPs/byte算术强度：工作负载的 FLOPs ÷ 字节数
判断标准：
- 算术强度 < 加速器强度 → 内存受限(memory-bound)
- 算术强度 > 加速器强度 → 计算受限(compute-bound)
5.3 各操作算术强度分析
| 操作 | 算术强度 | 类型 |
|---|---|---|
| ReLU | ~0.25 | 内存受限 |
| GELU | ~5 | 内存受限（即使计算量是 ReLU 的 20 倍） |
| 点积 | ~0.5 | 内存受限 |
| 矩阵-向量乘 | ~1 | 内存受限 |
| 矩阵-矩阵乘 | ~n/3 | 计算受限（n=1024 时约 340） |
关键洞察：
- ReLU 和 GELU 的隔离性能相同，因为瓶颈在内存，不在计算
- 矩阵乘法因数据量 O(n²) 但计算量 O(n³)，具有高算术强度
- Transformer 设计即保证大矩阵乘法主导，这是刻意的架构选择
5.4 推理 vs 训练
- 推理逐 token 生成 = 矩阵-向量乘 → 内存受限
- 训练处理整个序列 = 矩阵-矩阵乘 → 计算受限
6. 梯度计算与 6ND 公式
6.1 反向传播 FLOPs
对于一层 h2 = h1 @ w2：
- 前向：2 × B × D × D FLOPs
- 反向：需计算两个梯度
- h1.grad = h2.grad @ w2ᵀ → 2 × B × D × D
- w2.grad = h2.gradᵀ @ h1 → 2 × B × D × D
- 总计 4 × B × D × D
反向传播恰好是前向的 2 倍。
6.2 6ND 公式推导
对 L 层网络，每层 D² 参数：
- 前向：2 × B × (D² × L) = 2 × token数 × 参数量
- 反向：4 × token数 × 参数量
- 总计：6 × token数 × 参数量 FLOPs/步
此公式对短上下文的 Transformer 也是良好近似。上下文很长时，注意力 O(L²) 项会有额外贡献。
7. 训练内存核算
7.1 内存组成
以 Adam 优化器为例（每参数 4+4=8 字节状态）：
| 组件 | 每参数字节数 |
|---|---|
| 参数(bf16) | 2 |
| 梯度(bf16) | 2 |
| 优化器状态(fp32) | 8（Adam 需要一阶矩+二阶矩） |
| 总计 | 12 |
80GB H100 上最大模型规模：80e9 × 8 ÷ 12 ≈ 530 亿参数 注意：未计入激活值，实际取决于批大小与序列长度。
7.2 内存的双重作用
- 容量限制：决定能否容纳大模型
- 速度影响：数据传输耗时，优化器状态虽不主导计算，但影响能加载的模型大小
8. 内存优化技术
8.1 梯度累积 (Gradient Accumulation)
问题：大 batch 提高训练稳定性，但激活内存随 batch 线性增长。
方法：
- 在微批次上计算梯度并累积（不调 zero_grad）
- 每 batch_size/micro_batch_size 步更新参数并清零
激活内存降为：2 × micro_batch_size × D × L
8.2 激活检查点 (Activation Checkpointing)
思想：训练需要存储所有层激活值（内存 O(L)），推理不需梯度仅存当前层。
实现：前向时只保留部分层的激活值；反向时从最近检查点重算缺失激活。
- 不检查点：内存 O(L)，无需重算
- 完全不存储：内存 O(1)，计算 O(L²)——每层都从头算
- 平衡方案：每 √L 层存一个检查点 → 内存 O(√L)，计算 O(√L)
9. 课程作业关联
- 作业 1 涉及 tokenization 和 Transformer 架构实现
- 需要自己完成 Transformer 的 FLOPs 核算（比简单 MLP 更复杂）
- einops 在作业中可能要求使用
- 最后需实现 Adam 优化器及其内存核算
常见误区
- 混淆 FLOPs 与 FLOP/s：前者是工作总量，后者是速率
- 忘记 H100 规格要除以 2：稀疏性虚高
- 盲目追求大 batch：忽略激活内存线性增长
- 认为 ReLU 比 GELU 快：隔离性能相同，因为内存带宽是瓶颈
- 忽略 GPU 异步性：不调用
cuda.synchronize()的计时结果虚低 - 训练推理混淆：推理是矩阵-向量乘（内存受限），训练是矩阵-矩阵乘（计算受限）
复习清单
- [ ] 解释 fp16 为何易下溢，bf16 如何解决
- [ ] 写出混合精度训练的参数精度分配
- [ ] 计算 B×D @ D×K 的 FLOPs
- [ ] 定义 MFU 并解释其意义
- [ ] 计算 H100 的加速器强度
- [ ] 判断 ReLU/GELU/点积/矩阵-向量乘/矩阵乘法是内存受限还是计算受限
- [ ] 推导反向传播 FLOPs = 2 × 前向 FLOPs
- [ ] 解释 6ND 公式的来历
- [ ] 描述 Adam 每参数 12 字节的来源
- [ ] 简述梯度累积和激活检查点的工作原理
- [ ] 用 einops 完成一个矩阵乘法、reduction 和 reshape 操作
CS336 Lecture 2 深度解析：从 PyTorch 到资源核算——训练大模型前的必修课
在动手训练大模型之前，你必须回答两个问题：这个模型需要多少计算量？它需要多少内存？本讲将为你建立一套完整的"资源核算"方法论，让你在写下第一行训练代码之前，就能对成本心中有数。
导语：为什么需要资源核算？
在上一讲中，我们了解了语言模型（Language Model）的整体蓝图：从分词（Tokenization）到模型架构，再到训练策略。但有一个关键问题被搁置了——在给定的硬件资源下，我们究竟能训练多大的模型？需要多长时间？
本讲的核心目标是培养一种"资源核算"（Resource Accounting）的思维模式。这不是一个可选的优化技巧，而是训练现代大模型的必要前提。当你面对一个 700 亿参数的模型时，你不可能凭借直觉去猜测训练时间；你必须能够精确计算出所需的 FLOPs（浮点运算数）、内存占用，并理解硬件特性如何影响实际性能。
课程一开始，讲师就展示了两个极具实战性的问题：
- 在 1024 块 H100 上训练 700 亿参数、1.5 万亿 token 的模型需要多久？答案是约 143 天。
- 在 8 块 H100 上用 AdamW 优化器能训练的最大模型是多大？答案是约 530 亿参数。
这些数字不是凭空产生的，它们来自一个简单的公式：FLOPs = 6 × 参数量 × token 数。这个公式将成为你进行任何规模估算的起点。
一、张量：一切的基础
在深度学习的世界里，所有东西归根结底都是张量（Tensor）。参数、梯度、优化器状态、激活值、输入数据——它们都是不同形状和精度的张量。即使是像 DeepSeek 这样的大模型，其本质也就是一组具有特定形状的张量集合。
理解张量存储的内存占用非常简单：内存 = 元素个数 × 每个元素的字节数。例如，一个 4×8 的 float32 矩阵占用 128 字节（4×8×4）。
但现代深度学习早已超越了 float32 的范畴。让我们看看各种精度类型如何影响内存和性能：
精度选择：一场内存与稳定性的权衡
float32（fp32）：这是默认的精度，每个数占 4 字节。在深度学习中，我们很少使用更高精度的 float64，因为训练过程的随机性意味着我们不需要那么高的数值精度。
float16（fp16）：占 2 字节，看似是节省内存的好选择。但它的动态范围太窄——指数位只有 5 位（fp32 有 8 位）。如果你尝试用 fp16 表示 1e-8，它直接变成 0（下溢）。这会导致训练不稳定、出现 NaN 等问题。
bfloat16（bf16）：这是 Google Brain 在 2018 年开发的创新。它同样占 2 字节，但拥有和 fp32 相同的动态范围！秘诀在于它牺牲了尾数位来换取指数位。虽然分辨率（precision）变差了，但事实证明，在深度学习中这个权衡非常值得——毕竟训练过程本身带有随机性，不需要那么精细的粒度和精度。
混合精度训练：现代训练的标准做法是混合精度（Mixed Precision）——用 bf16 存储参数、激活值和梯度，用 fp32 存储优化器状态。PyTorch 的 AMP（Automatic Mixed Precision）库可以自动处理这个转换，在安全的情况下（如矩阵乘法）使用低精度，在危险的地方（如指数运算）保持高精度。
讲师还提到了更极端的选项：fp8 和 fp4。fp8 有两种变体，分别优化动态范围和分辨率；fp4 则每个值只有 4 比特，可能取值的范围从 -6 到 6。这些极端量化通常使用分块缩放（block scaling）来弥补精度的不足，但主要用于推理而非训练。
二、einops：让张量操作变得清晰
在进入 FLOPs 计算之前，讲师介绍了 einops 库——一个让张量操作变得直观的工具。传统的 PyTorch 操作如 x @ y.transpose(-2, -1) 很容易让人困惑：-2 和 -1 到底指什么维度？
einops 的核心理念是给维度命名。例如：
z = einsum(x, y, "batch seq1 hidden, batch seq2 hidden -> batch seq1 seq2")这个表达式清晰说明了：我们有两个输入张量，第一个的维度是 (batch, seq1, hidden)，第二个是 (batch, seq2, hidden)，输出是 (batch, seq1, seq2) —— 同时自动在 hidden 维度上求和（该维度未出现在输出中）。
einops 还提供 reduce（用于求和、求均值等操作）和 rearrange（用于重塑维度）。这些工具让代码更具可读性，也减少了人为错误。讲师特别强调：习惯 einops 后，你会以不同的方式思考张量操作——不再纠结于转置和索引，而是聚焦于数据流本身的逻辑。
三、FLOPs：衡量计算量的标准
**FLOPs（浮点运算数）**是衡量计算量的核心指标。它记录的是浮点运算的次数，如加法和乘法。与之容易混淆的是 FLOP/s（每秒浮点运算次数），后者衡量硬件速度。
计算一个操作的 FLOPs 通常很简单：
矩阵乘法（MatMul）：对于一个 B×D 的矩阵乘以 D×K 的矩阵，FLOPs = 2 × B × D × K。这个 2 来自每次乘法和加法。这在大多数情况下是准确量级（忽略了 D-1 次加法而非 D 次这种细微差别，对于大矩阵几乎无影响）。
逐元素操作：FLOPs = 元素个数。例如，ReLU 每次比较算一个 FLOP。
关键洞察：对于足够大的矩阵，矩阵乘法的 FLOPs 远远大于其他操作。这就是为什么我们在设计模型时，要尽量将计算重负载集中在矩阵乘法上。
6ND 公式的推导
现在让我们推导那个神奇的公式：训练一个模型的 FLOPs ≈ 6 × 数据点数 × 参数量。
假设一个简单的深度网络，每层是矩阵乘法加 ReLU：
- 前向传播：2 × B × D × D（这是矩阵乘法的 FLOPs，B 是数据点数，D 是参数维度）
- 反向传播：需要计算两个梯度——对输入的梯度（即回传的消息）和对参数的梯度。这两个操作各需要 2 × B × D × D 的 FLOPs，所以反向传播总共是 4 × B × D × D。
总 FLOPs = 前向 (2BD²) + 反向 (4BD²) = 6BD² = 6 × 数据点数 × 参数量（这里参数量 = D²）。
这个公式对于 Transformer 也是一个很好的近似，只要上下文长度不太长（如果上下文长度过长，注意力机制会引入额外的二次项）。
四、从理论到实践：MFU 与算术强度
知道模型的 FLOPs 只是第一步。真正的问题在于：你的硬件能否高效地执行这些操作？
MFU：模型 FLOPs 利用率
MFU（Model FLOPs Utilization） 衡量实际达到的 FLOP/s 与理论峰值的比值：
MFU = 实际 FLOP/s ÷ 宣传 FLOP/s问题在于，实际 FLOP/s 往往远低于规格说明书上的峰值。H100 的 bf16 峰值大约是 989 TFLOP/s（1979 除以 2，因为 1979 是稀疏模式的数值，稠密模式减半）。但在实际训练中，MFU 达到 50% 就已经很不错了。
为什么 MFU 这么低？答案在于内存带宽。
算术强度：计算与内存的平衡
这就是为什么我们需要理解算术强度（Arithmetic Intensity）：
加速器强度 = FLOP/s ÷ bytes/s（H100 大约是 989e12 ÷ 3.35e12 ≈ 295 FLOPs/byte）
算法算术强度 = FLOPs ÷ bytes（特定工作负载）核心规则：
- 如果算法算术强度 < 加速器强度 → 内存受限（Memory-bound）：瓶颈在数据传输
- 如果算法算术强度 > 加速器强度 → 计算受限（Compute-bound）：瓶颈在计算本身
让我们看几个具体例子（假设 bf16，2 字节/数）：
ReLU：读 x（2n 字节）+ 写 y（2n 字节），FLOPs = n（一次比较）。算术强度 ≈ 0.25 → 严重内存受限。这意味着 ReLU 的性能实际上与计算无关，完全取决于传输速度。
GELU：虽然计算量更大（约 20n FLOPs），算术强度 ≈ 5，但仍然远低于 295 → 还是内存受限！所以 GELU 和 ReLU 相比并没有显著的性能劣势——瓶颈根本不在于计算。
点积：读两个向量（2n + 2n 字节），FLOPs = 2n。算术强度 ≈ 0.5 → 内存受限。
矩阵向量乘法：读 x（2n）+ 读 w（2n²）+ 写 y（2n），FLOPs = 2n²。算术强度 ≈ 1 → 内存受限。这正是推理（单个 token 逐个生成）的场景。
矩阵乘法：三个 n×n 矩阵参与，FLOPs = 2n³，传输字节数 = 6n²。算术强度 ≈ n/3 → 计算受限！矩阵越大，这个比值越大，硬件利用率越高。
这个分析解释了为什么训练（使用大批量、大矩阵）比推理（单 token 逐步生成）的硬件利用率要高得多。也解释了为什么总有人强调"扩大批量大小"——更大的矩阵不仅数据效率更好，计算效率也更高。
为什么 MFU 只有 50%？
现在我们可以回答这个问题了。即使训练过程充斥着计算受限的矩阵乘法，仍然会有许多内存受限的操作（如激活函数）穿插其间。此外，GPU 的异步执行、通信开销等因素都会拉低实际利用率。所以，MFU 达到 50% 是常态，70-80% 则非常优秀。
五、训练的内存清单
现在让我们转向训练的实际内存需求。以一个深度网络为例：
参数：对于 L 层、每层 D×D 的矩阵，参数量 = D²L。用 bf16 存储：2 × 参数量 字节。
梯度：每个参数对应一个梯度，用 bf16 存储：2 × 参数量 字节。
优化器状态：这里通常使用 fp32（4 字节/参数）以保证稳定性：
- AdaGrad：存储梯度平方，4 字节/参数
- AdamW：存储一阶矩 + 二阶矩，8 字节/参数！
激活值：训练需要存储前向传播的中间结果用于反向传播。对于每层：2 × B × D × L 字节（bf16）。这解释了为什么批量大小（B）和层数（L）会显著影响内存需求。
关键洞察：现代深度学习训练中，优化器状态往往消耗最多内存，远超模型参数本身。这导致了一个悖论：16 位模型参数确实节省了内存，但为了训练稳定性，优化器状态仍然需要 32 位精度，这构成了主要的存储开销。
六、内存优化策略
既然内存如此宝贵，讲师介绍了两种重要的优化方法：
梯度累积（Gradient Accumulation）
当我们想使用大批量（如 64）但内存无法容纳时，可以将其拆分为多个微批次（如 4 个 16 的微批次）：
- 在每个微批次上计算梯度
- 不清除梯度，而是累加
- 每积累 4 个微批次后，执行一次参数更新
通过这种方式，我们可以享受大批量的训练稳定性，同时将激活内存降低了 4 倍。代码改动仅几行，却带来了显著的内存节省。
激活检查点（Activation Checkpointing）
默认情况下，训练需要存储所有层的激活值（O(L) 内存）。激活检查点（也称梯度检查点或重计算）通过用计算换内存：
核心思想：
- 在前向传播中，只存储某些"检查点层"的激活值（比如每 sqrt(L) 层存一个）
- 在反向传播中，当需要某个层的梯度时，从最近的检查点重新计算缺失的激活值
权衡分析：
- 不检查点：内存 O(L)，计算 O(L)
- 完全检查点（不存任何激活）：内存 O(1)，但计算 O(L²) —— 因为每一层都要从头重新计算
- 每 sqrt(L) 层设检查点：内存 O(sqrt(L))，重计算开销 O(sqrt(L))——这是在内存和计算之间的一个理想平衡
结论与展望
本讲的核心收获可以概括为以下几点：
- 一切皆为张量：模型、数据、优化器状态，本质上都是张量，内存计算是资源核算的基础。
- einops 让张量操作清晰明了：通过命名维度，大幅减少出错概率。
- 6ND 公式：训练一次迭代的 FLOPs ≈ 6 × 数据点数 × 参数量，让你能快速估算训练成本。
- 算术强度与 Roofline 分析：理解一个计算是内存受限还是计算受限，是优化性能的前提。
- 内存优化：梯度累积和激活检查点是降低训练内存需求的两大实用工具。
这些基础为下一讲 Tatsu 老师将要讲解的 Transformer 架构铺平了道路。当你理解了资源核算的方法论，你才能真正回答最初的两个问题：需要多少时间？需要多少内存？
作业预告：你将亲自参与实现 AdamW 优化器，为 Transformer 做详尽的 FLOPs 和内存核算——那时你会发现，本讲的知识将直接派上用场。
课程资源：
- 官方视频：CS336 2026 Lecture 2
- 课程主页：Stanford CS336
- 讲义代码：包含本讲所有示例和实验的完整 notebook
CS336 2026 Lecture 2：完整中文逐字稿
- 视频：官方视频
- 来源：英文人工字幕提取后翻译为中文；保留段落起始时间戳，只修正明显的断句、重复和语序。
- 说明：这是逐字稿，不是讲义摘要；相邻且语义连续的片段已合并为阅读段落；术语保留必要英文。
正文
[00:05.040] 希望大家都没被雨淋到。我反正被淋了。如我上次提到的，Marin 项目那次 1e23 FLOPs 的训练一直在跑，现在它跑完了。而且结果确实符合我们的预测。还记得我们在跑这个实验——这些曲线中的每一条本质上都是 IsoFLOPs 曲线，每条曲线都对应一批较小模型的运行结果。你要做的是找到计算最优的点，拟合出一条缩放定律。这就是我们预测损失的那个点，然后我们实际运行了模型，得到的损失值在 0.05 以内。我觉得这挺酷的。如果你把这个外推到 GPT-5 级别的性能，你就会得到这样的损失值。
[00:50.220] 当然，具体结果可能会因为缩放定律的不同而有所差异。好了，我就是想分享一下这个消息。上一讲我给了大家整个课程的一个概览。我们当时提到了分词（tokenization），那会是在第一个作业里要做的事情。今天我想讲讲资源核算（resource accounting），这更多偏向系统层面的内容。回顾一下，我们最核心的目标是在给定有限资源的情况下训练出最好的模型，这些资源可能是算力、内存，有时还有数据。但数据在本课程中不会成为真正的瓶颈。我们的目标就是最大化训练的计算效率。但是，在你能优化计算效率之前，你需要先理解一个给定计算本身的效率。为此，我们需要理解它的计算特性和内存特性。为了让大家感受一下，课程结束时你们将能够回答什么样的问题，这里有一个问题：
[02:05.750] 在 1024 块——实际上应该是 H100 上——训练一个 700 亿参数、1.5 万亿 token 的模型需要多长时间？你会怎么回答？其实有一个公式，我们会讲到如何得出 FLOPs 数量等于 6 倍参数量乘以 token 数量。我们可以查阅规格表，看看 H100 有多快。我们有一个叫 MFU 的指标，后面会讲，取 0.5。然后你可以估算硬件每天能提供多少 FLOPs。接着就能算出需要多少天。答案是 143 天。这里还有一个问题：在 H100 上用 AdamW 优化器能训练的最大模型是多大？你可以看：H100 每张卡有 80GB 显存，每个参数需要 2 加 2 加 4 加 4 字节。我们会解释这个数字是怎么来的。
[03:19.510] 然后你能得到的参数量大约是 530 亿。这里有一些注意事项，比如我们没有计算激活值，这取决于批大小和序列长度。所以这些都是非常粗略的快速估算。但希望大家在课程结束时能理解这些数字的来源。重点不是精确计算每一个细节，而是对大致的数量级心里有数。好，上次我讲了知识层面，以及你能从这门课中收获什么——也就是机制（mechanics），即事物如何运作。那么今天这部分会比较直白。这里的机制就是 PyTorch 如何工作、张量如何工作。应该相当简单——这里没有什么魔法。我想灌输给你们的心态是：资源核算将至关重要。我希望大家养成一个习惯，每当你写出这行代码时，都要想想它的性能特征。
[04:30.050] 然后是直觉（intuitions）部分。在这里我们只是要感受一下资源是什么样的，它们是如何被消耗的。今天不会有什么机器学习魔法。那些留给 Tatsu 下次课讲。好，那我们开始吧。让我们从底层出发，一步步构建。最底层是什么？最底层是张量。张量是存储一切的构建块。无论你是参数、梯度、优化器状态、数据还是激活值，本质上全都是张量。举个例子，你可以看看 DeepSeek 3.2 模型。你会发现模型本身就是一坨不同形状的张量。每个张量都有一些形状，以及某种精度，这个我稍后会讲。如你所知，张量包含了向量和矩阵的概念，并且可以推广到任意维度。好，我们来谈谈张量存储要占用多少空间。
[05:42.420] 这取决于张量的类型。一般来说，我们处理的是存储浮点数的张量，但张量也可以存储整数和其他类型。对于浮点数，通常你提到 float 的时候，我认为大家默认指的是 float32。一个 float32，你把它拆开来看，有 32 位。其中一位是符号位，8 位是指数位，这决定了动态范围。剩下的是尾数（或叫小数部分），它决定了精度变化。这也被称为 fp32 或单精度（single precision）。“单精度”这个术语来源于过去做科学计算的时候。float32 就像是基准。如果有人给你一个 float，你会预期它至少是单精度的。如果你需要更高精度，那就有双精度，也就是 float64。但在深度学习里，我们反其道而行之，因为即便是 32 位也太多了。
[07:02.250] 而且我们想要做的计算类型并不需要某些数值模拟那样高的精度。好，在讲其他类型之前，我们先看看 float32。让我们构造一个 4 乘 8 的矩阵。默认情况下，你创建的张量类型就是 float32。如果你想要别的类型，你需要显式声明。内存占用就是元素个数乘以元素大小，对 32 位数来说就是 4 字节。这样算下来是 128 字节。给大家一个直观感受，在 GPT-3 中——那已经是个相当老的模型了——前馈层中的一个矩阵大约是 2.3 GB。所以这些张量可以变得非常大。而且这还不是能想象到的最大规模。好，既然我们关心效率，我们通常会想要减少存储量。我们会看到，当你降低精度时，实际上就在节省内存。
[08:20.870] 同时你也在节省时间，因为处理 16 位数据会更快，比如说快一倍，但也不总是这样，要看情况。另外，通过减少内存——我们后面会看到——减少内存其实也能节省时间，这一点可能不那么直观。但希望到时候会变得很清晰。最明显的做法是，你想，好，那我们拿走一半的位数。于是你得到了 float16。float16 的构成是：一个符号位，只有 5 个指数位，其余都是尾数。float16 其实挺好的，就是动态范围比较差。比如说，你想构造一个 1e 减 8 的张量，那实际上直接就变成 0 了。你不能表示非常大的数，也不能表示非常小的数。原因就在于指数位——它只有 5 位，而 float32 有 8 位。
[09:35.690] 所以如果你用 fp16 训练——以前人们确实这么干过——你会遇到不稳定问题。会出现下溢、上溢，会出现 NaN。非常麻烦。于是 bfloat16 被发明了出来。这实际上是在 2018 年开发的，就是为了解决这个问题。当时的观察是：好吧，我们不想在总位数上妥协——总位数跟 fp16 一样。但我们要把一些位从尾数挪到指数上。这意味着它的动态范围比 float16 更大。而且它的动态范围实际上和 float32 是一样的。但当然，分辨率会更差，因为天下没有免费的午餐。不过事实证明，在很多深度学习应用里，这个权衡是非常值得的。你需要动态范围足够大，避免溢出和下溢。
[10:44.680] 而且因为训练过程本身比较粗糙、带有随机性，你并不需要那么高的分辨率。所以这部分我就跳过去不细讲了。那么好的，总结一下，这对训练意味着什么？你完全可以用 float32 来训练。如果你训练的是小模型，可能你根本不想操心这些，直接用 float32 就完全没问题。但它每个浮点数需要 4 字节——抱歉，4 字节的内存。这会占用很多内存。如果你用 float16 训练，那风险又太大了。所以 bf16 是个最佳平衡点。不过即使是 bf16 也可能有风险，可能稍微需要注意一下。现在大家普遍采用的一种做法是混合精度训练。让我——好，让我重新编译一下这个。我觉得这些运行结果有点过时了。混合精度训练是指某些计算使用一种精度，另外一些计算使用另一种精度。
[12:18.080] 一般来说，作为一个通用规则，bf16 用于参数、激活值和梯度。而优化器状态用 fp32。我们后面会再详细讨论这一点。要使用混合精度训练，PyTorch 有一个 AMP 库。我们不会讲太多细节，但基本上你只需要包装一下你的代码，PyTorch 的库会自动帮你处理，它会在安全的情况下尽量把计算转换成 bf16。比如说，矩阵乘法一般是安全的。但如果你做指数运算这类操作，它就会尽量保持 fp32。好的。所以我觉得 bf16 大概就是这门课最终会用到的地方。但如果你很想冒险尝试，你还可以走得更远。比如 fp8，这是四年前引入的，现在已经标准化了。
[13:22.610] 如果你看 fp8，实际上有两个版本，因为取决于你需要更大的动态范围还是更高的分辨率，有两种版本。这个我们就不讲了。但 NVIDIA 的 Transformer Engine 支持 fp8。最近，你甚至可以用到 fp4。去年，NVIDIA 开发了 nvfp4，每个值只有 4 比特。为了让大家有个概念，4 比特真不多。我可以把所有可能的值写在屏幕上这一行里，从负 6 到 6 之间。所以精度非常有限。不过这中间有个小技巧。因为如果你只是单纯用这些值，你根本没法好好训练。所以这里的实际含义是，每个值你有这 4 比特的自由度，但是是有分块的。
[14:23.670] 每个块可以相应地整体放大或缩小。所以实际上你可以表示更多数值，但你不能让每个单独的值都覆盖完整的动态范围。今年实际上发布了一个模型，Nemotron 3 Super，就是用 fp4 训练的，我觉得这挺酷的。好的。所以这些东西有些了解一下就好，有些你甚至碰不到。并不是说你创建一个张量，然后叫它 fp——你确实可以把它设为 fp4。但很多这种操作都是在底层由 NVIDIA 的软件栈自动完成的。好的。这个题外话讲得有点长了，但了解一下精度的复杂性可能还是有帮助的。在我继续之前，有什么问题吗？好的。抱歉，有同学问，当你有一个块的时候，你是把张量里所有的那些值都按这个块来统一缩放。那么所有这些值的比例关系仍然会保持在 fp4 的精度内。对。所以这个问题就是说——[听不清]最大值是为了获得更高的精度。
[15:36.220] 好的，这个问题是，我来把分块再解释一下。你有一个，比如一个块。在这个块内部，你可以在 4 比特范围内变化。另外，所有这些值还可以根据你的缩放因子有多少比特来整体放大或缩小。所以如果你看单个值，你实际上获得的不只是 4 比特的动态范围。但问题是，你不能让这个值在很高的位置，而相邻的值在很低的位。还有一个问题问的是，1 比特呢？[听不清] 1 比特会不会有影响？[听不清]是的，显然有很多训练相关的东西。所以也许这里是个好时机。这个问题是关于 1 比特的，因为你确实没法再比这更低了。所以训练和推理之间是有区别的。很多低比特的事情是这样，如果我们要讲——我觉得我们后面会讲量化——就是你用一个模型训练，可能是 bf16 精度。
[16:41.820] 然后你把它量化到，比如说 1 比特或 2 比特。这比直接训练一个 1 比特语言模型要容易得多，我觉得没有人真正做过后者。也许有可能，但我不认为有人训练出过什么可信的成果。好的，那我们继续吧。我们讲过了张量。内存的计算很简单，就是元素个数乘以每个元素占用的内存。默认情况下，你在 PyTorch 里创建的张量是在 CPU 上的。当然，如果你想让速度快起来，你得把它们移到 GPU 上。实际上，我有个小问题，这些幻灯片是在我笔记本电脑上运行的，这意味着我没有 GPU。所以有些代码我只会展示，但不会实际运行。好的，我们看看。这个可能不太有意思，但我觉得大家都知道怎么把张量移到 GPU 上。
[17:52.382] 但记得一定要这么做，否则你得不到加速效果。好的，我们讲完了张量的内存，这非常直接。现在我们来谈谈用张量做计算。在讲张量操作的 FLOPs 计算之前，我们先稍微岔开一下，聊聊 einops。你们有多少人之前用过 einops？好的，大概三分之二。好的，很好。那么 einops 背后的动机，对于那些可能还没被"洗脑"的人来说，就是这种代码非常容易出错。或者说我觉得看这样的代码很困惑。你有 x 和 y。然后你有 transpose 负 2 和负 1。你试着搞清楚负 2 和负 1 到底是什么意思。所以这可能就是使用变量名而不是索引的动机。einops 是一个用于操作张量的库，其中的维度是命名的。
[19:06.590] 这是受爱因斯坦求和约定的启发。有一个很好的教程你可以去学习。我就只讲一些基础知识。基本上，把 einsum 理解为一种带有良好记账机制的广义矩阵乘法。这里有一个例子。我有一个 3 乘 4 的矩阵，还有一个 4 乘 3 的矩阵。如果做矩阵乘法，这其实看起来很清晰，很容易理解。在 einops 中，基本上你说 x 有两个维度，行和列，我把它们命名为 seq1 和 hidden。然后我有 y。这也是一个矩阵，也有行和列，我把它们命名为 hidden 和 seq2。然后我生成一个张量，或者这里说的是矩阵，维度由 seq1 和 seq2 索引。任何没有在这里出现的维度，也就是 hidden，就会被求和消掉。
[20:14.330] 它的工作方式是，我要对所有出现在这里的变量——也就是 seq1、hidden 和 seq2——的所有可能取值进行枚举求和。我基本上就是索引到 x 里，索引到 y 里，把它们相乘，然后累加到结果 z 的 seq1 和 seq2 位置上。好的。所以你说这是——得了吧。这比那个简单多了。对吧？好，我们试一个更复杂的例子。这是我之前展示过的例子。现在我们有一个张量，它是 2×3×4，另一个张量也是 2×3×4。如果我按照老办法来做，基本上我要做的是转置最后这两个维度。然后因为 einsum 隐式地批量处理那些不是矩阵乘法维度的维度，你就得到答案了。所以你得稍微推理一下。Einops 让这变得非常清晰。它说有一个批量维度，有 seq1 和 hidden，有 seq2 和 hidden。我只需要生成 batch、seq1 和 seq2。
[21:38.640] 注意这里没有转置，因为我只是——从某种意义上说，我通过命名就完成了转置。如果我把维度写成 hidden 和 seq2，那就不需要转置了。我总是被转置搞糊涂。而我不需要去思考转置这件事，这让我很开心。好。如果你想要更灵活一点，你可以说，好吧，batch——我只需用点点点来代替。这意味着如果我有，比如说，一个秩为 10 的张量，有八个不同的批量维度，我就可以直接写点点点，而不需要把它们全部列出来。这在语言建模中会遇到，因为你可能有一个批量维度，可能有一个序列维度，可能有一个注意力头维度。你正试图对所有这些维度执行这个矩阵运算。你可能不想被这些搞得太烦心。
[22:36.902] 好的一点是你可以编写模块化代码，你可以写这个点点点，甚至不需要关心传入张量的形状。好了，那就是 einsum，它在 einops 库中。还有一个 reduce 操作。这是 sum、mean、max 和 min 的推广。比如说，如果你有这个张量，你想根据维度 dim 减 1 来求和，这意味着沿着最后一个维度求和。再说一次，我不喜欢这种表示法。但你可以做的是调用 reduce。基本上你要说的是有一些批量维度。在这种情况下，就是前两个维度。然后有一些隐藏维度，它没有出现在右侧，这意味着它会被求和。这里我写了 sum，意思是聚合归约操作是求和。但你可以把它换成 mean、max 或 min。
[23:47.470] 嗯，这样会有某种加速吗？这基本上会归结为相同类型的基础运算。你可以把它看作只是语法糖。所以它应该是一样的。好。最后我要讲的是 rearrange。我认为这是一个非常强大的工具。我觉得它会在某次作业中出现一次。有时候你会有一个维度实际上代表两个维度，而你想对其中一个进行操作。发生这种情况的原因是有时你有一个矩阵，你把它展平了。然后你可能想再解展平并重新展平。这就是它的工作方式。想象我有一个 3×8 的矩阵，但其中这个维度 8 实际上代表一个 2×4 的矩阵。所以我想把这个 2×4 的矩阵乘以这个 4×4 的矩阵。那么我要做的实际上是调用 rearrange。
[24:59.590] 我在这里做的是说，看，这是若干个批量维度，在这里只对应于第一个元素。然后在这里，我用括号来表示这个 h 实际上代表 heads 和 hidden1 的乘积，我必须——很明显有多种方式可以分解它。可以是 2×4、4×2。所以我说 head 的数量是 2，这意味着 hidden1 是 4。然后我可以把它分成两个维度，heads 和 hidden。好。所以这创建了这样的结果。之前 x 看起来是这样的，现在 x 看起来是这样的。我想这可能有点难以看清。然后你可以对 w 执行你的操作变换。这就是我们之前见过的，你有一个标准的矩阵乘法，其中 x 有一些批量维度。这里是 hidden×hidden 乘以 hidden2。然后一旦你完成了这个变换，你可以把它 rearrange 回去。
[26:17.660] 这很直接。你基本上看两个维度，然后把它们合并成一个维度。是的。你能把一个二维的东西重塑成一维吗？难道不是有两种方式吗？比如你可以按行优先或列优先来做？是的，所以问题是如果你有一个一维的东西然后把它变成——抱歉，你有一个二维的东西然后把它变成一维的，你按什么顺序做？嗯，你做它的顺序是由这里的顺序指定的。好。是的。所以有时候我觉得需要花一点时间来适应，但这是非常值得的，因为一旦你掌握了 einsum，你就会以不同的方式思考。所有那些转置和归约，所有那些，都变得更加流畅。你需要深入思考这些更专门的原语。
[27:19.740] 好了。所以我们打算用 einops。我们稍后会看到一些。现在让我们回到资源核算的问题。我有张量，我们已经讨论了它们如何占用内存。那么它们需要多少计算量呢？我们要用来衡量计算成本的东西是 FLOP 的数量。FLOP 是一个浮点运算。我们假设它是一个基本操作，比如加法或乘法。现在 GPU 还能做其他事情。但在大多数情况下，我们只是忽略它们，因为加法乘法这些是核心操作，会占用你大部分时间。有一点让我很恼火的是，如果我说 FLOPs 这个词，它实际上是有歧义的。有一种是 FLOPs，表示浮点运算的数量，通常写作小写 s 的 FLOPs。这是已完成计算量的度量。
[28:22.490] 还有一种叫 FLOP/s，是每秒浮点运算次数。有时它也会让人非常困惑，写成大写 S 的 FLOPS，但我总是会写 /s 来明确这是衡量硬件的速度。所以如果你看到 H100 有 800、989 万亿次浮点运算，那是后面那种。而当我提到 GPT-3 花了 325 或是多少 FLOPs 时，那是前面那种。只是先把这个说清楚。那么只是给你一个数量级的概念。当我谈论 1e22、1e23 或 1e25 这样的 FLOP 数量时，这些隐式地指的是某些模型的计算量或规模。如果你看 H100——如果你看这份光鲜的规格表——实际上这个数字不在这页上。好吧，算了。有一份规格表，我告诉你对于 bf16，FLOP/s 的数量是 1979。然后你去跑 benchmark，你会说等等，我实际拿到的并不是这个数。
[29:45.790] 然后你去看小字。有一个脚注说这是利用稀疏性的。稀疏矩阵，而稠密矩阵要除以 2。所以你必须总是把这些数字除以 2。这就是为什么你看到这些除以 2 的地方。好。这让我们可以——只是为了直觉——如果你有 8 个 H100，就是一个节点跑两周。那就是 8 乘以两周内的秒数。实际上，这看起来像是一周。好，无所谓，是一周的时间乘以你每秒能得到的 FLOP 数。大约是 5e21。好。这只是为了建立一种直觉，让你了解某些类型硬件有多少 FLOP，以及某些类型模型需要多少 FLOP。这没什么高深的。就是数学，就是粗略估算。好，现在我们来做一些更偏操作层面的东西。假设你有一个线性模型。事实证明，这种计算 FLOPs 的许多演算，核心其实都围绕着类似线性矩阵乘法（MatMul）展开。所以这样讨论其实并不会损失太多一般性。
[31:11.740] 假设你有 n 个点。每个点都是 d 维的。我们要把这些 d 维向量中的每一个都映射到一个 k 维的输出。所以 B 是点的数量，D 是输入维度，K 是输出维度。好，那我们构建一个 x，它是一个 B 行 D 列的数据矩阵。权重矩阵是 D 行 K 列的。当你做这个矩阵乘法时，问题在于这需要多少 FLOPs。事实证明，这基本上就是 2 乘以这三个维度的乘积。这样理解这个算法：对于每一个三元组，我们做一次乘法，然后还要做一次加法。所以有个减 1，因为在这种情况下，你实际上不需要做 D-1 次加法。但咱们先忽略这个。好的，我稍后再回到这一点。如果刚才讲得有点快，我觉得还有另一种推导方法。
[32:31.600] 那其他操作的 FLOPs 是多少呢？逐元素操作就是矩阵的大小。我觉得这很清楚。所以加法也需要 m 乘 n 的 FLOPs。总的来说，对于足够大的矩阵，你遇到的任何其他操作，其计算量都比不上矩阵乘法。所以总的来说，我们主要关注矩阵乘法（MatMul）在做什么，但在讨论内存时有一个重要的注意事项。是的。只是出于兴趣提一下，还有一些其他的算法可以做矩阵乘法。这仅仅是指常规的乘法方法吗？问题是在于，做矩阵乘法还有其他算法。你是指像亚三次方（sub-cubic）算法那样的？是的。所以总的来说，那是——人们为多个矩阵乘法探索的优化算法，更多是关于如何与系统协同设计，而不是这些更偏向渐进复杂度的算法。
[33:32.080] 是的。我们考虑的是同样的乘法运算，认为不可能像让除法比乘法更高效那样，让加法比乘法更高效。是的，所以我认为按照硬件构建的方式，这两者基本上是一样的。但确实，直觉上似乎加法可以比乘法做得更快。但硬件的方式是差不多的。好的。所以你可以把这个矩阵乘法（MatMul）理解为，B 是数据点的数量。而 D 乘以 K 是参数量。记住，x 是 B 行 D 列，w 是 D 行 K 列。所以思考这个公式的另一个方式是，做这个线性矩阵前向传播所需的 FLOPs 数，实际上是 2 乘以 token 数或数据点的数量乘以参数量。所以事实证明，这实际上可以推广到 Transformer，如果你还记得那个 6 乘以 n 乘以 d 的公式，我们可以看到那个雏形正在形成。好的。不幸的是，这些计算实际上意义不大，因为我是在 CPU 上做这些。
[34:54.050] 但我还是会在这里过一遍代码。我们到目前为止做的是测量 FLOPs。这是与硬件无关的。它就像你的模型需要做的计算数量。那么现在的问题是，在硬件上实际需要多长时间？找出答案的一种方法就是直接计时。在这门课里，我觉得在接下来的几讲中，我们会更多地讨论基准测试（benchmarking）。但这里先做个小预览。总的来说，当你在计时的时，特别是在 GPU 上，你必须调用 cuda synchronize 来确保——因为 GPU 是异步运行的，你需要确保有这个同步点。然后你执行操作。操作完成后，你还必须有一个同步屏障。如果你省略了这点，你会发现，哇，你的时间测出来真快。那是因为这是一个非阻塞调用，它直接就返回了。
[36:02.830] 而且通常比较好的做法是多次运行并取平均值。所以实际的每秒 FLOPs 基本上就是你做的 FLOPs 数乘以你在硬件上记录的时间。所以记住，还有另一个数字，就是 GPU 的规格说明书上的数字。让我看看能不能在这个链接里找到。我感觉可能链接搞错了。我下课后再修正一下。——它给出了某个 FLOPs 数字，是 989 teraflops。所以总的来说，实际的每秒 FLOPs 数通常会与宣传的每秒 FLOPs 数不同。好的。哎呀，抱歉。直接思考这个差异的方式，是叫做模型 FLOPs 利用率（Model FLOPs Utilization），简称 MFU。MFU 的定义是实际的每秒 FLOPs 除以宣传的每秒 FLOPs。这里忽略了通信和其他开销。
[37:34.220] 所以你基本上就是用实际值除以宣传值。总的来说，很少能获得超过——我是说，是的，你基本上永远得不到超过宣传值的数字。通常你会得到更少。一般来说，对于现代模型，如果你的 MFU 能达到 0.5 左右，你就应该对自己相当满意了。如果你只是一个直接的矩阵乘法，你也许能得到大概 0.8 甚至更高的值，但通常你达不到那么高。而有时候，如果你的情况真的很糟糕，你会得到类似 0.1 的值，这意味着你应该采取些措施了。好的。所以每当你编写你的模型时，你都可以计算。你现在知道如何通过结合两方面来计算 MFU 了：一是计算你的模型需要做的逻辑 FLOPs 数量，二是看墙上时钟时间，基本上就是做个除法。是的。
[38:42.497] 那边有问题吗？哦。[听不清] 是的。所以这个宣传值——什么是宣传的 FLOPs 数？这是指规格说明书里的数字，它已经是 989 这个数字除以 2 之后的结果了。然后在此基础上，根据你的计算情况，你通常只能达到其中的 0.5。[听不清] 你说你只得到了 50% 的 MFU？是的，那为什么你只得到了 50% 的 MFU？实际上，这是个好问题。当我们讨论内存瓶颈时，我会再回到这个话题。好的。那么总结一下，矩阵乘法通常主导计算。这是设计使然。每秒的 FLOPs 数取决于硬件。所以更好的硬件会带来更多的 FLOPs，它也取决于数据类型。这意味着，如果你看规格说明书，不同的数据类型会有不同的 FLOPs 数值。如今如果你尝试用 float32，速度会非常非常慢，因为硬件并不是为那种工作负载优化的。
[40:07.750] 而现在 bf16 或 fp8 会快得多。还有 MFU——现在你知道什么是 MFU 了。它就是实际 FLOPs 除以宣传的 FLOPs。好了。那么回到为什么 MFU 是 0.5 的问题。为了理解这个，我需要引入算术强度（arithmetic intensity）这个概念。原因在于，这不仅仅是做一堆矩阵乘法，然后看看矩阵乘法花了多少时间就完事了。这是我画的一个非常粗略的硬件示意图。你有高带宽内存（HBM）。然后你有计算核心所在的地方，也就是加速器芯片。然后你怎么计算呢？你得把你的输入、你的矩阵——这些张量位于下面这里——从内存发送到加速器。你进行计算，然后再把结果发送回去。所以如果你想衡量这需要多长时间，这取决于两件事。
[41:27.790] 一是加速器速度，也就是我们刚才讨论的内容。但另一个同样重要的因素是硬件的内存带宽，这一点我们还没谈到。如果你查看规格表，我们之前提到每秒的FLOPs是1979e12除以2。而每秒字节数，也就是内存带宽，这里是3.3太字节每秒。还记得我们之前为什么关注内存，以及存储数据需要多少空间吗？最明显的原因是，如果你的模型太大，放不进内存里，那可就麻烦了。但事实还表明，内存——你需要移动这些内存中的数据，这需要时间。所以实际上，数据量的大小也会影响速度。好的，接下来我要讲解一些操作，基本上会计算这些操作需要多长时间，并引入算术强度这个概念。
[42:38.040] 好的，假设我有一个百万维的bf16向量。我要对它计算ReLU。记住ReLU就是对整个向量逐元素取x和0的最大值。这里我统计两件事，一是移动的字节数。我需要读入x——把它复制到加速器中。每个元素是2字节，因为bf16每个浮点数占2字节，乘以n个浮点数。那就是2n。然后我要写回y，这又是2n。好的，这就是需要移动的字节数。那么做了多少次FLOPs呢？每个元素，我只是和零比较一下，就这样。所以是n次比较。现在看通信时间，即需要移动的字节数除以移动的速度。这得到时间，是1e减6秒。那计算时间呢？就是FLOPs除以每秒的FLOPs。所以是1减9。好的，这里还有一个重要的假设，我们通常尽量保持，那就是通信和计算是重叠进行的。
[44:19.453] 当我们更深入地讨论GPU时会详细讲这点。但核心思想是，在这种情况下，我们不会干坐着等数据移动过来。数据一到我们就开始计算，然后移回去。所以数据移动和计算是同时进行的。因此从数学上讲，我们假设总时间就是两者的最大值，因为我们假设它们可以完美重叠。实际上，不可能做到完美重叠，会有一些开销。但目前这样已经够用了。好的，所以总时间，正如我们这里看到的，是1e减6秒。好的，那如果问这里的瓶颈是什么？当通信时间大于计算时间时，我们称这个算法为内存受限，因为你大部分时间都花在等待数据比特传过来。而当计算时间大于通信时间时，就是计算受限，因为那时候你的瓶颈实际上是在做计算。那么在这个例子中，ReLU是什么？Rectified Linear Unit（线性修正单元）。哦，抱歉。它是内存受限还是计算受限？
[45:39.590] 内存受限。对。这很明显，因为计算量远小于通信量。那么这里还有另一种理解方式。这就是我要定义强度的部分。加速器的强度本质上就是加速器每传输一个字节能做多少工作。对于任何给定的加速器，根据规格表，你基本上就是用每秒FLOPs除以每秒字节数。那么每移动一个字节你能做多少有用的工作？对每个100来说是295。这意味着每个字节——你能做295次浮点运算。所以这是一个很好记的数字，大约300。好的。那么算法的算术强度就是对于这个工作负载，每个字节实际做了多少工作。如果你看ReLU计算的话，是FLOPs除以字节数，实际上——所以这实际上是四分之一，我猜，不是一半。
[47:02.390] 关键在于这个值非常小，是0.25。好的。现在我们可以用强度的语言来讨论瓶颈了。如果算术强度小于加速器强度，那就是内存受限；如果大于加速器强度，就是计算受限。这两个是等价的。如果你看代数部分，基本上就是两个分数。然后你乘一乘除一除——基本上就是把项调换一下。好的。所以在这个例子中，我们是内存受限的。总的来说，我们经常发现自己处于内存受限的情况，因为数据移动很昂贵。所以如果你能获得更高的算术强度，那就很好。0.25这个数字，如果有人告诉你算术密度是0.25，你应该说，哦，这真的很糟糕。对。那一些Transformer的典型算术强度是多少呢？
[48:08.300] 对，我们之后会讲到。好的。好了。那么提高算术强度的一个思路就是，让我们尽量在每个移动的字节上做更多的事情。GELU是另一种激活函数，看起来像这样，公式更加——不是那种带零的。如果做这个计算，来回移动的字节数仍然是2n加2n。这里的FLOPs大约每个逐元素标量操作是20次FLOPs，所以是20n。那么算术强度，我们说是5。这是一个粗略的估计。在这种情况下，我们是内存受限还是计算受限？内存受限。内存受限，因为5仍然小于295，远远小于。所以即使GELU做了很多工作，比ReLU多得多，但从结构上看，它仍然是内存受限的。
[49:30.020] 这意味着，如果你只是在计算ReLU和GELU，你会想，GELU这么复杂，肯定非常昂贵。但实际上它们完全一样，因为瓶颈根本不在那里。好的，现在让我们看一些线性操作，比如点积。点积，你有一个向量x，有一个向量w，大小为n，然后你做点积。那么移动了多少字节？你读入x，是2n；读入w，是2n；然后写入y，这是一个标量，是2。FLOPs的数量是你做了n次乘法和n减1次加法。所以是2n减1。好的。那算术强度是多少？哦，对于这个来说，大约是二分之一，这也相当糟糕。这意味着——希望你们已经明白了——这是内存受限的。那么矩阵向量乘法呢？x是一个向量，w是一个n乘n的矩阵。
[50:41.090] 你做这个乘积。你们有多少人认为这会是计算受限的？有多少人认为是内存受限的？好的，我们来看看。我要读入x，是2n；读入w，是2n平方；写入y，是2n。FLOPs的数量基本上就是做n个点积。所以是n乘以做一次点积的代价。算术强度几乎，几乎没有提高。所以这也是内存受限的。好的。现在我们来谈谈矩阵乘法。这时候事情就有趣了。我有一个n乘m的矩阵，另一个m乘n的矩阵。我把它们相乘。移动的字节数是2n平方加2n平方。然后我还要写回2n平方字节到y。这里的FLOPs是n平方个点积。那算术强度是多少？是300。呼。好的，所以是340。
[51:56.010] 一般来说，这个比值大约是 n 除以 3。从直觉上也说得通，因为你发送的是 n 平方量级的数据，但计算的是 n 三次方量级的操作。所以计算量除以发送量的比值是 n 的量级。矩阵越大，这个比值就越好。所以总的来说，这就是为什么你听到人们说“我们需要更大的批量大小”或“更大的矩阵”，原因就在于此。如果你的操作低于加速器的强度阈值，把东西做小并不会真正加速，效果是一样的。而如果你超过了这个阈值，那你实际上就在让 GPU 饱和运行了。好的，所以最后，这就是计算受限。好的，只要我们有大矩阵，我们的状态其实相当不错。我们处于计算受限状态，让加速器饱和运转。
[52:59.600] 之前那个关于 Transformer 的问题呢？其实你会发现，不仅在你们的作业里，而且在下一讲中也会看到，Transformer 本质上就是大矩阵乘法，中间夹杂一些其他操作。所以从算术强度角度来看，这是个好消息。这是有意设计的——Transformer 的设计方式就是为了具有高算术强度。这里有一点要提一下，算是为推理那一讲做个铺垫：矩阵向量乘积基本上就是你在做 Transformer 推理时发生的事情。因为推理时，你是一个 token 一个 token 地生成。所以你只能——这就像是一个向量和矩阵做点积。正如我们之前看到的，这是内存受限的。而在训练时，你拿到的是整个序列，一次性处理完。
[53:56.150] 好的。另外要注意，算术强度取决于精度。默认情况下，我们这里做的所有事情都是 fp16。好的。为了把它和之前关于 MFU 的问题联系起来——你可能得到较低 MFU 的原因是，虽然你可能在做——MFU 是实际 FLOPs 除以——抱歉，实际值除以理论峰值。所以如果你有非常大的内存瓶颈，那么你实际上不会获得很好的吞吐量，尽管——理论峰值的意思是，如果你没有内存瓶颈，你就只管从头到尾把模型的所有计算都跑完。好的，关于算术强度 Roofline 图的最后一点。有一种很好的方式——是的。[听不清] 一般来说，这些模型往往偏向内存受限，正如他们所说，在大部分计算中都是如此。但即便如此，加速器的利用率 50% 并不算好。也许你能到 70，也许到 80 那就算非常好了，也就是说加速器的计算能力相对于内存带宽来说太大了。
[55:18.990] 为什么会这样？为什么这些加速器——它们体量很大，却在那里闲着或者等内存？所以问题是，你能设计出具有更好特性的加速器吗？是的。也许当我们讨论 GPU 并更深入地了解它们的工作原理时，我们能聊聊为什么会这样。但如果你有答案，你应该告诉黄仁勋。也许你能设计出更好的硬件。好的，让我们来可视化算术强度和性能之间的关系。这个图基本上把算术强度放在 x 轴上。这里的每个切片对应一个特定的，比如说，算法。然后这里有这些线条。每条线对应一个特定的加速器，比如说 H100 或 B200 之类的。这个图展示的是——然后 y 轴是实际达到的每秒 FLOPs。如果你的算法算术强度很低，比如 ReLU 或点积，你会在这一侧，这意味着实际达到的每秒 FLOPs
[56:37.630] 不会像你的加速器峰值那么高。随着你提高算术强度，情况会变得更好。你会能够让你的硬件饱和运行，直到某个点。超过某个点之后，你就处于计算受限状态。显然你不可能超过峰值 FLOPs。好吧，那我们继续。好的，现在我要回到训练中所需操作的内存和计算的话题。到目前为止我们已经做了张量操作，主要是矩阵乘法。我们看到了它们占用多少内存、需要多少计算，以及两者之间的相互作用。现在让我们真正想一想训练需要什么。这是我们这里的运行示例。实际上这不是线性网络，只是一个深层网络。我要考虑的情况是：有一个输入，是 B 乘 D 的输入，然后有若干层，每层是一个 D 乘 D 的矩阵乘法。
[57:59.790] 然后产生一组预激活值，接着逐元素的 ReLU 产生第一层的激活值。然后这个过程反复进行。输出的大小是一样的。所以这个深层网络——我就是想看看它在 PyTorch 里长什么样，基本上就是一组模块。每个模块有一个权重向量——抱歉，一个权重矩阵。所以参数量是 D 的平方乘以层数 L。当你在一个批次的数据上运行模型时，我们做的就是逐层遍历，每一层基本上就是应用线性变换，然后做逐点的 ReLU 激活。然后对所有层都这样处理。这应该是一个很简单的模型，希望大家都清楚。好的，现在我们来谈谈梯度。让我们用一个更简单的例子。这就是一个简单的线性回归模型。我们有一个向量 1、2、3，一个权重向量 1、1、1。
[59:33.420] 然后我们做点积，形成这个 MSE 损失。好的。当我们做反向传播时，计算图中涉及到的每个变量都有一个梯度，要么被设置了，要么没有被设置。w.grad 被设置为 1、2、3。这只是 PyTorch 的基本机制。我觉得大家应该都很熟悉了。那么问题是，计算梯度需要多少计算量？好的，让我们来数一下计算梯度的 FLOPs。让我们用一个简化模型：有一个 x，是 B 乘 D 的，乘以 w1 矩阵。我们暂时忽略这个 ReLU 以简化问题，然后乘以 w2，它是同样的 D 乘 D——同样形状的 D 乘 D 矩阵。好的。我这里要用 einsum，原因稍后会说明。在这个线性网络中，你做的第一件事是拿 x 和 w1。x 是 batch 乘输入维度。
[01:00:55.860] w1 是输入维度乘输出维度，你得到一个 batch 乘输出维度。然后 h2 是拿 h1 和 w2，同样的道理。这些都只是矩阵乘法。然后你形成一个损失。我就是把它送到某个任意的操作，只是为了得到一个数值。好的。那么在反向传播中会发生什么——我会称这些为保留梯度，用于调试目的，你们后面会看到。你对损失做反向传播。问题是，这做了多少工作——多少浮点运算？让我们放大到一层。让我们聚焦在第二层，就是这一层。第二层拿 h1，乘以一个矩阵 w2，得到 h2。就这些。所以如果你看前向传播中的 FLOPs，这只是一个矩阵乘法。我们取三个维度，把它们乘在一起，然后乘以 2，对应 bf16。那是 2 字节。抱歉，那无关紧要。这个 2 是因为有一次加法和一次乘法。
[01:02:08.480] 那么在反向传播中，这是什么样子的？在反向传播中，如果你记得链式法则和反向传播算法，你需要计算两个东西。你需要计算反向消息，即损失对输入的梯度。而且你还必须计算相对于参数（parameters）的梯度。这些梯度长什么样呢？这其实就是根据这些的定义——我想这本质上就是链式法则，只是用 einsum 符号写出来了。你取 d loss d h2，也就是 h2 的梯度。那是一个 batch 乘以 Alt 的矩阵，然后你把它乘以 w2，w2 是 in 乘以 out 的维度。这样就会得到 h1 的梯度，也就是 d loss d h1，维度是 batch 乘以 n。所以我总是搞混这个，如果你学过微积分就会发现，其中一个需要转置，而且我总是忘记它们相乘的顺序。
[01:03:25.700] 我认为 einsum 让这个变得非常清晰，因为即使只看标量情况，你也能记住 h1 的梯度必须等于 h2 的梯度乘以 w2。唯一的问题就是，你如何确定索引？你基本上是通过查看这些张量的形状和命名维度来确定索引的。在这种情况下，它就是一个矩阵乘法，你要在输出维度上进行求和。当你这样做的时候，我们可以验证一下。所以 h1.grad 就是由 loss.backward() 计算出来的东西。实际上，这确实就是我在这里写出来的同一个东西，仅作为验证。所以你要计算的另一件事是 w，即关于参数的梯度。那就是 d loss d h2 乘以 h1。所以这和回传的信息是同一类东西，但是要与你不计算其梯度的另一个东西相乘。
[01:04:27.680] 你只需要写出各个维度。在这种情况下，你是在 batch 维度上求和。但好处是，你只要看一下这个，我们就能知道它的计算成本。FLOPs 的数量基本上就是所有维度的乘积。无论你在哪些维度上进行批处理，这都没关系。基本上就是对 i、j、k 进行枚举，然后进行聚合。聚合的方式不同，但 FLOPs 是相同的。好的。所以请注意，反向传播的代价恰好是前向传播的两倍。这是因为你必须计算两个梯度，一个是对每个参数的梯度，另一个是对输入的梯度。也就是那个不是参数的另一个东西。也许我停一下，关于这部分有什么问题吗？好的。那么这只是针对 w2 的情况。你只需要对网络中的所有参数都应用这个操作。然后把所有东西放在一起，你会看到，对于这个网络，前向传播的 FLOPs 是 2 乘以数据点数量（即 b）乘以参数量。
[01:05:59.430] 反向传播是这个的两倍，也就是 4 乘以数据点数量乘以参数量的 FLOPs。所以总计就是 2 加 4 等于 6，即 6 乘以数据点数量和参数量。这就是你可能在各种地方见过的 6nd 这个数字的来源。这只是通过统计前向和反向传播得到的。我们只是在讨论这些深度网络的时候这么做了。但事实证明，对于 transformer 来说，这也是一个很好的近似——只要上下文长度（context length）不太大。如果上下文长度太大，你就会有上下文长度的平方项，这会增加更多的 FLOPs，而这部分并没有包含在这种计算方式中。
[01:06:47.490] 好的。让我们稍微谈谈——现在我们有了梯度。接下来我们训练时需要做的另一件事是优化。所以这是我们的深度网络。为了不只是重复作业 1 里的内容，我将使用 AdaGrad 优化器，这是 2011 年的成果。这比 Adam 要早。你可以把它看作是介于 SGD 和 Adam 之间的东西。基本上它就是在 SGD 的基础上，考虑了梯度的二阶矩。动量（Momentum）是关注梯度的一阶矩。而 Adam 则是将两者结合起来。我在这里没有时间深入讲解优化器的细节，因为我们主要关注的是用法。所以我们可以定义一个优化器。当你计算出梯度，然后当你执行优化器步骤时，如果你是在定义一个新的优化器——比如在作业 1 中，你将实现 Adam——
[01:08:06.710] 对于每个参数组（在这个例子中是 w1 或 w2），我们会查看一个叫做优化器状态（optimizer state）的东西，这是优化器在运行过程中可以使用的一些存储空间。我们在 AdaGrad 中计算的是——就是梯度的平方，或者说是梯度平方的累积和。所以我们在这里从优化器状态中获取它，用当前的梯度更新它，然后再把它存回去。在你更新了 g2 之后，你就可以更新参数了。所以对于 AdaGrad，基本上就是除以平均梯度平方或者梯度平方和的平方根。好的。让我们看看。实际上我可能会跳过训练循环。其实，等等，等一下。我想我实际上跳过了一些我本不想跳过的东西。好的。那部分是关于优化器状态的，现在让我们看看优化器使用了多少内存。或者更广泛地说，内存使用情况是怎样的？
[01:09:29.840] 对于这个网络中的参数，每一层 L 有 D 平方个参数。如果我们用 fp16 存储，每个参数占用 2 字节。所以这就是参数量。激活值（Activations）——这是 2 乘以——其中 2 是 bf16 的字节数——batch 乘以 D 再乘以层数。每一层你都有激活值。你还有梯度，这基本上就是所有参数的一个副本。这也是 bf16。而优化器状态是对于每个——实际上这里应该是——我想这里应该是——抱歉，我会修正这个。这是个笔误。这里应该是参数量，而不是参数内存。这里也应该是参数量，所以是参数量乘以 2，然后这是参数量的 4 倍。这样做的原因是，为了稳定性，通常习惯使用 fp32 来存储优化器状态。
[01:10:31.980] 显然有人尝试过使用 bf16。最终发生的情况是，你需要对梯度取平方，并且要在多个步骤上求平均，这样操作很不稳定。对于 AdaGrad 来说，这里——我们每个参数需要 4 字节来存储优化器状态。对于 Adam——你需要存储一阶矩和二阶矩，所以是 8 字节。所以如果你想想的话，优化器状态实际上占用了很多内存。有一点需要注意的是，记住内存有两个作用。一个是——你必须在你的 HBM 中存储这些东西。但另一件事是，它还必须要传输到加速器上。一般来说，优化器状态并不是计算方面的瓶颈。所以这里的内存量对于速度性能来说其实不是那么重要，但这意味着你无法在内存中容纳大型模型。
[01:11:42.250] 好的。总结一下，正如我们提到的，这里的参数量是 D 乘以 D 乘以 L，而 FLOPs 数量是 D 乘以 token 数量或数据点数量再乘以参数量。对于 transformer 来说，这会稍微复杂一些。但你们会在作业 1 中处理这个问题。而且你们会更仔细地去做。好的。我将跳过训练循环。这只是一个总的回顾。在我们结束之前，我想快速谈两件事。一件是，正如我们所见，内存对于存储大型模型的能力有重要影响，但有时候也会影响你的速度。所以总的来说，你希望减少内存使用。通常人们会做两件事。一个是梯度累积（gradient accumulation）。一般来说，你会希望使用足够大的 batch size 来提高稳定性，直到达到一个临界 batch size，Tatsu 稍后会详细讨论这个。
[01:12:45.730] 正如我们所见，激活内存会随批次大小扩展。因此，如果你使用的批次大小过大，有时可能会遇到内存不足的情况。而梯度累积的做法是，你进行计算，你有这些微批次，然后你在微批次上计算梯度。然后你只需累积梯度，不要将梯度清零。每经过“批次大小除以微批次大小”这么多个步骤，你才更新一次参数并清零梯度。好的，所以这实际上是一个非常简单的代码改动，却能让你节省计算资源。另外一件我想简要提及的事情是激活检查点。在训练中，我们通常需要存储所有层的激活值。这是默认做法。实际上，有趣的是，对于推理来说，我们不需要计算梯度，因此我们只需要存储当前层的激活值。
[01:13:40.760] 但对于训练来说，内存使用量是 B 乘以 D 乘以 2 乘以 L。问题在于，你能否减少这个量呢？激活检查点，也称为梯度检查点或重计算，其核心思想是：在前向传播过程中，你只为部分层保留激活值。而在反向传播过程中，你从最后一个检查点重新计算缺失的激活值。这就是为什么它被称为检查点。这是一个你在系统中常见的通用技巧，本质上是一种权衡。如果你想减少内存，你可以选择重新计算。所以在这里，我们要定义一下具体做法，以便实际操作。这其实相当简单。实际上，我没有——好吧，也许我没有看到——好的，基本上你要做的就是使用相同的模型，只不过你在某一层上加上 Torch 的 utils.checkpoint。
[01:14:47.960] 这意味着进行这个计算，但不要存储任何中间激活值，只存储必要的内容。这意味着如果你存储了所有的激活值，记住在你的深度网络中，你有 ReLU 之前的东西和 ReLU 之后的东西，对吧。所以那是很大的存储量。相反，如果你对那些每个块都包含一个线性和一个 ReLU 的块进行激活检查点操作，你就不需要存储 ReLU 之前的值，基本上就能省下一半的激活内存。你可以用一半的内存实现同样的效果。然后当你在进行反向传播时，你需要 g3，但你可以很容易地从 h2 计算出 g3。好的。你还可以更进一步。你可以说，好吧，我可以存储所有层，但极端情况下，我可以任何层都不存储。这样内存效率就会达到最大化。唯一的缺点是计算量会变成 L 的平方。因为对于每一层，你都必须从头开始计算。
[01:15:57.080] 也许一个比较理想的平衡点是，如果你将检查点存储在 sqrt(L) 层处，那意味着你的激活内存是 sqrt(L)，而你的重计算开销也是 sqrt(L)。这样就能达到平衡。好的，总结一下今天的讲座，所有操作都作用于张量：参数、梯度、激活值、优化器状态和数据。我们介绍了 einops，希望你们能够接受它作为一种思考张量操作的方式。6 乘以数据点数量再乘以参数数量这个公式，我们现在已经解开了它的神秘面纱，它就是一个训练步骤的 FLOPs 数量。实际上，如果这是一个训练步骤，那么这里的批量大小应该是 2。
[01:16:46.790] 然后我们讨论了算术强度和屋顶线分析，这使我们能够判断一个计算是内存密集型还是计算密集型。矩阵乘法是计算密集型的。基本上其他所有操作都是内存密集型的。最后，梯度累积和激活检查点是减少内存使用的两种方法。通过减少内存使用，你可以使用更大的批次大小。好的，今天的讲座就到这里。下周，Tatsu 将会讲解架构方面的内容。
CS336 2026 · Lecture 2 · PyTorch (einops), resource accounting (FLOPs, memory, arithmetic intensity)
- 视频：官方视频
- 字幕：en-US · manual captions
- 处理：连续字幕合并为自然段；每段保留首条字幕的起始时间戳
- 状态：原语言字幕稿，待校对
正文
[00:05.040] I hope everyone is staying dry. I'm not. So as I mentioned last time, the Marin project had a 1e23 FLOPs which was running, and it finished. And it actually matched the forecast. So remember we were running this. Each of these curves is essentially IsoFLOPs curve, which is a bunch of smaller model runs. And you try to find the compute optimal point. You fit a scaling law. And this was the point where we predicted the loss, and we ran the model and it got loss within 0.05. So I thought that was pretty cool. And if you extrapolate out to GPT-5 level performance, this is the loss you get.
[00:50.220] Of course, your mileage might vary depending on how these scaling loss are. OK, I just wanted to share that news. So last lecture I gave an overview of the entire class. And we are talking about tokenization, which is going to be on the first assignment. Today, I want to talk about resource accounting, which is going to be more on the systems side of things. So to recall, the main thing we're trying to do is train the best model we can given a finite set of resources, which could be compute, memory, sometimes data. But that's not really going to be a limiting factor for us in this class. And our goal is to simply to maximize the computational efficiency of our training. So before you can optimize the computational efficiency, we need to understand the efficiency of a given computation. And for that, we need to understand the compute and memory characteristics. Just to give you a taste of the type of questions you will hopefully be able to answer by the end of the class, so here's a question.
[02:05.750] How long would it take to train a 70 billion parameter model on 15 trillion tokens on 1024-- actually, this should be H100. So how do you answer that? Well, there's a formula, which we'll talk about how you can get the number of FLOPs to be 6 times the number of parameters times the number of tokens. We can look up the spec sheet to see how fast the H100 is. We have this thing which called MFU, which we'll talk about 0.5. Then you can estimate the number of FLOPs that you need that the hardware gives you per day. And then you can compute the number of days. So the number of days is 143. Here's another question. What's the largest model you can train on H100s using AdamW? So you can look at, well, H100s have 80 gigabytes of each VM memory, the number of bytes per parameter, which is 2 plus 2 plus 4 plus 4. We'll explain where that comes from.
[03:19.510] And then the number of parameters is that you can get is going to be about 53 billion. So there's some caveats here that we don't count the activations which depends on batch size and the sequence length. So this is all very rough back-of-the-envelope calculations. But hopefully by the end of this class you'll understand where these come from. And the point is not to precisely calculate every single thing, but just get the rough shape of things. OK. So last time I talked about knowledge and what you can take away from this class-- mechanics, which are how things work. So today, that would be pretty straightforward. The mechanics are just how PyTorch work, how tensors work. It should be fairly-- there's no magic here. The mindset I want to impart on you is that resource accounting is going to be very crucial. And I want everyone to get in the habit whenever you write this line of code, think about the performance characteristics.
[04:30.050] And then finally intuitions. Here we're just going to get a sense of the resources, how they're spent. There's going to be no ML magic today. I'll leave that to Tatsu for the next lecture. OK. So let's get into things. So let's start bottom up and start building up. So what is at the bottom? At the bottom are tensor. So tensors are the building block of storing everything. If you have parameters gradients, optimizer states, data activation, everything essentially is a tensor. So for example, you can take a look at the DeepSeek 3.2 model. And you see that the model itself is a bunch of different tensors. Each tensor has some shapes and also some precision, which I'll talk about later. And so as you know, tensor is subsume vectors, matrices, and can generalize to any number of entities. OK. So let's talk about how much tensors take to store.
[05:42.420] So it depends on the type of tensor. So in general, we're going to be dealing with tensors that store floating point, but tensors can also store integers and other types. So for floating point, typically whenever you talk about float, I think the standard people refer to as float what is called float32. So a float32, if you break it down, has 32 bits. One of the bits is a sign, 8 bits are the exponent, which gives you dynamic range. And the rest is the mantissa or the fraction, which gives you variation. This is also known as fp32 or single precision. And the term single precision comes from the fact that back in the day when you were doing scientific computing, float32 was like a baseline. It was just like you would expect if someone gave you a float, you would expect it to be at least single precision. And if you want more precision, then you can get double precision as that's float64. But in deep learning, we're going the other way, because even 32 is a lot.
[07:02.250] And the types of computations that we want to do don't demand the high precision that some kind of numeric simulations do. OK. So before we get to other types, let's just look at float32. So let's construct a 4 by 8 matrix. By default, the type of tensor you create is float32. So if you want something else, you should declare it. And the memory usage is just the number of elements times the element size here, which is 4 bytes for a 32-bit number. And that's going to be 128 bytes. So just to give you a kind of perspective, so in GPT-3, which is a fairly old model one of the matrix in the feedforward layer is about 2.3 gigabytes. So these tensors can get quite big. And this is not even the biggest one that one can imagine. OK. So since we're interested in efficiency, we want to generally reduce the amount of storage. And we'll see that as you reduce the precision, you actually save memory.
[08:20.870] And you also save time because operating on 16 bits is going to be faster, let's say, twice as fast, but not always. It depends. And then memory, by reducing memory, we'll see later that actually reducing memory can save time as well, which is maybe less obvious. But it will hopefully become clear. So the obvious thing is you say, OK, let's take away half the bits. Now you have float16. So float16 says you have a sign. You have only 5 bits of exponent, and then the rest is the mantissa. So float16 is good, except for its dynamic range is poor. So even if you have, let's say, try to construct 1e minus 8 tensor, then that is actually just 0. So you can't really represent very big numbers, and you can't represent very small numbers. And the reason is that this exponent it's only 5 bits of exponent compared to 8.
[09:35.690] So if you train with fp16, which people did back in the day, you will get instability. You will get underflow. You get overflow. You'll get NaNs. It's pretty challenging. So bfloat16 was invented. So this was developed in actually 2018 to address this issue. And the observation was that well, let's not compromise on the number of bits. The number of bits is going to be the same as fp16. But we're going to shift some of the bits from the mantissa to the exponent. So that means it has more dynamic range than float16. And it actually has the same dynamic range as float32. But of course, the resolution is worse because there's no free lunch here. But it turns out that in a lot of deep learning applications, this is well worth the trade-off. You want the dynamic range to not overflow and underflow.
[10:44.680] And because things are kind of sloppy and stochastic anyway, you don't need that much resolution. So let me actually skip over this part. So OK. So to summarize, what are the implications for training? So you can absolutely train with float32. And if you're training a small model, you probably just-- and you don't want to worry about it, just float32 is fine. But it requires 4 bytes of-- sorry, 4 bytes of memory per float. And that can take up a lot of memory. And if you train with float16, then that's going to be too risky. So bf16 is the sweet spot. Even bf16 can be risky as well. Maybe [INAUDIBLE] a little bit. One thing that people have-- has become common practice is to use mixed precision training. And the actually let me-- OK, let me compile this again. I think these are stale. So mixed precision training is where some of the computations use some precision, and for other computations use other precision.
[12:18.080] So in general, as a general rule, bf16 is what you would use for parameters, activations, and gradients. And for optimizer states, you would use fp32. And we'll discuss a little bit more about that later. And to do invoke mixed precision training, so PyTorch has an AMP library. We're not going to talk too much about this, but you basically wrap your code, and PyTorch has-- the library takes care of it automatically for you in the sense that it tries to cast things into a bf16 when it's safe. So, for example, MatMuls are generally safe. But if you try to do exponentiation, then it will try to leave things as fp32. OK. So we could do bf16 I think is probably where this class will end. But if you're feeling very adventurous, you can go farther. So fp8, this was introduced four years ago. It actually has been standardized.
[13:22.610] So if you look at fp8, there's actually two versions because depending on if you need more dynamic range or more resolution. And there's two versions of them. And so we're not going to talk about that. But NVIDIA has some-- the transformer engine supports fp8. And most recently, you can actually go down to fp4. So last year, NVIDIA developed nvfp4, and there's only 4 bits per value. So if you just to-- so we're on the same page. 4 bits is not a lot. I can write all the values down on a single line here between minus 6 and 6. So that's not very much precision. Now there's a little bit of a cheat here. Because if you just naively only use these values, you're not going to be able to train very well. So what actually this means is that every value you can have this 4 bits of freedom, but there are blocks.
[14:23.670] And each block can be scaled up and down accordingly. So you can actually represent more values but you can't represent the full dynamic range for every single value. And there is a model released actually this year, the Nemotron 3 Super, which was trained in fp4, which I think is pretty cool. OK. So some of this it's just good to know. And some of this you can't even touch. It's not like you create a tensor and you call it fp-- you can make it fp4. A lot of this is done under the hood by NVIDIA's software stack. OK. This was a bit of a long digression, but it's maybe helpful to appreciate that the intricacies of precision here. Any questions before I move on? Yeah. Sorry. It was a question like, when you have a block, you're scaling all those same things in tensor by that block. So then the ratios of all that will still be in fp4. Yeah. So the question is to just-- [INAUDIBLE] maximum value is to have more specificity.
[15:36.220] Yeah, so the question is, OK, just to explain the block a bit more. So you have, let's say a block. Within that block, you can vary up within the 4 bits. And in addition, all those values can be scaled up and down according to how many bits you're scaling factor is. So if you look at an individual value, you actually get more than 4 bits of dynamic range. But it's just like you can't have this value be way over here and the neighboring value way down here. One question about, what about 1 bits [INAUDIBLE] affect 1 bit [INAUDIBLE]? Yeah, so there's a lot of training obviously. So maybe this is a good point. So the question is about 1 bit because you can't really go lower than that. So there's a difference between training and inference. So a lot of the low bit stuff, if we talk about-- I think we'll talk about quantization later, is that you train a model on and maybe even bf16.
[16:41.820] And then you quantize into, let's say, 1 or 2 bits. And that is much easier than training a 1-bit language model, which I don't think anyone has done. Maybe it's possible, but I don't think anyone has trained anything credible there. OK. Let us go on then. So we talked about tensors. And the memory calculus is pretty simple, just the number of elements times however much memory each element takes up. So by default, the tensors you create in PyTorch are going to be on CPU. And of course, you have to-- if you want things to go fast, you want to move them to GPU. And actually, so there's a bit of a-- I have a slight issue with the slides that they were executed on my laptop, which means that I don't have GPU. So some of the code I'll just show but not execute. OK. Let's see. This is maybe not that interesting, but I think everyone knows how to move tensors to GPU.
[17:52.382] But just remember to do that. Otherwise you won't get your speed-ups. OK, so we talked about memory of tensors, which is very straightforward. Now let's talk about computing with tensors. So before talking about FLOPs accounting for the tensor operations, let's take a little bit of a digression to talk about einops. So how many of you are familiar used einops before? OK. So maybe 2/3. OK. Good. So the motivation behind einops, for those of you who might not be indoctrinated is, it's very easy to mess up. Or I find it very confusing to look at code such as this. And you have x and y. And then you have transpose minus 2 and minus 1. And you're trying to figure out what minus 2 and minus 1 is. And so this is maybe the motivation for using variable and names rather than indices. So einops is a library for manipulating tensors where the dimensions are named.
[19:06.590] And this is inspired by Einstein summation notation. There's a nice tutorial which you can go through. I'm just going to cover some of the basics. So basically, the way to think about einsum is a generalized matrix multiplication with good bookkeeping. So here's an example. So I have a matrix 3 by 4. I have a 4 by 3 matrix. And if I do the MatMul, this is actually pretty nice. It's pretty easy to understand. In einops, basically you say x has two dimensions, the row and the column, which I'm going to name seq1 and hidden. I'm going to have y. That's a matrix which has also rows and columns, which I'm going to name hidden and seq2. And I'm going to produce tensor or here a matrix, where the dimensions are indexed by seq1 and seq2. And anything that is not mentioned here, the hidden gets summed out.
[20:14.330] So the way this works is I'm going to sum over-- enumerate over all possible values of all the variables that occur here, which are seq1's hidden and seq2. And I'm going to basically index into x, index into y, multiply them and dump them, accumulate them into the result z sub seq1 and seq2. All right. So you say this is-- come on. This is much easier than this, Right OK, let's try a more complicated example. So here's the example I showed before. Now we have a tensor. It's 2 by 3 by 4, another tensor 2 by 3 by 4. And if I were doing things the old way, basically what I'm doing is transposing these last two. And then because at implicitly batches the dimensions that are not the MatMul dimensions, then you get the answer. So you have to reason about this a bit. Einops makes this very clear. It says there's a batch dimension. There's seq1 hidden, seq2 hidden. And I'm just going to produce batch seq1 and seq2.
[21:38.640] And notice that there's no transpose because I just-- in some sense I've done the transpose by the naming. If I had hidden and seq2, then it would be no transpose. I always get confused by transposes. And the fact I don't have to think about transposing it makes me happy. OK. So if you want to get fancy, you can say, well, batch-- I'm just going to replace with dot, dot, dot. And this means that if I had, let's say, rank 10 tensor with eight different batching dimensions, I can just write dot dot, dot without enumerating all of them. And this comes up in language modeling because you might have a batch dimension. You might have a sequence dimension. You might have a head dimension. And you're trying to do this matrix operation for all of them. And you might not want to have to just worry about this.
[22:36.902] And the nice thing is that you can write modular code where you can write this dot, dot, dot without worrying about even the shape of the tensor that comes in. All right. So that's eisum which is in the einops library. There's a reduce. So this is a generalization of sum, mean, max, and min. So for example, if you have, let's say, this tensor and you want to sum according to dimension of dim minus 1, which means sum along the last dimension here. So again, I don't like this notation. But what you can do is call reduce. And basically what you say is that there's some batching dimensions. In this case, it would be these first two. And then there's some hidden dimension which doesn't appear on the right side, which means that it gets summed. And here, I put sum, which means that the aggregation reduction operation is sum. But you can replace it with mean or max or min.
[23:47.470] Yeah. Is there some speed-up to this? This basically reduces to the same type of primitive operations. You can think about it as just sugar. So it should be the same. OK. So the final thing I'll talk about is rearrange. So this is, I think, a pretty powerful tool. I think it will come up in an assignment once. So sometimes you have a dimension that actually represents two dimensions, and you want to operate on one of them. And the reason this happens is that sometimes you have a matrix and you flatten it. And then you want to maybe unflatten and flatten. So this is the way that it works here. So imagine I have a matrix 3 by 8 but where this dimension, 8 actually represents a 2 by 4 matrix. So I want to multiply that 2 by 4 matrix by this 4 by 4 matrix. So what I'm going to do, actually, is to call rearrange.
[24:59.590] And what I am doing here is saying, look, this is some number of batch dimensions, which here just corresponds to the first element here. And then here, I use parentheses to say that this h actually represents the product of heads and hidden1, where I have to-- obviously there's multiple ways to decompose this. It could be 2 by 4, 4 by 2. So I said the number of heads is 2, which means that hidden1 is 4. And then I can break that up into two dimensions, heads and hidden. OK. So this creates this. So before x looked like this, and now x looks like this. So I guess this might be a little bit hard to see. So then you can perform your operation transformation on w. And this is what we've seen before where you have just a standard MatMul where there's some number of batching dimensions for x here. And this is hidden times hidden by hidden2. And then once you've done that transformation, you can rearrange it back.
[26:17.660] And this is straightforward. You basically look at two dimensions, and then you can group them into one dimension. Yeah. Could you take a two-dimensional thing and shape it into one dimension, aren't there two ways? Like, you can do it row major or column major? Yeah, so the question is if you have a one-dimensional thing and you shift it into-- sorry, you have a two-dimensional thing and you shift it into one-dimensional thing, which way do you do it? Well, the order you do it is specified in the order here. OK. Yeah. So sometimes I find it takes a bit of time to get it used to, but it's well worth it because once you have an einsum, you just think in a different way. And all the transposes and reductions, all that, it's just it becomes more fluid. You have to think through these more bespoke primitives.
[27:19.740] All right. So OK. We're going to use einops. We'll see a little bit later. So now let's go return back to the resource accounting question. So I have tensors. We've talked about how they take memory. So how much compute do they take? So the thing we're going to use to measure computation cost is the number of FLOPs. A FLOP is a floating-point operation. And we're going to assume it's a basic operation like addition or multiplication. So now there's other things that GPUs can do. But for the most part, we're just going to ignore them because these are the bread and butter, and are going to eat up most of your time. So one thing that is a pet peeve is that if I say the word FLOPs, it's actually ambiguous what I mean. So there's FLOPs, which is saying the number of floating-point operations, usually written FLOPs with a lowercase s. This is a measured amount of computation done.
[28:22.490] And then there's FLOP/s, which is floating-point operations per second. Sometimes it's also very confusing and written as FLOPS with uppercase S, which-- but I'm going to always write /s to make it clear that this is measuring the speed of hardware. So if you go and see that H100s have 800, 989 teraflops, it's the second. And when I say that GPT-3 took 325 or whatever it is FLOPs, that's the former. So just to get that out of the way. So just to give you an order of magnitude. So the number of FLOPs when I talk about 1e22 or 23 or 25, these are referring implicitly to the amount of compute or the scale of some of these models. And so if you look at H100s-- if you look at this glossy spec sheet-- actually it's not on this page. OK. Forget it. There is a spec sheet and I'll tell you that for bf16, the number of FLOP/s is 1979. And then you go and you benchmark it, and it's like wait a minute, that's not actually what I'm getting.
[29:45.790] And then you go read the fine print. And there's a footnote that says this is with sparsity. So sparse matrix and for dense over 2. So you always have to take these numbers divide by 2. So that's why you see these divide by 2. OK. So this allows us to-- just for intuition, so if you have 8 H100s so there's one node for two weeks. That's 8 times the number of seconds per in two weeks. Actually, this looks like it's one week. OK fine. It's one week times the number of FLOPs you get per second. So that's about 5e 21. OK. So this is just of building intuition for the number of FLOPs that certain types of hardware have and how many FLOPs certain types of models require. It's nothing fancy. It's just math and napkin math. OK. So now let's do something more mechanical. So suppose you have a linear model. It turns out that a lot of this calculus of counting FLOPs is actually going to be at the core it's like linear MatMuls. So this is actually not without too much loss of generality.
[31:11.740] So you have n points. Each point is d-dimensional. And we're going to map each of these d-dimensional vectors to a k-dimensional output. So B is going to be the number of points. D is the number of dimension, input dimensions, and K is the number of output dimensions. OK. So let's construct some x, which is the data matrix B by D. The weight matrix is D by K. And when you do the MatMul, the question is how many FLOPs that is. And it turns out that this is going to be 2 times basically the product of all the three dimensions. And the way to see that is that we have one multiplication for each triple and then also one addition. So there's a minus 1, because if you don't have to actually add like in this case D minus 1 times. But let's ignore that. OK. I'll come back to this. If this was a bit fast, I think there's another way to derive this.
[32:31.600] So what about the FLOPs of other operations? So elementwise operations are just the size of a matrix. I think that's fairly clear. So addition also requires m n FLOPs. So in general, no other operation you'll encounter as expensive as matrix multiplication for large enough matrices. So in general, we're just going to focus on what MatMuls are doing, with important caveat of when we talk about memory. Yeah. Just from interest, there are some other algorithms for doing matrix multiplication. Is this only for just doing it normally? So the question is that there are other algorithms of doing matrix multiplication. You mean like sub cubic algorithms? Yes. So in general, that's-- the optimization, the algorithms that people are going to explore for multiple matrix multiplications are going to be much more about how you co-design with the systems, rather than these more asymptotic algorithms.
[33:32.080] Yeah. We're considering additional multiplication in the same way is not possible to do division more efficiently than multiplication. Yeah, so I think the way the hardware is built, the two are basically the same. But yeah, intuitively it seems like I can do addition faster than I do multiplication. But the way the hardware is kind of the same. OK. So you can think about this MatMul as B is the number of data points. And D K is the number of parameters. So remember x is B by D and w is D by K. So another way to think about this formula is that the number of FLOPs in order to do a forward pass of this linear matrix is actually 2 times the number of tokens or data points times the number of parameters. So it turns out that this actually generalizes to transformers, which if you remember the 6 times n times d formula, we can see the shape of that forming. OK. So unfortunately these calculations are not going to be very meaningful because I'm doing this on CPU.
[34:54.050] But I'll just walk through the code here. So what we've so far done is measured FLOPs. So this is independent of hardware. It's just like the number of calculations you need to do for your model. So now the question is, how long does it actually take on hardware? So one way to find out is you just time it. So in this class, I think in the few lectures, we're going to talk more about benchmarking. But here's a little preview. So in general, when you time especially on GPU, you have to call cuda synchronize to make sure that-- because the GPU is running asynchronously, you want to make sure that you have this synchronization point. And then you perform the operation. And after operation, then you also have to have the synchronization barrier. If you omit this, you're going to find that wow, your timings are really fast. And that's because this is a non-blocking call. It just returns.
[36:02.830] And often it's general good practice to try this multiple times and take the average. So the actual FLOPs per second is basically the number of FLOPs you did times the time that you recorded on your hardware. And so remember, there's also another number, which is the GPU has a spec sheet. Let's see if I can pull it up on this link. I feel like maybe I have the wrong link. I'll have to fix that after class. --that it gives you some number of FLOPs, which was 989 teraflops. And so in general, the number of actual FLOPs per second is going to be different from the promised FLOPs per second. OK. Oops, sorry. And the way to directly think about the discrepancy is something called Model FLOPs Utilization, or MFU. And the definition of MFU is the actual FLOPs per second divided by the promise of FLOPs per second. And here this is ignoring the communication and other overhead.
[37:34.220] So you basically take the actual and divide by the promise. So in general, it's rare that you get more than was-- I mean, it's-- yeah, you just never get anything more than what you were promised. Often you get less. And in general, if you get about MFU of 0.5 for modern models, you should be pretty happy with yourself. If you have just like a straight up MatMul, you can get maybe potentially like 80.8 even, but you usually can't get that high. And sometimes if you have really-- something's really wrong, you'll get something like 0.1, which means that you should do something. OK. So whenever you write your model, you can calculate. You now know how to calculate the MFU through a combination of counting the number of logical FLOPs that your model needs to do, and then looking at the wall clock time and essentially dividing. Yeah.
[38:42.497] Was there a question over there? Oh. [INAUDIBLE] Yeah. So is the promise-- what is the promise FLOPs number? That is in the spec sheet, it is already divided by 2 of the 989 number. And then on top of that, you only get 0.5 of that in general depending on your computation. [INAUDIBLE] you're getting 50% MFU? Yeah, so why are you getting only 50% MFU? Actually, that's a good question. I'll come back to that when we talk about memory bottlenecks. OK. So to summarize here, matrix multiplications dominate generally the computations. And that's by design. And the number of FLOPs per second depends on the hardware. So better hardware leads to more FLOPs and also it depends on the data type. Which means that if you look at the spec sheet, different data types will have different FLOPs. If you try to do float32 nowadays, it's going to be really, really slow because they're not really optimizing for that workload.
[40:07.750] Whereas now bf16 or fp8 are going to be much faster. And MFU-- now you know what MFU is. It's the actual FLOPs divided by promised FLOPs. All right. So to go back to the question of why is MFU 0.5. And to understand that, I'm going to have to introduce this idea of arithmetic intensity. And the reason is that, well, it's not just doing a bunch of MatMuls and then you're done and looking at how long the MatMuls take. This is my very cartoon version of what hardware looks like. You have high bandwidth memory. And then you have where the compute cores are, the accelerators chips are. And then how do you compute? Well, you have to send your inputs, your matrices-- the tensors are sitting down here-- from the memory to the accelerator. You do the computation, and then you send it back. So if you want to measure how long this takes, this depends on two things.
[41:27.790] One is the accelerator speed, which is what we've talked about just now. But the other thing that matters is the memory bandwidth of your hardware, which we haven't talked about. And if you look at the spec sheet, we talked about how the FLOPs per second was 1979e12 divided by 2. And the bytes per second, which is the memory bandwidth, this is 3.3 terabytes per second. And remember why we were looking at memory, how much things took to store. It's most obviously that, well, if you have a model that's too big it doesn't fit in your memory, that's not going to be fun. But also it turns out that memory-- you need to move this memory, which takes time. So actually the size of how large things actually influences speed as well. OK. So I'm going to talk through some operations, and I'm going to basically compute how long things are going to take and introduce this idea of arithmetic intensity.
[42:38.040] OK, so suppose I have a million dimensional vector of bf16. And I'm going to just compute a ReLU on this. So remember ReLU is just max of x and 0 done elementwise on the entire vector here. So I count two things, one is the number of bytes that were moved. So I have to read x in-- copy it into the accelerators. And each of this is going to be 2 because bf16 is 2 bytes per float times n floats. So that's 2n. And then I'm going to write y back. So that's another 2n. OK. So that's the number of bytes that have to be moved. And then how many FLOPs were done? Well, each of these elements, I'm just comparing it with zero, and that's it. So that's n comparisons. So now I look at the communication time, which is the number of bytes that I needed to move divided by the speed of that movement. And that gives me the time, which is 1e minus 6 seconds. And what about the computation time? That's the FLOPs divided by FLOPs per second. So that's 1 minus 9. OK. So there's also going to be another important assumption which generally, we try to hold is that we overlap communication and computation.
[44:19.453] We're going to talk more about that when we talk more deeply about GPUs. But the idea is that in this case as we don't sit here waiting for the things to move. As soon as they're there, we start computing them, and then we move them back. So this movement and also the compute is happening at the same time. So mathematically, we're just going to assume that the total time is the max of the two, because we're going to assume that we can perfectly overlap them. In practice, it's not going to be perfectly overlapped. There's going to be some overhead. But this is good enough for now. OK. So the total time, as we see here, is 1e minus 6. OK. So if you ask what is the bottleneck here? So when the communication time is greater than the computation time, then we call it the algorithm memory bound, because you're spending most of your time just waiting for bits to show up. And when the computation time is greater and the communication time, that's compute bound, because then you're actually-- your bottleneck is actually doing the compute. And so in this case, what is ReLU? Rectified Linear Unit. Oh, sorry. Is it memory bound or compute bound?
[45:39.590] Memory bound. Yeah. And it's clear, because the compute is way less than the communication. So here's another way to see it. And this is where I'm going to define the intensity. So what is the intensity of an accelerator is essentially how much work can the accelerator do per byte transferred. And for any given accelerator based on the spec sheet, you basically have the FLOPs per second divided by the bytes per second. So how much useful work can you do per byte that's moved for each 100? That's 295. So that means for every byte of-- you can do 295 floating point operations. So that's an intuitive number to have in your head, about 300. OK. So now the arithmetic intensity of algorithm is how much actual work was done per byte for this workload. And if you look at it for the ReLU computation, it's FLOPs over bytes and it's actually-- so this is actually a quarter, I guess, not half.
[47:02.390] So the point is that it's very small, 0.25. OK. So now we can talk about bottlenecks through the language of intensity. So something is memory bound if the arithmetic intensity is smaller than the accelerator intensity and compute bound if it's greater than accelerator intensity. So these are equivalent. And if you look at the algebra, it's basically you have two fractions. And then you just multiply and divide to-- you basically switch the terms around. OK. So in this case, we're memory bound. So in general, we're going to find ourselves in a situation where we're memory bound because data movement is expensive. And so if you can get higher arithmetic intensity, that's good. So 0.25, if you see that number, if someone tells you arithmetic density is 0.25, you should say, oh, this is really bad. Yeah. What's some typical arithmetic intensity for some transformer [INAUDIBLE]?
[48:08.300] Yeah, so we'll get to that. OK. All right. So OK. So one way to think about increasing arithmetic intensity is that let's just try to do more stuff per unit of byte moved. So the GELU is another activation. It looks like this, some formula that is more-- doesn't have zeros. And if you do this calculation, the number of bytes are moved back and forth is still 2n plus 2n. And the FLOPs here it's about 20 FLOPs per elementwise scalar operation, so it's 20n. So the arithmetic intensity is, let's say, 5. This is a crude estimate. And in this case, are we memory bound or compute bound? Memory bound. Memory bound, because 5 is still smaller than 295, way smaller. So even though GELU does a lot of work, more work than ReLU, in the way that things are structured, it's still memory bound.
[49:30.020] Which means that if you were just computing ReLU and GELU, you think that, well, GELU is so complicated, it must be really expensive. But actually it's exactly the same, because that's not where the bottleneck is. OK, so now let's look at some linear operations, so dot product. So dot product, you have a vector x. You have a vector w, size n, and you take the dot product. So how many bytes are moved? So you read x, which is 2n; read w, which is 2n; and then you write y, which is a scalar, which is 2. And the number of FLOPs is you do the n multiplications and n minus 1 additions. So that's 2n minus 1. OK. So what's the arithmetic intensity? Oh, for this one, it's about half, which is also pretty bad. Which means that-- hopefully you get the idea-- it's memory bound. So what about a matrix vector product? So I have this x is a vector. w is an n by n matrix.
[50:41.090] And you form this product. How many of you think this will be compute bound? How many of you think memory bound? OK. So let's see. So I'm going to read x read, which is 2n; read w, which is 2n squared; and write y, which is 2n. And there's the number of FLOPs is basically you're doing n dot products. So that's n times the cost of doing a dot product. And the arithmetic intensity is barely, barely higher. So this is also memory bound. OK. So now let's talk about matrix multiplication. So this is where things get interesting. So I have an n by m matrix, another m by n matrix. And I multiply them. And the number of bytes that were moved was 2n squared plus 2n squared. And then I have to write 2n squared bytes back to y. And the number of FLOPs here is n squared dot products. And then what's arithmetic intensity? It's 300. Whew. OK, so 340.
[51:56.010] And in general, it's roughly n over 3. And intuitively this makes sense because you're sending n squared things, but you're computing n cubed things. So the number of things you're computing over a number of things you're sending is order n. And this gets better the larger you make the matrices. So in general, this is why when you hear people talk about, oh, we need to make large batch sizes or have large matrices, it's exactly this. If you're under the accelerator intensity, making things smaller doesn't actually speed things up. It's all the same. Whereas if you get to a point where you're over this, then you're actually saturating your GPUs. OK. So finally, this is compute bound. OK. So as long as we have large matrices, we're actually pretty good. We're in a compute bound saturating the accelerator.
[52:59.600] And the question earlier about what about transformers? Well, it turns out we'll see in both your assignment, but also in the next lecture that transformers are essentially big matrix multiplications with some things sprinkled in between. So that's good news from arithmetic intensity. And this is, by design, transformer is designed in a certain way to have high arithmetic intensity. And one comment here, just to foreshadow the inference lecture, is that matrix vector products is essentially what goes on when you're doing transformer inference. Because inference, you're generating one token at a time. And so you only get to-- it's like a vector that you're trying to dot product with a matrix. And that's as we saw with memory bound. Whereas at training time, you get this whole sequence and you're processing all at once.
[53:56.150] OK. So note also that the intensity depends on the precision. So by default everything we're doing here is fp16. OK. So also to tie it to the other question about MFU, and the reason you might be getting low MFU is that while you might be doing-- so MFU is a promise FLOPs-- sorry, actual over promise. So if you have really large memory bottlenecks, then you're not actually going to get very good throughput, even though you-- the promise is if you didn't have memory bottlenecks, you were just like going through and doing all the computations of your model. OK. Final thing on arithmetic intensity roofline plots. So there's a nice way to-- yeah. [INAUDIBLE] generally speaking, these models tend to be memory bound, as they say for most of the computation. But yet the accelerators are 50% is not as good. And maybe you get 70, maybe you get 80 that's really good, meaning that the accelerator is outsized compared to memory bandwidth.
[55:18.990] Why is it like that? Why do these accelerators-- are just big while they're just idling or waiting for memory? So the question is, could you design maybe accelerators that had better characteristics? Yeah. Maybe when we talk about GPUs and understand a bit more how they work, we can talk about why this is. But if you have an answer, you should tell Jensen. And maybe you can design a better hardware. OK, so let's visualize the relationship between arithmetic intensity and performance. So this plot basically plots the arithmetic intensity on the x-axis. So every slice here corresponds to a particular, let's say, algorithm. And then we have these lines here. And each of these lines corresponds to, let's say, a particular accelerator. Maybe H100 or B200 or so. And so what this shows is-- and then on the y-axis is the FLOPs per second that are realized. So if your algorithm has low arithmetic intensity, like, ReLU or dot products, you're going to be over here, which means that the FLOPs per second realize
[56:37.630] is going to be not as high as your peak accelerator. And as you increase the memory arithmetic intensity, that's going to-- things are going to be better. You're going to be able to saturate your hardware up until a certain point. And after a certain point your compute bound. And obviously you can't exceed the peak FLOPs. All right. So let's go on here. OK. So now I'm going to go back and talk about memory and compute for the operations that we need in training. So far we've done tensor operations, basically MatMuls. And we saw basically how much memory it took and how much compute it took and the interaction between them. So now let's actually think about what it takes to train. So here is our running example here. Actually it's not a linear network, just a deep network. So I'm going to consider a case where I have an input, which is a B by D input, and a number of layers where each layer is a D by D MatMul.
[57:59.790] And then that produces some set of preactivations and then element wise ReLU that produces some activations of the first layer. And then this is repeated again and again. And the output is just going to be the same size. So the deep network, I mean, just to see what it looks like in PyTorch, just so it's basically a set of blocks. Each block has a weight vector-- sorry, a weight matrix associated with it. And so when you the number of parameters here is D squared times the number of layers L. And when you run the model on the batch of data, what we're doing is we're going through the layers, and each layer we basically apply the linear transformation and then a pointwise ReLU activation. And then we do that for all the layers. So this is hopefully straightforward very simple model. OK. So now let's talk about gradients. So let's use actually even simpler example. So this is just a simple linear model regression. So we have a vector 1, 2, 3, a weight vector 1, 1, 1.
[59:33.420] And we take the dot product, and we form this MSE loss. OK. And what happens when we take the-- do the backward pass is that each of the variables involved in this computation graph has a gradient that is either set or not set. So the w.grad gets set to 1, 2, 3. So this is just basic mechanics of PyTorch. I think it should be familiar to all of you. So now the question is, how much compute does gradients take? OK. So let's count FLOPs for computing gradients. So let's take a simplified model where you have an x, which is BID and times w1 matrix. We're going to ignore this ReLU for now just for simplicity times w2, which is the same D by D-- same shape D by D matrix. OK. I'm going to use einsum here for reasons that will become clear. So the first thing you do in this linear network is you take x and you take w1. x is batch by input dimension.
[01:00:55.860] w1 is as in buy out, and you get a batch buy out. And h2 takes h1 and w2, and it's the same story. These are just MatMuls. And then you form a loss. I'm just sending it to some arbitrary just to get a number out. OK. So what happens in the backward pass, I'm going to call these retain gradients for debugging purposes, which you'll see later. You take the backward pass on the loss. And the question is, how much work did that-- how many floating point operations was that? So let's zoom in on one layer. Let's focus on the second layer here, so this layer. And so the second layer takes h1 and just multiplies it by a matrix w2 to get h2. And that's it. So if you look at the FLOPs in the forward pass, this is just a MatMul. And we just take the three dimensions, and we multiply them together and bf16. That's 2 bytes. Sorry, that's irrelevant. This is just 2 because it's an addition and a multiplication.
[01:02:08.480] So in the backward pass, what does this look like? So in the backward pass, if you remember your chain rule and your backprop algorithm, you have to compute two things. You have to compute the backward message, the gradient of the loss with respect to your input. And you also have to compute the gradient with respect to your parameters. And what do those gradients look like? This is just by the definition of these-- I guess this is just the chain rule here written in einsum notation. You take the d loss d h2, which is h2 grad. That's a batch by Alt matrix, and you multiply it by w2, which is in by out. And then that gives you h1 grad, which is d loss d h1 and which is batch n. So I always get this confused whenever if you see-- learn calculus, there's like, one of them has a transpose, and I forget which order to put them in.
[01:03:25.700] And einsum I think makes it very clear, because you can remember even just looking at the scalar case that it has to look like h1 grad is h2 grad times w2. And the only question is, how do you index things? And you index things by basically looking at the shape of these-- and the named dimensions. And in this case, it's just a matrix multiplication where you are summing over the output dimension. And when you do that, we can check that. So h1.grad was the thing that was computed by loss.backward. And this is actually indeed the same thing that I wrote out here just as a sanity check. So the other thing you have to compute is w, the gradient with respect to the parameters. And that's d loss d h2 times h1. So it's the same kind of backward message that's coming back but multiplying with the other thing that you're not taking the gradient with respect to.
[01:04:27.680] And you just write out the dimensions. And in this case, you're summing over the batch dimension. But the nice thing is that if you just look at this, we know how expensive this is. The number of FLOPs is essentially the product of all the dimensions. It doesn't matter which ones you're batching or not. It's basically you enumerate over the i, j, k, and you aggregate. The way that you aggregate is different, but the FLOPs is the same. OK. So notice that the backward pass is exactly twice as expensive as the forward pass. And this is because you have to take-- compute two gradients, one with respect to the parameters for each parameter, and then one with respect to the input, the other thing that's not the parameter. Maybe I'll pause. Any questions about that. OK. So now that was just for w2. And you just need to apply this to all the parameters in the network. And you put it all together, you will see that for this network, the forward pass is 2 times the number of data points, which is b times the number of parameters FLOPs.
[01:05:59.430] And the backward pass is twice of that, which is 4 times the number of data points times the number of parameters FLOPs. And so the grand total is 2 plus 4 is 6 and 6 times the number of data points and parameters. So this is where the 6 nd comes from that you might have seen in various places. It's just by counting forward and backward. So we did this for just these deep networks. But it turns out that this is actually a good approximation for transformers as well-- as long as the context length isn't too large. If the context length is too large, then you get the context length squared, and that's more FLOPs that isn't in this kind of accounting.
[01:06:47.490] OK. So let's talk a bit about-- so now we have gradients. Now we have to-- the other piece when you do training is we have to do optimization. So here's our deep network. Just so that I'm not just giving you what's in assignment 1, I'm going to use the AdaGrad optimizer, which this from-- this is from 2011. This is predates Adam. It's where you can think about it as somewhere in between SGD and Adam. And basically it's SGD where you look at the second moments of the gradient. Momentum is what you look at the first moment gradients. And Adam is where combine the two of them. I'm not going to have time to go into the optimizer details here, since we're focusing mostly on the usage. So we can define an optimizer. And when you compute the gradients and you compute the gradients and then when you take an optimizer step, if you're defining a new optimizer, so in the assignment 1, you're going to implement Adam.
[01:08:06.710] For each of the parameter groups which are, in this case w1 or w2, we look at-- there's something called optimizer state, which is of storage that you can-- the optimizer uses while it's running. And what we're computing in AdaGrad is something-- is the squared gradients, the sum of the squared gradients. So here we're getting it from the optimizer state. We're updating it with the current gradient, and we're storing it back. And then after you update the g2, then you update the parameters. So in AdaGrad, it's you basically divide by the square root of the average gradient squared or sum of gradient squared, rather. OK. So let's see. So I'm actually going to maybe skip the training loop. Actually, wait. Hold on. I think I actually skipped something I didn't mean to. OK. So that was the optimizer state, and now let's look at how much memory the optimizer is using. Or in general what is the memory usage?
[01:09:29.840] So for parameters in this network, there's D squared parameters per each of the L layers. And each parameter takes 2 bytes if we're storing it in fp16. So that's the number of parameters. Activations-- this is 2 times-- which is for bf16 batch times D times the number of layers. For every layer you have activation. You have gradients which is basically a copy of all the parameters. And this is also bf16. And the optimizer state is for every-- actually this should be the-- I think this should be-- sorry, I'll fix this. This is a typo. This should be the number of parameters, not parameter memory. And this should be number of parameters, so number of parameters times 2, and then this is 4 times the number of parameters. And the reason for this is that it's customary to use fp32 for the optimizer states for stability reasons.
[01:10:31.980] Obviously people have tried using bf16. And what ends up happening is that you're taking squares, and you're averaging multiple steps, and it doesn't-- it's not very stable. And for AdaGrad, it's here-- we have 4 bytes per parameter for storing the optimizer state. Adam-- you store the first order bit and the second order moments, so that's 8 bytes. So if you think about the optimizer states is actually a lot of the memory used. One note is that remember memory serves two purposes. One is how much-- you have to store this thing in your HBM. But the other thing is that it has to be shipped to the accelerators. And in general, the optimizer state is not really the bottleneck for compute. So the amount of memory here is not really so important for performance in terms of speed, but it means that you can't fit large models in your memory.
[01:11:42.250] OK. So to put it together, as we mentioned, the number of parameters here is D times D times L, and the number of FLOPs is D times number of tokens or number of data points times number of parameters. So for transformers, this is going to be a bit more complicated. But you're going to do that in assignment 1. And you do it more carefully. OK. I'm going to skip the training loop. This is just a general review. Two things I want to quickly touch on before we conclude. One is that as we see, memory does have an important effect both on the ability to store large models, but also sometimes on your speed. And so in general, you want to reduce your memory usage. So there's two things that people typically do. One is gradient accumulation. So in general, you want to use batch sizes that are large enough to improve stability up to a critical batch size, which Tatsu will talk about later.
[01:12:45.730] And as we saw that activation memory scales with batch size. So you might want to run out of memory at some point if you have two large batch sizes. So gradient accumulation says, well, you compute. You have these micro batches, and you compute the gradient on the micro batches. And then you just accumulate the gradients. You don't zero out the gradients. And every batch size over microbatch size steps, you update the parameters and zero out the gradients. OK, so this is actually a very simple code change which allows you to save on compute. So the other thing I'll quickly mention is activation checkpointing. So in training, we need to store in general the activations of all the layers. This is done by default. Actually, interestingly, for inference we don't need to compute the gradient, so we only need to store the current layer's activations.
[01:13:40.760] But for training, the memory usage is B times D times 2 times L. And the question is, can you reduce this? So activation checkpointing, also known as gradient checkpointing or rematerialization-- the key idea here is that in the forward pass, you just keep activations for only a subset of the layers. And in the backward pass, you recompute the missing activations from the last checkpoint. That's why it's called checkpointing. And this is a general trick that you see in systems which is trading off. If you want to reduce memory, you can just recompute things. So here we're going to define this to operationalize this. This is actually fairly easy. Actually, I didn't-- OK, maybe I didn't see the-- OK, so basically what you do here is you have the same model, except for you just add Torch utils at checkpoint on a layer.
[01:14:47.960] So that means do this computation, but don't store any of the intermediate activations. Only store what is needed. So this means that if you store all the activations, remember in your deep network, you have the thing that's pre the ReLU and the thing that's after the ReLU right. So that's a lot of storage. And instead if you do activation checkpointing on those blocks where each block has a linear and ReLU, you don't store the pre ReLU, and you save basically half. You can get away with half the memory. And then when you're doing the backward pass, you need g3, but you can compute g3 easily from h2. OK. You can go further than this. You can say, well, I can store all the layers, but in the extreme case, I can just store no layers. And so that will be maximally memory efficient. The only thing is that the compute is going to be L squared. Because for every one of these layers, you have to start from the beginning.
[01:15:57.080] And maybe a sweet spot is if you store the checkpoints at square root L layers, that means your activation memory is square root of L, and your recomputation overhead is also square root of L. So that's balanced. OK. So to summarize this lecture, everything is operating on tensors-- parameters, gradients, activations, optimization states, data. We introduced this einops, which hopefully you guys can embrace as a way to think about tensor operations. 6 times number of data points times number of parameters is a formula which now we have demystified as the number of FLOPs per-- for training step. And actually if this is training step, this should be batch size.
[01:16:46.790] And then we talked about arithmetic intensity and roofline analysis, which allows us to diagnose whether a computation is memory bound or compute bound. Matrix multiplications are compute bound. Basically everything else is memory bound. And then finally, gradient accumulation, activation, checkpointing are ways to reduce the memory. And by reducing the memory, that allows you to use bigger batch sizes. OK. So that's it for today's lecture. So next week, Tatsu will talk about architectures.