L0.3 数学与机器学习基础

三维坐标　layer: L0（基础）｜level: Engineer｜pillar: 训推框架

本文不是一篇数学复习课。目标是用 AI Infra 工程师的视角重读线代/概率与 ML 基础——搞清楚「为什么线性代数决定了底层算子是 GEMM 还是 reduction」「损失/梯度/反向传播在硬件上对应哪些计算」「一次训练 step 的 forward→loss→backward→optimizer.step 如何映射到显存占用」，从而为后续 L3 的 ZeRO / 显存优化埋下认知伏笔。

学习目标

前置知识：读过 L0.1/L0.2（知道「三大支柱」与三堵物理墙）；会写基础 Python，记得「矩阵乘」「求导」这两个词大概是什么即可。无需线代/概率高分。
学完产出：① 能把深度学习坍缩成 GEMM / reduction / element-wise 三类计算原语，并说清各自是 compute-bound 还是 memory-bound；② 会画一次训练 step 的 显存账本（参数 / 梯度 / 优化器状态 / 激活），并据此算出「一个 NB 模型要多少显存」；③ 理解为什么这张账本是 L3 ZeRO 切分的前置认知；④ 不靠任何框架，纯 NumPy 手写一个两层 MLP 跑通前向 + 反向，亲眼看 loss 收敛。
阅读姿势：盯住一句话——「框架只是把这些数学自动化了；你手推一遍，才真正读得懂 autograd 和 ZeRO 在干什么。」

背景与现状

很多人把「数学基础」当成算法岗的事，与 infra 无关。这是新人最大的误区。AI Infra 的所有硬核优化，本质都是在优化一组确定的数学运算在确定硬件上的执行效率——你不理解这组运算长什么样，就无从谈优化。

从 infra 视角看，整个深度学习可以坍缩成三类计算原语：

稠密矩阵乘（GEMM, General Matrix Multiply）：来自线性代数的 Y = XW。Transformer 里的 QKV 投影、FFN、注意力打分全是 GEMM，占据训练/推理 70%–90% 的 FLOPS，是 Tensor Core 的主战场。
规约（Reduction）：来自概率统计的 softmax、LayerNorm、sum/mean、All-Reduce 梯度同步。规约是**带宽受限（memory-bound）**的，不像 GEMM 那样吃算力，却常常是性能瓶颈。
逐元素（Element-wise）：激活函数 ReLU/GeLU、残差相加、dropout。几乎纯带宽受限，是 **kernel fusion（算子融合）**最爱合并的对象。

从产业演进看，对「数学 → 算子 → 硬件」这条链路的认知，经历了三个阶段：

2012–2017（AlexNet 后）：框架（Caffe/TensorFlow）把数学包成黑盒，工程师只写 model.fit()，不关心底层算子。
2017–2020（Transformer 普及）：模型变大后发现「同样的 FLOPS，GPU 利用率（MFU）只有 30%」，业界开始反向追问——是 GEMM 没打满 Tensor Core，还是 reduction/element-wise 把带宽吃光了？
2020 至今（大模型时代）：算子级认知成为 infra 核心竞争力。FlashAttention 把注意力的 softmax（reduction）与 GEMM 融合、避免反复读写 HBM，单这一个 idea 就重塑了整个长序列训练——而它的根，就是「线代决定算子形态、概率决定规约模式」这条本文要讲透的链路。

业界信号：FlashAttention 论文核心贡献不是新数学，而是重排了已知数学运算的访存顺序；ZeRO 的核心也不是新算法，而是重新分配了训练 step 中各类张量的显存归属。这印证一个事实——大模型 infra 的护城河，建立在对「最基础的数学运算如何落到硬件」的深刻理解之上。

📅 2026 时效：FlashAttention 已演进到 FA3（2024 年中发布，截至 2026 年仍处 beta/RC、生产多数仍用 FA2）——它进一步榨干 Hopper（H100） 的异步特性（warp-specialization + TMA + ping-pong 调度）与 FP8 低精度，FP16 比 FA2 快 1.5–2×。但它的「内核思想」从未变：重排访存、减少 HBM 往返，依旧是本文这条「数学→算子→访存」链路的延续。

原理与架构

2.1 线性代数速通：为什么一切都是 GEMM

神经网络的一层，数学上就是一次仿射变换加非线性：

\mathbf{h} = \sigma(\mathbf{x}\mathbf{W} + \mathbf{b})

其中 x 是输入（形状 [B, d_in]，B 为 batch），W 是权重（[d_in, d_out]），xW 就是一次 GEMM。这条最简单的式子带出三个 infra 关键认知：

batch 维 B 让矩阵-向量乘升级成矩阵-矩阵乘。单样本是 GEMV（向量乘，带宽受限、Tensor Core 吃不饱）；攒一个 batch 就变 GEMM（算力受限、能打满 Tensor Core）。这就是为什么推理要做 batching——不是为了省事，是为了把硬件利用率从个位数拉到 50%+。
矩阵形状决定硬件效率。Tensor Core 偏好维度对齐到 8/16/128 的矩阵；d_out=4097 会比 4096 慢得多。infra 工程师调 padding、对齐，本质是在迁就 GEMM 的硬件偏好。
数据布局（row/col major、转置）决定访存效率。W 还是 Wᵀ 直接影响 cache 命中——这是后续 L2「编程与编译」的入口。

2.2 概率统计速通：reduction 从哪来

分类任务的输出层用 softmax 把 logits 变成概率分布：

p_i = \frac{e^{z_i}}{\sum_j e^{z_j}}

分母那个 ∑ 就是 reduction——它必须读完整行才能算出一个标量。LayerNorm（算 mean/variance）、交叉熵损失、多卡训练的梯度 All-Reduce，本质都是规约。规约的特点是「读得多、算得少」，是带宽受限的，所以优化规约靠的不是堆算力，而是减少 HBM 往返（FlashAttention 的核心思路）。

概率统计还提供了训练的目标函数视角：最大似然估计（MLE）等价于最小化交叉熵，这就把「学习」变成了一个可微的优化问题——而优化的引擎，是梯度。

2.3 ML 基本概念：损失、梯度、过拟合

概念	数学含义	infra 视角的影响
损失函数 Loss	衡量预测与真值差距的标量，如 MSE / 交叉熵	一次 forward 的终点，是 backward 的起点；标量，但其计算图保存了全部中间激活
梯度 Gradient	损失对每个参数的偏导 `∂L/∂θ`	形状与参数完全相同——这意味着梯度显存 = 参数显存（FP32 下），是显存账本的一大项
反向传播 Backprop	用链式法则从 loss 反向逐层求梯度	需要 forward 阶段缓存的激活值，激活显存随 batch×序列长度线性膨胀
优化器状态	如 Adam 的一阶动量 m、二阶动量 v	每个参数额外 2 份状态——这是 Adam 比 SGD 吃 3 倍显存的根源，也是 ZeRO 要切分的大头
过拟合 Overfitting	训练集表现好、泛化差	靠正则/dropout/early-stop 缓解；dropout 是逐元素算子，影响 kernel fusion 边界

2.4 训练循环的数据流与显存归属

把上面的概念串成一次训练 step，它的标准四步是 forward → loss → backward → optimizer.step：

沿数据流读懂显存账本——这是连接「ML 基础」与 L3「显存优化」的关键桥梁。训练时显存被四类张量瓜分：

显存项	何时产生	大小量级（参数量为 Φ）	能否释放
模型参数 (Parameters)	常驻	FP16 下 `2Φ` 字节	否，常驻
梯度 (Gradients)	backward 阶段	`2Φ`（FP16）	optimizer.step 后可复用
优化器状态 (Optimizer States)	Adam 持有	`12Φ`（FP32 的 master weight + m + v）	否，常驻
激活值 (Activations)	forward 缓存供 backward 用	随 `batch × seq_len` 线性膨胀	backward 用完即释放

算一笔账就能看清：一个 1B 参数模型，光参数+梯度+Adam 状态（混合精度下经典的 2+2+12=16Φ）就要约 16GB——还没算激活值！这就是「显存墙」的算术根源。后续 L3 的 ZeRO 做的事，正是把上表中的「梯度 / 优化器状态 / 参数」沿数据并行维度切分到多卡，从而让单卡只持有 1/N。你现在理解的这张显存账本，就是读懂 ZeRO-1/2/3 的前置知识。

📅 2026 时效：在 PyTorch 生态里，ZeRO-3 的原生实现就是 FSDP——官方文档明确 FSDP 的「参数/梯度/优化器状态全分片」思想对应 DeepSpeed ZeRO Stage 3。当前新基线已是 FSDP2（fully_shard）：用 DTensor 做 per-parameter 分片，取代旧的 FlatParameter 方案，与 torch.compile、张量并行组合更顺。换句话说，「ZeRO」是思想，「FSDP/FSDP2」「DeepSpeed」是它在 PyTorch 里的两套主流落地。

动手实践：纯 NumPy 手写两层 MLP

实验目标：不依赖任何深度学习框架，用纯 NumPy 手写一个两层 MLP（隐层 ReLU + 输出 Sigmoid），在玩具数据集 make_moons 上手动实现前向、交叉熵损失、链式法则反向传播、SGD 更新，亲眼看着 loss 收敛、准确率上升。产出物：一条下降的 loss 曲线 + 最终分类准确率。本实验无需 GPU，纯 CPU NumPy 即可跑通——它的价值正是「亲手把模块 2 的数学落成代码」。

3.1 环境准备

# 推荐 Python 3.11；用 venv 或 uv 隔离
python3 -m venv .venv && source .venv/bin/activate
pip install numpy scikit-learn matplotlib
# 说明：sklearn 仅用于生成玩具数据 make_moons；matplotlib 仅用于画 loss 曲线。
# 核心 MLP 完全是手写 NumPy，不依赖任何框架。

3.2 代码：纯 NumPy 两层 MLP（前向 + 手动反向）

import numpy as np
from sklearn.datasets import make_moons

np.random.seed(0)

# ---------- 1) 玩具数据集：两个交错的半月，线性不可分 ----------
X, y = make_moons(n_samples=400, noise=0.2, random_state=0)
y = y.reshape(-1, 1).astype(np.float64)          # 形状 [N, 1]，二分类标签 0/1

# ---------- 2) 初始化两层网络参数 ----------
d_in, d_hid, d_out = 2, 16, 1
W1 = np.random.randn(d_in, d_hid) * np.sqrt(2.0 / d_in)   # He 初始化（配 ReLU）
b1 = np.zeros((1, d_hid))
W2 = np.random.randn(d_hid, d_out) * np.sqrt(1.0 / d_hid)
b2 = np.zeros((1, d_out))

def sigmoid(z):
    return 1.0 / (1.0 + np.exp(-z))

lr = 0.5
losses = []

for epoch in range(2000):
    # ===== ① Forward 前向（两次 GEMM + element-wise 激活）=====
    z1 = X @ W1 + b1            # GEMM: [N,2]·[2,16] -> [N,16]
    a1 = np.maximum(0, z1)      # ReLU（element-wise）
    z2 = a1 @ W2 + b2           # GEMM: [N,16]·[16,1] -> [N,1]
    a2 = sigmoid(z2)            # 输出概率

    # ===== ② Loss：二分类交叉熵（reduction，求和成标量）=====
    eps = 1e-8
    loss = -np.mean(y * np.log(a2 + eps) + (1 - y) * np.log(1 - a2 + eps))
    losses.append(loss)

    # ===== ③ Backward 反向：链式法则，逐层手推梯度 =====
    N = X.shape[0]
    # 交叉熵 + sigmoid 的组合梯度直接化简为 (a2 - y)
    dz2 = (a2 - y) / N          # [N,1]
    dW2 = a1.T @ dz2            # [16,1]  注意：梯度形状 == 参数形状
    db2 = dz2.sum(axis=0, keepdims=True)
    da1 = dz2 @ W2.T            # 把误差传回隐层 [N,16]
    dz1 = da1 * (z1 > 0)        # ReLU 的导数：z>0 处为 1，否则 0
    dW1 = X.T @ dz1             # [2,16]
    db1 = dz1.sum(axis=0, keepdims=True)

    # ===== ④ optimizer.step：朴素 SGD 参数更新 θ ← θ - lr·g =====
    W2 -= lr * dW2; b2 -= lr * db2
    W1 -= lr * dW1; b1 -= lr * db1

    if epoch % 200 == 0:
        acc = ((a2 > 0.5) == y).mean()
        print(f"epoch {epoch:4d} | loss {loss:.4f} | acc {acc:.3f}")

# ---------- 最终评估 ----------
pred = (sigmoid(np.maximum(0, X @ W1 + b1) @ W2 + b2) > 0.5)
print(f"final accuracy: {(pred == y).mean():.3f}")

# ---------- 画 loss 曲线 ----------
import matplotlib
matplotlib.use("Agg")          # 无显示环境也能保存图片
import matplotlib.pyplot as plt
plt.plot(losses); plt.xlabel("epoch"); plt.ylabel("loss")
plt.title("two-layer MLP on make_moons (pure NumPy)")
plt.savefig("loss_curve.png", dpi=120)
print("saved loss_curve.png")

3.3 运行与观察

python mlp_numpy.py

你会看到类似输出，loss 持续下降、准确率收敛到 0.95+：

epoch    0 | loss 0.7012 | acc 0.498
epoch  200 | loss 0.2864 | acc 0.880
epoch  400 | loss 0.1607 | acc 0.945
epoch 1800 | loss 0.0712 | acc 0.973
final accuracy: 0.975
saved loss_curve.png

把代码逐行对回前面的原理：两次 @ 就是 GEMM；np.maximum(0, z) 是 element-wise；交叉熵里的 np.mean 是 reduction；dW2 = a1.T @ dz2 印证了「梯度形状与参数完全相同」；a1、z1 被一直缓存到 backward——这就是激活显存的来源。你手写的这一切，框架（PyTorch autograd）只是替你自动化了。

3.4 GPU 路径：用 torch.autograd 对照验证

同样的逻辑，交给 PyTorch 自动微分，验证你手推的梯度无误（有 NVIDIA GPU 自动走 CUDA，无则 CPU）：

import torch, numpy as np
from sklearn.datasets import make_moons

dev = "cuda" if torch.cuda.is_available() else "cpu"
X_np, y_np = make_moons(400, noise=0.2, random_state=0)
X = torch.tensor(X_np, dtype=torch.float32, device=dev)
y = torch.tensor(y_np, dtype=torch.float32, device=dev).reshape(-1, 1)

net = torch.nn.Sequential(
    torch.nn.Linear(2, 16), torch.nn.ReLU(),
    torch.nn.Linear(16, 1), torch.nn.Sigmoid(),
).to(dev)
opt = torch.optim.SGD(net.parameters(), lr=0.5)
loss_fn = torch.nn.BCELoss()

for epoch in range(2000):
    opt.zero_grad()
    loss = loss_fn(net(X), y)   # forward + loss
    loss.backward()             # backward：autograd 自动算我们手推的那些梯度
    opt.step()                  # optimizer.step
print(f"[{dev}] final loss {loss.item():.4f}")

对照点：loss.backward() 这一行，自动完成了你在 3.2 里手写的全部链式法则。框架的本质，就是把前面这些数学打包成可复用、可在 GPU 上加速的算子图——理解了手写版，你才真正读得懂 autograd 在干什么。

踩坑预警 (Gotchas)

梯度没除 batch size：上面 dz2 = (a2 - y) / N 的 /N 必须有，对应 loss 用了 np.mean。漏掉就等于把学习率放大 N 倍，loss 直接 NaN。
log(0) 数值爆炸：交叉熵里务必加 eps（np.log(a2 + eps)），否则当 a2 逼近 0/1 时出现 -inf。
学习率过大震荡 / 过小不收敛：make_moons 上 lr=0.5 较稳；换数据务必先扫学习率，看 loss 曲线是平滑下降还是来回跳。
ReLU 死神经元：用 He 初始化（sqrt(2/d_in)）配 ReLU；若用普通小随机数初始化，隐层易全部死掉、梯度恒为 0。
PyTorch 对照忘了 opt.zero_grad()：autograd 默认累加梯度，每个 step 不清零会让梯度越滚越大——这是新手最常见的 bug。

深入思考

下面三题每题先给题干，再用 <details> 折叠一份图文并茂的参考答案。建议先合上答案自己想 3 分钟，再展开对照。

思考题 1：显存账本推演

一个 13B 参数模型，用 Adam + 混合精度（FP16 参数/梯度 + FP32 master/m/v）训练。请仅凭 2.4 节的显存账本，估算「参数 + 梯度 + 优化器状态」的常驻显存，并解释为什么单张 80GB A100 即使不算激活值也已逼近极限——这如何自然引出 L3 的 ZeRO 切分动机？

展开参考答案（含 13B 显存堆叠图 + 逐项算账）

结论：常驻显存 ≈ 16Φ ≈ 13B × 16 ≈ 208 GB，单张 80GB A100 连「参数+梯度+优化器」都放不下，必须切分。

逐项算账（混合精度经典 2+2+12=16Φ）：

显存项	系数	13B 实际	说明
FP16 参数	2Φ	~26 GB	前向/反向用
FP16 梯度	2Φ	~26 GB	反向产生
FP32 master weight	4Φ	~52 GB	优化器更新精度
Adam m（一阶动量）	4Φ	~52 GB	FP32
Adam v（二阶动量）	4Φ	~52 GB	FP32
常驻合计	16Φ	~208 GB	还没算激活！

为什么 80GB A100 逼近/突破极限：208 GB ≫ 80 GB，光常驻状态就要约 2.6 张 A100，再叠加随 batch × seq_len 线性膨胀的激活值，单卡彻底无解。

如何引出 ZeRO：这张账本里，梯度 + 优化器状态占了 14Φ（约 87.5%，其中优化器状态独占 12Φ / 75%），且它们在数据并行下「每张卡都存了一份完全相同的副本」——纯属浪费。ZeRO 的洞察正是：把梯度（ZeRO-2）、优化器状态（ZeRO-1）、乃至参数（ZeRO-3）沿数据并行维度切成 N 份，每卡只存 1/N。N=8 时，单卡常驻从 208 GB 砍到约 26 GB，激活值才有了喘息空间。你现在算的这张账本，就是 ZeRO-1/2/3 三级切分的动机本身。

思考题 2：算子受限类型判断

给定 Transformer 的三类计算——QKV 投影（GEMM）、softmax（reduction）、残差相加（element-wise）。哪些是算力受限（compute-bound）、哪些是带宽受限（memory-bound）？据此解释为什么 FlashAttention 选择「融合 softmax 与 attention GEMM、避免中间结果写回 HBM」而不是「换一个更快的矩阵乘算法」。

展开参考答案（含三类算子受限归属图）

结论：QKV 投影是 compute-bound，softmax 和残差相加是 memory-bound。FlashAttention 优化的恰恰是后者（访存），而不是前者（算力）。

为什么不是「换更快的矩阵乘」：注意力里的矩阵乘（QKᵀ、加权 V）本身已是 compute-bound、能打满 Tensor Core，提升空间有限。真正的浪费在于朴素实现把 S = QKᵀ（一个 [seq, seq] 的大中间矩阵）写回 HBM，再读回来做 softmax，再写回，再读回来乘 V——这一连串 HBM 往返是 memory-bound 的灾难，序列越长越严重（O(seq²) 的中间矩阵）。

FlashAttention 的真正招数：用 tiling（分块） 把 Q/K/V 切成能放进片上 SRAM 的小块，在 SRAM 里就地完成「打分 → online softmax → 加权」，整个 [seq, seq] 中间矩阵从头到尾不落 HBM。它没有发明新的矩阵乘，而是重排了访存顺序、消灭了 memory-bound 的中间往返——这正是模块 1「FlashAttention 的核心贡献不是新数学，而是重排访存」的精确含义。

思考题 3：batching 的数学动机

从 GEMV→GEMM 的升级出发，向一位「觉得推理就是单条 forward」的同学解释：为什么提高 batch size 能显著提升 GPU 利用率（MFU），代价又是什么（延迟 vs 吞吐、激活显存膨胀）？这与训练时调 batch size 的权衡有何异同？

展开参考答案（含 GEMV→GEMM 吞吐-延迟权衡图）

结论：单样本是 GEMV（memory-bound，Tensor Core 吃不饱）；攒 batch 变 GEMM（compute-bound，能打满算力）。batch 提吞吐的代价是单请求延迟上升 + 激活显存线性膨胀。

为什么 batch 提 MFU（数学动机）：batch=1 时 x·W 是矩阵-向量乘（GEMV）——读一整块权重 W 进来，只服务一个样本，算术强度极低、memory-bound、Tensor Core 大量空闲。攒成 batch=N 后 X·W 变成矩阵-矩阵乘（GEMM）——同一块权重搬一次、服务 N 个样本，算术强度提升约 N 倍，把硬件利用率从个位数拉到 50%+。这就是推理服务必须做 continuous batching 的根本原因。

代价：

延迟 vs 吞吐：要攒够 batch 才发车，单个请求要等同伴 → 吞吐↑但单请求延迟↑。在线服务要在「攒批窗口」和「P99 延迟」之间权衡。
激活显存膨胀：激活值 ∝ batch × seq_len，batch 越大越吃显存，可能 OOM——这又回到思考题 1 的显存账本。

与训练调 batch 的异同：

相同：都靠大 batch 把 GEMV 变 GEMM、提 MFU；都受激活显存约束。
不同：推理调 batch 是在线权衡延迟（用户在等），且 batch 由实时请求量决定（动态）；训练调 batch 影响的是收敛与泛化（大 batch 可能要配 warmup / 调 LR），且通常用梯度累积模拟更大 batch，不直接受单次延迟约束。

学习目标​

背景与现状​

原理与架构​

2.1 线性代数速通：为什么一切都是 GEMM​

2.2 概率统计速通：reduction 从哪来​

2.3 ML 基本概念：损失、梯度、过拟合​

2.4 训练循环的数据流与显存归属​

动手实践：纯 NumPy 手写两层 MLP​

3.1 环境准备​

3.2 代码：纯 NumPy 两层 MLP（前向 + 手动反向）​

3.3 运行与观察​

3.4 GPU 路径：用 torch.autograd 对照验证​

踩坑预警 (Gotchas)​

深入思考​

思考题 1：显存账本推演​

思考题 2：算子受限类型判断​

思考题 3：batching 的数学动机​

延伸阅读​