CS336 · 第 2 课

PyTorch 构建块与资源核算

训练一个 70B 模型到底要花多少钱、多长时间？答案可以用"餐巾纸数学"（napkin math）算出来——你只需要会乘法和理解几个关键数字。这一课建立整个课程的方法论：从第一性原理出发，用资源核算驱动每一个工程决策。

这一讲到底解决什么问题？

上一课建立了全课框架：理解源于构建，效率驱动一切。从这一讲开始，每个工程选择都要回答两个问题：模型状态能否放进显存，完成训练需要多少计算时间？PyTorch tensor、operation、module 和 training loop 是实施这些选择的工具，不是彼此孤立的 API 清单。

PyTorch 工具层

贯穿始终的是资源核算——对两种核心资源保持持续追踪：

• 内存（GB）：你的模型能放进去吗？
• 算力（FLOPs）：训练要多久？

我们不会在这课讲 Transformer 架构本身（Assignment 1 会深入），而是用更简单的模型来演示原理——这些原理在任何规模都适用。

先建立账本，再做餐巾纸估算

Training
tokens → forward → loss → backward → gradients → optimizer → new weights

显存
├─ model states（通常随 P 缩放）
│  ├─ parameters / master parameters
│  ├─ gradients
│  └─ optimizer moments m, v
├─ activations（随 batch、sequence、layers、hidden shape 缩放）
├─ temporary tensors / CUDA workspace
└─ communication buffers / allocator fragmentation

Inference
weights + KV cache + temporary activations
（没有 backward、gradients、Adam states）

import torch

# 创建方式
x = torch.tensor([[1., 2, 3], [4, 5, 6]])  # 从数据
x = torch.zeros(4, 8)                        # 全零
x = torch.ones(4, 8)                         # 全一
x = torch.randn(4, 8)                        # 随机初始化
x = torch.empty(4, 8)                        # 未初始化（最快）

x = torch.tensor([
    [0., 1, 2, 3],
    [4, 5, 6, 7],
    [8, 9, 10, 11],
])

# stride 描述在存储中跳过多少元素到达下一个维度
x.stride(0)  # = 4  （到下一行，跳过 4 个元素）
x.stride(1)  # = 1  （到下一列，跳过 1 个元素）

# 元素 [r, c] 的索引 = r * stride(0) + c * stride(1)
# 例如 [1, 2] → 1*4 + 2*1 = 6

x = torch.tensor([[1., 2, 3], [4, 5, 6]])

# 切片 - view，共享存储
y = x[0]       # [1., 2, 3]，共享存储
y = x[:, 1]    # [2, 5]，共享存储

# reshape（可能 view，可能拷贝）
y = x.view(3, 2)  # view，共享存储

# transpose - view，共享存储
y = x.t()
y.is_contiguous()  # False！非连续

# 非连续 tensor 不能再 view，需要先 contiguous()
y.view(2, 3)  # RuntimeError!
y.contiguous().view(2, 3)  # OK，但会拷贝

内存核算：数据类型

餐巾纸估算要先把模型中的对象换算成 bytes。参数、梯度、optimizer state 和 activation 都是 tensor；dtype 决定每个元素占多少空间，也决定动态范围和可使用的硬件路径。

Adam 训练到底保存了哪些对象

因此 16 bytes/param 或 18 bytes/param 都只是某个 precision + optimizer + framework recipe 的账本。真正方法是逐项写“元素数 × dtype bytes × 是否复制/分片”，不能背一个常数。

aggregate HBM = 8 × 80 GB = 640 GB
upper bound = 640×10^9 / 18 ≈ 35.6B parameters

精度与动态范围

为什么 bfloat16 是深度学习的宠儿？

# float16 的问题：小数值下溢
x = torch.tensor([1e-8], dtype=torch.float16)
print(x)  # tensor([0.])  ← 下溢了！

# bfloat16 没问题
x = torch.tensor([1e-8], dtype=torch.bfloat16)
print(x)  # tensor([1.0000e-8])  ← 保留了

直觉：bfloat16 把有效精度从 float16 的 10 位尾数降到 7 位，但保留了与 fp32 相同的 8 位指数，因此动态范围更适合训练。代价是舍入误差更大；“低精度一定有正则化效果”并不是通用保证，数值稳定性仍取决于算子、归一化、累加精度和优化器实现。

FP8：H100 的新武器

H100 支持两种 FP8 格式：

• E4M3：4 位指数 + 3 位尾数，范围约 ±448，精度稍高，常用于前向的权重与激活。
• E5M2：5 位指数 + 2 位尾数，范围更大、精度更低，常用于反向梯度。

这是 NVIDIA Transformer Engine 常见的 hybrid recipe，不是硬性规则；实际训练还需要缩放、amax 统计，并通常以更高精度累加。

FP4 / NVFP4：Spring 2026 新增的极低精度视角

Spring 2026 官方 trace 进一步讨论了 NVIDIA 在 2025 年公开的 NVFP4：每个值只有 4 bit，并为小块数据保存独立 scale。这样扩大了局部动态范围，但同一块中的值共享尺度，因此表示能力与误差会强烈依赖分块方式。

# 官方 trace 给出的 NVFP4 离散值直觉
values = [-6, -4, -3, -2, -1.5, -1, -0.5,
           0, 0.5, 1, 1.5, 2, 3, 4, 6]

# 真实存储成本不能只写成 0.5 byte/value：
# 还要计入每个 block 的 scale，以及可能保留的高精度副本/累加器。

GPT-3 的 d_model = 12288, d_ff = 4 × 12288 = 49152
一个矩阵: 12288 × 49152 × 4 bytes (fp32)
     = 2,415,919,104 bytes ≈ 2.3 GB

GPT-3 有 96 层 → 一个模型光参数就 175 GB（fp32）

计算核算：FLOPs

内存核算回答“能否放下”，但不能回答“需要多久”。接下来把 operation shape 转成 FLOPs，再把整次训练归约到参数量和 token 数。

什么是 FLOP？

FLOP（Floating-Point Operation）是一次浮点加法或乘法。注意区分：

• FLOPs（小写 s）：Floating-Point Operations，计算量的度量
• FLOP/s（大写 S）：Floating-Point Operations per Second，硬件速度

矩阵乘法的 FLOPs

在 dense Transformer 中，参数相关计算主要来自矩阵乘法。先看具体 shape：

X:  [100, 256]
W:  [256, 512]
Y = X @ W
Y:  [100, 512]       # 内维 256 消去，外维保留

给定 dY: [100, 512]
dW = X.T @ dY
   = [256,100] @ [100,512]
   = [256,512]       # 必须与 W 同 shape

dX = dY @ W.T
   = [100,512] @ [512,256]
   = [100,256]       # 必须与 X 同 shape

forward Y: 2 × 100 × 256 × 512 = 26,214,400 FLOPs
backward dW: 2 × 256 × 100 × 512 = 26,214,400 FLOPs
backward dX: 2 × 100 × 512 × 256 = 26,214,400 FLOPs

所以这一线性层：backward ≈ 2 × forward
不是经验魔法，而是 backward 有两个同量级 matmul。

Exact identity矩阵相乘的合法 shape 与 \(dW\)、\(dX\) 的 shape。

Napkin-math approximation把长度 \(D\) 的 dot product 记作 \(2D\) FLOPs，以及忽略 bias、activation、normalization 等较小项。

从矩阵乘法到 6PD 公式

现在只看一个 token穿过所有参数主导的线性层。令 \(P\) 为被该 token 使用的 dense parameter count：

one token forward  ≈ 2P     # 每个权重参与一次 multiply-add
one token backward ≈ 4P     # weight-gradient 与 input-gradient 两条 matmul
one token training ≈ 6P

D 个实际处理的 training tokens:
training compute ≈ 6P × D = 6PD

\(D\) 是训练过程中实际 processed token count。它不是数据文件的 GB 数，也不必等于 unique dataset token count。100B-token dataset 训练 3 epochs，若每个 token 每轮都使用，则 \(D=300B\)。这里用 \(D\) 而不是 \(T\)，避免把总数据 token 数与单个 sequence length 混淆。

为什么 batch size 没出现在 6PD

固定 D = 1,000,000 processed tokens

方案 A: 1,000 tokens/step × 1,000 steps = 1,000,000
方案 B: 2,000 tokens/step ×   500 steps = 1,000,000

两者一阶 parameter-related compute 都约为 6P × 1,000,000。

Batch size 改变的是如何把同一批 token 分组，而不是 token 总数。因此它不单独出现在总 FLOPs 的一阶式中；但它会改变 GPU utilization、activation memory、optimizer step 次数、优化动力学与多卡通信形态，所以“总 FLOPs 相同”绝不等于“训练时间和最终模型相同”。

P = 70×10^9 parameters
D = 15×10^12 processed tokens

C ≈ 6PD
  = 6 × 70×10^9 × 15×10^12
  = 6.3×10^24 FLOPs

其他操作的 FLOPs

Model FLOPs Utilization (MFU)

总 FLOPs 除以 GPU 峰值只能得到理想下界。真实 kernel 会受到数据搬运、shape、通信和软件开销影响，因此需要 MFU 把理论峰值与端到端训练吞吐连接起来。

定义

MFU = 实际 FLOP/s / 硬件峰值 FLOP/s

更准确地说，MFU 是模型按约定 FLOPs 口径得到的有效 FLOP/s占对应 dtype、dense/sparse 模式下理论峰值的比例。分子不是 GPU 上发生过的所有指令。

GPU utilization = 100%
MFU             = 40%

MFU 的损失可能来自：

• 内存读写等待（memory-bound 操作）
• 通信开销（多 GPU 时）
• Kernel launch 开销
• Python/PyTorch 开销

硬件峰值参考

实测 MFU

import timeit
import torch

# 大矩阵乘法，bf16
B, D, K = 16384, 32768, 8192
x = torch.ones(B, D, device='cuda', dtype=torch.bfloat16)
w = torch.randn(D, K, device='cuda', dtype=torch.bfloat16)

# 预热（JIT 编译、缓存预热）
for _ in range(5):
    _ = x @ w
torch.cuda.synchronize()

# 正式计时
num_trials = 5
start = timeit.default_timer()
for _ in range(num_trials):
    y = x @ w
torch.cuda.synchronize()
elapsed = (timeit.default_timer() - start) / num_trials

flops = 2 * B * D * K
actual_flops_per_sec = flops / elapsed
promised_flops_per_sec = 989e12  # H100 bf16
mfu = actual_flops_per_sec / promised_flops_per_sec
print(f"MFU: {mfu:.2%}")  # 通常 40-60%

算术强度与 Ridge Point

MFU 告诉我们“离峰值还有多远”，arithmetic intensity 进一步解释“为什么”：同样的 GPU，有些 operation 受计算单元限制，有些 operation 先耗尽 HBM bandwidth。

算术强度（Arithmetic Intensity）

对于矩阵乘法 Y = XW（X: B×D, W: D×F）：

读 X: B × D × 2 bytes (bf16)
读 W: D × F × 2 bytes
写 Y: B × F × 2 bytes
总传输: 2BD + 2DF + 2BF bytes

FLOPs: 2BDF

算术强度 = 2BDF / (2BD + 2DF + 2BF)
         ≈ B （当 D, F >> B 时）

Ridge Point

每个硬件有一个ridge point——计算和内存带宽的平衡点：

H100 ridge point = 989 TFLOP/s / 3.35 TB/s ≈ 295

如果算术强度 > 295 → 在该简化 roofline 下更可能 compute-bound
如果算术强度 < 295 → 在该简化 roofline 下更可能 memory-bound

在 D,F >> B 且 traffic 模型成立时，B > 295 才可能跨过 ridge point

Compute-bound 与 memory-bound 不是“好/坏”标签，而是当前 workload 的性能上限分别由计算峰值还是带宽决定。优化方向也不同：前者减少 FLOPs 或提高矩阵单元效率，后者减少 bytes 或增加数据复用。

逐 token decode 为什么容易受带宽限制？

当并发规模很小、每次只生成一个 token 时，线性层接近矩阵—向量乘法：读取大量权重，却只为少量 token 做计算，因此权重复用不足，往往更容易落入 memory-bound 区域。

• Decode MLP：增大 batch 可以让多个请求复用同一批权重，提高算术强度，因此它并非永远 memory-bound。
• Decode attention：每一步都要读取随 context 增长的 KV cache，而单步 query 很小，通常对 memory bandwidth 特别敏感。

因此推理优化既要减少权重与 KV 的数据移动，也要通过 batching 提高复用；具体瓶颈仍需结合 batch、context、dtype、kernel 与硬件实测。

# 传统写法——你能一眼看出在算什么吗？
scores = query @ key.transpose(-2, -1)  # -2, -1 是什么维度？
x = x.view(x.size(0), x.size(1), num_heads, -1)  # -1 是多少？

from einops import einsum

# attention scores: Q @ K^T
scores = einsum(query, key,
    'batch seq1 head dim, batch seq2 head dim -> batch head seq1 seq2'
)
# 不再需要 transpose，维度命名让意图一目了然

from einops import rearrange

# 把 hidden 拆成 heads × head_dim
x = rearrange(x, 'batch seq (heads dim) -> batch heads seq dim', heads=8)

# 合并回去
x = rearrange(x, 'batch heads seq dim -> batch seq (heads dim)')

from jaxtyping import Float
import torch

# 默认主要用于类型标注；配合 jaxtyped + typeguard/beartype 可运行时检查
x: Float[torch.Tensor, "batch seq heads dim"] = torch.ones(2, 3, 8, 64)

# 简单模型: y = 0.5 * (x @ w - 5)^2
x = torch.tensor([1., 2, 3])
w = torch.tensor([1., 1, 1], requires_grad=True)

# 前向
pred_y = x @ w           # = 6
loss = 0.5 * (pred_y - 5).pow(2)  # = 0.5

# 反向——自动计算 d(loss)/d(w)
loss.backward()
print(w.grad)  # [1, 2, 3]  (= x * (pred_y - 5))

前向: Y = X @ W                    → 2BDF FLOPs
反向:
  dL/dX = dL/dY @ W^T              → 2BDF FLOPs
  dL/dW = X^T @ dL/dY              → 2BDF FLOPs

总计: 6BDF = 3 × 前向

# 不稳定的初始化
w = nn.Parameter(torch.randn(input_dim, output_dim))
# output 的元素大小 ~ sqrt(input_dim) → 随维度增长！

# Xavier/Glorot 初始化
w = nn.Parameter(torch.randn(input_dim, output_dim) / np.sqrt(input_dim))
# output 的元素大小 ~ 1（与维度无关）

# 更安全的截断版本
nn.init.trunc_normal_(w, std=1/np.sqrt(input_dim), a=-3, b=3)

class Linear(nn.Module):
    def __init__(self, input_dim, output_dim):
        super().__init__()
        self.weight = nn.Parameter(
            torch.randn(input_dim, output_dim) / np.sqrt(input_dim)
        )

    def forward(self, x):
        return x @ self.weight

class Cruncher(nn.Module):
    """简单深度线性模型，用于演示。"""
    def __init__(self, dim, num_layers):
        super().__init__()
        self.layers = nn.ModuleList([Linear(dim, dim) for _ in range(num_layers)])
        self.final = Linear(dim, 1)

    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        return self.final(x).squeeze(-1)

model = Cruncher(dim=64, num_layers=2)
print(f"参数数: {sum(p.numel() for p in model.parameters())}")
# 64*64 + 64*64 + 64 = 8256

for step in range(num_steps):
    # 1. 获取数据
    x, y = get_batch(batch_size=B)

    # 2. 前向传播
    pred_y = model(x)
    loss = F.mse_loss(pred_y, y)

    # 3. 反向传播
    loss.backward()

    # 4. 参数更新
    optimizer.step()

    # 5. 清除梯度（重要！否则梯度会累积）
    optimizer.zero_grad()  # 当前 PyTorch 默认 set_to_none=True

# 对于参数量 P，batch 大小 B，序列长度 L，维度 D，层数 L_num
# 假设 float32 + AdamW

parameters:       4P  (fp32 参数)
gradients:         4P  (fp32 梯度)
optimizer_state:   8P  (m + v, 各 4P)
activations:       ~4 × B × D × L_num × L_seq  (取决于具体结构)

总计（不含 activations）: 16P bytes

# 常见 bf16 参数 + fp32 主参数/梯度/Adam 状态核算：
mixed_precision_total: 2P + 4P + 4P + 8P = 18P bytes

# 注意：框架可能不保留独立 master 参数，梯度 dtype 也可能不同，
# 所以 18P 是预算模型，不是所有实现的固定常数。

# 1. 对标准 .npy 文件使用 mmap_mode 懒加载
data = np.load('data.npy', mmap_mode='r')

# 2. pin_memory → 异步 CPU→GPU 传输
x = x.pin_memory().to(device, non_blocking=True)

# 3. 多个 worker 预取
loader = DataLoader(dataset, num_workers=4, pin_memory=True,
                    prefetch_factor=2, persistent_workers=True)

# 保存
checkpoint = {
    'model': model.state_dict(),
    'optimizer': optimizer.state_dict(),
    'step': step,
    'loss': loss.item(),
}
torch.save(checkpoint, f'checkpoint_{step}.pt')

# 加载
checkpoint = torch.load('checkpoint.pt')
model.load_state_dict(checkpoint['model'])
optimizer.load_state_dict(checkpoint['optimizer'])

混合精度训练

核心思想

import torch

for x, y in loader:
    # bf16 指数范围与 fp32 相同，通常不需要 GradScaler
    with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
        # 前向用 bf16——激活值只占一半内存
        output = model(x)
        loss = criterion(output, y)

    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

如果使用 float16，梯度更容易下溢，通常应使用 torch.amp.GradScaler("cuda")。对 bf16 无条件套用 GradScaler 会掩盖两种格式在动态范围上的关键差异。

内存节省

# 无混合精度 (fp32):
参数: 4P, 梯度: 4P, 优化器: 8P → 16P bytes

# 混合精度 (bf16 前向 + fp32 主参数):
bf16 参数: 2P
fp32 主参数: 4P
fp32 梯度: 4P
fp32 优化器: 8P
bf16 激活值: ~half → 显著节省！

总计参数相关: ~18P（稍多，但激活值节省巨大）

知识检查

如果只能记住 5 件事

1. 先把参数、梯度、optimizer state、activation 和临时 buffer 都换算成 bytes，再判断模型能否放下。
2. Transformer dense training 的一阶计算量可用 \(6PD\) 估算，但它不包含通信、数据管道、故障和评估时间。
3. MFU 把端到端吞吐与理论模型 FLOPs、硬件峰值连接起来；比较时必须说明模型、精度和硬件。
4. Arithmetic intensity 是 FLOPs / bytes；与 ridge point 的相对位置决定当前更可能受计算还是带宽限制。
5. PyTorch 的 tensor layout、view/copy、autograd 和 mixed precision 都会改变真实资源账本，不能只看数学公式。

延伸阅读