第2章：数学基础 - AIInfraGuide

AI Infra 不要求每天手推复杂定理，但要求看到公式时能立刻回答三个工程问题：张量是什么形状、需要多少计算和存储、怎样实现才稳定。

本章围绕这三个问题建立数学直觉。我们会从向量和矩阵开始，连接到 GEMM 与分块；从概率分布开始，连接到 Softmax、交叉熵和采样；从链式法则开始，连接到反向传播；最后讨论浮点数、误差和混合精度。目标不是证明所有结论，而是做到“公式能读、维度能推、数量级能估、数值风险能判断”。

📑 目录

1. 学习目标、符号与三类问题
2. 线性代数：张量、变换与分解
3. 分块矩阵与 GEMM 工程直觉
4. 概率论：从随机变量到模型分布
5. Softmax、交叉熵与 KL 散度
6. 微积分与反向传播
7. 优化、梯度稳定性与归一化
8. 数值计算与混合精度
9. AI Infra 综合算例
总结
自我检验清单
练习题
参考资料

1. 学习目标、符号与三类问题

完成本章后，你应该能够：

区分标量、向量、矩阵和高阶张量，并熟练推导运算后的 shape；
解释点积、范数、矩阵乘法、秩、特征分解与 SVD 的直觉；
把矩阵乘法拆成 block，理解 GEMM tiling 为什么数学等价；
计算期望、方差、协方差，理解常见概率分布与条件概率；
从 logits 得到稳定的 Softmax，理解温度、交叉熵、熵和 KL 散度；
用链式法则读懂计算图，手推线性层和 Softmax 交叉熵的梯度；
解释梯度消失/爆炸、残差连接和归一化对训练的影响；
区分 FP32、TF32、FP16、BF16 与 FP8 的精度和动态范围；
识别溢出、下溢、灾难性消减、非结合性等数值风险；
估算一个算子的参数量、激活量、FLOPs、访存量与算术强度。

1.1 本章符号

符号	含义
$x$	标量，或语境明确时的变量
$\mathbf{x}$	向量
$\mathbf{X}$	矩阵或高阶张量
$X_{ij}$	矩阵第 $i$ 行第 $j$ 列元素
$\mathbb{R}^{m \times n}$	$m$ 行、 $n$ 列的实矩阵集合
$\mathbf{X}^\top$	转置
$\mathbf{X}^{-1}$	逆矩阵（若存在）
$\lVert \mathbf{x} \rVert_2$	向量的 L2 范数
$\mathbb{E}[X]$	期望
$\operatorname{Var}(X)$	方差
$\nabla_x f$	$f$ 对 $x$ 的梯度
$\partial f / \partial x$	偏导数
$\odot$	逐元素乘法

深度学习代码常用 batch-first 记法：

$B$ ：batch size；
$S$ ：sequence length；
$H$ ：hidden size；
$V$ ：vocabulary size；
$N_h$ ：attention head 数；
$D_h = H / N_h$ ：每个 head 的维度。

1.2 读公式时先问三个问题

以语言模型输出投影为例：

\mathbf{Y} = \mathbf{X}\mathbf{W}, \qquad \mathbf{X} \in \mathbb{R}^{(BS) \times H},\quad \mathbf{W} \in \mathbb{R}^{H \times V}

第一问，形状是否匹配？内维 $H$ 相同，输出是 $(BS) \times V$ 。

第二问，代价是多少？每个输出元素做 $H$ 次乘加，总量约为：

2BSHV \quad \text{FLOPs}

这里把一次乘法和一次加法各计一个浮点操作。

第三问，实现风险是什么？权重很大、输出 logits 很大；需要考虑数据类型、内存布局、分块、并行策略和 Softmax 的数值稳定性。

这三问贯穿整个 AI Infra 技术栈。

2. 线性代数：张量、变换与分解

2.1 标量、向量、矩阵和张量

标量是一个数，如学习率 $\eta = 10^{-4}$ ；
向量是一维有序数组，如一个 token 的隐藏状态 $\mathbf{x} \in \mathbb{R}^{H}$ ；
矩阵是二维数组，如线性层权重 $\mathbf{W} \in \mathbb{R}^{H \times 4H}$ ；
张量是多维数组，如一批序列的隐藏状态 $\mathbf{X} \in \mathbb{R}^{B \times S \times H}$ 。

工程中的“张量”还带有数学 shape 之外的信息：

数据地址 + shape + stride + dtype + device + layout

两个张量 shape 相同，不代表内存布局相同；数值相同，也不代表 dtype 和误差特性相同。

2.2 Shape、索引和 stride

考虑一个 shape 为 $(2, 3)$ 的行主序矩阵：

\mathbf{X} = \begin{bmatrix} 1 & 2 & 3 \\ 4 & 5 & 6 \end{bmatrix}

底层连续存储为 [1, 2, 3, 4, 5, 6]，以元素为单位的 stride 是 $(3, 1)$ 。元素 $X_{ij}$ 的线性偏移为：

\text{offset}(i,j) = 3i + j

转置视图 $\mathbf{X}^\top$ 的 shape 是 $(3,2)$ ，stride 变成 $(1,3)$ ，通常无需复制数据。但按转置后的最后一维扫描时，访问不再连续。

reshape 是否复制取决于原 layout 能否用新 shape/stride 表示。看到 view、reshape、transpose、contiguous 时，都要同时思考逻辑维度和物理布局。

2.3 逐元素运算与广播

逐元素加法要求对应元素能配对：

\mathbf{C} = \mathbf{A} + \mathbf{B}, \qquad C_{ij} = A_{ij} + B_{ij}

广播允许某些长度为 1 或缺失的维度被逻辑扩展。例如：

\mathbf{X} \in \mathbb{R}^{B \times S \times H}, \qquad \mathbf{b} \in \mathbb{R}^{H}

$\mathbf{X}+\mathbf{b}$ 会在 batch 和 sequence 维复用同一个 bias，结果 shape 仍为 $(B,S,H)$ 。

从尾部维度向前对齐时，每对维度必须满足：

两者大小相等；或
至少一个大小为 1；或
某一方不存在该维度。

广播通常是逻辑视图，不会先复制出完整张量。但后续 Kernel 的访存和归约模式仍会受影响。更危险的是，错误 shape 也可能“恰好能广播”，代码不报错却算错语义。

2.4 点积：相似度与加权求和

两个 $n$ 维向量的点积：

\mathbf{x}^\top \mathbf{y} = \sum_{i=1}^{n}x_i y_i

它有两种常用直觉。

加权求和： $\mathbf{x}$ 是权重， $\mathbf{y}$ 是值。

方向相似性：

\mathbf{x}^\top \mathbf{y} = \lVert \mathbf{x}\rVert_2 \lVert \mathbf{y}\rVert_2 \cos\theta

因此余弦相似度为：

\cos\theta = \frac{\mathbf{x}^\top\mathbf{y}} {\lVert\mathbf{x}\rVert_2\lVert\mathbf{y}\rVert_2}

Attention 用 Query 与 Key 的点积衡量匹配程度；Embedding 检索常用归一化后的点积近似余弦相似度。

如果向量范数接近 0，分母会不稳定，工程实现通常加一个很小的 $\epsilon$ 。

2.5 范数：长度、误差与约束

常见向量范数：

\lVert \mathbf{x} \rVert_1 = \sum_i |x_i|

\lVert \mathbf{x} \rVert_2 = \sqrt{\sum_i x_i^2}

\lVert \mathbf{x} \rVert_\infty = \max_i |x_i|

矩阵的 Frobenius 范数把所有元素视为一个长向量：

\lVert \mathbf{A} \rVert_F = \sqrt{\sum_i\sum_j A_{ij}^2}

范数在工程中用于：

衡量参数、激活或梯度规模；
梯度裁剪；
比较参考实现与优化实现的误差；
正则化；
向量归一化。

比较浮点结果时，绝对误差和相对误差各有作用：

e_{abs} = |\hat{x}-x|,\qquad e_{rel} = \frac{|\hat{x}-x|}{|x|}

当参考值 $x$ 接近 0 时，相对误差会被放大，所以测试库通常组合绝对容差 atol 和相对容差 rtol：

|\hat{x}-x| \le \text{atol} + \text{rtol}\cdot |x|

2.6 矩阵乘法是一组点积

若：

\mathbf{A} \in \mathbb{R}^{M \times K}, \qquad \mathbf{B} \in \mathbb{R}^{K \times N}

则：

\mathbf{C} = \mathbf{A}\mathbf{B} \in \mathbb{R}^{M \times N}

每个元素为：

C_{ij} = \sum_{k=1}^{K} A_{ik}B_{kj}

把形状当作接口检查：

(M, K) @ (K, N) -> (M, N)

中间的 $K$ 被归约，外侧的 $M,N$ 保留下来。标准 GEMM 常写成：

\mathbf{C} \leftarrow \alpha\mathbf{A}\mathbf{B} + \beta\mathbf{C}

矩阵乘法一般不满足交换律： $\mathbf{A}\mathbf{B}$ 与 $\mathbf{B}\mathbf{A}$ 可能 shape 不同，或数值不同。但满足结合律：

(\mathbf{A}\mathbf{B})\mathbf{C} = \mathbf{A}(\mathbf{B}\mathbf{C})

数学上相等不代表浮点结果逐位相同，也不代表计算代价相同。选择不同结合顺序可以显著改变中间张量大小和 FLOPs。

2.7 Batched Matrix Multiplication

多头 Attention 中常见：

\mathbf{Q} \in \mathbb{R}^{B \times N_h \times S_q \times D_h}

\mathbf{K} \in \mathbb{R}^{B \times N_h \times S_k \times D_h}

转置最后两维后：

\mathbf{K}^\top \in \mathbb{R}^{B \times N_h \times D_h \times S_k}

批量矩阵乘法：

\mathbf{Q}\mathbf{K}^\top \in \mathbb{R}^{B \times N_h \times S_q \times S_k}

前两维 $(B,N_h)$ 是 batch 维，每个 (batch, head) 独立执行一个矩阵乘法。不要把“高阶张量乘法”想得过于神秘，它通常只是对最后两维做矩阵乘法，前面的维度负责批处理和广播。

2.8 线性变换与仿射变换

严格的线性变换满足：

f(a\mathbf{x}+b\mathbf{y}) = af(\mathbf{x})+bf(\mathbf{y})

矩阵乘法 $f(\mathbf{x})=\mathbf{W}\mathbf{x}$ 是线性变换。神经网络中所谓“Linear 层”通常还包含 bias：

\mathbf{y}=\mathbf{W}\mathbf{x}+\mathbf{b}

这在数学上是仿射变换。多个没有激活函数的线性/仿射层可以合并成一个，因此非线性激活是深层网络表达复杂函数的关键。

2.9 转置、单位矩阵与逆矩阵

转置交换行列：

(\mathbf{A}^\top)_{ij} = A_{ji}

常用规则：

(\mathbf{A}\mathbf{B})^\top = \mathbf{B}^\top\mathbf{A}^\top

单位矩阵 $\mathbf{I}$ 满足：

\mathbf{I}\mathbf{A}=\mathbf{A}\mathbf{I}=\mathbf{A}

方阵 $\mathbf{A}$ 若存在逆矩阵，则：

\mathbf{A}^{-1}\mathbf{A}= \mathbf{A}\mathbf{A}^{-1}=\mathbf{I}

求解线性方程 $\mathbf{A}\mathbf{x}=\mathbf{b}$ 时，数学上可写 $\mathbf{x}=\mathbf{A}^{-1}\mathbf{b}$ ，但数值实现通常不显式求逆，而使用 LU、QR、Cholesky 等分解直接求解。显式求逆往往计算更多、误差更大。

2.10 线性相关、张成空间与秩

如果一组向量中某个向量能由其他向量线性组合得到，它们线性相关。矩阵的秩（rank）可以理解为独立方向的数量：

\operatorname{rank}(\mathbf{A}) \le \min(M,N)

秩低意味着矩阵中的信息存在冗余，可以由更小的两个矩阵近似：

\mathbf{A} \approx \mathbf{U}\mathbf{V},\qquad \mathbf{U}\in\mathbb{R}^{M\times r}, \mathbf{V}\in\mathbb{R}^{r\times N}, \quad r \ll \min(M,N)

这会把参数量从 $MN$ 降到 $r(M+N)$ ，也是低秩适配和模型压缩的数学基础。

2.11 特征值与特征向量

对方阵 $\mathbf{A}$ ，如果存在非零向量 $\mathbf{v}$ 和标量 $\lambda$ 使：

\mathbf{A}\mathbf{v}=\lambda\mathbf{v}

则 $\mathbf{v}$ 是特征向量， $\lambda$ 是对应特征值。直觉上，矩阵变换在这个方向上只缩放，不改变方向。

特征值可帮助理解：

线性动态系统是否放大/衰减；
Hessian 在不同方向的曲率；
梯度经过多层 Jacobian 连乘时为什么爆炸或消失；
协方差矩阵的主要变化方向（PCA）。

不是所有矩阵都能在实数域下良好特征分解；工程中应根据矩阵结构选择合适算法，而不是把 eig 当万能工具。

2.12 奇异值分解（SVD）

任意实矩阵 $\mathbf{A}\in\mathbb{R}^{M\times N}$ 都可以写成：

\mathbf{A} = \mathbf{U}\mathbf{\Sigma}\mathbf{V}^\top

其中：

$\mathbf{U}$ 的列是左奇异向量；
$\mathbf{V}$ 的列是右奇异向量；
$\mathbf{\Sigma}$ 对角线上的 $\sigma_1\ge\sigma_2\ge\cdots\ge0$ 是奇异值。

保留最大的前 $r$ 个奇异值：

\mathbf{A}_r = \mathbf{U}_{:,1:r} \mathbf{\Sigma}_{1:r,1:r} \mathbf{V}_{:,1:r}^\top

$\mathbf{A}_r$ 是一种最优的 rank- $r$ 近似（在常见的 2-范数和 Frobenius 范数意义下）。直觉上，小奇异值对应影响较弱的方向，可以舍弃以换取压缩。

SVD 本身也有计算和内存成本。大型模型权重的实际压缩通常使用截断、随机化算法或直接训练低秩因子，而不是对所有巨大矩阵做完整 SVD。

2.13 LoRA 的低秩更新

对预训练权重 $\mathbf{W}_0\in\mathbb{R}^{d_{out}\times d_{in}}$ ，LoRA 冻结 $\mathbf{W}_0$ ，只训练低秩更新：

\mathbf{W} = \mathbf{W}_0 + \Delta\mathbf{W}

\Delta\mathbf{W} = \frac{\alpha}{r}\mathbf{B}\mathbf{A}

其中：

\mathbf{A}\in\mathbb{R}^{r\times d_{in}},\qquad \mathbf{B}\in\mathbb{R}^{d_{out}\times r}

新增参数量从完整更新的 $d_{out}d_{in}$ 变为：

r(d_{in}+d_{out})

例如 $d_{in}=d_{out}=4096$ 、 $r=16$ ：

\frac{16(4096+4096)}{4096^2} = 0.0078125 \approx 0.78\%

低秩只减少可训练参数与相应优化器状态，不自动消除基座权重、激活和前向计算成本。部署时可选择合并权重，或保留 adapter 动态切换，两者有不同的显存、延迟和灵活性权衡。

3. 分块矩阵与 GEMM 工程直觉

3.1 为什么可以分块计算

把矩阵沿 $M$ 、 $N$ 和归约维 $K$ 切成块。以 $2\times2$ block 为例：

\mathbf{A}= \begin{bmatrix} \mathbf{A}_{11} & \mathbf{A}_{12}\\ \mathbf{A}_{21} & \mathbf{A}_{22} \end{bmatrix},\qquad \mathbf{B}= \begin{bmatrix} \mathbf{B}_{11} & \mathbf{B}_{12}\\ \mathbf{B}_{21} & \mathbf{B}_{22} \end{bmatrix}

则：

\mathbf{C}_{11} = \mathbf{A}_{11}\mathbf{B}_{11} + \mathbf{A}_{12}\mathbf{B}_{21}

\mathbf{C}_{12} = \mathbf{A}_{11}\mathbf{B}_{12} + \mathbf{A}_{12}\mathbf{B}_{22}

其他块同理。每个输出块沿 $K$ 方向累加若干局部乘积。因此 GPU Kernel 可以：

取一小块 $\mathbf{A}$ 和一小块 $\mathbf{B}$ ；
搬到片上共享内存/寄存器；
重复使用这些数据计算多个输出元素；
沿 $K$ 方向迭代并累加；
最后把输出块写回 HBM。

这就是 tiling 的数学依据。它改变计算顺序和数据搬运，不改变目标公式。

3.2 朴素 GEMM 的计算与访存

对于 $(M,K)@(K,N)$ ：

输出元素数： $MN$ ；
每个元素执行 $K$ 次乘加；
FLOPs 约为 $2MKN$ 。

朴素实现若每次乘加都从全局内存加载 $A_{ik}$ 和 $B_{kj}$ ，同一个元素会被大量重复读取。GEMM 的高性能来自数据复用：

$A$ 的一个元素被同一输出行的多个列复用；
$B$ 的一个元素被同一输出列的多个行复用；
累加器在寄存器中复用 $K$ 次。

3.3 一个简化的算术强度估算

算术强度定义为：

\text{Arithmetic Intensity} = \frac{\text{FLOPs}}{\text{bytes transferred from target memory}}

理想情况下，FP16 GEMM 每个输入矩阵只从 HBM 读一次，输出写一次，粗略字节数为：

2(MK+KN+MN)

于是：

AI \approx \frac{2MKN}{2(MK+KN+MN)}

实际还有 cache、对齐、读改写、临时结果和融合操作。这个估算不是精确性能模型，但能帮助判断算子更可能受算力还是带宽限制。

3.4 tile 不是越大越好

更大的 tile 提高数据复用，却会消耗更多：

共享内存；
寄存器；
每个 block 的线程和同步；
边界处理成本。

资源占用过高会降低一个 SM 同时驻留的 block/warp 数量，影响延迟隐藏。tile 选择是在数据复用、占用率、指令效率和形状适配之间折中。

3.5 尾块与 padding

当 $M,N,K$ 不是 tile size 的整数倍时，最后一块会越过有效边界。常见做法：

加边界判断/掩码；
把输入 padding 到对齐尺寸；
为常见整齐 shape 提供快路径，其他 shape 走通用路径。

padding 会增加计算和存储，分支会增加控制开销。哪种方式更快取决于 shape 和硬件。

3.6 浮点分块结果为什么可能不同

实数加法满足结合律，但浮点加法不严格满足：

(a+b)+c \ne a+(b+c)

不同 tile、线程归约树、Tensor Core 路径会改变累加顺序，所以优化前后结果可能不逐位一致。正确性验证应使用适合 dtype 和问题规模的误差容限，并关注误差是否随归约长度系统性放大。

4. 概率论：从随机变量到模型分布

4.1 随机变量和概率分布

随机变量把随机结果映射为数。离散随机变量用概率质量函数：

p_X(x)=P(X=x),\qquad \sum_x p_X(x)=1

连续随机变量用概率密度函数：

p_X(x)\ge0,\qquad \int_{-\infty}^{\infty}p_X(x)\,dx=1

连续变量在单点的概率通常为 0，区间概率由密度积分得到。模型里常见分布：

分布	取值	典型场景
Bernoulli	$0/1$	二分类、Dropout mask
Categorical	$1,\ldots,V$	下一个 token 分布
Uniform	区间/有限集合	初始化、随机采样
Gaussian	实数	初始化、噪声、近似分析

4.2 联合概率、边缘概率与条件概率

联合概率描述多个事件同时发生：

P(X=x,Y=y)

对另一个变量求和或积分得到边缘概率：

P(X=x)=\sum_y P(X=x,Y=y)

条件概率：

P(X=x\mid Y=y) = \frac{P(X=x,Y=y)}{P(Y=y)}

自回归语言模型用链式法则分解序列概率：

P(x_1,x_2,\ldots,x_T) = \prod_{t=1}^{T}P(x_t\mid x_1,\ldots,x_{t-1})

这解释了训练时的 next-token prediction 和推理时逐 token 生成为什么是同一概率模型的两个过程。

4.3 独立与条件独立

若：

P(X,Y)=P(X)P(Y)

则 $X,Y$ 独立。独立比“不相关”更强；协方差为 0 不一定独立。

给定 $Z$ 后条件独立写作：

P(X,Y\mid Z)=P(X\mid Z)P(Y\mid Z)

分布式采样、数据并行和统计估计常隐含独立同分布（i.i.d.）假设。真实数据中的重复样本、分片偏差或序列相关性可能破坏该假设。

4.4 Bayes 公式

P(X\mid Y) = \frac{P(Y\mid X)P(X)}{P(Y)}

其中 $P(X)$ 是先验， $P(Y\mid X)$ 是似然， $P(X\mid Y)$ 是后验。虽然常规 LLM 训练不直接手算 Bayes 公式，但它是概率推断、参数估计和不确定性建模的基础。

4.5 期望：概率加权平均

离散变量：

\mathbb{E}[X]=\sum_x xP(X=x)

连续变量：

\mathbb{E}[X]=\int x p(x)\,dx

期望具有线性性质，不要求变量独立：

\mathbb{E}[aX+bY] = a\mathbb{E}[X]+b\mathbb{E}[Y]

训练目标通常是对数据分布的期望损失：

\min_\theta \mathbb{E}_{(x,y)\sim p_{data}} [\mathcal{L}(f_\theta(x),y)]

实际无法遍历真实分布，于是用 mini-batch 样本均值近似。

4.6 方差与标准差

方差衡量围绕均值的波动：

\operatorname{Var}(X) = \mathbb{E}[(X-\mu)^2] = \mathbb{E}[X^2]-\mu^2

标准差与原变量单位相同：

\sigma = \sqrt{\operatorname{Var}(X)}

若 $X,Y$ 独立：

\operatorname{Var}(X+Y) = \operatorname{Var}(X)+\operatorname{Var}(Y)

训练中的梯度噪声、初始化尺度、归一化和量化误差都与方差有关。

4.7 协方差与相关系数

协方差衡量两个变量共同变化：

\operatorname{Cov}(X,Y) = \mathbb{E}[(X-\mu_X)(Y-\mu_Y)]

标准化得到相关系数：

\rho_{X,Y} = \frac{\operatorname{Cov}(X,Y)}{\sigma_X\sigma_Y}

对向量随机变量，协方差矩阵为：

\mathbf{\Sigma} = \mathbb{E}[(\mathbf{x}-\boldsymbol{\mu}) (\mathbf{x}-\boldsymbol{\mu})^\top]

它是对称半正定矩阵。PCA 对协方差矩阵做特征分解，找出数据方差最大的正交方向。

4.8 样本均值、方差与 mini-batch

给定样本 $x_1,\ldots,x_n$ ，样本均值：

\bar{x}=\frac{1}{n}\sum_{i=1}^{n}x_i

常用无偏样本方差：

s^2=\frac{1}{n-1}\sum_{i=1}^{n}(x_i-\bar{x})^2

深度学习算子中的“方差”是否除以 $n$ 还是 $n-1$ 取决于具体定义。LayerNorm 通常使用总体方差形式（除以元素数量），统计库的默认值可能不同。复现算子时要查 API 语义，不要只看名字。

4.9 最大似然与负对数似然

给定数据 $D=\{(x_i,y_i)\}$ ，最大似然估计选择让观测数据最可能的参数：

\theta^* = \arg\max_\theta \prod_i p_\theta(y_i\mid x_i)

乘积容易数值下溢，且不便求导。取对数把乘积变成求和：

\theta^* = \arg\max_\theta \sum_i \log p_\theta(y_i\mid x_i)

等价地最小化负对数似然（NLL）：

\mathcal{L}_{NLL} = -\sum_i \log p_\theta(y_i\mid x_i)

对数是数值计算和概率建模之间的重要桥梁。

5. Softmax、交叉熵与 KL 散度

5.1 Logits 不是概率

模型最后一层输出 $V$ 个任意实数：

\mathbf{z}=(z_1,z_2,\ldots,z_V)

这些值叫 logits。它们可以为负，也不要求和为 1。Softmax 把它们转换为分类分布：

p_i = \operatorname{softmax}(\mathbf{z})_i = \frac{e^{z_i}}{\sum_{j=1}^{V}e^{z_j}}

显然 $p_i>0$ 且 $\sum_i p_i=1$ 。

Softmax 对统一平移不敏感：

\operatorname{softmax}(\mathbf{z}+c) = \operatorname{softmax}(\mathbf{z})

因为分子分母都会乘以 $e^c$ 。这个性质正是稳定实现的依据。

5.2 数值稳定的 Softmax

若直接计算 $e^{1000}$ ，有限精度浮点数会溢出。令：

m=\max_j z_j

稳定形式为：

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

最大的指数输入为 0，因此最大指数值为 1；其他值不大于 1。伪代码：

def stable_softmax(x):
    maximum = max(x)
    exps = [exp(value - maximum) for value in x]
    denominator = sum(exps)
    return [value / denominator for value in exps]

减最大值防止正向溢出，但很小的项仍可能下溢到 0。这通常代表它相对最大项确实可以忽略；如果后续还要取对数，则应直接使用稳定的 log_softmax，避免先变成 0 再 log(0)。

5.3 LogSumExp

定义：

\operatorname{LSE}(\mathbf{z}) = \log\sum_j e^{z_j}

稳定形式：

\operatorname{LSE}(\mathbf{z}) = m + \log\sum_j e^{z_j-m}

于是：

\log\operatorname{softmax}(\mathbf{z})_i = z_i - \operatorname{LSE}(\mathbf{z})

交叉熵实现通常融合 log_softmax + NLLLoss，既减少中间张量和访存，也避免不稳定的“先 Softmax 再取 log”。

5.4 温度参数

带温度 $T>0$ 的 Softmax：

p_i(T)= \frac{e^{z_i/T}}{\sum_j e^{z_j/T}}

$T<1$ ：放大 logit 差异，分布更尖锐；
$T>1$ ：缩小差异，分布更平坦；
$T\to0^+$ ：逐渐接近只选择最大 logit；
$T\to\infty$ ：逐渐接近均匀分布。

温度改变的是采样分布，不等于修改模型权重。工程实现应先缩放 logits，再应用稳定 Softmax，并谨慎处理非常小的温度。

5.5 交叉熵

真实分布 $\mathbf{q}$ 与模型分布 $\mathbf{p}$ 的交叉熵：

H(\mathbf{q},\mathbf{p}) = -\sum_i q_i\log p_i

若真实标签是 one-hot，正确类别为 $y$ ：

q_y=1,\quad q_{i\ne y}=0

则：

\mathcal{L} = -\log p_y

模型给正确类别越高概率，损失越小。若 $p_y=1$ ，损失为 0；若 $p_y$ 接近 0，损失很大。

5.6 信息熵

分布自身的熵：

H(\mathbf{p}) = -\sum_i p_i\log p_i

熵衡量不确定性：

全部概率集中在一个类别时，熵最小；
$V$ 个类别均匀分布时，熵最大，为 $\log V$ 。

对数底决定单位：自然对数对应 nat，以 2 为底对应 bit。机器学习损失通常使用自然对数。

5.7 KL 散度

从分布 $\mathbf{q}$ 到 $\mathbf{p}$ 的 KL 散度：

D_{KL}(\mathbf{q}\lVert\mathbf{p}) = \sum_i q_i\log\frac{q_i}{p_i}

它满足：

D_{KL}(\mathbf{q}\lVert\mathbf{p})\ge0

且两分布相同时为 0。但它通常不对称：

D_{KL}(\mathbf{q}\lVert\mathbf{p}) \ne D_{KL}(\mathbf{p}\lVert\mathbf{q})

所以 KL 散度不是严格意义的距离。

交叉熵可分解为：

H(\mathbf{q},\mathbf{p}) = H(\mathbf{q}) + D_{KL}(\mathbf{q}\lVert\mathbf{p})

训练时真实分布 $\mathbf{q}$ 固定， $H(\mathbf{q})$ 与模型参数无关，所以最小化交叉熵等价于最小化相应 KL 散度。

KL 散度会出现在知识蒸馏、分布匹配、RLHF/PPO 约束和推测解码分析中。实现时要明确 API 接收的是概率、log 概率还是 logits，以及 KL 的方向。

5.8 Perplexity

若平均 token 负对数似然为 $\bar{L}$ ，困惑度：

\operatorname{PPL}=e^{\bar{L}}

它可直觉理解为模型在每一步面对的“有效候选数”。但不同 tokenizer、数据预处理、上下文长度和是否忽略特殊 token 都会影响 PPL，不能脱离评测设置直接横比。

5.9 Top-k 与 Top-p 采样

模型输出分布后，解码策略决定如何选 token：

Greedy：选择最大概率 token；
Top-k：只在概率最高的 $k$ 个 token 中采样；
Top-p（nucleus）：选择累计概率至少达到 $p$ 的最小候选集合，再归一化采样。

Top-p 集合大小会随分布尖锐程度动态变化。实现时通常先过滤 logits，把被排除项设为负无穷，再执行稳定 Softmax 和采样。

分布式推理若 vocabulary 被张量并行切分，global top-k/top-p 需要跨设备聚合局部候选或使用等价的分布式算法。

6. 微积分与反向传播

6.1 导数：局部变化率

单变量函数的导数：

f'(x)= \lim_{h\to0}\frac{f(x+h)-f(x)}{h}

导数表示在 $x$ 附近，输入微小变化如何影响输出：

f(x+\Delta x) \approx f(x)+f'(x)\Delta x

常用导数：

\frac{d}{dx}x^n=nx^{n-1}

\frac{d}{dx}e^x=e^x

\frac{d}{dx}\log x=\frac{1}{x}

6.2 偏导数与梯度

多变量标量函数 $f(x_1,\ldots,x_n)$ 对单个变量的变化率是偏导。把所有偏导排成向量得到梯度：

\nabla_{\mathbf{x}}f= \begin{bmatrix} \frac{\partial f}{\partial x_1}\\ \vdots\\ \frac{\partial f}{\partial x_n} \end{bmatrix}

梯度指向函数上升最快的方向，因此梯度下降沿负梯度更新：

\mathbf{x}_{t+1} = \mathbf{x}_t-\eta\nabla f(\mathbf{x}_t)

$\eta$ 是学习率。

6.3 Jacobian 与 Hessian

向量函数 $\mathbf{y}=f(\mathbf{x})$ 的 Jacobian：

\mathbf{J}_{ij} = \frac{\partial y_i}{\partial x_j}

它描述每个输入分量对每个输出分量的局部影响。

标量函数的 Hessian 是二阶偏导矩阵：

\mathbf{H}_{ij} = \frac{\partial^2 f} {\partial x_i\partial x_j}

Hessian 描述局部曲率。深度学习参数量太大，通常不会显式构造完整 Hessian，但 Hessian-vector product、曲率近似和特征值分析仍用于优化与稳定性研究。

6.4 链式法则

若：

y=f(u),\qquad u=g(x)

则：

\frac{dy}{dx} = \frac{dy}{du}\frac{du}{dx}

多层网络就是复合函数。对：

\mathbf{h}_1=f_1(\mathbf{x}),\quad \mathbf{h}_2=f_2(\mathbf{h}_1),\quad L=f_3(\mathbf{h}_2)

梯度从损失反向传播：

\frac{\partial L}{\partial \mathbf{x}} = \frac{\partial L}{\partial \mathbf{h}_2} \frac{\partial \mathbf{h}_2}{\partial \mathbf{h}_1} \frac{\partial \mathbf{h}_1}{\partial \mathbf{x}}

概念上是 Jacobian 连乘；实际自动微分不会显式创建巨大 Jacobian，而是高效计算 vector-Jacobian product（反向模式）。

6.5 计算图与反向模式自动微分

设：

a=xy,\qquad b=a+x,\qquad L=b^2

前向阶段保存必要中间量。反向从：

\frac{\partial L}{\partial L}=1

开始：

\frac{\partial L}{\partial b}=2b

因为 $b=a+x$ ：

\frac{\partial L}{\partial a}=2b

$x$ 有两条路径影响 $L$ ，梯度需要相加：

\frac{\partial L}{\partial x} = \frac{\partial L}{\partial b}\frac{\partial b}{\partial x} + \frac{\partial L}{\partial a}\frac{\partial a}{\partial x} = 2b + 2by

\frac{\partial L}{\partial y}=2bx

这解释了计算图中一个张量被多处使用时为什么要累积梯度。

6.6 线性层的反向传播

采用 batch 行向量约定：

\mathbf{Y}=\mathbf{X}\mathbf{W}+\mathbf{b}

其中：

\mathbf{X}\in\mathbb{R}^{B\times d_{in}},\quad \mathbf{W}\in\mathbb{R}^{d_{in}\times d_{out}},\quad \mathbf{Y}\in\mathbb{R}^{B\times d_{out}}

给定上游梯度：

\mathbf{G}=\frac{\partial L}{\partial \mathbf{Y}} \in\mathbb{R}^{B\times d_{out}}

则：

\frac{\partial L}{\partial \mathbf{X}} = \mathbf{G}\mathbf{W}^\top \in\mathbb{R}^{B\times d_{in}}

\frac{\partial L}{\partial \mathbf{W}} = \mathbf{X}^\top\mathbf{G} \in\mathbb{R}^{d_{in}\times d_{out}}

\frac{\partial L}{\partial \mathbf{b}} = \sum_{i=1}^{B}\mathbf{G}_{i,:} \in\mathbb{R}^{d_{out}}

只看 shape 就能快速检查公式：

dX: (B, dout) @ (dout, din) -> (B, din)
dW: (din, B) @ (B, dout) -> (din, dout)
db: 对 batch 维归约 -> (dout,)

一个线性层的反向通常包含两次 GEMM 和一次归约，这也是训练算力显著高于单次前向的原因之一。

6.7 Softmax 的 Jacobian

Softmax：

p_i=\frac{e^{z_i}}{\sum_k e^{z_k}}

其偏导：

\frac{\partial p_i}{\partial z_j} = p_i(\delta_{ij}-p_j)

写成矩阵：

\mathbf{J}_{softmax} = \operatorname{diag}(\mathbf{p}) - \mathbf{p}\mathbf{p}^\top

Softmax 每个输出依赖所有输入，因此 Jacobian 不是对角矩阵。但实际反向仍可在 $O(V)$ 时间内计算，无需物化 $V\times V$ 矩阵。

6.8 Softmax + 交叉熵的简洁梯度

one-hot 标签 $\mathbf{y}$ ，损失：

L=-\sum_i y_i\log p_i

与 Softmax 联合求导后得到：

\frac{\partial L}{\partial z_i}=p_i-y_i

这是深度学习最重要的简化之一。对正确类别，梯度是 $p_y-1$ ；对其他类别，梯度是 $p_i$ 。预测越错，推动 logits 修正的力度越大。

融合交叉熵 Kernel 可以共享最大值和归约结果、减少中间概率张量，并直接生成所需梯度。

6.9 梯度检查

自动微分实现可用有限差分做小规模检查：

\frac{\partial f}{\partial x_i} \approx \frac{f(\mathbf{x}+\epsilon\mathbf{e}_i) -f(\mathbf{x}-\epsilon\mathbf{e}_i)}{2\epsilon}

中心差分通常比单边差分准确，但 $\epsilon$ 不能太大或太小：太大有截断误差，太小会被浮点舍入误差淹没。梯度检查适合小尺寸、双精度和确定性函数，不适合直接检查带随机性或巨大张量的完整训练任务。

7. 优化、梯度稳定性与归一化

7.1 全量梯度、SGD 与 mini-batch

经验风险：

L(\theta)=\frac{1}{N}\sum_{i=1}^{N}L_i(\theta)

全量梯度每步使用所有样本，成本高。mini-batch 梯度：

\hat{g} = \frac{1}{B}\sum_{i\in\mathcal{B}} \nabla_\theta L_i(\theta)

是全量梯度的随机估计。基本 SGD 更新：

\theta_{t+1}=\theta_t-\eta\hat{g}_t

更大 batch 往往降低梯度估计方差并提高硬件利用率，但需要更多激活内存，且优化超参数可能要相应调整。

7.2 动量

一种常见动量形式：

\mathbf{v}_t = \beta\mathbf{v}_{t-1}+(1-\beta)\mathbf{g}_t

\theta_{t+1} = \theta_t-\eta\mathbf{v}_t

动量对历史梯度做指数加权平均，减小短期噪声，在持续一致的方向上加速。

7.3 Adam 的尺度自适应直觉

Adam 同时维护一阶矩和二阶矩的指数移动平均：

\mathbf{m}_t = \beta_1\mathbf{m}_{t-1} +(1-\beta_1)\mathbf{g}_t

\mathbf{v}_t = \beta_2\mathbf{v}_{t-1} +(1-\beta_2)\mathbf{g}_t^2

偏差修正后：

\theta_{t+1} = \theta_t -\eta\frac{\hat{\mathbf{m}}_t} {\sqrt{\hat{\mathbf{v}}_t}+\epsilon}

平方和开方均逐元素进行。Adam 为每个参数保存额外状态；若状态使用 FP32，内存开销往往是大模型训练的重要部分，也因此催生了 ZeRO 等分片技术。

7.4 梯度消失与爆炸

深层复合函数的梯度包含许多 Jacobian 连乘：

\frac{\partial L}{\partial \mathbf{h}_0} = \frac{\partial L}{\partial \mathbf{h}_L} \prod_{l=1}^{L} \frac{\partial \mathbf{h}_l} {\partial \mathbf{h}_{l-1}}

若这些变换在相关方向上的尺度长期小于 1，梯度指数衰减；长期大于 1，梯度指数增长。

工程症状包括：

梯度范数趋近 0，早期层几乎不更新；
梯度范数突然巨大，loss 变为 inf/nan；
混合精度下更早触发下溢或溢出。

缓解手段包括合理初始化、残差连接、归一化、梯度裁剪、合适学习率和稳定的数据类型。

7.5 残差连接为什么帮助梯度传播

残差块：

\mathbf{y}=\mathbf{x}+F(\mathbf{x})

其 Jacobian：

\frac{\partial \mathbf{y}}{\partial \mathbf{x}} = \mathbf{I}+\frac{\partial F}{\partial \mathbf{x}}

恒等路径提供了一条不完全依赖 $F$ 的梯度通路。它不保证永远不会梯度爆炸/消失，但显著改善深层网络的优化条件。

7.6 初始化与方差传播

考虑：

y=\sum_{i=1}^{n}w_i x_i

若 $w_i,x_i$ 独立、均值为 0：

\operatorname{Var}(y) = n\operatorname{Var}(w)\operatorname{Var}(x)

为了让各层激活方差大致稳定，可令 $\operatorname{Var}(w)$ 与 $1/n$ 同阶。Xavier、Kaiming 初始化会根据 fan-in、fan-out 和激活函数选择尺度。

初始化不是越小越稳定：过小会让信号和梯度衰减，过大则导致饱和或爆炸。

7.7 LayerNorm

对一个 token 的隐藏向量 $\mathbf{x}\in\mathbb{R}^{H}$ ：

\mu=\frac{1}{H}\sum_{i=1}^{H}x_i

\sigma^2=\frac{1}{H} \sum_{i=1}^{H}(x_i-\mu)^2

\operatorname{LayerNorm}(x_i) = \gamma_i \frac{x_i-\mu}{\sqrt{\sigma^2+\epsilon}} +\beta_i

LayerNorm 沿隐藏维归约，不依赖 batch 中其他样本，因此适合可变序列和自回归模型。

从 Kernel 角度，它至少涉及均值归约、方差归约和逐元素变换。朴素实现多次读写 HBM；优化实现会融合归约与仿射变换，并处理长向量的并行归约。

7.8 RMSNorm

RMSNorm 不减均值，只按均方根缩放：

\operatorname{RMS}(\mathbf{x}) = \sqrt{\frac{1}{H}\sum_i x_i^2+\epsilon}

\operatorname{RMSNorm}(x_i) = \gamma_i\frac{x_i}{\operatorname{RMS}(\mathbf{x})}

它减少了均值相关计算，但“公式更少”不必然按同比例提升端到端速度，实际收益取决于融合、访存和模型整体瓶颈。

7.9 梯度裁剪

全局 L2 范数裁剪：

\mathbf{g}\leftarrow \mathbf{g}\cdot \min\left(1,\frac{c}{\lVert\mathbf{g}\rVert_2+\epsilon}\right)

当梯度范数超过阈值 $c$ 时，整体按比例缩小，方向不变。分布式训练中，若梯度被分片，计算全局范数需要跨 rank 汇总局部平方和：

\lVert\mathbf{g}\rVert_2 = \sqrt{\sum_r\sum_{i\in r}g_i^2}

因此看似简单的数学操作也会引入集合通信。

8. 数值计算与混合精度

8.1 浮点数是有限集合

典型二进制浮点数表示为：

(-1)^s \times \text{significand} \times 2^{\text{exponent}}

位数被分给符号、指数和尾数（有效数字）。指数位决定动态范围，尾数位决定相邻可表示数的精细程度。

格式	总位数	指数位	尾数字段位	主要特点
FP32	32	8	23	范围和精度较均衡
TF32	19（乘法输入语义）	8	10	保留 FP32 范围，降低乘法精度，常用 FP32 累加
FP16	16	5	10	精度尚可但动态范围较小
BF16	16	8	7	接近 FP32 动态范围，精度低于 FP16
FP8 E4M3	8	4	3	精度优先、范围较小，具体编码依实现规范
FP8 E5M2	8	5	2	范围优先、精度更低，具体编码依实现规范

表中不把隐含前导位计入尾数字段。TF32 通常是 NVIDIA Tensor Core 的计算模式，不是常规内存存储 dtype。

8.2 精度和动态范围是两件事

FP16 尾数字段比 BF16 多，因此同一数量级下通常能表示更细的差别；但 FP16 指数位少，最大有限值只有 65504，非零数的范围也窄得多。

BF16 牺牲一部分精度换取与 FP32 相同的指数位数，因此训练中通常更不易溢出/下溢。它仍然会有明显舍入误差，尤其在把很小增量加到很大数上时。

8.3 舍入与机器精度

很多十进制小数不能被二进制浮点精确表示：

>>> 0.1 + 0.2 == 0.3
False

浮点数间距随数值绝对值增大而增大。若 $|\delta|$ 远小于当前 $x$ 附近的可表示间距：

\operatorname{fl}(x+\delta)=x

这就是为何优化器常保留 FP32 master weights：低精度权重上过小的更新可能被直接舍掉。

8.4 非结合性和归约顺序

由于每步都舍入：

\operatorname{fl}(\operatorname{fl}(a+b)+c) \ne \operatorname{fl}(a+\operatorname{fl}(b+c))

并行归约使用树形加法，不同线程数、block 划分和原子操作顺序可能产生略不同结果。非确定性不一定代表实现错误，但必须在可接受误差内，并符合任务对复现性的要求。

提高归约精度的手段包括：

用更高精度累加；
成对/树形求和，避免极端尺度长期相加；
Kahan 等补偿求和（有额外指令成本）；
先局部归约，再合并部分和。

8.5 溢出、下溢和非有限值

overflow：结果绝对值超过最大有限值，常变成 $\pm\infty$ ；
underflow：结果太接近 0，进入 subnormal 或舍入为 0；
NaN：如 $0/0$ 、 $\infty-\infty$ 或非法运算产生；
inf 参与后续运算可能迅速传播为 NaN。

排查 loss 变成 NaN 时，应找到第一个非有限张量，而不是只盯最终 loss。常见源头有过大 logits、除零、负数开方、无效 mask、梯度溢出和错误的自定义 Kernel 边界。

8.6 灾难性消减

两个非常接近的大数相减时，高位有效数字抵消，剩余结果的相对误差可能巨大。例如用：

\operatorname{Var}(X) = \mathbb{E}[X^2]-\mathbb{E}[X]^2

计算方差时，两项可能都很大且接近，相减后精度损失严重。

更稳定的方法包括 two-pass variance 或 Welford 在线算法。后者逐步更新均值和平方离差：

n \leftarrow n+1

\delta = x_n-\mu_{n-1}

\mu_n = \mu_{n-1}+\frac{\delta}{n}

M_{2,n}=M_{2,n-1}+\delta(x_n-\mu_n)

总体方差为 $M_{2,n}/n$ 。Welford 还能合并不同分块的统计量，适合并行归约。

8.7 条件数：输入误差会被放大多少

问题的条件数描述输入微小扰动对输出的影响。对可逆矩阵，在某个一致范数下：

\kappa(\mathbf{A}) = \lVert\mathbf{A}\rVert \lVert\mathbf{A}^{-1}\rVert

条件数大表示问题病态：即使算法实现正确、输入只含很小舍入误差，输出也可能变化很大。

要区分：

问题本身是否病态；
所选算法是否数值稳定。

更多位数能缓解舍入误差，但不能从根本上消除病态问题的敏感性。

8.8 混合精度训练的基本模式

混合精度不是“把所有东西都换成 FP16”。常见模式是：

大型矩阵乘使用 FP16/BF16 输入，Tensor Core 加速；
乘加在 FP32 或更高的指定累加精度中完成；
归约、归一化、Softmax 等敏感操作保留或内部提升到 FP32；
优化器状态和/或 master weights 保留 FP32；
根据硬件和算子选择输出 dtype。

每个框架和硬件的具体策略不同，应查实际算子文档和 profiler，而不是从输入 dtype 猜内部计算路径。

8.9 Loss Scaling 为什么有用

反向传播中的 FP16 小梯度可能下溢为 0。令缩放因子为 $S$ ：

L' = S L

则：

\nabla L'=S\nabla L

梯度先被放大到 FP16 可表示范围，反向完成后再除以 $S$ ：

\nabla L=\frac{\nabla L'}{S}

动态 Loss Scaling 会：

用当前 scale 计算反向；
检查梯度是否含 inf/NaN；
若溢出，跳过更新并减小 scale；
连续稳定若干步后尝试增大 scale。

Loss Scaling 主要解决小梯度下溢，scale 过大反而会造成溢出。BF16 动态范围更大，通常不依赖 Loss Scaling，但仍可能因模型本身不稳定产生非有限值。

8.10 稳定算法通常也更适合融合

数值稳定与性能优化并不总是冲突：

Stable Softmax 先做 max 归约，再做 exp/sum 归约；Online Softmax 可合并分块状态；
log_softmax + NLL 融合减少中间概率写回；
Welford 同时维护统计量，适合 LayerNorm 的块级归约；
GEMM 用低精度输入、高精度累加兼顾吞吐与误差。

关键是选择具有可组合状态的数学形式，让 Kernel 能在有限片上存储中分块处理。

9. AI Infra 综合算例

9.1 算例一：Transformer 线性层的 shape 与代价

输入：

\mathbf{X}\in\mathbb{R}^{B\times S\times H}

权重：

\mathbf{W}\in\mathbb{R}^{H\times 4H}

把前两维展平：

\mathbf{X}'\in\mathbb{R}^{(BS)\times H}

输出：

\mathbf{Y}'=\mathbf{X}'\mathbf{W} \in\mathbb{R}^{(BS)\times 4H}

再 reshape 为 $(B,S,4H)$ 。FLOPs：

2\cdot(BS)\cdot H\cdot4H =8BSH^2

若 $B=8,S=2048,H=4096$ ：

8\times8\times2048\times4096^2 \approx 2.20\times10^{12}\ \text{FLOPs}

这只是一次前向线性投影，说明大模型为何高度依赖 Tensor Core GEMM。

9.2 算例二：多头 Attention 的完整维度

从：

\mathbf{X}\in\mathbb{R}^{B\times S\times H}

生成 Q/K/V 后 reshape 并转置：

\mathbf{Q},\mathbf{K},\mathbf{V} \in\mathbb{R}^{B\times N_h\times S\times D_h}

其中 $H=N_hD_h$ 。

分数矩阵：

\mathbf{S} =\frac{\mathbf{Q}\mathbf{K}^\top}{\sqrt{D_h}} \in\mathbb{R}^{B\times N_h\times S\times S}

应用因果 mask 和 Softmax：

\mathbf{P}=\operatorname{softmax}(\mathbf{S}+\mathbf{M})

其中未来位置的 mask 逻辑上为 $-\infty$ ，使其 Softmax 概率为 0。

输出：

\mathbf{O}=\mathbf{P}\mathbf{V} \in\mathbb{R}^{B\times N_h\times S\times D_h}

合并 head 后回到 $(B,S,H)$ 。

若显式物化 FP16 的 $\mathbf{S}$ ，仅它就需要：

2BN_hS^2\ \text{bytes}

当 $B=1,N_h=32,S=32768$ 时：

2\times1\times32\times32768^2 =64\ \text{GiB}

这还没算概率、反向中间量和其他层。FlashAttention 的关键就是不把完整 $S\times S$ 中间矩阵写回 HBM，而在片上分块完成稳定 Softmax 与 $PV$ 累积。

9.3 算例三：为什么除以 $\sqrt{D_h}$

假设 $q_i,k_i$ 独立、均值 0、方差 1：

\mathbf{q}^\top\mathbf{k} =\sum_{i=1}^{D_h}q_i k_i

每项 $q_i k_i$ 期望约为 0、方差约为 1，则和的方差约为 $D_h$ ，标准差约为 $\sqrt{D_h}$ 。

除以 $\sqrt{D_h}$ 后，分数方差恢复到约 1：

\operatorname{Var}\left( \frac{\mathbf{q}^\top\mathbf{k}}{\sqrt{D_h}} \right)\approx1

避免维度增大时 logits 过度扩张，Softmax 进入极端饱和区。

9.4 算例四：Online Softmax 的分块合并

一行 logits 被切成多个 block。已处理部分的最大值和指数和为 $(m,l)$ ：

m=\max_i x_i,\qquad l=\sum_i e^{x_i-m}

新 block 的状态为 $(m_b,l_b)$ 。合并最大值：

m_{new}=\max(m,m_b)

调整到同一基准后合并指数和：

l_{new} = e^{m-m_{new}}l +e^{m_b-m_{new}}l_b

这使 Softmax 可以分块、流式且数值稳定地计算。若同时维护加权 Value 累加器，还能避免保存完整注意力矩阵，这是 FlashAttention 的数学核心之一。

9.5 算例五：LoRA 的参数与计算

完整线性层：

\mathbf{y}=\mathbf{x}\mathbf{W}_0^\top

LoRA：

\mathbf{y} =\mathbf{x}\mathbf{W}_0^\top +\frac{\alpha}{r} (\mathbf{x}\mathbf{A}^\top)\mathbf{B}^\top

形状：

x:     (..., din)
A^T:   (din, r)       -> (..., r)
B^T:   (r, dout)      -> (..., dout)

在线 adapter 会多做两个小 GEMM；若提前把 $\mathbf{B}\mathbf{A}$ 合并进基座权重，推理仍是一个大 GEMM，但失去低成本动态切换 adapter 的便利。

9.6 算例六：分布式均值为什么要加权

两个 rank 的局部样本数分别为 $n_1,n_2$ ，局部均值为 $\mu_1,\mu_2$ 。全局均值不是简单的 $(\mu_1+\mu_2)/2$ ，除非 $n_1=n_2$ 。正确值：

\mu= \frac{n_1\mu_1+n_2\mu_2}{n_1+n_2}

更一般地，每个 rank 先汇总局部 sum 和 count，AllReduce 后再相除。分布式 loss、指标和吞吐统计都要明确分母，尤其是最后一个不完整 batch 或序列 mask 不同的场景。

9.7 算例七：显存估算

若一个张量 shape 为 $(B,S,H)$ ，元素字节数为 $b$ ，其连续存储大小：

M=B\times S\times H\times b

例如 BF16 的 $(8,4096,8192)$ ：

8\times4096\times8192\times2 =536{,}870{,}912\ \text{bytes} =512\ \text{MiB}

训练峰值显存还包括参数、梯度、优化器状态、保存的激活、临时 workspace、通信 buffer、内存碎片和框架上下文。单个张量估算只是起点。

9.8 从数学式到 Kernel 的固定检查表

看到一个新算子时，按以下顺序拆解：

输入、输出和中间量的 shape 是什么？
哪些维度保留，哪些维度归约？
能否分块？分块状态如何合并？
FLOPs 和理论最小访存量是多少？
是否存在广播、转置或非连续 stride？
哪些操作对精度敏感，需要高精度累加？
是否会出现 exp 溢出、除零、消减或长归约误差？
中间张量能否融合消除？
边界 shape 和尾块如何处理？
用什么参考实现、dtype 容差和极端输入验证？

这份检查表把抽象数学转换成可执行的 Infra 工作流。

总结

本章建立了四条相互连接的主线：

线性代数描述张量如何变换：shape 决定合法性，秩和分解决定压缩可能，分块决定数据如何复用；
概率论描述模型输出的含义：logits 经 Softmax 变成分布，交叉熵对应负对数似然，KL 衡量分布差异；
微积分描述参数如何学习：链式法则沿计算图反向传播，线性层反向仍由 GEMM 和归约构成；
数值计算决定公式在硬件上是否可靠：有限精度会舍入、溢出和下溢，稳定公式与混合精度策略同样重要。

对 AI Infra 工程师而言，“懂数学”最终体现在：能把公式翻译成 shape、归约、分块、数据移动、精度选择和验证标准。

自我检验清单

练习题

推导 (B, Nh, Sq, Dh) @ (B, Nh, Dh, Sk) 的输出 shape，并标出 batch 维与归约维。
给定 $M=4096,K=4096,N=16384$ 的 FP16 GEMM，估算 FLOPs 和理想最小输入输出字节数。
写程序验证 $\operatorname{softmax}(\mathbf{z}+c)=\operatorname{softmax}(\mathbf{z})$ ，并比较朴素实现与稳定实现处理大 logits 的结果。
从 Softmax Jacobian 和交叉熵开始，手推 $\partial L/\partial z_i=p_i-y_i$ 。
推导 Y = GELU(XW + b) 的反向 shape，不要求展开 GELU 的具体导数。
比较 FP16、BF16 和 FP32 对序列 [1, 1e-3, 1e-3, ...] 求和的误差，并尝试改变求和顺序。
实现 Welford 方差，与 E[x²] - E[x]² 在“大均值、小方差”数据上的结果比较。
对一个 $4096\times4096$ 权重计算 rank 8、16、64 LoRA 的参数比例和额外前向 FLOPs。
计算 $B=2,N_h=32,S=8192,D_h=128$ 时显式 FP16 Attention score 矩阵的大小。
两个 rank 分别处理 7 个和 5 个有效 token，局部 loss 均值分别为 2.0 和 3.0；计算正确全局均值，并说明简单平均为什么错。
用有限差分检查一个小线性层手写反向，尝试不同 $\epsilon$ 并观察误差曲线。
选择一个 PyTorch 算子，按“shape、归约、FLOPs、访存、数值风险、验证方法”写一页分析。

搜索